Enable Indexer cache for DS v3.2 decoding

[RissyRan]

Description

Enable Indexer cache for DS v3.2 decoding, to unblock the eval benchmark for DS v3.2 model with sparse attention bringup.

• DS reference implementation for Indexer is here
• Add init_indexer_cache & update_indexer_cache for indexer cache
• Other small changes

Tests

• All runners are green
• Training end-to-end (no impact) with a smaller model version: link
• Test against reference implementation still green: link
• Decoding gives reasonable
  • small seq len to skip the indexer: link
  • large seq len to process the indexer max_prefill_predict_length=3072 max_target_length=4096: link

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.

[github-actions]

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

[codecov]

Codecov Report

❌ Patch coverage is 86.95652% with 6 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_mla.py	84.21%	2 Missing and 4 partials ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

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

[shuningjin]

Thank you! There were some attention autoregressive/decoding tests:

• mha:

  maxtext/tests/unit/attention_test.py

  Line 371 in 3fbe3be

  def test_autoregression(self):

• mla:

  maxtext/tests/unit/attention_test.py

  Line 1195 in 3fbe3be

  def test_autoregression(self, rope_type):

Shall we add a similar one to test indexer cache in mla? example from gemini

[RissyRan]

Thank you! There were some attention autoregressive/decoding tests:

• mha:

  maxtext/tests/unit/attention_test.py

  Line 371 in 3fbe3be

  def test_autoregression(self):

• mla:

  maxtext/tests/unit/attention_test.py

  Line 1195 in 3fbe3be

  def test_autoregression(self, rope_type):

Shall we add a similar one to test indexer cache in mla? example from gemini

Sounds good. Also enabled the mla assertion, and did some sanity check for decoding: long seq, short seq