Cast attention inputs to the autocast dtype in dispatch_attention_fn

src/diffusers/models/attention_dispatch.py

@@ -408,6 +408,22 @@ def dispatch_attention_fn(
 ) -> torch.Tensor:
     attention_kwargs = attention_kwargs or {}
 
+    # Under autocast, RMSNorm QK norms return fp32 while value stays bf16/fp16. Torch SDPA casts
+    # its inputs to the autocast dtype internally, but backends calling external kernels
+    # (flash-attn, sage, ...) receive the mix as-is and crash, so do the same cast here.
+    # float64 is not autocast-eligible. See https://github.com/huggingface/diffusers/issues/14104.
+    # is_compiling() short-circuits the is_autocast_available guard (it protects eager runs on
+    # devices without autocast, e.g. "meta"; Dynamo can't trace it before torch 2.12).
+    device_type = query.device.type
+    if (torch.compiler.is_compiling() or torch.amp.is_autocast_available(device_type)) and torch.is_autocast_enabled(
+        device_type
+    ):
+        autocast_dtype = torch.get_autocast_dtype(device_type)
+        if query.dtype != torch.float64:
+            query, key, value = query.to(autocast_dtype), key.to(autocast_dtype), value.to(autocast_dtype)
+        if attn_mask is not None and torch.is_floating_point(attn_mask) and attn_mask.dtype != torch.float64:
+            attn_mask = attn_mask.to(autocast_dtype)
+
     if backend is None:
         # If no backend is specified, we either use the default backend (set via the DIFFUSERS_ATTN_BACKEND environment
         # variable), or we use a custom backend based on whether user is using the `attention_backend` context manager

tests/models/test_attention_processor.py

@@ -1,3 +1,4 @@
+import contextlib
 import importlib.metadata
 import tempfile
 import unittest
@@ -8,9 +9,14 @@
 from packaging import version
 
 from diffusers import DiffusionPipeline
+from diffusers.models.attention_dispatch import (
+    AttentionBackendName,
+    _AttentionBackendRegistry,
+    dispatch_attention_fn,
+)
 from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor
 
-from ..testing_utils import torch_device
+from ..testing_utils import is_torch_compile, require_torch_accelerator, torch_device
 
 
 class AttnAddedKVProcessorTests(unittest.TestCase):
@@ -133,3 +139,130 @@ def test_conversion_when_using_device_map(self):
 
         self.assertTrue(np.allclose(pre_conversion, conversion, atol=1e-3))
         self.assertTrue(np.allclose(conversion, after_conversion, atol=1e-3))
