feat: 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.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.