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.

• Baseline xprof Step time: 77 s

• After proposed implementation xprof. Step time 73 s

  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!

[github-actions]

🤖 Hi @JHCuc3m, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This PR introduces an elegant compiler optimization by inserting a jax.lax.optimization_barrier after the head aggregation step in the indexer loss calculation of DeepSeek-V3.2. This forces the XLA compiler to materialize the smaller head-aggregated tensor, preventing redundant recomputation of the memory-intensive softmax pipeline from raw QK scores. The proposed change is highly effective, leading to a measurable step-time reduction for long context sequences (128K) while preserving full numerical and behavioral correctness.

🔍 General Feedback

• High-Quality Optimization: Preventing compiler fusion that triggers redundant softmax recomputations is an excellent way to speed up step times in long-context models without increasing memory peak requirements.
• Robustness and Safety: The implementation of jax.lax.optimization_barrier is perfectly safe, device-agnostic, and works seamlessly with the standard JAX AD pipeline since the gradients are already detached at this point.
• Code Cleanliness: The codebase standards are well-maintained, and the added comment provides great clarity on the intent behind this optimization barrier.

[github-actions]

🤖 Hi @JHCuc3m, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @JHCuc3m, but I was unable to process your request. Please see the logs for more details.

[gobbleturk]

Amazing find!