diff --git a/primus/backends/megatron/core/models/diffusion/flux/attention.py b/primus/backends/megatron/core/models/diffusion/flux/attention.py
index b444b527c..e10b52eb0 100644
--- a/primus/backends/megatron/core/models/diffusion/flux/attention.py
+++ b/primus/backends/megatron/core/models/diffusion/flux/attention.py
@@ -390,6 +390,12 @@ def forward(
             query = apply_rotary_pos_emb(query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q)
             key = apply_rotary_pos_emb(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv)
 
+        # Q/K RMSNorm (for_qk) can leave Q/K in fp32 while V stays bf16; AITER
+        # FlashAttention requires Q, K, V to share one dtype. Align Q/K to V at the
+        # attention-kernel boundary (covers both RMSNorm and RoPE promotion).
+        query = query.to(value.dtype)
+        key = key.to(value.dtype)
+
         # Core attention computation
         if self.checkpoint_core_attention and self.training:
             core_attn_out = self._checkpointed_attention_forward(
@@ -566,6 +572,12 @@ def forward(
             query = apply_rotary_pos_emb(query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q)
             key = apply_rotary_pos_emb(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv)
 
+        # Q/K RMSNorm (for_qk) can leave Q/K in fp32 while V stays bf16; AITER
+        # FlashAttention requires Q, K, V to share one dtype. Align Q/K to V at the
+        # attention-kernel boundary (covers both RMSNorm and RoPE promotion).
+        query = query.to(value.dtype)
+        key = key.to(value.dtype)
+
         # Core attention computation
         if self.checkpoint_core_attention and self.training:
             core_attn_out = self._checkpointed_attention_forward(