Optimize MLA generate_mask with TPU-native vectorized threshold cutoff and configurable exact-k pruning

Replace the broadcast comparison in `generate_mask` with a vectorized threshold cutoff comparison, adding the `indexer_mask_exact_topk` configuration flag to dynamically support raw thresholding or exact-k prefix-sum pruning while retaining `generate_mask_broadcast` for cross-checking in unit tests.

- Two options:
  - Option 1 (`indexer_mask_exact_topk = True`, default): Runs a cumulative prefix-sum pruning scan (`jnp.cumsum`) over active boundaries, keeping exactly the first k selected tokens and discarding tied overflows.
  - Option 2 (`indexer_mask_exact_topk = False`): Executes raw vectorized threshold comparison. Highly performant (enables speedups), but may unmask > k tokens under boundary ties.
- Verification:
  - Retained original broadcast implementation as `generate_mask_broadcast` and verified bit-for-bit equivalence in `attention_test.py` across equivalence, exact-k tie, and unsorted tests (Passed).
  - Added unit tests for equivalence, raw thresholding ties, exact-k pruning ties, and sequence-length bypass in `attention_test.py` (Passed).

TAG=agy
CONV=859f8f98-4c27-4a27-8df2-9fd1c7fa5d61

Author: zcjhao

src/maxtext/configs/base.yml

@@ -422,6 +422,7 @@ indexer_n_heads: 64
 indexer_topk: 2048
 indexer_use_approx_top_k: false
 indexer_approx_top_k_recall: 0.95
+indexer_mask_exact_topk: true
 # Determines the training strategy for the indexer:
 # - false (Dense Warm-up): Computes indexer loss over all tokens. Used with `trainable_parameters_mask` to freeze other model parameters.
 # - true (Sparse Training): Computes indexer loss over top-k tokens only and detaches the indexer input for independent optimization.

src/maxtext/configs/types.py

@@ -687,6 +687,14 @@ class AttentionIndexer(BaseModel):
       False, description="Whether to use approximate top-k selection for the indexer on TPU."
   )
   indexer_approx_top_k_recall: float = Field(0.95, description="Recall target for approximate top-k selection.")
+  indexer_mask_exact_topk: bool = Field(
+      True,
+      description=(
+          "When True, enforces that exactly k elements are unmasked by the indexer under boundary ties."
+          " When False, uses raw thresholding which is faster but may unmask > k elements"
+          " during ties."
+      ),
+  )
 
 
 class Llama4Attention(BaseModel):

src/maxtext/layers/attention_mla.py

