Eliminate duplicated softmax recomputation in indexer loss.

[JHCuc3m]

Description

This PR eliminates a duplicated softmax computation in the indexer loss calculation of DeepSeek-V3.2.

Why is this change being made?

In the original implementation, the compiler (XLA) failed to reuse the head-aggregated attention probabilities for the subsequent sequence reduction (L1 normalization). Instead, it recomputed the entire softmax pipeline from the raw QK scores redundantly.

Solution

By inserting jax.lax.optimization_barrier after the head aggregation step in calculate_indexer_loss, we force the compiler to materialize and reuse the head-aggregated intermediate tensor for the subsequent sequence reduction instead of recomputing the entire softmax from raw QK scores.

It was observed that this reduces execution time for 128K sequence length long context, therefore it should be less of an issue for storing softmax tensor in shorter sequences.

Implementation Details

• Modified calculate_indexer_loss in src/maxtext/layers/attention_mla.py to insert the barrier.

Tests

All tests passed.

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](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-
  documentation-files).

TAG=agy
CONV=5ff94b54-4171-4309-8704-6046df05eb13

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!