Merge pull request #4454 from AI-Hypercomputer:zjiahao/DSA3.2-duplicated-softmax-computation

PiperOrigin-RevId: 949271937

Author: Google-ML-Automation

src/maxtext/layers/attention_mla.py

@@ -1084,6 +1084,8 @@ def calculate_indexer_loss(
 
     # Aggregate heads: [b, h, t, s] -> [b, t, s]
     attention_probs = jnp.sum(attention_probs, axis=1)
+    # Force materialization and prevent fusion across this point to reuse the intermediate tensor
+    attention_probs = jax.lax.optimization_barrier(attention_probs)
     # L1 normalize aggregated target distribution
     attention_probs = attention_probs / (jnp.sum(attention_probs, axis=-1, keepdims=True) + EPS)