Cast attention inputs to the autocast dtype in dispatch_attention_fn

src/diffusers/models/attention_dispatch.py

@@ -408,6 +408,19 @@ def dispatch_attention_fn(
 ) -> torch.Tensor:
     attention_kwargs = attention_kwargs or {}
 
+    # Under autocast, fp32-cast-policy ops (like the RMSNorm QK norms in Flux models) return fp32
+    # while `value` stays bf16/fp16. SDPA casts its inputs inside its C++ autocast kernel; backends
+    # calling external kernels (flash-attn, sage, ...) crash on the mix, so do the same cast here.
+    # float64 is not autocast-eligible; "meta" has no autocast key and `is_autocast_enabled` raises.
+    # See https://github.com/huggingface/diffusers/issues/14104.
+    device_type = query.device.type
+    if device_type != "meta" 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,125 @@ 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))
+
+
+@require_torch_accelerator
+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):
+        # The fp32 q/k + bf16 v defaults reproduce what the fp32 autocast policy for `aten::rms_norm`
+        # does to the QK-norm pattern: q/k leave the norm upcast, value never passes through it.
+        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, monkeypatch):
+        # Checks on: the opt-in DIFFUSERS_ATTN_CHECKS validation sees the tensors the backend will
+        # actually receive, so it must not raise on mixed autocast inputs (on main it false-raised).
+        monkeypatch.setattr(_AttentionBackendRegistry, "_checks_enabled", True)
+        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_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, match="Expected query, key, and value to have the same dtype"):
+            dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE)
+
+    @pytest.mark.skipif(
+        torch_device not in ["cuda", "xpu"], reason="aten::rms_norm has an fp32 autocast policy for CUDA/XPU only"
+    )
+    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 policy only exists on recent torch
+        # (and is debated upstream), so probe for it instead of assuming.
+        device_type = torch.device(torch_device).type
+
+        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