@@ -221,26 +221,45 @@ def apply_partial_rope(
     x = jnp.concatenate([x_pe, x_nope], axis=-1)
     return x
 
-  def generate_mask(self, topk_indices, s):
-    """
-    Creates a mask for top-k indices.
+  def generate_mask(self, indexer_score: Array, topk_values: Array) -> Array:
+    """Creates a mask for top-k indices
 
     Args:
-        topk_indices: [b, t, k] int - The indices to keep.
-        s: int - The total size to select from.
+        indexer_score: [b, t, s] float - The computed relevance scores for all tokens.
+        topk_values: [b, t, k] float - The top-k selected score values from jax.lax.top_k.
 
     Returns:
-        mask: [b, t, s] - `0.0` at topk_indices, `DEFAULT_MASK_VALUE` (large negative) elsewhere.
+        mask: [b, t, s] - `0.0` for top-k selected elements, `DEFAULT_MASK_VALUE` elsewhere.
     """
-    # 1. Create a range [0, 1, ..., s-1]
-    # 2. Broadcast compare against [b, t, k] to get [b, t, k, s]
-    # 3. Use .any() to see if a s-index is present in any of the k slots
-    is_topk = (jnp.arange(s) == topk_indices[..., None]).any(axis=-2)
-    # 4. Use where to select between 0.0 and the mask value
-    # cast values to dtype
+    cutoff_threshold = topk_values[..., -1:]
+
     val_true = jnp.array(0.0, dtype=self.dtype)
     val_false = jnp.array(DEFAULT_MASK_VALUE, dtype=self.dtype)
-    return jnp.where(is_topk, val_true, val_false)
+
+    if self.config.indexer_mask_exact_topk:
+      # Prune ties by keeping only the first k unmasked tokens along sequence dimension
+      k = topk_values.shape[-1]
+      is_strictly_greater = indexer_score > cutoff_threshold
+      is_equal = indexer_score == cutoff_threshold
+
+      # Use cumsum to rank both strictly greater and equal tokens (XLA fuses these scans)
+      sg_rank = jnp.cumsum(is_strictly_greater.astype(jnp.int32), axis=-1)
+      eq_rank = jnp.cumsum(is_equal.astype(jnp.int32), axis=-1)
+
+      # Cap strictly greater elements to exactly k
+      sg_kept = is_strictly_greater & (sg_rank <= k)
+
+      # Calculate how many equals we still have room for
+      num_sg_kept = jnp.minimum(sg_rank[..., -1:], k)
+      num_eq_to_keep = k - num_sg_kept
+
+      selected = sg_kept | (is_equal & (eq_rank <= num_eq_to_keep))
+
+      return jnp.where(selected, val_true, val_false)
+    else:
+      # Raw threshold cutoff masking (optional: enables speedups, but may unmask > k tokens under indexer scores ties)
+      raw_mask = indexer_score >= cutoff_threshold
+      return jnp.where(raw_mask, val_true, val_false)
 
   def __call__(
       self,
@@ -367,16 +386,16 @@ def __call__(
 
     # TopK selection based on index score
     if self.config.indexer_use_approx_top_k:
-      _, topk_indices = jax.lax.approx_max_k(
+      topk_values, topk_indices = jax.lax.approx_max_k(
           indexer_score,
           k=self.indexer_topk,
           recall_target=self.config.indexer_approx_top_k_recall,
       )
     else:
-      _, topk_indices = jax.lax.top_k(indexer_score, k=self.indexer_topk)  # topk_indices [b, t, k]
+      topk_values, topk_indices = jax.lax.top_k(indexer_score, k=self.indexer_topk)  # [b, t, k]
 
     # Create Sparse Index Mask: 0 and large negatives
-    indexer_mask = self.generate_mask(topk_indices, k.shape[1])  # [b, t, s]
+    indexer_mask = self.generate_mask(indexer_score, topk_values)  # [b, t, s]
 
     # Re-apply attention mask after TopK: in case number of unmasked tokens < TopK
     if attention_mask is not None:

tests/unit/attention_test.py

@@ -2011,8 +2011,8 @@ def test_indexer_loss(self):
     attention_mask = self.get_causal_mask_for_indexer(batch_size, q_len, kv_len)
     indexer_score += attention_mask
 
-    topk_indices = jnp.array([[[0, 1], [0, 1], [0, 1]], [[0, 1], [0, 1], [0, 1]]])
-    indexer_mask = mla.indexer.generate_mask(topk_indices, kv_len) + attention_mask
+    topk_values, _ = jax.lax.top_k(indexer_score, k=2)
+    indexer_mask = mla.indexer.generate_mask(indexer_score, topk_values) + attention_mask
 
     loss_dense = mla.calculate_indexer_loss(
         indexer_score=indexer_score,
@@ -2061,8 +2061,8 @@ def test_indexer_loss_kl_divergence_zero(self):
     # Indexer score matches the shape and is uniform
     indexer_score = jnp.zeros((batch_size, q_len, kv_len)) + attention_mask
 
-    topk_indices = jnp.array([[[0, 1], [0, 1], [0, 1]], [[0, 1], [0, 1], [0, 1]]])
-    indexer_mask = mla.indexer.generate_mask(topk_indices, kv_len) + attention_mask
+    topk_values, _ = jax.lax.top_k(indexer_score, k=2)
+    indexer_mask = mla.indexer.generate_mask(indexer_score, topk_values) + attention_mask
 
     loss = mla.calculate_indexer_loss(
         indexer_score=indexer_score,
@@ -2230,6 +2230,220 @@ def full_indexer_loss_fn(inputs_q, inputs_kv, low_rank_q, mla, sparse_training=s
         self.assertTrue(jnp.all(grad_kv == 0.0))
         self.assertTrue(jnp.all(grad_low_rank_q == 0.0))
 
+  def old_generate_mask(self, topk_indices, s, dtype=jnp.float32):
+    """Old baseline implementation using pairwise broadcast comparison.
+
+    Retained exclusively in unit tests as a ground-truth reference for cross-checking mathematical equivalence.
+    """
+    is_topk = (jnp.arange(s) == topk_indices[..., None]).any(axis=-2)
+    val_true = jnp.array(0.0, dtype=dtype)
+    val_false = jnp.array(DEFAULT_MASK_VALUE, dtype=dtype)
+    return jnp.where(is_topk, val_true, val_false)
+
+  def test_generate_mask_threshold_equivalence(self):
+    """Verifies that TPU-native threshold cutoff masking matches exact top-k selection when scores are unique."""
+    mla_config_args = self.config_arguments.copy()
+    mla_config_args["use_indexer"] = True
+    mla_config_args["indexer_topk"] = 64
+    mla_config_args["attention"] = "dot_product"
+
+    _, mla = self.init_mla(mla_config_args, rope_type="default")
+
+    jax_rng = jax.random.PRNGKey(0)
+    b, t, s, k = 2, 128, 1024, 64
+    dtype = jnp.float32
+
+    scores = jax.random.normal(jax_rng, (b, t, s), dtype=dtype)
+    # Add tiny position-dependent epsilon so all scores are strictly unique (preventing sorting tie-breaker divergence)
+    scores = scores + jnp.arange(s, dtype=dtype) * (1e-6 / s)
+    topk_values, topk_indices = jax.lax.top_k(scores, k=k)
+
+    # Call original broadcast Indexer logic for cross-checking
+    mask_original = self.old_generate_mask(topk_indices, s, dtype=dtype)
+
+    # Call actual optimized Indexer logic
+    mask_threshold = mla.indexer.generate_mask(scores, topk_values)
+
+    self.assertTrue(jnp.allclose(mask_original, mask_threshold, atol=1e-5))
+
+  def test_generate_mask_threshold_ties_exact_k(self):
+    """Verifies that prefix-sum pruning guarantees exactly k unmasked tokens even with boundary ties."""
+    mla_config_args = self.config_arguments.copy()
+    mla_config_args["use_indexer"] = True
+    mla_config_args["indexer_topk"] = 3
+    mla_config_args["attention"] = "dot_product"
+    mla_config_args["indexer_mask_exact_topk"] = True
+
+    _, mla = self.init_mla(mla_config_args, rope_type="default")
+
+    k = 3
+    dtype = jnp.float32
+    scores = jnp.array(
+        [
+            [
+                [0.9, 0.8, 0.5, 0.5, 0.1, 0.0, -1.0, -2.0, -3.0, -4.0],
+                [0.5, 0.5, 0.9, 0.8, 0.1, 0.0, -1.0, -2.0, -3.0, -4.0],
+            ]
+        ],
+        dtype=dtype,
+    )
+
+    topk_values, topk_indices = jax.lax.top_k(scores, k=k)
+    mask = mla.indexer.generate_mask(scores, topk_values)
+    mask_original = self.old_generate_mask(topk_indices, s=scores.shape[-1], dtype=dtype)
+
+    val_true = jnp.array(0.0, dtype=dtype)
+
+    self.assertFalse(jnp.isnan(mask).any())
+    self.assertEqual(jnp.sum(mask[0, 0] == val_true), 3)  # Exactly 3 tokens (exact k) unmasked
+    self.assertEqual(jnp.sum(mask[0, 1] == val_true), 3)  # Exactly 3 tokens (exact k) unmasked
+
+    # Assert equivalence to original broadcast baseline
+    self.assertTrue(jnp.allclose(mask_original, mask, atol=1e-5))
+
+    # Assert exact unmasked elements
+    np.testing.assert_array_equal(
+        mask[0, 0] == val_true,
+        [True, True, True, False, False, False, False, False, False, False],
+    )
+    np.testing.assert_array_equal(
+        mask[0, 1] == val_true,
+        [True, False, True, True, False, False, False, False, False, False],
+    )
+
+  def test_generate_mask_threshold_ties_unsorted(self):
+    """Verifies that elements strictly greater than cutoff are preserved even if they appear after ties."""
+    mla_config_args = self.config_arguments.copy()
+    mla_config_args["use_indexer"] = True
+    mla_config_args["indexer_topk"] = 3
+    mla_config_args["attention"] = "dot_product"
+    mla_config_args["indexer_mask_exact_topk"] = True
+
+    _, mla = self.init_mla(mla_config_args, rope_type="default")
+
+    k = 3
+    dtype = jnp.float32
+    scores = jnp.array(
+        [
+            [
+                # 0.9 is strictly greater but appears after three 0.5s.
+                # If cumsum was used unconditionally on (score >= cutoff), 0.9 would get rank 4 and be masked out!
+                # Correct behavior: keep 0.9, and the first two 0.5s to reach exactly k=3.
+                [0.5, 0.5, 0.5, 0.9, 0.1, 0.0, -1.0, -2.0, -3.0, -4.0],
+            ]
+        ],
+        dtype=dtype,
+    )
+
+    topk_values, topk_indices = jax.lax.top_k(scores, k=k)
+    mask = mla.indexer.generate_mask(scores, topk_values)
+    mask_original = self.old_generate_mask(topk_indices, s=scores.shape[-1], dtype=dtype)
+
+    val_true = jnp.array(0.0, dtype=dtype)
+
+    self.assertFalse(jnp.isnan(mask).any())
+    self.assertEqual(jnp.sum(mask[0, 0] == val_true), 3)  # Exactly 3 tokens unmasked
+
+    # Assert equivalence to original broadcast baseline
+    self.assertTrue(jnp.allclose(mask_original, mask, atol=1e-5))
+
+    np.testing.assert_array_equal(
+        mask[0, 0] == val_true,
+        [True, True, False, True, False, False, False, False, False, False],
+    )
+
+  def test_generate_mask_approx_k_overflow(self):
+    """Verifies exact-k guarantee when approx_top_k underestimates the threshold."""
+    mla_config_args = self.config_arguments.copy()
+    mla_config_args["use_indexer"] = True
+    mla_config_args["indexer_topk"] = 3
+    mla_config_args["attention"] = "dot_product"
+    mla_config_args["indexer_mask_exact_topk"] = True
+
+    _, mla = self.init_mla(mla_config_args, rope_type="default")
+
+    dtype = jnp.float32
+    scores = jnp.array(
+        [
+            [
+                # Simulating approx_max_k returning an underestimated threshold of 0.5.
+                # However, 0.9, 0.8, 0.7, 0.6 (4 elements) are strictly > 0.5.
+                [0.9, 0.8, 0.7, 0.6, 0.5, 0.0, -1.0, -2.0, -3.0, -4.0],
+            ]
+        ],
+        dtype=dtype,
+    )
+
+    # Artificially supply a threshold of 0.5 at the end
+    topk_values = jnp.array([[[1.0, 1.0, 0.5]]], dtype=dtype)
+    mask = mla.indexer.generate_mask(scores, topk_values)
+
+    val_true = jnp.array(0.0, dtype=dtype)
+
+    self.assertFalse(jnp.isnan(mask).any())
+    self.assertEqual(jnp.sum(mask[0, 0] == val_true), 3)  # Exactly 3 tokens unmasked
+
+    # It should keep the first 3 elements that are > 0.5 which are 0.9, 0.8, 0.7
+    np.testing.assert_array_equal(
+        mask[0, 0] == val_true,
+        [True, True, True, False, False, False, False, False, False, False],
+    )
+
+  def test_generate_mask_threshold_ties_raw(self):
+    """Verifies that raw thresholding allows more than k unmasked tokens under boundary ties."""
+    mla_config_args = self.config_arguments.copy()
+    mla_config_args["use_indexer"] = True
+    mla_config_args["indexer_topk"] = 3
+    mla_config_args["attention"] = "dot_product"
+    mla_config_args["indexer_mask_exact_topk"] = False
+
+    _, mla = self.init_mla(mla_config_args, rope_type="default")
+
+    k = 3
+    dtype = jnp.float32
+    scores = jnp.array(
+        [
+            [
+                [0.9, 0.8, 0.5, 0.5, 0.1, 0.0, -1.0, -2.0, -3.0, -4.0],
+                [0.7, 0.6, 0.5, 0.2, 0.1, 0.0, -1.0, -2.0, -3.0, -4.0],
+            ]
+        ],
+        dtype=dtype,
+    )
+
+    topk_values, _ = jax.lax.top_k(scores, k=k)
+    mask = mla.indexer.generate_mask(scores, topk_values)
+
+    self.assertFalse(jnp.isnan(mask).any())
+    self.assertEqual(jnp.sum(mask[0, 0] == 0.0), 4)  # 4 tokens unmasked (tied >= 0.5)
+    self.assertEqual(jnp.sum(mask[0, 1] == 0.0), 3)  # Exactly 3 tokens unmasked (no boundary ties)
+
+  def test_generate_mask_sequence_smaller_than_k(self):
+    """Verifies that the indexer handles sequence length smaller than or equal to k by returning None."""
+    mla_config_args = self.config_arguments.copy()
+    mla_config_args["use_indexer"] = True
+    mla_config_args["indexer_topk"] = 10  # k = 10
+    mla_config_args["attention"] = "dot_product"
+
+    cfg, mla = self.init_mla(mla_config_args, rope_type="default")
+
+    dtype = jnp.float32
+    inputs_q = jnp.zeros((1, 5, cfg.emb_dim), dtype=dtype)
+    inputs_kv = jnp.zeros((1, 5, cfg.emb_dim), dtype=dtype)  # s = 5 <= k
+    low_rank_q = jnp.zeros((1, 5, cfg.q_lora_rank), dtype=dtype)
+    inputs_positions = jnp.zeros((1, 5), dtype=jnp.int32)
+
+    mask, indices, score = mla.indexer(
+        inputs_q=inputs_q,
+        low_rank_q=low_rank_q,
+        inputs_kv=inputs_kv,
+        inputs_positions=inputs_positions,
+    )
+
+    self.assertIsNone(mask)
+    self.assertIsNone(indices)
+    self.assertIsNone(score)
+
 
 class Qwen3NextGatedDeltaNetTest(unittest.TestCase):
   """Test for the Gated Delta Net in Qwen3-Next"""