Support Ring Attention with DeepSeek DSA Sparse Indexer

[zcjhao]

Description

This PR adds end-to-end support for DeepSeek Sparse Attention (DSA) Indexer with Tokamax Ring Context Parallelism during training. Previously, use_indexer=True was prohibited with Ring Attention due to a lack of dynamic per-ring-step mask extraction. This PR also solve 2 bugs concerning load balancing with Indexer and indexer losses.

Problems Solved & Key Design Choices:

1. Dynamic Per-Ring-Step Mask Extraction & Grid Scheduling:
   • Forward Pass: Slices the global dynamic indexer_mask into per-step KV-shard blocks (ring_axis_idx - step) % ring_axis_size and tiles them into hardware blocks (block_q, block_kv).
   • Backward Pass ($dK/dV$ Transposition): Forward Splash Attention iterates $Q$-major (q_blocks, kv_blocks), whereas the backward $dK/dV$ pass iterates $KV$-major (kv_blocks, q_blocks). We apply .swapaxes(0, 1) and . swapaxes(-1, -2) when is_dkv=True, aligning the mask blocks with the hardware execution schedule.

Forward Flow: From Global Mask to TPU Register Tile

1. Global Indexer Mask Tensor [Batch, S_local, S_global]
                           │
                           ▼ (Reshape to 4D [Batch, S_local, N_chips, S_local])
2. Ring Step i (Device r) ───► Dynamic Slice for current chip:
                               local_idx_mask = mask_4d[:, :, (r - i) % N, :]
                           │
                           ▼ (_inject_local_indexer_mask)
3. Combine with Causal Mask: combined_mask = (q_pos >= kv_pos) & local_idx_mask
                           │
                           ▼ (Tile into [128, 128] blocks)
4. Package into MaskInfo.partial_mask_blocks: [Num_Tiles, sa_block_q, sa_block_kv]
                           │
                           ▼ (Pallas TPU Kernel)
5. Inside VMEM Register: logits = where(tile_mask == 1, (Q @ K.T) * scale, -1e30)

Backward Flow: From Saved Mask to Transposed Gradient Tiles

  1. Saved Indexer Mask Tensor from Forward Residuals [Batch, S_local, S_global]
                                    │
                                    ▼ (Reshape to 4D [Batch, S_local, N_chips, S_local])
  2. Ring Step i (Device r) ────────► Dynamic Slice for current rotating KV chip:
                                      local_idx_mask = mask_4d[:, :, (r - i) % N, :]
                                    │
                                    ▼ (_inject_local_indexer_mask with is_dkv=True)
  3. Combine with Causal Mask:      combined_mask = (q_pos >= kv_pos) & local_idx_mask
                                    │
                                    ▼ (Tile into [sa_block_q_dkv, sa_block_kv_dkv] blocks)
  4. Transpose for KV-Major Grid:   blocks = blocks.swapaxes(0, 1)    # [kv_blocks, q_blocks]
                                    blocks = blocks.swapaxes(-1, -2)  # [sa_block_kv_dkv, sa_block_q_dkv]
                                    │
                                    ▼ (Package into MaskInfo.partial_mask_blocks)
  5. Package into MaskInfo:         [Num_Tiles, sa_block_kv_dkv, sa_block_q_dkv]
                                    │
                                    ▼ (Pallas TPU Backward Kernel)
  6. Inside VMEM Registers:         dQ = dP @ K  (masked by Q-KV tile)
                                    dK = dP.T @ Q (masked by transposed KV-Q tile)
                                    dV = P.T @ dO (masked by transposed KV-Q tile)

2. Isolated Flax NNX Auxiliary Loss:

   • In Flax NNX, _apply_layers_sequentially filters scanned layer states via nnx.filter_state(..., nnx.Not((nnx.RngState, nnx.Intermediate))). Subclassing nnx.Intermediate caused jax.lax.scan to prune the loss variable during execution, resulting in indexer_loss: 0.000 and broken backward gradients.
   • We define class indexer_losses(nnx.Variable): which inherits directly from nnx.Variable rather than nnx Intermediate. This allows the variable to pass through jax.lax.scan unharmed without modifying core decoder scanning logic.

