Enable Indexer cache for DS v3.2 decoding

[RissyRan]

Description

This is a clean version of previous PR after rebase.

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
• Add encoding_dsv32.py from HF files to enable specific encoding for V3.2
• Other small changes for benchmark time out setting

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
• Reasonable MMLU pro valuation in b/469547559, and sanity check on decoding - 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.

[codecov]

Codecov Report

❌ Patch coverage is 95.00000% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attention_mla.py	93.75%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[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.

[github-actions]

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

[shuningjin]

Thanks for enabling decoding for DeepSeek Sparse Attention, and adding specialized encoding!

Really appreciate the comprehensive testing for pre-filling and generation (across unit tests, standalone decoding, and the API server). It is amazing that you reproduce reasonable MMLU-pro result for large scale with limited resources.

[github-actions]

🤖 Hi @shuningjin, 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 pull request enables the Indexer cache for DeepSeek V3.2 decoding, which is a critical component for sparse attention bringup in MaxText. It also introduces necessary encoding/decoding logic for DeepSeek-V3.2's specialized output format, including thinking/reasoning blocks, and updates the API server to handle these new fields.

🔍 General Feedback

• Reasoning Content: The addition of reasoning_content to ChatMessage is a great improvement for supporting modern LLMs that use separate thinking blocks.
• Cache Implementation: The Indexer KV cache implementation follows the existing MLA patterns well and correctly handles masking for uninitialized slots during decoding.
• Model Compatibility: The model name checks in the API server might be too restrictive, potentially missing the official Hugging Face model names. Consider using more robust string matching as suggested.
• Testing: The new unit tests effectively verify both prefill and autoregressive modes for the indexer, covering different sequence length scenarios.

[Rohan-Bierneni]

Thank you for making this change! As we have discussed in previous pr, there can be some other optimizations made to the cache, but for functionality reusing kvcache for indexer was a good choice.

Also thank you for adding tests and getting MMLU-Pro to work with api_server.

LGTM!