Support Ring Attention with DeepSeek DSA Sparse Indexer

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/types.py

@@ -3497,6 +3497,11 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
             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"
@@ -3699,8 +3704,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:
@@ -3738,6 +3741,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:
@@ -3750,8 +3755,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
+  )