Flux-family attention: flash-attn and other backends fail in an autocast context

### Describe the bug

- q and k are normed before attention in flux-style attention.
- in an autocast context, norms are calculated in float32 no matter their input - and are returned as float32
- the result gets passed on to attention
- no problem for torch scaled_dot_product_attention, it can handle float32 - doesn't fail, and still executes in low precision in an autocast context
- other attention backends don't though. even though there are checks (if activated), it still fails with a check then.

Not sure about the solution.

1. avoid propagation to float32 for norms, as some other functions do?
2. autocast-like behaviour in the attention backend to accept float32 there, even though flash-attn itself doesn't?

Details:

`FluxAttnProcessor.__call__` (and the equivalent processors for Flux2, Chroma, Krea2, etc., which share the same `norm_q`/`norm_k` + `apply_rotary_emb` + `dispatch_attention_fn` pattern) can hand `dispatch_attention_fn` a `query`/`key` that ended up in `float32` even though the model is meant to run in `bfloat16`, whenever the whole forward pass runs inside an active `torch.autocast` context (e.g. this happens routinely in training code that mixes weight dtypes across parts of the model).

Root cause, in order:

1. `attn.norm_q` / `attn.norm_k` are `torch.nn.RMSNorm`. Under an active `torch.autocast`, `RMSNorm` is on PyTorch's "compute in fp32" autocast policy list, so it returns an `fp32` tensor even when its input was `bf16`. This part is intentional, documented PyTorch autocast behavior, not itself a bug.
2. `apply_rotary_emb` [...] is a no-op in terms of dtype
3. `value` never goes through a norm or RoPE, so it stays `bf16`.
4. `dispatch_attention_fn` is called with `query`/`key` in `fp32` and `value` in `bf16`.

For the `native`/`_native_cudnn`/`_native_efficient`/`_native_flash` (torch-SDPA-based) backends this is harmless: `torch.nn.functional.scaled_dot_product_attention` is itself a PyTorch op with a registered autocast policy, so under an active autocast context it downcasts **all** of its inputs to the autocast dtype before running, regardless of what dtype they arrive in (verified below — this holds even under an explicit `sdpa_kernel()` backend restriction).

But several other backends call an external kernel function directly (`flash_attn_func` for `AttentionBackendName.FLASH`, and similarly for the hub/sage/xformers/aiter/npu/xla variants) — these are not PyTorch ops with an autocast policy, so no such implicit cast happens. They receive the raw fp32 query/key + bf16 value:
crash inside the external kernel's own dtype assertion (e.g. flash-attn's `"FlashAttention only support fp16 and bf16 data type"`), or

---

the latter part *Drafted by Claude* but shortened

### Reproduction

for example, run Chroma in an autocast context, install flash-attn, set the backend

### Logs