3. Loss Harvesting & Objective Injection:

   • In train.py's loss_fn, indexer_losses is popped before generic intermediates (mirroring the upstream MTP loss pattern).
   • The loss is extracted across scanned transformer layers and injected into the scalar optimization objective (loss += indexer_loss), restoring full automatic differentiation VJP gradient flow to indexer Query/Key projection weights (wq_b, wkv_b).

4. Fix Indexer Load Balancing Mask Bug in Splash Attention

   • During Context Parallel load balancing, the indexer_mask was not being restored to chronological order before passing it to Splash Attention (unlike the Key and Value tensors, which were reordered).
   • We verified that without explicitly un-permuting the mask in wrap_flash_attention():

   indexer_mask = max_utils.reorder_sequence(tensor=indexer_mask, cp_size=cp_size, seq_dim=2, to_contiguous=True)
   

The newly added unit test test_tpu_flash_attention_context_parallel_with_indexer comparing multi-chip Flash Attention load-balancing against a single-chip baseline fails without it.

screenshot

Tests

• Verified Ring operations in trace xprof and non-zero indexer_loss from logs

Image

HLO

• configs_value_test.py

  • Added: test_tpu_tokamax_ring_config_validation_accepts_indexer to verify that pyconfig.initialize successfully accepts the combination of MLA, the DSA Sparse Indexer, and Tokamax Ring Attention.
  • Updated: Removed the outdated indexer rejection parameterization from test_tpu_tokamax_ring_config_validation_rejects_unsupported_configs.

• tokamax_ring_attention_test.py

  • Added: test_call_ring_attention_threads_indexer_mask_without_segment_ids to test batch vectorization and mapping threads on the dynamic indexer_mask without segmentation IDs.
  • Added: test_call_ring_attention_threads_indexer_mask_with_segment_ids to test batch vectorization and mapping threads on the dynamic indexer_mask when utilizing segmentation IDs.

• attention_test.py

  • Added: test_tpu_dot_product_context_parallel_with_indexer — Parameterized forward equivalence test verifying that standard dot-product MLA + Indexer calculates logits successfully and accurately across standard context parallelism permutations (testing combinations of matrix sizes and load_balance=True/False).
  • Added: test_tpu_flash_attention_context_parallel_with_indexer — Parameterized forward equivalence test verifying that standard Flash Attention MLA + Indexer matches reference outputs properly across standard context parallelism combinations.
  • Added: test_tpu_flash_attention_ring_context_parallel_with_indexer — Parameterized forward equivalence test verifying that MLA + Indexer under Tokamax Ring Attention identically matches the logits of single-device generic dot-product MLA + Indexer.
  • Added: test_tpu_flash_attention_ring_context_parallel_grad_with_indexer — Parameterized backward gradient equivalence test verifying that the backward input gradients and auxiliary indexer_losses match bit-for-bit with the reference gradients when utilizing is_dkv=True transposition.

• train_nnx_test.py

  • Added: test_indexer_losses_harvested_and_injected_into_loss to simulate multi-layer Transformer intermediate loss states and accurately validate the extraction and calculation of the expected_indexer_loss scalar objective.

  • Ran deepseek32_vs_reference_test.py (we werify MaxText Dot Product Baseline matches Official PyTorch version, and our previous tests show that Ring Attention version matches Dot product baseline. Because Ring Attention = Dot Product Baseline and Dot Product Baseline = PyTorch Reference, this shows that distributed Ring Attention with Indexer matches the reference implementation):

