Add optional Approximate Top-K configuration for MLA Indexer

[JHCuc3m]

Description

This PR adds an optional configuration parameter indexer_use_approx_top_k to the Multi-head Latent Attention (MLA) Indexer, allowing users to enable JAX's TPU-optimized approx_max_k primitive instead of the default exact top_k selection.

Why is this change being made?

During performance investigations of DeepSeek-V3.2 in long-context mode (128K sequence length) with 128-way Context Parallelism (CP=128), the top-k was identified a major bottleneck in the MLA Indexer forward pass:

• The indexer computes a local score matrix of shape [1, 1024, 131072] on each device.
• The default exact top_k selection is implemented on TPU via an expensive Bitonic Sort along the 131K sequence dimension.
• Sorting 131K elements per layer across 58 layers introduces massive step-time overhead.

Benchmarks show a 4x speedup on f32[1,1024,65536] tensors when tested with the approximate path enabled using a recall target of 0.95.

Why this is a good solution

JAX's approx_max_k employs block-based reduction optimized for TPU Matrix Units (MXU). It avoids full bitonic sorting, reducing complexity from $\mathcal{O}(N \log^2 N)$ to $\approx \mathcal{O}(N + K \log K)$ with a significantly smaller constant factor. Paper: https://arxiv.org/pdf/2206.14286.

Workload show a ~4x speedup on f32[1, 1024, 65536] tensors when tested on TPU with the approximate path enabled using a recall target of 0.95.

Specific Implementation Details

1. Configuration Schema (types.py): Added indexer_use_approx_top_k and indexer_approx_top_k_recall to the AttentionIndexer Pydantic class to pass configuration validation.
2. Default Config (base.yml): Exposed the parameters with safe defaults (indexer_use_approx_top_k: false, indexer_approx_top_k_recall: 0.95).
3. Model Architecture (attention_mla.py): Updated Indexer.__call__ to conditionally route the selection to jax.lax.approx_max_k when enabled.

Shortcomings & Future Improvements

• There is not systematic study on how the accuracy lost of using indexer_use_approx_top_k instead of top_k might affect downstream model performance, while it is expected to be minimal when a high recall rate is used.

Tests

1. Regression Guard (Default Path)

We ran the attention unit test suite with the default configuration (indexer_use_approx_top_k=false) to ensure no regressions:

• Command: pytest tests/unit/attention_test.py
• Result: PASSED (20 passed, 32 skipped).

2. Compilation & Tracing Safety

We added a new unit test, test_indexer_with_approx_top_k, to verify that the new path compiles and traces successfully in JAX:

• Command: pytest tests/unit/attention_test.py -k test_indexer_with_approx_top_k
• Result: PASSED.

3. Mathematical Correctness & Recall Tracking

We added a correctness test, test_approx_top_k_recall, which generates random scores of shape [4, 16, 1024], runs both exact and approximate top-K ($K=64$), and calculates the actual recall:

• Command: pytest tests/unit/attention_test.py -k test_approx_top_k_recall -s
• Result: PASSED (Achieved 100% recall on CPU).

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 and provided workload links above if applicable. https://xprof.corp.google.com/trace_viewer/zjiahao-9369408277271509466?view_start=23269.070&view_end=23384.958
• [] I have made or will make corresponding changes to the doc if needed.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[JHCuc3m]

Code quality checker failed from existing file content

[NuojCheng]

Is 0.05 too large? What about making it 0.01?

[JHCuc3m]

updated!

[NuojCheng]

I think it is a very cool optimization! Thank you Jiahao