```
Traceback (most recent call last):
File "\OneTrainer\modules\ui\TrainUI.py", line 716, in __training_thread_function
trainer.train()
File "\OneTrainer\modules\trainer\GenericTrainer.py", line 744, in train
model_output_data = self.model_setup.predict(self.model, batch, self.config, train_progress)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\modules\modelSetup\BaseChromaSetup.py", line 241, in predict
packed_predicted_flow = model.transformer(
^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\nn\modules\module.py", line 1773, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\nn\modules\module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\expermential guff\venv\src\diffusers\src\diffusers\models\transformers\transformer_chroma.py", line 577, in forward
encoder_hidden_states, hidden_states = block(
^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\nn\modules\module.py", line 1771, in _wrapped_call_impl
return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_dynamo\eval_frame.py", line 736, in compile_wrapper
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\nn\modules\module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\modules\util\checkpointing_util.py", line 90, in forward
def forward(self, *args, **kwargs):
File "\OneTrainer\venv\Lib\site-packages\torch\_dynamo\eval_frame.py", line 929, in _fn
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\aot_autograd.py", line 1241, in forward
return compiled_fn(full_args)
^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\runtime_wrappers.py", line 370, in runtime_wrapper
all_outs = call_func_at_runtime_with_args(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\utils.py", line 126, in call_func_at_runtime_with_args
out = normalize_as_list(f(args))
^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\utils.py", line 100, in g
return f(*args)
^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\autograd\function.py", line 576, in apply
return super().apply(*args, **kwargs) # type: ignore[misc]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\runtime_wrappers.py", line 2074, in forward
fw_outs = call_func_at_runtime_with_args(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\utils.py", line 126, in call_func_at_runtime_with_args
out = normalize_as_list(f(args))
^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\runtime_wrappers.py", line 556, in wrapper
return compiled_fn(runtime_args)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\runtime_wrappers.py", line 750, in inner_fn
outs = compiled_fn(args)
^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_inductor\output_code.py", line 584, in __call__
return self.current_callable(inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_inductor\utils.py", line 2716, in run
out = model(new_inputs)
^^^^^^^^^^^^^^^^^
File "\AppData\Local\Temp\torchinductor\bq\cbqrfpkp4z55klkpln3fescrhmpvo7357m3j6qqgugdzy6ldyfze.py", line 3114, in call
buf137 = torch.ops.flash_attn._flash_attn_forward.default(buf135, buf136, buf134, 0.0, 0.08838834764831845, False, -1, -1, 0.0, None, False)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_ops.py", line 829, in __call__
return self._op(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_library\autograd.py", line 111, in autograd_impl
result = forward_no_grad(*args, Metadata(keyset, keyword_only_args))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_library\autograd.py", line 40, in forward_no_grad
result = op.redispatch(keyset & _C._after_autograd_keyset, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_ops.py", line 836, in redispatch
return self._handle.redispatch_boxed(keyset, *args, **kwargs) # type: ignore[return-value]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_library\custom_ops.py", line 344, in backend_impl
result = self._backend_fns[device_type](*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_compile.py", line 53, in inner
return disable_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_dynamo\eval_frame.py", line 929, in _fn
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\torch\_library\custom_ops.py", line 377, in wrapped_fn
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "\OneTrainer\venv\Lib\site-packages\flash_attn\flash_attn_interface.py", line 91, in _flash_attn_forward
out, softmax_lse, S_dmask, rng_state = flash_attn_gpu.fwd(
^^^^^^^^^^^^^^^^^^^
RuntimeError: FlashAttention only support fp16 and bf16 data type
Exception raised from mha_fwd at C:\Users\junya\actions-runner\_work\flash-attention-prebuild-wheels\flash-attention-prebuild-wheels\flash-attention\csrc\flash_attn\flash_api.cpp:374 (most recent call first):
00007FFC8050199400007FFC805018F0 c10.dll!c10::Error::Error [ @ ]
00007FFC8050040A00007FFC805003B0 c10.dll!c10::detail::torchCheckFail [ @ ]
00007FFAF25B42FD00007FFAF25A5EA0 flash_attn_2_cuda.cp312-win_amd64.pyd!c10::ivalue::Object::operator= [ @ ]
00007FFAF25C3C9B00007FFAF25BD050 flash_attn_2_cuda.cp312-win_amd64.pyd!PyInit_flash_attn_2_cuda [ @ ]
00007FFAF25C3DB400007FFAF25BD050 flash_attn_2_cuda.cp312-win_amd64.pyd!PyInit_flash_attn_2_cuda [ @ ]
00007FFAF25AF2BC00007FFAF25A5EA0 flash_attn_2_cuda.cp312-win_amd64.pyd!c10::ivalue::Object::operator= [ @ ]
00007FFC780B79E400007FFC780B73F4 python312.dll!PyThread_acquire_lock_timed [ @ ]
00007FFC7810201800007FFC78101990 python312.dll!PyObject_Vectorcall [ @ ]
00007FFC781019C500007FFC78101990 python312.dll!PyObject_Vectorcall [ @ ]
00007FFC78102EA500007FFC78102610 python312.dll!PyEval_EvalFrameDefault [ @ ]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [ @ ]
00007FFC7813281500007FFC781326F0 python312.dll!PyObject_Call [ @ ]
00007FFC7813275F00007FFC781326F0 python312.dll!PyObject_Call [ @ ]
00007FFC60AF7BBA00007FFC60ACFF50 torch_python.dll!torch::FunctionParameter::type_name [ @ ]
00007FFC60AEDAC800007FFC60ACFF50 torch_python.dll!torch::FunctionParameter::type_name [ @ ]
00007FFC60AF15EB00007FFC60ACFF50 torch_python.dll!torch::FunctionParameter::type_name [ @ ]
00007FFC60AD17D600007FFC60ACFF50 torch_python.dll!torch::FunctionParameter::type_name [ @ ]
00007FFC601E25F600007FFC601E1730 torch_python.dll!c10::ivalue::Future::devices [ @ ]
00007FFC780B79E400007FFC780B73F4 python312.dll!PyThread_acquire_lock_timed [ @ ]
00007FFC780E560E00007FFC780E556C python312.dll!PyObject_MakeTpCall [ @ ]
00007FFC7824673800007FFC78223794 python312.dll!PyThread_tss_is_created [ @ ]
00007FFC7813281500007FFC781326F0 python312.dll!PyObject_Call [ @ ]
00007FFC7813275F00007FFC781326F0 python312.dll!PyObject_Call [ @ ]
00007FFC781066EE00007FFC78102610 python312.dll!PyEval_EvalFrameDefault [ @ ]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [ @ ]
00007FFC780FBD8E00007FFC780FB9BC python312.dll!PyArg_CheckPositional [ @ ]
00007FFC7813281500007FFC781326F0 python312.dll!PyObject_Call [ @ ]
00007FFC7813275F00007FFC781326F0 python312.dll!PyObject_Call [ @ ]
00007FFC781066EE00007FFC78102610 python312.dll!PyEval_EvalFrameDefault [ @ ]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [ @ ]
00007FFC7813281500007FFC781326F0 python312.dll!PyObject_Call [ @ ]
00007FFC7813275F00007FFC781326F0 python312.dll!PyObject_Call [ @ ]
00007FFC60AF7B2F00007FFC60ACFF50 torch_python.dll!torch::FunctionParameter::type_name [ @ ]
00007FFB5B08924100007FFB5B089080 torch_cpu.dll!c10::Dispatcher::callBoxed [ @ ]
00007FFC608BF02500007FFC608BEF60 torch_python.dll!torch::jit::invokeOperatorFromPython [ @ ]
00007FFC608BC08D00007FFC608BBF50 torch_python.dll!torch::jit::_get_operation_for_overload_or_packet [ @ ]
00007FFC6081AFCA00007FFC608153F0 torch_python.dll!c10d::PythonOnCompletionHook::PythonOnCompletionHook [ @ ]
00007FFC6082F6C100007FFC608153F0 torch_python.dll!c10d::PythonOnCompletionHook::PythonOnCompletionHook [ @ ]
00007FFC607C820600007FFC6074CED0 torch_python.dll!THPPointer<_frame>::dup [ @ ]
00007FFC601E25F600007FFC601E1730 torch_python.dll!c10::ivalue::Future::devices [ @ ]
00007FFC780B79E400007FFC780B73F4 python312.dll!PyThread_acquire_lock_timed [ @ ]
00007FFC781327A600007FFC781326F0 python312.dll!PyObject_Call [ @ ]
00007FFC781066EE00007FFC78102610 python312.dll!PyEval_EvalFrameDefault [ @ ]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [ @ ]
00007FFC78133E8B00007FFC78133E18 python312.dll!PyObject_FastCallDictTstate [ @ ]
00007FFC781D319B00007FFC781D311C python312.dll!PyObject_Call_Prepend [ @ ]
00007FFC781D30C600007FFC781D2DF0 python312.dll!PyDictProxy_New [ @ ]
00007FFC7810201800007FFC78101990 python312.dll!PyObject_Vectorcall [ @ ]
00007FFC781019C500007FFC78101990 python312.dll!PyObject_Vectorcall [ @ ]
00007FFC78102EA500007FFC78102610 python312.dll!PyEval_EvalFrameDefault [ @ ]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [ @ ]
00007FFC78133E8B00007FFC78133E18 python312.dll!PyObject_FastCallDictTstate [ @ ]
00007FFC781D319B00007FFC781D311C python312.dll!PyObject_Call_Prepend [ @ ]
00007FFC781D30C600007FFC781D2DF0 python312.dll!PyDictProxy_New [ @ ]
00007FFC7810201800007FFC78101990 python312.dll!PyObject_Vectorcall [ @ ]
00007FFC781019C500007FFC78101990 python312.dll!PyObject_Vectorcall [ @ ]
00007FFC78102EA500007FFC78102610 python312.dll!PyEval_EvalFrameDefault [ @ ]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [ @ ]
00007FFC7813281500007FFC781326F0 python312.dll!PyObject_Call [ @ ]
00007FFC780893B700007FFC7808931C python312.dll!PyObject_Call [ @ ]
00007FFC6069275F00007FFC6068D170 torch_python.dll!torch::autograd::registerFunctionPreHook [ @ ]
00007FFC780B7A1600007FFC780B73F4 python312.dll!PyThread_acquire_lock_timed [ @ ]
00007FFC781327A600007FFC781326F0 python312.dll!PyObject_Call [ @ ]
00007FFC781066EE00007FFC78102610 python312.dll!PyEval_EvalFrameDefault [ @ ]
```

### System Info

6d71b76aceff935192e58fee38c5cc5d8d227cf0

### Who can help?

_No response_

Start with FluxAttnProcessor.__call__ and the equivalent Flux2, Chroma, and Krea2 processors, tracing norm_q/norm_k, apply_rotary_emb, and dispatch_attention_fn. Reproduce Chroma under torch.autocast with flash-attn enabled, then verify the affected non-SDPA backends no longer fail on the resulting dtype mismatch.

Written by the indexing model from the issue text.