tests/unit/deepseek32_vs_reference_test.py::DeepseekV32IndexerTest::test_indexer_match0
  PASSED [ 12%]
    tests/unit/deepseek32_vs_reference_test.py::DeepseekV32IndexerTest::test_indexer_match1
  PASSED [ 25%]
    tests/unit/deepseek32_vs_reference_test.py::DeepseekV32IndexerTest::test_indexer_match2
  PASSED [ 37%]
    tests/unit/deepseek32_vs_reference_test.py::DeepseekV32MLATest::test_mla_parity_dot_product_s128_k128
  PASSED [ 50%]
    tests/unit/deepseek32_vs_reference_test.
  py::DeepseekV32MLATest::test_mla_parity_dot_product_s128_k128_c4 PASSED [ 62%]
    tests/unit/deepseek32_vs_reference_test.py::DeepseekV32MLATest::test_mla_parity_dot_product_s128_k4
  PASSED [ 75%]
    tests/unit/deepseek32_vs_reference_test.py::DeepseekV32MLATest::test_mla_parity_dot_product_s2_k4
  PASSED [ 87%]
    tests/unit/deepseek32_vs_reference_test.py::DeepseekV32MLATest::test_mla_parity_dot_product_s8_k4
  PASSED [100%]

    ======================== 8 passed in 56.02s ========================

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.

[gemini-code-assist]

Code Review

This pull request enables support for sparse indexer masks within the TPU Tokamax ring attention kernel, integrating indexer loss logging and auxiliary loss calculations into the training loop. The review feedback highlights a few critical runtime issues: a mismatch in the evaluation metric dictionary key for logging indexer loss, an incorrect keyword argument (_shape instead of shape) when instantiating FullMask, and a potential ValueError when concatenating zero-dimensional arrays in the loss calculation when scan_layers=False.

[zcjhao]

LGTM at high level! One question about indexer loss. When onboarded, we were testing Linen instead of NNX. Could you have a run with old version to see if your new changes align with previous runs? Thanks!

enable_nnx: false
pure_nnx_decoder: false
pure_nnx: false

It seems Linen is currently broken for DSA, Line, and the same configs succeed in nnx (Should I add a check that rejects linen + Indexer?). What happened was that I suddenly stopped seeing Indexer loss calculation in my traces, and I found that it seems that due to some recent changes of NNX, indexer_losses silently defaulted to 0 (breaking loss function and backward for Indexer), just like in #4525 for mtp_loss because

maxtext/src/maxtext/layers/nnx_decoders.py

Line 1092 in 5e89a48

clean_state = nnx.filter_state(scanned_state, nnx.Not((nnx.RngState, nnx.Intermediate)))

that dropped all nnx.Intermediate after a layer pass. I followed I similar approach by creating a custom subclass for indexer_losses, making sure it does not get dropped by the nnx.Intermediate filter and manually pop it in train.py. I believe the only part this would affect is how indexer_losses is handled, and I have checked that it correctly appears again in the logs as shown in the PR description, all other components remain the same, I believe comparison with Linen is not required.

[github-actions]

🤖 Hi @zcjhao, 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 introduces comprehensive, end-to-end support for DeepSeek Sparse Attention (DSA) Indexer with Tokamax Ring Context Parallelism during training. The implementation includes dynamic per-ring-step mask extraction for both the forward and backward passes, correct handling of the auxiliary Flax NNX loss variables under scanned layers, and load balancing mask reordering fixes. The codebase is well-structured, follows established project conventions, and includes solid test coverage for all new functionality.

🔍 General Feedback

• Excellent Architecture & Design: The slicing and tiling logic of the global indexer mask into per-step KV-shard blocks, and aligning transposition schedules for the backward pass with the hardware scheduler, is extremely well thought out and executed.
• Robust Integration with NNX: Defining class indexer_losses directly as an nnx.Variable instead of inheriting from nnx.Intermediate perfectly solves the state-pruning issue during scanned layer scans.
• Thorough Test Coverage: The addition of parameterized unit tests verifying standard dot-product, Flash Attention, and Ring Context Parallelism equivalence is excellent.
• Areas for Improvement: Suggested minor fixes to prevent test hardcoding (so parameterized test coverage of higher context parallelism isn't overridden) and a potential performance optimization to move mask boolean conversion out of scanned hot loops.

[huytransformer]

Can u please address these 2 then this LGTM

[huytransformer]

Can u check if dq_reduction_step=3 give correct dQ when ring attention use indexer mask? If not perhaps reject this combination for now?

[huytransformer]

can u add indexer_losses to Decoder.scan_decoder_layers() and test with pure_nnx=False and scan_layers=True