Fix TPU kernel OOB memory access

[darisoy]

Description

This PR fixes a critical TPU hardware crash (Denial of Service / Machine Check Exception) caused by a negative prefetch index in MaxText ragged attention kernels, and establishes full CPU-only CI test coverage using JAX Pallas Interpret mode.

Proposed Changes & Implementation:

1. Negative Prefetch Index Prevention: Clamped last_good_block in ragged_attention.py to 0 using jnp.maximum(0, ...) to ensure empty sequence paddings safely access valid bounds, preventing out-of-bounds TPU hardware crashes (CWE-125).
2. Preserved Intermediate Scaling & Documented: Retained the intermediate o = o * l unnormalization inside ragged_mha and ragged_gqa because the outer MaxText Attention layer (attention_op.py) explicitly relies on unnormalized outputs to perform multi-chunk prefill/decode cache merging and final normalization. Added clear explanatory inline comments to document this architectural requirement.
3. Unified CPU Interpret Mode: Exposed interpret parameter in MHA/GQA wrappers to support CPU-based testing and added full unit test coverage on CPU.

BUG=510376806

---

Tests

The changes were verified using both live TPU hardware (tpu7x-8 VM) and CPU-based Pallas Interpret mode.

1. TPU-Based Verification

The unit tests in tests/unit/kernels_test.py correctly align scales using the division workaround (ragged_out / ragged_denom) to verify unnormalized outputs against reference implementations.

Command run inside the virtual environment on TPU VM:

source maxtext_venv/bin/activate
pytest tests/unit/kernels_test.py -k "not Cpu"

Result: 3 passed (TPU execution verified).

2. CPU-Based Verification & Coverage (Ragged Attention)

To ensure code coverage in CPU-only CI environments, we added a new test suite RaggedAttentionCpuTest in tests/unit/kernels_test.py. This suite forces JAX Pallas to run in Interpret Mode on CPU (bypassing TPU compilation).

Command run locally or in CPU-only CI:

JAX_PLATFORMS=cpu pytest tests/unit/kernels_test.py -k "Cpu"

Result: 3 passed (CPU interpret mode verified). This provides full test coverage for ragged_attention.py in standard GitHub CI runners.

Checklist

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

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• 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, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[darisoy]

Note on Code Coverage (SparseCore Kernels)

You will notice that Codecov reports 0% Patch Coverage for the newly added security and bounds-checking lines in the following two files:

• src/maxtext/kernels/gather_reduce_sc.py (Mosaic SparseCore)
• src/maxtext/kernels/gather_reduce_pallas.py (Pallas SparseCore)

Why these lines are not covered in the CI:

1. Hardware Requirements: These SparseCore gather-reduce kernels rely on low-level Mosaic MLIR dialects and SparseCore ASIC hardware, which are only supported on actual TPU instances (TPU v5e/v6e/v7+).
2. CPU CI Limitation: The GitHub CI Code Coverage job runs on standard CPU-only runners. Mosaic MLIR has no CPU interpreter, and Pallas CPU interpret mode does not support or emulate SparseCore hardware features.
3. Skipped Tests: Consequently, the SparseCore unit tests (tests/gather_reduce_sc_test.py) have hardware checks that automatically skip them on CPU, meaning these files are never executed during the CI coverage runs.

Manual Verification:
All SparseCore clamping, off-by-one, and underflow bounds-checking fixes have been manually verified and compiled successfully on a live tpu7x-8 VM.

(Note: The Ragged Attention fixes in ragged_attention.py have been fully covered and verified in CI using JAX Pallas's CPU Interpret Mode in tests/unit/kernels_test.py).

[clee1994]

LGTM, thanks for fixing this!