Merge pull request #2716 from AI-Hypercomputer:tokamax_splash_sink

PiperOrigin-RevId: 834512715

Author: Google-ML-Automation

src/MaxText/layers/attention_op.py

@@ -1291,12 +1291,10 @@ def wrap_flash_attention(
         decoder_segment_ids_tuple = None
 
       if self.config.use_tokamax_splash:
-        if max_logit_value is not None:
-          attention_output = jax.vmap(partial(splash_kernel, max_logit_value=max_logit_value))(
-              query, key, value, decoder_segment_ids_tuple
-          )
-        else:
-          attention_output = jax.vmap(splash_kernel)(query, key, value, decoder_segment_ids_tuple)
+        kernel = partial(splash_kernel, max_logit_value=max_logit_value)
+        attention_output = jax.vmap(lambda q, k, v, d, s: kernel(q, k, v, d, sinks=s), in_axes=(0, 0, 0, 0, None))(
+            query, key, value, decoder_segment_ids_tuple, sinks
+        )
       else:
         attention_output = jax.vmap(splash_kernel, in_axes=(0, 0, 0, 0, None))(
             query, key, value, decoder_segment_ids_tuple, sinks