Fix soft cap in JAX Splash attention

Author: huytransformer

src/maxtext/kernels/attention/jax_flash_attention.py

@@ -159,10 +159,10 @@ def compute_attention_block(output, l, m):
             (batch_size, num_kv_heads, q_groups, block_q, block_kv),
         )
 
-        s_i_j = jnp.where(broadcasted_mask, s_i_j, mask_value)
         if cap is not None:
           s_i_j = jnp.tanh(s_i_j / cap)
           s_i_j = s_i_j * cap
+        s_i_j = jnp.where(broadcasted_mask, s_i_j, mask_value)
         m_i_j = s_i_j.max(axis=-1)
         p_i_j = jnp.exp(s_i_j - m_i_j[..., None])
         l_i_j = p_i_j.sum(axis=-1)

src/maxtext/layers/attention_op.py

@@ -1500,6 +1500,7 @@ def kernel_fn(q, k, v, d, s):
             block_q=self.block_q,
             mask=materialized_mask,
             mask_value=DEFAULT_MASK_VALUE,
+            cap=attn_logits_soft_cap,
         )
         if record_max_logits:
           # The native JAX splash attention implementation does not currently expose the softmax statistics

tests/unit/attention_test.py

@@ -47,6 +47,7 @@
 )
 from maxtext.layers.attentions import Attention
 from maxtext.layers import embeddings
+from maxtext.kernels.attention import jax_flash_attention
 from maxtext.configs import pyconfig
 from maxtext.models.qwen3 import Qwen3NextGatedDeltaNet
 import numpy as np
@@ -56,6 +57,41 @@
 from tests.utils.test_helpers import get_test_config_path
 
 
+class JaxFlashAttentionTest(unittest.TestCase):
+  """Tests for JAX flash attention."""
+
+  def test_flash_attention_block_masked_soft_cap(self):
+    cap = 1.0
+    mask_value = -1.0e9
+    query = jnp.array([[[[10.0], [10.0]]]], dtype=jnp.float32)
+    key = jnp.array([[[[10.0], [0.0]]]], dtype=jnp.float32)
+    value = jnp.array([[[[1.0], [3.0]]]], dtype=jnp.float32)
+    mask = jnp.array([[True, True], [True, False]])
+
+    output = jax_flash_attention.flash_attention_block_masked(
+        query,
+        key,
+        value,
+        segment_ids=None,
+        block_kv=2,
+        block_q=1,
+        mask=mask,
+        mask_value=mask_value,
+        cap=cap,
+    )
+
+    logits = jnp.einsum("bhqd,bhkd->bhqk", query, key)
+    logits = jnp.tanh(logits / cap) * cap
+    logits = jnp.where(mask[None, None, :, :], logits, mask_value)
+    expected = jnp.einsum("bhqk,bhkd->bhqd", jax.nn.softmax(logits, axis=-1), value)
+    np.testing.assert_allclose(
+        np.asarray(output),
+        np.asarray(expected),
+        rtol=1e-6,
+        atol=1e-6,
+    )
+
+
 class BidirectionalBlockMaskTest(unittest.TestCase):
   """Test for make_bidirectional_block_mask."""