Fix scale to 1.0 for cudnn_flash_jax

Author: Angelogeb

src/maxtext/layers/attention_op.py

@@ -1678,7 +1678,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,