From 359d3e9041bb2a10b4f89d36fcdc5bb5251131ea Mon Sep 17 00:00:00 2001 From: Jiahao Chen Zhou Date: Thu, 16 Jul 2026 22:33:41 +0000 Subject: [PATCH] Eliminate duplicated softmax recomputation in indexer loss. By inserting `jax.lax.optimization_barrier` after the head aggregation step in `calculate_indexer_loss`, we force the compiler to reuse the head-aggregated intermediate tensor for the subsequent sequence reduction instead of recomputing the entire softmax pipeline from raw QK scores. TAG=agy CONV=5ff94b54-4171-4309-8704-6046df05eb13 --- src/maxtext/layers/attention_mla.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index d5f79a1a4b..e81b0a9b73 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -1077,6 +1077,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)