Support packed TPU Tokamax ring attention

[huytransformer]

Description

Enable sequence packing for TPU Tokamax ring attention.

Packing was previously rejected by config validation when combined with the TPU ring. Ring itself already supports packed sequences so this pr lifts the rej from the TPU Tokamax ring validation in configs/types.py and from AttentionOp, and resolves context_parallel_reorder_strategy=AUTO with packing and ring to DUAL_CHUNK_SWAP on TPU.

This is important for ring because ring CP targets long contexts, where real documents rarely fill the window on their own.

Tests

CPU:
python3 -m pytest tests/unit/configs_value_test.py -k tokamax_ring
5 passed (18 subtests), including the new accepts_packing and accepts_packed_load_balance cases.

TPU (v5p, 4 chips):
python3 -m pytest tests/unit/attention_test.py -k ring_context_parallel
8 passed. Forward and gradient equivalence against dot product attention for packed and unpacked, with and without load balancing.

E2E (v5p-128)
20 step training run with packing=True on c4 (tfds), context_parallel_strategy=ring, context_parallel_load_balance=True, ici_context_parallelism=4, reorder strategy AUTO. Loss decreased from 10.84 to 8.06 over 20 steps

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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.