Fix soft cap in JAX Splash attention

[huytransformer]

Description

Fixes attention logit soft cap in the JAX Splash attention path. Example, gemma2 configs set attn_logits_soft_cap, but the use_jax_splash path did not pass that value into flash_attention_block_masked.

Tests

Ran:

python3 -m py_compile \
  src/maxtext/layers/attention_op.py \
  src/maxtext/kernels/attention/jax_flash_attention.py \
  tests/unit/attention_test.py

python3 -m pytest tests/unit/attention_test.py::JaxFlashAttentionTest::test_flash_attention_block_masked_soft_cap

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!