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.

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

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 accepts the combination of MLA, Sparse Indexer, and Tokamax Ring Attention.
  • Updated: Removed the old indexer rejection case 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: Tests batch vectorization and threading of dynamic indexer_mask without segmentation IDs.
  • Added: test_call_ring_attention_threads_indexer_mask_with_segment_ids: Tests batch vectorization and threading of dynamic indexer_mask with segmentation IDs.
• attention_test.py:
  • Added: test_tpu_flash_attention_ring_context_parallel_with_indexer: Parameterized forward equivalence test (load_balance=True/False) verifying that MLA + Indexer under Tokamax Ring Attention matches single-device generic dot-product MLA + Indexer.
  • Added: test_tpu_flash_attention_ring_context_parallel_grad_with_indexer: Parameterized backward gradient equivalence test (load_balance=True/False) verifying that backward input gradients and auxiliary indexer_losses match bit-for-bit with is_dkv=True transposition.
• Added: test_indexer_losses_harvested_and_injected_into_loss to train_nnx_test.py: simulates multi-layer transformer intermediate loss and calculates expected_indexer_loss.
• 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.

[codecov]

Codecov Report

❌ Patch coverage is 23.59551% with 68 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
.../tokamax_splash_attention/ring_attention_kernel.py	1.81%	54 Missing ⚠️
...axtext/kernels/attention/tokamax_ring_attention.py	46.15%	6 Missing and 1 partial ⚠️
src/maxtext/common/metric_logger.py	0.00%	4 Missing and 2 partials ⚠️
src/maxtext/layers/attention_op.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[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 end-to-end integration for the DeepSeek Sparse Attention (DSA) Indexer under Tokamax Ring Context Parallelism during training and evaluation. The design correctly implements dynamic per-ring-step slicing and rotation of indexer masks in both forward and backward passes, and addresses Flax NNX auxiliary loss harvesting by defining a dedicated indexer_losses variable.

🔍 General Feedback

• Excellent Architectural Alignment: Defining class indexer_losses(nnx.Variable) to bypass scanned layer intermediate filters is a brilliant and clean solution that integrates perfectly with upstream patterns like MTP losses.
• Robust Integration Testing: The additions of parameterized equivalence tests comparing distributed Ring Context Parallel MLA+Indexer outputs/gradients against single-device dot-product baselines are outstanding and provide high assurance of numerical correctness.
• Clear Documentation: The inclusion of ASCII forward and backward data flow diagrams in the PR description is extremely helpful for understanding the complex grid scheduling and transposition logic.

[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 end-to-end support for DeepSeek Sparse Attention (DSA) Indexer with Tokamax Ring Context Parallelism during training. The implementation includes crucial enhancements such as dynamic per-ring-step mask extraction, transposition for backward pass schedule alignment, and utilizing a custom Flax NNX Variable class to prevent auxiliary loss pruning in scanned JAX loops.

🔍 General Feedback

• High Quality Architecture: The design choices are exceptional, particularly the transposition of blocks for $dK/dV$ backward-pass alignment and utilizing a custom nnx.Variable class to elegant bypass scan filters.
• Robust and Comprehensive Testing: The PR is extremely well-tested with parameterized forward/backward gradient equivalence tests and unit tests covering dynamic mask batch-vectorization.
• Backwards Compatibility: Modifications are backwards-compatible and preserve original behavior when the sparse indexer or Tokamax ring context parallelism is not enabled.

[github-actions]

🟢 Minor style note: The leading newline `\n` in the log message can make the log output formatting less consistent. Removing it is recommended.

Suggested change

	moe_lb_loss = 0.0
	max_logging.debug("No Indexer loss found. Defaulting to 0.0.")

[RissyRan]

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

[RissyRan]

cc @huytransformer helps take a review on kernel part