From 31f876cd56635745af250390ea4246b9dde1da0c Mon Sep 17 00:00:00 2001 From: Daniel Medeiros Date: Sun, 5 Jul 2026 08:24:00 -0300 Subject: [PATCH 1/5] Cast attention inputs to the autocast dtype in dispatch_attention_fn 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 via its own autocast registration, but backends that call external kernels (flash-attn, sage, xformers, ...) receive the tensors as-is and fail, e.g. "FlashAttention only support fp16 and bf16 data type". Mirror SDPA's autocast behavior at the dispatch boundary so every backend sees the same dtypes as the native one. Fixes #14104. Co-Authored-By: Claude Fable 5 --- src/diffusers/models/attention_dispatch.py | 20 +++ tests/models/test_attention_processor.py | 139 ++++++++++++++++++++- 2 files changed, 158 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 9414c151fd67..209e27bba275 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -408,6 +408,26 @@ def dispatch_attention_fn( ) -> torch.Tensor: attention_kwargs = attention_kwargs or {} + # Match the autocast behavior of `torch.nn.functional.scaled_dot_product_attention`: its autocast + # registration casts eligible floating-point inputs (autocast skips float64) to the autocast dtype + # before running. Under an active autocast context, ops with an fp32 cast policy (such as the + # `torch.nn.RMSNorm` QK norms in Flux-family models) return float32 even for bf16/fp16 inputs while + # `value` keeps the autocast dtype, and backends that call external kernels (flash-attn, sage, + # xformers, ...) bypass autocast's dispatch handling and fail on such inputs. Normalizing here keeps + # every backend consistent with the native backend. See https://github.com/huggingface/diffusers/issues/14104. + # `is_autocast_available` guards device types without an autocast dispatch key (e.g. "meta"), for + # which `is_autocast_enabled` raises. It is skipped when compiling: Dynamo cannot trace it before + # torch 2.12, and compiled graphs never run on such devices anyway. + 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 diff --git a/tests/models/test_attention_processor.py b/tests/models/test_attention_processor.py index 8b45c2148504..9b8749a46f1b 100644 --- a/tests/models/test_attention_processor.py +++ b/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,134 @@ 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 AttentionDispatchAutocastTests(unittest.TestCase): + """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) + + self.assertEqual(received["query"], torch.bfloat16) + self.assertEqual(received["key"], torch.bfloat16) + self.assertEqual(received["value"], torch.bfloat16) + self.assertEqual(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) + self.assertEqual(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) + self.assertEqual(received["attn_mask"], torch.bool) + + def test_autocast_passes_backend_dtype_checks(self): + # 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 + + original = _AttentionBackendRegistry._checks_enabled + _AttentionBackendRegistry._checks_enabled = True + try: + with torch.autocast(device_type, dtype=torch.bfloat16): + out = dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE) + finally: + _AttentionBackendRegistry._checks_enabled = original + self.assertEqual(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 self.assertRaises(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"): + self.skipTest("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: + self.skipTest("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) + + self.assertEqual(received["query"], torch.bfloat16) + self.assertEqual(received["key"], torch.bfloat16) + self.assertEqual(received["value"], torch.bfloat16) + self.assertEqual(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) + self.assertEqual(out.dtype, torch.bfloat16) + + torch.compiler.reset() + compiled = torch.compile(dispatch_attention_fn, fullgraph=True) + out = compiled(query, key, value) + self.assertEqual(out.dtype, torch.bfloat16) From d26452b84794da762a37791e0b1814a35d995aba Mon Sep 17 00:00:00 2001 From: Daniel Medeiros Date: Fri, 10 Jul 2026 09:41:22 -0300 Subject: [PATCH 2/5] Use plain pytest style for the autocast dispatch tests Requested in review: drop unittest.TestCase in favor of plain asserts, pytest.raises/pytest.skip, and the monkeypatch fixture. Renamed to TestAttentionDispatchAutocast so default pytest collection picks it up. Co-Authored-By: Claude Fable 5 --- tests/models/test_attention_processor.py | 46 +++++++++++------------- 1 file changed, 21 insertions(+), 25 deletions(-) diff --git a/tests/models/test_attention_processor.py b/tests/models/test_attention_processor.py index 9b8749a46f1b..5c20fcbe27ab 100644 --- a/tests/models/test_attention_processor.py +++ b/tests/models/test_attention_processor.py @@ -141,7 +141,7 @@ def test_conversion_when_using_device_map(self): self.assertTrue(np.allclose(conversion, after_conversion, atol=1e-3)) -class AttentionDispatchAutocastTests(unittest.TestCase): +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` @@ -182,10 +182,10 @@ def test_autocast_casts_inputs_to_autocast_dtype(self): with torch.autocast(device_type, dtype=torch.bfloat16): out = dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE) - self.assertEqual(received["query"], torch.bfloat16) - self.assertEqual(received["key"], torch.bfloat16) - self.assertEqual(received["value"], torch.bfloat16) - self.assertEqual(out.dtype, torch.bfloat16) + 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() @@ -196,33 +196,29 @@ def test_autocast_casts_floating_point_mask_but_not_bool_mask(self): 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) - self.assertEqual(received["attn_mask"], torch.bfloat16) + 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) - self.assertEqual(received["attn_mask"], torch.bool) + assert received["attn_mask"] == torch.bool - def test_autocast_passes_backend_dtype_checks(self): + 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 - original = _AttentionBackendRegistry._checks_enabled - _AttentionBackendRegistry._checks_enabled = True - try: - with torch.autocast(device_type, dtype=torch.bfloat16): - out = dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE) - finally: - _AttentionBackendRegistry._checks_enabled = original - self.assertEqual(out.dtype, torch.bfloat16) + 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 self.assertRaises(RuntimeError): + with pytest.raises(RuntimeError): dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE) @require_torch_accelerator @@ -232,24 +228,24 @@ def test_rms_norm_qk_pattern_under_autocast(self): # 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"): - self.skipTest("aten::rms_norm only registers an fp32 autocast policy for 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: - self.skipTest("this torch version has no fp32 autocast policy for aten::rms_norm.") + 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) - self.assertEqual(received["query"], torch.bfloat16) - self.assertEqual(received["key"], torch.bfloat16) - self.assertEqual(received["value"], torch.bfloat16) - self.assertEqual(out.dtype, torch.bfloat16) + 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): @@ -264,9 +260,9 @@ def test_torch_compile_fullgraph_under_autocast(self): compiled = torch.compile(dispatch_attention_fn, fullgraph=True) with torch.autocast(device_type, dtype=torch.bfloat16): out = compiled(query, key, value) - self.assertEqual(out.dtype, torch.bfloat16) + assert out.dtype == torch.bfloat16 torch.compiler.reset() compiled = torch.compile(dispatch_attention_fn, fullgraph=True) out = compiled(query, key, value) - self.assertEqual(out.dtype, torch.bfloat16) + assert out.dtype == torch.bfloat16 From 049feefb343d64777c71dce961e3040dd11a8fab Mon Sep 17 00:00:00 2001 From: Daniel Medeiros Date: Fri, 10 Jul 2026 10:00:09 -0300 Subject: [PATCH 3/5] Tighten the autocast comment in dispatch_attention_fn Co-Authored-By: Claude Fable 5 --- src/diffusers/models/attention_dispatch.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 209e27bba275..dbc6454aa9f0 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -408,16 +408,12 @@ def dispatch_attention_fn( ) -> torch.Tensor: attention_kwargs = attention_kwargs or {} - # Match the autocast behavior of `torch.nn.functional.scaled_dot_product_attention`: its autocast - # registration casts eligible floating-point inputs (autocast skips float64) to the autocast dtype - # before running. Under an active autocast context, ops with an fp32 cast policy (such as the - # `torch.nn.RMSNorm` QK norms in Flux-family models) return float32 even for bf16/fp16 inputs while - # `value` keeps the autocast dtype, and backends that call external kernels (flash-attn, sage, - # xformers, ...) bypass autocast's dispatch handling and fail on such inputs. Normalizing here keeps - # every backend consistent with the native backend. See https://github.com/huggingface/diffusers/issues/14104. - # `is_autocast_available` guards device types without an autocast dispatch key (e.g. "meta"), for - # which `is_autocast_enabled` raises. It is skipped when compiling: Dynamo cannot trace it before - # torch 2.12, and compiled graphs never run on such devices anyway. + # 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 From 38a870dfc2ddaddd8e675e3836fa91628cc09ac1 Mon Sep 17 00:00:00 2001 From: Daniel Medeiros Date: Fri, 10 Jul 2026 10:56:57 -0300 Subject: [PATCH 4/5] Simplify the autocast guard to a literal meta-device check --- src/diffusers/models/attention_dispatch.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index dbc6454aa9f0..f1969ccc0b0e 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -408,16 +408,13 @@ 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). + # 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 (torch.compiler.is_compiling() or torch.amp.is_autocast_available(device_type)) and torch.is_autocast_enabled( - 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) From e864233a94e4a11c34eafe0006cedb56dea5e90b Mon Sep 17 00:00:00 2001 From: Daniel Medeiros Date: Fri, 10 Jul 2026 10:57:00 -0300 Subject: [PATCH 5/5] Gate the dispatcher autocast tests behind require_torch_accelerator and tighten them --- tests/models/test_attention_processor.py | 35 ++++++++++-------------- 1 file changed, 15 insertions(+), 20 deletions(-) diff --git a/tests/models/test_attention_processor.py b/tests/models/test_attention_processor.py index 5c20fcbe27ab..6fc7fb0a5357 100644 --- a/tests/models/test_attention_processor.py +++ b/tests/models/test_attention_processor.py @@ -141,17 +141,20 @@ def test_conversion_when_using_device_map(self): 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 + 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) @@ -174,7 +177,10 @@ def recording_backend(query, key, value, **kwargs): finally: _AttentionBackendRegistry._backends[AttentionBackendName.NATIVE] = original - def test_autocast_casts_inputs_to_autocast_dtype(self): + 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 @@ -203,32 +209,21 @@ def test_autocast_casts_floating_point_mask_but_not_bool_mask(self): 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): + with pytest.raises(RuntimeError, match="Expected query, key, and value to have the same dtype"): dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE) - @require_torch_accelerator + @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 fp32 autocast policy for `aten::rms_norm` - # only registers for CUDA/XPU and only on recent torch, so probe for it instead of assuming. + # 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 - 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) @@ -249,7 +244,7 @@ def test_rms_norm_qk_pattern_under_autocast(self): @is_torch_compile def test_torch_compile_fullgraph_under_autocast(self): - # The autocast branch must stay fullgraph-traceable — probing autocast state with APIs that + # 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,