Support Ring Attention with DeepSeek DSA Sparse Indexer and others

src/maxtext/common/metric_logger.py

@@ -223,6 +223,10 @@ def _log_training_metrics(self, metrics, step):
       log_parts.append(f"main_model_loss: {loss - mtp_loss:.3f}")
       log_parts.append(f"mtp_loss: {mtp_loss:.3f}")
 
+    if getattr(self.config, "use_indexer", False):
+      indexer_l = scalars.get("learning/indexer_loss", 0.0)
+      log_parts.append(f"indexer_loss: {float(indexer_l)}")
+
     max_logging.log(", ".join(log_parts))
 
   def _log_eval_metrics(self, metrics, step):
@@ -246,6 +250,11 @@ def _log_eval_metrics(self, metrics, step):
       )
     if "eval/avg_dpo_reward_accuracy" in scalars:
       log_parts.append(f"dpo_reward_accuracy={scalars['eval/avg_dpo_reward_accuracy']:.3f}")
+
+    if getattr(self.config, "use_indexer", False):
+      indexer_l = scalars.get("eval/avg_indexer_loss", 0.0)
+      log_parts.append(f"avg_indexer_loss={indexer_l:.3f}")
+
     max_logging.log(", ".join(log_parts))
 
   def _log_running_eval_metrics(self, metrics, step):

src/maxtext/configs/base.yml

@@ -193,7 +193,7 @@ logits_dot_in_fp32: false  # whether to use fp32 in logits_dense or shared_embed
 cast_logits_to_fp32: true # whether to cast the logits to fp32. the higher precision is generally beneficial, but it can vary slightly.
 float32_qk_product: false # in dot_product attention, whether to cast to fp32 the inputs to qk product
 float32_logits: false # in dot_product attention, whether to cast to fp32 the inputs to softmax
-mla_qk_head_chunk_size: 0 # Limits HBM footprint by sequentially evaluating the QK matrix across the unsharded local heads dimension natively.
+mla_qk_head_chunk_size: 0 # Limits HBM footprint by sequentially evaluating the QK matrix in the Indexer across the unsharded local heads dimension natively.
 float32_weight_sum: true # whether to use full fp32 precision to sum expert weights for numerical stability
 float32_gate_logits: false # whether to cast inputs to fp32 to compute MoE gate logits for numerical stability
 

src/maxtext/configs/types.py

@@ -664,7 +664,7 @@ class Attention(BaseModel):
       0,
       ge=0,
       description=(
-          "Chunk size over heads dimension for QK attention dot product in mla. "
+          "Chunk size over heads dimension for QK attention dot product in mla.  "
           "Default is 0 (no chunking). Reduces memory footprint at the cost of time."
       ),
   )
@@ -2942,6 +2942,8 @@ def _validate_usp_context_parallelism(self):
       raise ValueError("TPU USP attention requires use_tokamax_splash=True.")
     if self.use_jax_splash:
       raise ValueError("TPU USP attention requires use_jax_splash=False.")
+    if self.use_indexer:
+      raise ValueError("TPU USP attention does not support sparse indexer masks.")
     if self.attention_type != "global":
       raise ValueError("TPU USP attention is initially supported only for global causal attention.")
     if self.packing:
@@ -2952,8 +2954,6 @@ def _validate_usp_context_parallelism(self):
       raise ValueError("TPU USP attention does not support ragged attention.")
     if self.attention_sink:
       raise ValueError("TPU USP attention does not support attention sinks.")
-    if self.use_indexer:
-      raise ValueError("TPU USP attention does not support sparse indexer masks.")
     if self.use_chunked_prefill:
       raise ValueError("TPU USP attention does not support chunked prefill yet.")
     if self.use_multimodal:
@@ -3639,7 +3639,13 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
             f"`mla_qk_head_chunk_size` ({self.mla_qk_head_chunk_size}) must cleanly divide exactly into "
             f"`indexer_n_heads` ({self.indexer_n_heads})."
         )
+
     if self.use_indexer:
+      if self.attention_type != AttentionType.MLA.value:
+        raise ValueError(
+            f"`use_indexer=True` requires `attention_type='{AttentionType.MLA.value}'`, since only the "
+            "MLA indexer produces this mask."
+        )
       if self.q_lora_rank == 0:
         raise NotImplementedError("Sparse indexer has not implemented for q_lora_rank = 0.")
       supports_dot_product = self.attention == "dot_product"
@@ -3816,6 +3822,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         raise ValueError("TPU Tokamax ring attention requires context_parallel_size > 1.")
       if self.context_sharding != "context":
         raise ValueError("TPU Tokamax ring attention requires context_sharding='context'.")
+      if self.use_indexer and self.dq_reduction_steps != 0:
+        raise ValueError("TPU Tokamax ring attention with sparse indexer mask only supports dq_reduction_steps=0.")
       if self.dq_reduction_steps not in (0, 3):
         raise ValueError("TPU Tokamax ring attention requires dq_reduction_steps to be 0 or 3.")
       if self.max_target_length % (context_parallel_size * context_parallel_size) != 0:
@@ -3844,8 +3852,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         raise ValueError("TPU Tokamax ring attention does not support ragged attention.")
       if self.attention_sink:
         raise ValueError("TPU Tokamax ring attention does not support attention sinks.")
