Merge pull request #3312 from Angelogeb:cudnn_jax-scale-fix

PiperOrigin-RevId: 892411074

Author: Google-ML-Automation

src/maxtext/layers/attention_op.py

@@ -1672,7 +1672,7 @@ def cudnn_jax_flash_attention(
           key,
           value,
           mask_type=MaskType.CAUSAL,
-          scale=1.0 / math.sqrt(head_dim),
+          scale=1.0,
           dropout_rate=self.dropout_rate,
           qkv_layout="BTNH",
           return_residual=True,