Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -616,11 +616,10 @@ class GemmoeSdpaAttention(GemmoeAttention):
         causal_mask = attention_mask
         if attention_mask is not None and cache_position is not None:
             causal_mask = causal_mask[:, :, cache_position, : key_states.shape[-2]]
-            
-            # Cast query, key, and value states to the same dtype (bf16)
-            query_states = query_states.to(dtype=torch.bfloat16)
-            key_states = key_states.to(dtype=torch.bfloat16)
-            value_states = value_states.to(dtype=torch.bfloat16)
+
+        # Convert causal_mask to the same dtype as query_states
+        if causal_mask is not None:
+            causal_mask = causal_mask.to(dtype=query_states.dtype)
 
         # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
         # Reference: https://github.com/pytorch/pytorch/issues/112577.
@@ -629,7 +628,6 @@ class GemmoeSdpaAttention(GemmoeAttention):
             key_states = key_states.contiguous()
             value_states = value_states.contiguous()
 
-        
         attn_output = torch.nn.functional.scaled_dot_product_attention(
             query_states,
             key_states,