-      if self.use_indexer:
-        raise ValueError("TPU Tokamax ring attention does not support sparse indexer masks.")
       if self.use_chunked_prefill:
         raise ValueError("TPU Tokamax ring attention does not support chunked prefill yet.")
       if self.moba:
@@ -3883,6 +3889,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         raise ValueError("TPU Ulysses attention requires use_tokamax_splash=True.")
       if self.use_jax_splash:
         raise ValueError("TPU Ulysses attention requires use_jax_splash=False.")
+      if self.use_indexer:
+        raise ValueError("TPU Ulysses attention does not support sparse indexer masks.")
       if self.attention_type != "global":
         raise ValueError("TPU Ulysses attention is initially supported only for global causal attention.")
       if self.context_parallel_load_balance:
@@ -3895,8 +3903,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         raise ValueError("TPU Ulysses attention does not support ragged attention.")
       if self.attention_sink:
         raise ValueError("TPU Ulysses attention does not support attention sinks.")
-      if self.use_indexer:
-        raise ValueError("TPU Ulysses attention does not support sparse indexer masks.")
       if self.use_chunked_prefill:
         raise ValueError("TPU Ulysses attention does not support chunked prefill yet.")
       if self.use_multimodal:

src/maxtext/kernels/attention/tokamax_ring_attention.py

@@ -110,8 +110,6 @@ def validate_tokamax_ring_runtime(
     raise ValueError("TPU Tokamax ring attention does not support chunked prefill yet.")
   if sinks is not None:
     raise ValueError("TPU Tokamax ring attention does not support attention sinks.")
-  if indexer_mask is not None:
-    raise ValueError("TPU Tokamax ring attention does not support indexer masks.")
   if bidirectional_mask is not None:
     raise ValueError("TPU Tokamax ring attention does not support bidirectional masks.")
   if record_max_logits:
@@ -283,6 +281,7 @@ def make_sharded_ring_attention_kernel(
     ring_axis: str,
     attn_logits_soft_cap: float | None,
     maybe_shard_with_pspec: Any,
+    mask: Any = None,
 ):
   """Builds and shards the Tokamax ring attention kernel for MaxText."""
   splash_config = build_splash_config(
@@ -295,11 +294,17 @@ def make_sharded_ring_attention_kernel(
   if config.use_max_logit_estimate > 0:
     splash_config = dataclasses.replace(splash_config, max_logit_const=config.use_max_logit_estimate)
 
-  mask = _make_causal_mask(
-      (query.shape[2], key.shape[2]),
-      context_parallel_size,
-      load_balanced=config.context_parallel_load_balance,
-  )
+  if mask is None:
+    # When using the indexer, causal masking is unified into the dynamic indexer_mask
+    # and applied dynamically per block; use FullMask to avoid duplicate static masks.
+    if getattr(config, "use_indexer", False):
+      mask = tokamax_splash_mask.FullMask((query.shape[2], key.shape[2]))
+    else:
+      mask = _make_causal_mask(
+          (query.shape[2], key.shape[2]),
+          context_parallel_size,
+          load_balanced=config.context_parallel_load_balance,
+      )
 
   @functools.partial(jax.jit, static_argnames=["single_head_mask"])
   def wrap_ring_kernel(single_head_mask):
@@ -331,15 +336,27 @@ def call_ring_attention(
     decoder_segment_ids_q: Any,
     decoder_segment_ids_kv: Any,
     ring_kernel: Any,
+    indexer_mask: Any = None,
 ):
   """Calls a Tokamax ring attention kernel over the MaxText batch dimension."""
   if (decoder_segment_ids_q is None) != (decoder_segment_ids_kv is None):
     raise ValueError("decoder_segment_ids_q and decoder_segment_ids_kv must both be set or both be None.")
+  # Vectorize execution across batch dimension, threading indexer_mask when present.
+  # Note: ring_kernel expects positional arguments (q, k, v, segment_ids, sinks, indexer_mask).
   if decoder_segment_ids_q is None:
-    return jax.vmap(lambda q, k, v: ring_kernel(q, k, v, None), in_axes=(0, 0, 0))(query, key, value)
-
-  def call_one(q, k, v, q_segment_ids, kv_segment_ids):
+    if indexer_mask is None:
+      return jax.vmap(lambda q, k, v: ring_kernel(q, k, v, None, None, None), in_axes=(0, 0, 0))(query, key, value)
+    return jax.vmap(
+        lambda q, k, v, im: ring_kernel(q, k, v, None, None, im),
+        in_axes=(0, 0, 0, 0),
+    )(query, key, value, indexer_mask)
+
+  def call_one(q, k, v, q_segment_ids, kv_segment_ids, im=None):
     segment_ids = ring_attention_kernel.SegmentIds(q_segment_ids, kv_segment_ids)
-    return ring_kernel(q, k, v, segment_ids)
+    return ring_kernel(q, k, v, segment_ids, None, im)
 
-  return jax.vmap(call_one, in_axes=(0, 0, 0, 0, 0))(query, key, value, decoder_segment_ids_q, decoder_segment_ids_kv)
+  if indexer_mask is None:
+    return jax.vmap(call_one, in_axes=(0, 0, 0, 0, 0))(query, key, value, decoder_segment_ids_q, decoder_segment_ids_kv)
+  return jax.vmap(call_one, in_axes=(0, 0, 0, 0, 0, 0))(
+      query, key, value, decoder_segment_ids_q, decoder_segment_ids_kv, indexer_mask
+  )