fix: MTP CP-aware roll, segment-aware roll, and synthetic packing support

[chiaotung97]

Summary

Fixes 3 issues when combining Multi-Token Prediction (MTP) with
All-Gather Context Parallelism (AG-CP) and Packing.

Checklist

• CP-aware left shift (_shift_left_one_cp_aware) with ppermute backward ring
• Segment-aware roll (roll_and_mask_by_segment) to prevent cross-document leakage
• Synthetic data packed segment IDs (_make_packed_segment_ids)
• Unit tests: 30/30 passing on TPU v6e-4 (4 devices)
• Backward compatibility: all paths degrade gracefully
• Merge conflicts with main resolved

Problem

#	Issue	Root Cause	Impact
1	jnp.roll does not cross CP rank boundaries	MTP shifts left each iteration to predict the "next token", but jnp.roll only operates within the local shard	Cross-rank right-neighbor tokens are lost under CP
2	roll_and_mask is segment-unaware	Under packing, a left shift can cross document boundaries (e.g., DocA's last token shifts into DocB's first)	Cross-document target leakage contaminates gradient signals
3	Synthetic data has no real segment boundaries	segment_ids is always all-ones even though base.yml defaults to packing: true	Segment boundary logic is untestable on synthetic data; train_utils.py also blocks the combination outright

Solution

1. CP-Aware Left Shift (_shift_left_one_cp_aware)

Uses jax.lax.ppermute in a backward ring: rank r sends its first token to
rank r−1, receives rank r+1's first token, and places it in its last
position. Degrades to jnp.roll when no CP axis is in scope — zero overhead.

2. Segment-Aware Roll (roll_and_mask_by_segment)

Shifts both x and segment_ids via _shift_left_one_cp_aware, then masks
positions where seg_current != seg_next (document boundary) or
seg_current == 0 (padding). Falls back to roll_and_mask when
segment_ids=None. All rolling variables in
MultiTokenPredictionBlock.__call__ now use this function.

3. Synthetic Data with Packed Segment IDs

_make_packed_segment_ids splits each row into 2..N randomly-sized segments
with sequential integer IDs starting from 1. Removed the train_utils.py
guard that rejected synthetic + packing + CP.

Files Changed

File	Change	+/−
layers/multi_token_prediction.py	_shift_left_one_cp_aware, roll_and_mask_by_segment, wiring	+103/−4
input_pipeline/synthetic_data_processing.py	_make_packed_segment_ids	+43/−2
utils/train_utils.py	Remove synthetic+packing+CP guard	−5
tests/unit/multi_token_prediction_test.py	17 new tests (segment + CP + packed_ids)	+296
docs/design/ag_cp_mtp_packing_fix.md	Design document	+141

Unit Tests

30 tests, all passing (TPU v6e-4, 4 devices):

Test class	Tests	Description
TestRollAndMask	3	Existing — local roll and mask
TestRollAndMaskBySegment	8	Segment boundary masking (no CP)
TestMakePackedSegmentIds	6	Segment ID generation: shape, contiguity, determinism
TestShiftLeftOneCpAware	5	CP-aware shift: 1D/2D/3D, last-rank zero, rank-boundary non-zero
TestRollAndMaskBySegmentWithCp	4	CP + segment boundaries, padding mask, single-seg, 3D
Upstream tests (layer, block, MTP loss, quantize)	4	Existing — unchanged

Backward Compatibility

• _shift_left_one_cp_aware: degrades to jnp.roll when CP=1 or no "context" axis
• roll_and_mask_by_segment: degrades to roll_and_mask when segment_ids=None
• roll_and_mask(shift=-1): equivalent to original path when CP is off
• synthetic_data_processing: segment IDs remain jnp.ones when packing=False

Verification

#	Configuration	TPU v6e-4
1	MTP + packing	✅
2	MTP + CP=2	✅
3	MTP + CP=2 + packing	✅
4	Unit tests	✅ 30/30

[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

❌ Patch coverage is 90.90909% with 5 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/multi_token_prediction.py	91.66%	1 Missing and 2 partials ⚠️
...axtext/input_pipeline/synthetic_data_processing.py	89.47%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!