Fix fp32 dtype leak from QK-norm in Flux2 attention processors

[AL3708]

Problem

torch.nn.RMSNorm is on autocast's fp32 promotion list. Under
autocast(bfloat16) (or fp16), the QK-norm layers in Flux2Attention
therefore return fp32 tensors, while the value tensor stays in half
precision:

query = attn.norm_q(query)   # -> fp32 under autocast
key = attn.norm_k(key)       # -> fp32 under autocast
value = value.unflatten(...)  # stays bf16

apply_rotary_emb preserves the wider dtype, so query/key arrive at the
attention backend as fp32 alongside a half-precision value.

Attention backends that require a uniform half dtype reject this outright.
flash-attn raises on mismatched input dtypes rather than casting, so training
Flux.2 with mixed_precision=bf16 and a flash-attn backend fails inside the
attention call.

This affects both attention processors:

• Flux2AttnProcessor, including the added_kv_proj_dim encoder branch
  (norm_added_q / norm_added_k)
• Flux2ParallelSelfAttnProcessor

Fix

Cast the normalised query/key back to the dtype of the corresponding
value tensor, immediately after the norm.

Using value.dtype rather than a hardcoded bfloat16:

• stays correct under fp16 autocast,
• becomes a no-op in full-precision runs,
• keeps the encoder branch consistent by using encoder_value.dtype.

Placing the cast directly after the norm (rather than after
apply_rotary_emb) also covers the maybe_metal_flash_rope_attention path,
which applies rotary embeddings internally.

Reproduction

Self-contained: builds a Flux2Attention directly and calls it under
autocast(bfloat16). No checkpoint, no dataset, ~30 lines.

import torch
from simpletuner.helpers.models.flux2.transformer import Flux2Attention

attn = Flux2Attention(
    query_dim=256, heads=4, dim_head=64, added_kv_proj_dim=256
).to("cuda", dtype=torch.bfloat16)
attn.processor._attention_backend = "flash"

hidden = torch.randn(1, 256, 256, device="cuda", dtype=torch.bfloat16)
encoder_hidden = torch.randn(1, 64, 256, device="cuda", dtype=torch.bfloat16)

with torch.autocast("cuda", dtype=torch.bfloat16):
    attn(hidden_states=hidden, encoder_hidden_states=encoder_hidden)

On main (37d6fb2), A100-SXM4-80GB, torch 2.13.0+cu126:

norm_q(torch.bfloat16) -> torch.float32

  File "simpletuner/helpers/models/flux2/transformer.py", line 367, in forward
    return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb, **kwargs)
  File "simpletuner/helpers/models/flux2/transformer.py", line 269, in __call__
    hidden_states = dispatch_attention_fn(
  File ".../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

With this patch applied, the same script returns
[(1, 256, 256), (1, 64, 256)].

Verified alongside it:

• casting to value.dtype is bit-identical to casting explicitly to bf16
  (torch.equal → True), so the fix changes no numerics
• the same code path works under autocast(float16), which a hardcoded
  bfloat16 cast would have broken