Optimize MLA generate_mask with scatter update

[JHCuc3m]

Description

This PR optimizes the generate_mask function in the Multi-head Latent Attention (MLA) sparse indexer by replacing a high-complexity broadcasted comparison with a JAX scatter-based update.

• The Problem: The previous implementation (jnp.arange(s) == topk_indices[..., None]).any(axis=-2) had an algorithmic complexity of $O(b \cdot t \cdot k \cdot s)$. At a 128K context window, which theoretically creates a massive virtual intermediate tensor of ~274 Billion elements (~256 GB). On TPU, although the XLA compiler successfully fused the loop to avoid physical HBM allocation (reducing HBM traffic to ~142 MB), the hardware remained heavily bottlenecked by the sheer volume of element-wise ALU comparisons (274 Billion ops) executed in registers, taking ~153 ms per step.

• The Solution: The optimized version initializes the mask to DEFAULT_MASK_VALUE $O(b \cdot t \cdot s)$ and then uses indexing to scatter-write 0.0 at the selected topk_indices $O(b \cdot t \cdot k)$, where $k \ll s$ (2,048 vs 131,072).

• Implementation Details: Initializes the mask using jnp.full and performs the scatter update using broadcasted batch and time indices (batch_indices = jnp.arange(b)[:, None, None], time_indices = jnp.arange(t)[None, :, None]) via mask.at[batch_indices, time_indices, topk_indices].set(0.0).

Tests

1. Mathematical Equivalence Test

Added test_generate_mask_equivalence to tests/unit/attention_test.py to verify that the new scatter-based implementation produces identical results to the old broadcast-based implementation.

• Result: Passed (exact match within atol=1e-5).

2. Regression Testing

Ran the existing unit test suite to ensure no regressions in attention or model functionality:

• tests/unit/attention_test.py (Passed)
• tests/unit/model_test.py (Passed)

Checklist

• 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. (Validation run zj-scatter-mask-val is currently running on the cluster).
• 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!