Optimize MLA generate_mask

[JHCuc3m]

Description

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.

• Two options:
  • Option 1 (indexer_mask_exact_topk = False): Executes raw vectorized threshold comparison. Allows > k tokens under boundary ties.
  • Option 2 (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.

Context & Implementation Details

• The Problem: The baseline implementation (jnp.arange(s) == topk_indices[..., None]).any(axis=-2) executes an all-pairs broadcast index comparison.
• The Solution: We leverage the $k$-th score from top_k as a decision threshold (cutoff = topk_values[..., -1:]).
  Value-thresholding (score >= cutoff) converts the random index-matching query into an element-to a single value comparison.
• Configurable Tie Handling: Under boundary ties, raw thresholding unmasks all tied tokens (allows $>k$ tokens). To allow strict density control, Option 2 executes a sequence-wide cumsum over active threshold boundaries. This counts and filters the unmasked elements by selection rank.

Tests

Correctness & Boundary Verification

Added unit tests in tests/unit/attention_test.py asserting that vectorized threshold cutoff masking:

• Matches exact top-k selection when scores are unique (test_generate_mask_threshold_equivalence).
• Unmasks $k+x$ elements under ties when indexer_mask_exact_topk = False (raw cutoff,
  test_generate_mask_threshold_ties_raw).
• Unmasks exactly $k$ elements under ties when indexer_mask_exact_topk = True (prefix-sum pruning,
  test_generate_mask_threshold_ties_exact_k).
• Edge case when k is less than sequence length test_generate_mask_sequence_smaller_than_k.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

[✓] I have performed a self-review of my code.
[✓] 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.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @RissyRan, 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 PR introduces an optimized generate_mask implementation for Multi-head Latent Attention (MLA) indexers, replacing all-pairs broadcast comparisons with a TPU-native vectorized value thresholding. While the performance optimization and the configuration-driven tie-breaker are excellent, there is a critical logic error in the prefix-sum pruning logic which could result in masking out the top elements of the sequence when they appear after boundary ties.

🔍 General Feedback

• Smarter Vectorization: Moving from all-pairs broadcasting comparison (jnp.arange(s) == topk_indices[..., None]).any(...) to a single cutoff threshold score >= cutoff_threshold is a high-impact performance improvement on TPU hardware.
• Robustness in Tests: The added unit tests are well-structured, but they initially failed to catch the ordering dependency of the jnp.cumsum operation because the test inputs were sorted in descending order. Adding unsorted test inputs will prevent regressions.

[github-actions]

🔴 Critical - This implementation of exact-k pruning computes jnp.cumsum directly over raw_mask (which includes all elements >= cutoff_threshold). Because jnp.cumsum processes elements from left to right along the sequence dimension, it will assign smaller rank values to elements at the beginning of the sequence regardless of their actual score value.

If scores strictly greater than the cutoff threshold appear later in the sequence, their cumulative rank can exceed k and they will be incorrectly masked out. For example, if scores are [0.8, 0.8, 0.9, 0.8] with $k = 2$ and a cutoff of 0.8, raw_mask is [True, True, True, True], and cumsum selects the first two elements [0.8, 0.8] at indices 0 and 1, incorrectly masking out the maximum value 0.9 at index 2. This violates the fundamental top-k attention invariant.

To resolve this, we should always keep all elements strictly greater than the cutoff, and only prune elements that are exactly equal to the cutoff using a cumulative sum over the tie elements.

Suggested change

	selected_rank = jnp.cumsum(raw_mask.astype(jnp.int32), axis=-1)
	return jnp.where((raw_mask) & (selected_rank <= topk_values.shape[-1]), val_true, val_false)
	else:
	# Raw threshold cutoff masking (optional: enables speedups, but may unmask > k tokens under indexer scores ties)
	if self.config.indexer_mask_exact_topk:
	# Keep all tokens strictly greater than cutoff, and prune ties (equal to cutoff)
	is_strictly_greater = indexer_score > cutoff_threshold
	is_equal = indexer_score == cutoff_threshold
	num_strictly_greater = jnp.sum(is_strictly_greater, axis=-1, keepdims=True)
	num_equal_to_keep = topk_values.shape[-1] - num_strictly_greater
	equal_rank = jnp.cumsum(is_equal.astype(jnp.int32), axis=-1)
	selected = is_strictly_greater | (is_equal & (equal_rank <= num_equal_to_keep))
	return jnp.where(selected, val_true, val_false)

[RissyRan]

Could you cross check and also keep the original implementation if possible?

[github-actions]

🟡 Medium - The current unit test uses a score array sorted in descending order, which masked the fact that jnp.cumsum on the entire raw mask was ordering-sensitive. We should add an unsorted test case to verify that elements strictly greater than the cutoff threshold are correctly preserved even if they occur after boundary ties.

Suggested change

	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.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), 3) # Exactly 3 tokens (exact k) unmasked
	self.assertEqual(jnp.sum(mask[0, 1] == 0.0), 3) # Exactly 3 tokens (exact k) unmasked
	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.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), 3) # Exactly 3 tokens (exact k) unmasked
	self.assertEqual(jnp.sum(mask[0, 1] == 0.0), 3) # Exactly 3 tokens (exact k) unmasked
	# Verify correctness when higher scores are located after boundary ties (non-sorted sequence)
	scores_unsorted = jnp.array(
	[
	[
	[0.5, 0.5, 0.9, 0.8, 0.1, 0.0, -1.0, -2.0, -3.0, -4.0],
	[0.5, 0.5, 0.7, 0.6, 0.1, 0.0, -1.0, -2.0, -3.0, -4.0],
	]
	],
	dtype=dtype,
	)
	topk_values_unsorted, _ = jax.lax.top_k(scores_unsorted, k=k)
	mask_unsorted = mla.indexer.generate_mask(scores_unsorted, topk_values_unsorted)
	self.assertEqual(jnp.sum(mask_unsorted[0, 0] == 0.0), 3)
	self.assertEqual(jnp.sum(mask_unsorted[0, 1] == 0.0), 3)
	# Ensure that top elements are unmasked (mask value of 0.0)
	self.assertEqual(mask_unsorted[0, 0, 2], 0.0) # 0.9
	self.assertEqual(mask_unsorted[0, 0, 3], 0.0) # 0.8

[RissyRan]

Please consider the comment above.

[RissyRan]

Could you cross check and also keep the original implementation if possible?

[RissyRan]

Wondering if you see any perf regression due to this change, including throughput and memory usage?

[RissyRan]

Please consider the comment above.