+
+
+class TestAttentionDispatchAutocast:
+    """Regression tests for https://github.com/huggingface/diffusers/issues/14104.
+
+    Under an active torch.autocast context, ops with an fp32 cast policy (e.g. the `torch.nn.RMSNorm`
+    QK norms in Flux-family models) return float32 while `value` keeps the autocast dtype. Native SDPA
+    downcasts all of its inputs itself, but backends calling external kernels (flash-attn, sage, ...)
+    receive the tensors as-is and reject the mismatch — so `dispatch_attention_fn` must normalize
+    dtypes the same way SDPA's autocast registration does.
+    """
+
+    def _qkv(self, qk_dtype=torch.float32, v_dtype=torch.bfloat16):
+        query = torch.randn(1, 8, 2, 16, device=torch_device, dtype=qk_dtype)
+        key = torch.randn(1, 8, 2, 16, device=torch_device, dtype=qk_dtype)
+        value = torch.randn(1, 8, 2, 16, device=torch_device, dtype=v_dtype)
+        return query, key, value
+
+    @contextlib.contextmanager
+    def _record_native_backend_input_dtypes(self):
+        original = _AttentionBackendRegistry._backends[AttentionBackendName.NATIVE]
+        received = {}
+
+        def recording_backend(query, key, value, **kwargs):
+            received["query"], received["key"], received["value"] = query.dtype, key.dtype, value.dtype
+            if kwargs.get("attn_mask") is not None:
+                received["attn_mask"] = kwargs["attn_mask"].dtype
+            return original(query, key, value, **kwargs)
+
+        _AttentionBackendRegistry._backends[AttentionBackendName.NATIVE] = recording_backend
+        try:
+            yield received
+        finally:
+            _AttentionBackendRegistry._backends[AttentionBackendName.NATIVE] = original
+
+    def test_autocast_casts_inputs_to_autocast_dtype(self):
+        query, key, value = self._qkv()
+        device_type = torch.device(torch_device).type
+
+        with self._record_native_backend_input_dtypes() as received:
+            with torch.autocast(device_type, dtype=torch.bfloat16):
+                out = dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE)
+
+        assert received["query"] == torch.bfloat16
+        assert received["key"] == torch.bfloat16
+        assert received["value"] == torch.bfloat16
+        assert out.dtype == torch.bfloat16
+
+    def test_autocast_casts_floating_point_mask_but_not_bool_mask(self):
+        query, key, value = self._qkv()
+        device_type = torch.device(torch_device).type
+        float_mask = torch.zeros(1, 2, 8, 8, device=torch_device, dtype=torch.float32)
+        bool_mask = torch.ones(1, 2, 8, 8, device=torch_device, dtype=torch.bool)
+
+        with self._record_native_backend_input_dtypes() as received:
+            with torch.autocast(device_type, dtype=torch.bfloat16):
+                dispatch_attention_fn(query, key, value, attn_mask=float_mask, backend=AttentionBackendName.NATIVE)
+        assert received["attn_mask"] == torch.bfloat16
+
+        with self._record_native_backend_input_dtypes() as received:
+            with torch.autocast(device_type, dtype=torch.bfloat16):
+                dispatch_attention_fn(query, key, value, attn_mask=bool_mask, backend=AttentionBackendName.NATIVE)
+        assert received["attn_mask"] == torch.bool
+
+    def test_autocast_passes_backend_dtype_checks(self, monkeypatch):
+        # The opt-in debugging checks (DIFFUSERS_ATTN_CHECKS) validate the tensors the backend will
+        # actually receive, so mixed autocast inputs must not trip them.
+        query, key, value = self._qkv()
+        device_type = torch.device(torch_device).type
+
+        monkeypatch.setattr(_AttentionBackendRegistry, "_checks_enabled", True)
+        with torch.autocast(device_type, dtype=torch.bfloat16):
+            out = dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE)
+        assert out.dtype == torch.bfloat16
+
+    def test_no_autocast_dtypes_pass_through_unchanged(self):
+        # Outside autocast the dispatcher must not touch dtypes: mismatched inputs keep raising
+        # (from SDPA itself for the native backend), exactly as before.
+        query, key, value = self._qkv()
+        with pytest.raises(RuntimeError):
+            dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE)
+
+    @require_torch_accelerator
+    def test_rms_norm_qk_pattern_under_autocast(self):
+        # End-to-end shape of the Flux-family bug: bf16 q/k go through `torch.nn.RMSNorm` under
+        # autocast and come out fp32, value stays bf16. The fp32 autocast policy for `aten::rms_norm`
+        # only registers for CUDA/XPU and only on recent torch, so probe for it instead of assuming.
+        device_type = torch.device(torch_device).type
+        if device_type not in ("cuda", "xpu"):
+            pytest.skip("aten::rms_norm only registers an fp32 autocast policy for CUDA/XPU.")
+
+        norm = torch.nn.RMSNorm(16, eps=1e-6, device=torch_device, dtype=torch.bfloat16)
+        query, key, value = self._qkv(qk_dtype=torch.bfloat16)
+
+        with torch.autocast(device_type, dtype=torch.bfloat16):
+            if norm(query).dtype != torch.float32:
+                pytest.skip("this torch version has no fp32 autocast policy for aten::rms_norm.")
+
+        with self._record_native_backend_input_dtypes() as received:
+            with torch.autocast(device_type, dtype=torch.bfloat16):
+                query, key = norm(query), norm(key)
+                out = dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE)
+
+        assert received["query"] == torch.bfloat16
+        assert received["key"] == torch.bfloat16
+        assert received["value"] == torch.bfloat16
+        assert out.dtype == torch.bfloat16
+
+    @is_torch_compile
+    def test_torch_compile_fullgraph_under_autocast(self):
+        # The autocast branch must stay fullgraph-traceable — probing autocast state with APIs that
+        # Dynamo cannot trace would graph-break every compiled model at each attention call.
+        # `backend` is left unset (default active backend, i.e. native): passing the enum explicitly
+        # takes the `AttentionBackendName(backend)` path, which Dynamo cannot trace before torch 2.12,
+        # and model processors compile with `backend=None` anyway.
+        query, key, value = self._qkv(qk_dtype=torch.bfloat16)
+        device_type = torch.device(torch_device).type
+
+        compiled = torch.compile(dispatch_attention_fn, fullgraph=True)
+        with torch.autocast(device_type, dtype=torch.bfloat16):
+            out = compiled(query, key, value)
+        assert out.dtype == torch.bfloat16
+
+        torch.compiler.reset()
+        compiled = torch.compile(dispatch_attention_fn, fullgraph=True)
+        out = compiled(query, key, value)
+        assert out.dtype == torch.bfloat16