Eliminate redundant sorts in BWD pass by caching indexes between embedding chunks

[Shuwen-Fang]

Description

Optimizes token routing in ring_ragged_sort during embedding chunking by eliminating redundant jnp.argsort calls in both the forward and backward passes.

Context & Impact

Performance Impact

• Latency Reduction: Saves ~600 µs per chunk per layer in the backward pass (see Profile Breakdown Screenshot).

How & Why Sort Is Removed

All embedding chunks within a layer share identical token routing (topk_indices_2d and valid_rows_mask), but every embedding chunk repeatedly recomputed the SparseCore valid-row index permutation via jnp.argsort in _preprocess. This PR caches the sorted by validity array by computing it once during the forward pass, then passing it into future bwd passes via custom VJP residuals _ring_ragged_sort_bwd, so that jnp.argsort runs once per layer across all forward and backward chunk ops.

---

Tests

1. Unit Test Verification:

   • Executed test_moe_emb_chunking_random_routing in tests/unit/moe_test.py on TPU (tpu7x-128) on VM:

   tests/unit/moe_test.py::RoutedMoeTest::test_moe_emb_chunking_random_routing . [100%]
   ================ 1 passed, 41 deselected in 40.00s ================
   

2. Loss & Perplexity Correctness Run (num_moe_emb_chunks=4):

   • Executed 20-step pretraining correctness runs (tmp/run_correctness.sh) comparing main branch vs sort_index_once branch with 4 embedding chunks active:
     • Step 0 Loss: 12.258 (main) vs 12.258 (sort_index_once)
     • Step 10 Loss: 11.985 (main) vs 11.985 (sort_index_once)
     • Step 19 Loss: 11.919 (main) vs 11.919 (sort_index_once)
     • Step 19 Perplexity: 150053.750 (main) vs 150053.750 (sort_index_once)
   • Confirmed 100% numerical equivalence across all 20 steps under chunked execution.

3. Sort Elimination Verification via Profile:

   • Deployed XManager experiment XID 273744286 (deepseekv3-fsdp512-ep8-chunked-2) to verify via profile trace that argsort overhead in the backward pass was eliminated.

Checklist

• I have performed a self-review of my code.
• 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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

❌ Patch coverage is 93.84615% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
.../maxtext/kernels/ragged/ragged_gather_reduce_v2.py	85.18%	2 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[NuojCheng]

we need jnp.int32 for accumulation, otherwise got regression issue

[NuojCheng]

emmm not clear why make this change

[NuojCheng]

1. You could add a test protecting your feature.
2. Could you add a regression test using flags in here? I would like to make sure our performance is not impacted by this change. Feel free to reduce layers and chip sizes for easier repro.