Cast attention inputs to the autocast dtype in dispatch_attention_fn

[danielxmed]

What does this PR do?

Fixes #14104.

Starting with torch 2.12, aten::rms_norm registers an fp32 cast policy for AutocastCUDA/AutocastXPU. Bisected via torch._C._dispatch_dump("aten::rms_norm"): absent on 2.6 through 2.11, present on 2.12, and verified behaviorally on GPU for 2.8 vs 2.12. The policy is intentional on the torch side (pytorch/pytorch#174824, a numerics fix), currently debated as a memory/throughput regression in pytorch/pytorch#188955, but torch 2.12 stable ships with it either way. Under an active torch.autocast context the torch.nn.RMSNorm QK norms used by Flux-family attention processors therefore return float32 even for bf16/fp16 inputs, while value, which never passes through a norm, keeps the autocast dtype. dispatch_attention_fn forwarded that fp32/bf16 mix untouched:

• native backends worked: aten::scaled_dot_product_attention has its own autocast registration that downcasts its autocast-eligible floating inputs;
• external-kernel backends crashed: flash, flash_hub, sage, aiter, xformers, npu, xla call the kernel directly, outside autocast's dispatch handling. flash-attn dies with RuntimeError: FlashAttention only support fp16 and bf16 data type, and flex rejects the dtype mix as well;
• even the opt-in DIFFUSERS_ATTN_CHECKS dtype checks false-raised on the native backend under autocast, although the op itself would have run fine.

The affected pattern (torch.nn.RMSNorm QK-norm, then RoPE, then dispatch_attention_fn) appears in 18 model files: Flux, Flux2, Chroma, Bria, BriaFibo, OvisImage, LongCatImage, LTX, LTX2, Wan, WanAnimate, ChronoEdit, SkyReelsV2, Helios, JoyImage, JoyImage Edit Plus (merged July 8, a week after the issue was filed), MotifVideo, ErnieImage, plus the Anima condition embedder. So this fixes the dispatcher once rather than patching every processor: when autocast is active for the tensors' device type, dispatch_attention_fn now casts q/k/v and floating-point masks to torch.get_autocast_dtype(device_type), mirroring SDPA's own autocast registration, including its eligibility rule (float64 tensors and bool/int masks pass through untouched). Since dispatch_attention_fn is the documented drop-in for F.scaled_dot_product_attention, this is parity with the op it replaces, not new policy:

• numerically a no-op for the native backend (SDPA performed this exact cast internally);
• provably dead branch outside autocast: mismatched dtypes still raise exactly as before, covered by a test;
• also repairs the context-parallel _native_cudnn/_native_flash paths, which call the decomposed aten ops directly and bypassed the composite's autocast handling;
• anyone who wants fp32 attention inside an autocast region uses torch.autocast(..., enabled=False), same as with raw SDPA.

Two implementation details worth flagging for review:

• torch.is_autocast_enabled(device_type) raises for device types without an autocast dispatch key. The only such device type diffusers tensors realistically report is meta (meta-device forward works on main and keeps working), so the guard is a literal device_type != "meta" check. Everything left in the branch is Dynamo-traceable on torch 2.6, the repo floor: torch.compile(dispatch_attention_fn, fullgraph=True) traces with zero graph breaks, autocast on and off, verified on 2.6.0 and 2.12.1 and locked in by a @is_torch_compile regression test.
• Autocast is a first-party workflow (81 files under examples/ use it, including Flux LoRA training validation), and this failure hits any trainer that selects a non-native attention backend. The issue was reported from OneTrainer, which has since made mixed precision always-on (always enable mixed-precision / autocasting Nerogar/OneTrainer#1481), so their whole user base lands on this path once they move to torch 2.12.

Tests

New TestAttentionDispatchAutocast (plain pytest style, gated behind require_torch_accelerator so it stays off the fast PR CI) in tests/models/test_attention_processor.py, the first unit tests for dispatch_attention_fn:

• test_autocast_casts_inputs_to_autocast_dtype: mixed fp32 q/k + bf16 v under autocast, with the DIFFUSERS_ATTN_CHECKS validation enabled: the backend fn must receive uniform bf16 and the checks must not false-raise (both fail on main);
• test_autocast_casts_floating_point_mask_but_not_bool_mask (fails on main);
• test_no_autocast_dtypes_pass_through_unchanged: outside autocast the dispatcher still doesn't touch dtypes, the mismatch keeps raising SDPA's exact "same dtype" message;
• test_rms_norm_qk_pattern_under_autocast (skipif CUDA/XPU only): end-to-end shape of the bug through a real torch.nn.RMSNorm, asserting the fp32 upcast actually happens (guards against PyTorch policy changes hollowing the test) and that the backend receives bf16;
• test_torch_compile_fullgraph_under_autocast (@is_torch_compile): the autocast branch stays fullgraph-traceable.

Validation:

pytest tests/models/test_attention_processor.py::TestAttentionDispatchAutocast -v
  # CUDA (RTX A4000): 5 passed on torch 2.12.1+cu126, 5 passed on torch 2.13.0+cu126
  # CPU boxes: the class skips; with the accelerator gate lifted, the CPU-runnable
  # tests pass on torch 2.6.0 and 2.12.1, fail on main
pytest tests/models/transformers/test_models_transformer_flux.py \
       tests/models/transformers/test_models_transformer_flux2.py \
       tests/models/transformers/test_models_transformer_chroma.py  # 93 passed, 358 skipped
ruff check / ruff format --check / utils/check_copies.py  # clean

GPU validation (CUDA, real flash-attn2 kernels, RunPod RTX A4000, torch 2.12.1+cu126):

main:  FluxAttention forward under cuda autocast, backend=flash_hub
       -> RuntimeError: FlashAttention only support fp16 and bf16 data type
fixed: RMSNorm(bf16) under cuda autocast -> torch.float32 (policy confirmed on CUDA)
       FluxAttention forward with backend=flash_hub: OK, out=torch.bfloat16
       max |flash_hub - native| under autocast: < 1e-5
       TestAttentionDispatchAutocast on CUDA: 5 passed
       meta-device forward through the dispatcher: works, with and without autocast active

Limitations

• Backends that only support fp16/bf16 still (correctly) raise if the autocast dtype itself can't be handled by the kernel. Unchanged semantics, same as native.
• The rms_norm fp32 policy is CUDA/XPU-only, so the CPU regression tests construct the dtype mix explicitly; the accelerator-gated test covers the real policy path.

Before submitting

• Did you use an AI agent to help with this PR? Yes, Claude Code.
  • Read the "Coding with AI agents" guide.
  • Self-reviewed the diff against .ai/review-rules.md.
• Read the contributor guideline.
• Coordination: maintainer go-ahead in Flux-family attention: flash-attn and other backends fail in an autocast context #14104 (comment) (full analysis + proposal in Flux-family attention: flash-attn and other backends fail in an autocast context #14104 (comment))
• Wrote new necessary tests.
• Documentation: no doc changes needed. docs/source/en/optimization/attention_backends.md describes backend selection and opt-in checks, neither of which changes; autocast now simply behaves consistently across backends.

Who can review?

@sayakpaul (attention dispatcher) @DN6 (recent attention_dispatch.py fixes)

[sayakpaul]

Left two high-level comments.

[sayakpaul]

Why does it have to depend on torch.compiler.is_compiling()?

[danielxmed]

Two stacked constraints. In eager, torch.is_autocast_enabled(device_type) raises for device types without an autocast dispatch key ("meta" is the one diffusers actually hits, e.g. during meta-device init), so it needs the torch.amp.is_autocast_available(device_type) guard. But Dynamo cannot trace that guard before torch 2.12 (it only entered the trace rules there): on torch 2.6, the floor in setup.py, compiling this function with fullgraph=True fails with

Unsupported: Graph break due to unsupported builtin torch._C._is_autocast_available

torch.compiler.is_compiling() resolves both sides: it folds to a constant inside the graph, so compiled code skips the untraceable guard, and nothing is lost because compiled graphs never execute on meta anyway. Eager keeps the guard. test_torch_compile_fullgraph_under_autocast locks the fullgraph behavior in; verified on 2.6.0 and 2.12.1.

Once the minimum torch reaches 2.12 the short-circuit can be dropped and the condition becomes just is_autocast_available(...) and is_autocast_enabled(...). Happy to restructure if you would rather have it factored into a small helper.

[sayakpaul]

This all feels very very LLM-ish. Could you please cut the jargon and explain things out of your own understanding? It's just easier for everyone to follow. Otherwise, it's too much unnecessary cognitive load.

[danielxmed]

Sure, plain version: torch.is_autocast_enabled crashes on device types without autocast support ("meta" is the real case). The is_autocast_available check protects against that, but torch 2.6, our minimum, can't trace that check under torch.compile(fullgraph=True). So is_compiling() skips it inside compiled code, which is safe: compiled graphs never execute on meta tensors. Eager keeps the protection, compile skips the untraceable call. Once the minimum torch is 2.12, this collapses to just the two autocast calls.

I also tightened the comment block in the code itself (049feef), same story in fewer lines.

One note: happy to simplify whenever it helps, and a plain "please make this simpler" does the job. The "LLM-ish" framing wasn't necessary.

[sayakpaul]

So, we are saying:

Torch SDPA casts its inputs to the autocast dtype internally, but backends calling external kernels

How does SDPA handle it in a manner that doesn't offend Dynamo?

torch.is_autocast_enabled crashes on device types without autocast support ("meta" is the real case).

But is there any real merit to consider that? One would probably never initialize a model on a meta device and then switch between attention backends, I think. If we consider that, does the following still apply?

but torch 2.6, our minimum, can't trace that check under torch.compile(fullgraph=True).

[danielxmed]

Fair challenge, and it made the code simpler.

On SDPA: it never checks anything at the Python level. scaled_dot_product_attention has autocast kernels registered in C++ (torch._C._dispatch_dump("aten::scaled_dot_product_attention") lists AutocastCUDA/CPU/XPU/MPS entries from aten/src/ATen/autocast_mode.cpp), so the dispatcher casts the inputs before Dynamo ever sees a branch. When AOTAutograd traces through the dispatcher, the casts just show up in the graph as ordinary ops. Our cast lives in Python, which is why traceability became a topic at all.

On meta: agreed that nobody switches attention backends on a meta model. I swept the codebase to be sure, and every meta usage in src/ and tests/ is weight init or offload bookkeeping, none reaches attention. The catch is that torch.is_autocast_enabled("meta") raises whether autocast is active or not, so an unguarded call would break plain meta-device forwards (shape or memory tracing) that work on main today. That protection doesn't require is_autocast_available, though. A string check covers it, and the guard is now

if device_type != "meta" and torch.is_autocast_enabled(device_type):

Which also answers the last question: with is_autocast_available gone there is nothing left that Dynamo can't trace. is_autocast_enabled and get_autocast_dtype compile fine under fullgraph=True on torch 2.6 (zero graph breaks per torch._dynamo.explain, correct recompile when autocast toggles between calls), so is_compiling() is gone too. Done in 38a870d; re-validated on an A4000 (torch 2.12.1+cu126): flash_hub under cuda autocast passes, the full test class is green on CUDA, and the fullgraph test passes on 2.6.0 and 2.12.1.

[sayakpaul]

This should be implemented with pytest.

[danielxmed]

Done in d26452b: dropped unittest.TestCase for plain asserts, pytest.raises/pytest.skip, and the monkeypatch fixture for the checks toggle. Renamed the class to TestAttentionDispatchAutocast so default pytest collection picks it up. Re-ran on torch 2.6.0 and 2.12.1 (CPU): 5 passed, 1 GPU-skip on both.

[sergereview]

🤗 Serge says:

Clean, well-scoped fix. dispatch_attention_fn now mirrors SDPA's autocast input-casting so external-kernel backends (flash, sage, xformers, ...) stop receiving the fp32/bf16 mix produced by the torch.nn.RMSNorm QK-norm + autocast interaction on torch 2.12. The guards are sound: is_autocast_available correctly shields device types (e.g. meta) where is_autocast_enabled would raise, and the torch.compiler.is_compiling() short-circuit keeps the branch fullgraph-traceable on the 2.6 floor. All three autocast APIs accept the string device_type form since torch 2.4, so the floor is safe.

Tests

• The new AttentionDispatchAutocastTests cover the casting, mask eligibility (float vs bool), the opt-in dtype checks, the pass-through-outside-autocast invariant, the end-to-end RMSNorm pattern (accelerator-gated with a runtime probe), and fullgraph compile. Imports (is_torch_compile, require_torch_accelerator, registry internals) all resolve.

Minor (non-blocking)

• The float64 eligibility gate keys only off query.dtype != torch.float64 before casting all of q/k/v, whereas SDPA evaluates autocast eligibility per tensor. A genuinely mixed float64/bf16 q/k/v set would be handled slightly differently, but that combination doesn't arise in the models this targets, so it's not worth complicating the branch.

LGTM.

serge v0.1.0 · model: claude-opus-4-8 · 5 LLM turns · 5 tool calls · 36.3s · 75302 in / 2174 out tokens

[sayakpaul]

I would mark these tests with require_torch_accelerator because I don't think they need to run on our fast PR CI.

[danielxmed]

Done in e864233. The class is gated now, so all of it skips on the CPU-only fast PR CI and runs on the GPU workflows instead.

[sayakpaul]

Is the qk_type set to torch.float32 to emulate the QKnorm behaviour?

[danielxmed]

Yes, exactly. fp32 q/k with bf16 v is what the aten::rms_norm policy produces on the QK-norm path, and value keeps the autocast dtype because it never passes through the norm. The helper now says this in a comment (e864233).

[sayakpaul]

This is already tested in test_autocast_casts_inputs_to_autocast_dtype

[danielxmed]

They overlapped on the happy path, agreed. Merged in e864233: the checks toggle now lives inside test_autocast_casts_inputs_to_autocast_dtype and the standalone test is gone. What it guarded is the opt-in DIFFUSERS_ATTN_CHECKS validation, which runs before the backend is invoked and false-raised under autocast on main; keeping the toggle enabled in the surviving test covers that with one test fewer.

[sayakpaul]

I would even be specific and check that there's a dtype mismatch error in the RuntimeError message (assuming that's the case).

[danielxmed]

Done in e864233: match="Expected query, key, and value to have the same dtype". That's SDPA's shared input validation message, identical on torch 2.6.0 and 2.12.1, and the same on CUDA (checked on an A4000).

[sayakpaul]

(nit): Would use pytest.mark.skipif and use if torch_device not in ["cuda", "xpu"] as the reason.

[danielxmed]

Done in e864233. The runtime probe right below it stays: the fp32 policy also depends on the torch version (new in 2.12), so the version half can't be a static skipif.

[sayakpaul]

Cool test!

[sayakpaul]

Left further comments and apologies if my statement on LLM-ish framing came off as offensive. Thanks for speaking your mind!

[danielxmed]

No offense taken, and thanks for saying it. All good.

Round summary, everything pushed:

• The guard is now device_type != "meta" and torch.is_autocast_enabled(device_type). Both torch.compiler.is_compiling() and torch.amp.is_autocast_available are gone; reasoning in the thread.
• Tests: class gated behind require_torch_accelerator (fast PR CI skips it), the checks test merged into the first test, pytest.raises asserts the exact dtype-mismatch message, skipif nit applied.
• Re-validated on an RTX A4000 (torch 2.12.1+cu126): flash_hub forward under cuda autocast passes with max |flash_hub - native| < 1e-5, the full class is green on CUDA, meta-device forward keeps working, and the fullgraph compile test passes on torch 2.6.0 and 2.12.1.