[DO NOT MERGE] Compile compute_chunked_logps at chunk_size > 1024 under context parallelism

[parambole]

Do not merge. Not part of the 128-chip best run, which ships in #2878 and AI-Hypercomputer/maxtext#5718. Kept open for reference.

TL;DR

• Problem: with context parallelism, the GRPO log-prob head (compute_chunked_logps) does not compile on TPU at compute_logps_chunk_size 2048 or 4096 (cp=4): Reshape should have supported layout before reaching the emitter.
• Fix: when chunk_size > 1024, keep the 4-byte-per-token arrays that enter and leave the chunk loop (int32 targets in; f32 log-probs, and entropies when requested, out) replicated over the sequence axes. Chunk 4096 then compiles. This fixes the default path; [DO NOT MERGE] Opt-in sharded GRPO log-prob head: gather the vocab table once, scan local tokens (-11.8% step time, v7x 397B bench) #2873 is the opt-in alternative that bypasses it.
• Result: -9.1 s/step (~7%) for chunk 4096 with this PR vs chunk 1024 without it (121.6-121.8 vs 130.8 s/step), on the v7x MaxText RL trainer bench (128 chips, fsdp32 x cp4 x ep2). Cross-pod comparison.
• Default: on only when chunk_size > 1024 and a mesh is set with jax.set_mesh (as the MaxText training engine does). At 1024 and below the compiled program is unchanged. Override: TUNIX_LOGPS_REPLICATE_SEQUENCE=auto|always|never.
• Risk: low. Only data movement changes: on CPU, log-probs, entropies and gradients are bit-identical to the unconstrained program. TPU loss/gradient numerics are not checked yet (see Results).

How it works

flowchart LR
  subgraph S1["Today: chunk_size above 1024, sequence split over context"]
    direction TB
    A1["int32 targets [B, L]"] --> A2["all-to-all into the chunk loop"]
    A2 --> A3["lax.scan over chunks"]
    A3 --> A4["f32 log-probs: all-to-all back"]
    A4 --> A5["TPU compile error:<br/>Reshape should have supported layout"]
  end
  subgraph S2["This PR"]
    direction TB
    B1["int32 targets [B, L]"] --> B2["replicate sequence dims:<br/>one all-gather, 4 B/token"]
    B2 --> B3["lax.scan over chunks<br/>(bf16 hidden states unchanged)"]
    B3 --> B4["f32 log-probs and entropies:<br/>replicate sequence dims"]
    B4 --> B5["chunk 4096 compiles"]
  end

Loading

The scan takes one chunk per iteration, so XLA moves the context split into each chunk with an all-to-all on every array that crosses the loop boundary; for the two 32-bit arrays (and the log-prob cotangent in the backward pass) the TPU compiler folds a transpose into that all-to-all's reshape and cannot emit it. Replicating only those arrays over the sequence axes (256 KiB per device per array at the bench cell) turns each all-to-all into one all-gather plus a local slice, while the batch dim, the bf16 hidden states and the vocabulary matmuls are left alone.

Results

TPU rows: bench cell of maxtext#5701, 128 v7x chips, fsdp32 x cp4 x ep2, ragged_buffer_factor=2.75, Pathways raiden_20261009183300, 2026-10-14.

Measurement	Where	Result
Step time (s/step)	TPU; cross-pod (chunk 1024 ran on three other pods)	chunk 4096 + this PR, steps 1-4: 121.750 / 121.763 / 121.644 / 121.648. Chunk 1024 without this PR, one value per pod: 130.779 / 130.803 / 130.805. -9.1 s/step
Regular micro-batch	TPU, same runs	12.44-12.51 s vs 13.50 s (-1.0 s)
Compile at chunk 4096	TPU	with this PR: compiled in 628.0 s, 5/5 steps. Without it: the error above (2026-10-13 23:26Z run)
MoE routing at chunk 4096	TPU, same run	0 dropped and 0 replayed assignments on 45/45 micro-batches
Chunk 1024 program, with vs without this PR	full bench program, CPU AOT for tpu7x	0 non-metadata lines differ in the optimized jit_fwd_bwd HLO (127,300 lines)
Prediction written before the TPU run	from an older Pathways build where unfixed 4096 compiled	-9.0 s/step (band -11 to -6); measured -9.1

Caveats: the step-time delta is cross-pod (the same-pod chunk-1024 arm was cut for time; the 1024 baseline varied by 0.026 s/step across three pods). Not yet run on the newer Pathways build raiden_20261013212200. TPU loss/gradient numerics are not checked (the bench payload gives 0 loss by construction); they rest on the CPU tests, plus a forward probe within 1.8e-5 relative of chunk 1024.

Changes

• tunix/rl/common.py
  • _replicate_sequence_dims(x, seq_dims) (new): with_sharding_constraint with None on the sequence dims and UNCONSTRAINED on the rest, under the mesh set by jax.set_mesh. Returns x unchanged when no mesh is set.
  • _replicate_sequence_in_chunk_loop(chunk_size) (new): the on/off rule.
  • compute_chunked_logps: constrains the int32 targets before the chunk reshape, and the f32 log-probs (and entropies) right after the scan.
  • Env var TUNIX_LOGPS_REPLICATE_SEQUENCE: auto (default, on when chunk_size > 1024), always, never. Any other value raises. Read at trace time.
• tests/rl/chunked_logps_sequence_replication_test.py (new): 34 tests on an 8-device CPU mesh (batch over fsdp, sequence over context).

Testing

CPU: 34/34 new tests, and 83/83 in tests/rl/common_test.py with and without the change; five deliberately broken copies of the change are each caught by the intended tests. AOT (CPU, tpu7x target, libtpu 0.0.51.dev20261010) at cp=4 / 128 chips and cp=8 / 256 chips, with negative controls (never, no mesh, ids only, batch-dim constraint) that reproduce the failure. TPU: 5/5 steps on the bench cell.

python -m pytest tests/rl/chunked_logps_sequence_replication_test.py tests/rl/common_test.py
# TPU: the maxtext#5701 bench cell with compute_logps_chunk_size=4096.
# TUNIX_LOGPS_REPLICATE_SEQUENCE=never is the negative control (today's program).

Full tables and the per-test list are in the collapsed section below.

Risks / rollout

• At chunk 1024 and below the traced program is unchanged, so existing runs are bit-identical. Rollback: TUNIX_LOGPS_REPLICATE_SEQUENCE=never.
• The 1024 threshold is empirical for these compiler builds. The real trigger looks like chunks per device (8 or fewer fails): at cp=8 with one row per device, chunk 1024 already fails (CPU AOT for tpu7x, not run on TPU), so use always there (or a chunk above 1024).
• Engages only under jax.set_mesh. Under a legacy with mesh: block it stays a no-op, as today. tunix's own RLEngine (alias RLCluster) enters its mesh that way (rl_cluster.py:1310), so that path is unchanged.
• Does not remove the per-chunk vocabulary all-gather or the duplicated per-context-rank head compute; the opt-in head in [DO NOT MERGE] Opt-in sharded GRPO log-prob head: gather the vocab table once, scan local tokens (-11.8% step time, v7x 397B bench) #2873 does.

Related

• [DO NOT MERGE] Opt-in sharded GRPO log-prob head: gather the vocab table once, scan local tokens (-11.8% step time, v7x 397B bench) #2873: opt-in sharded log-prob head. When it engages (packed rows, TUNIX_LOGPS_SHARDED_HEAD=local), compute_per_token_logps returns before calling compute_chunked_logps, so this PR only changes the default path. The only textual conflict is one import line (import math / import os); with both applied, 19/19 + 34/34 + 83/83 CPU tests pass.
• maxtext#5701: the RL trainer bench used for the TPU and AOT numbers.
• feat(recipes): opt-in MOE_EXCLUDE_PADDING, LOG_MOE_DROPS and RETRY_DROPLESS_FINAL_MICRO_BATCH knobs #2841: recipe knobs, including an overridable COMPUTE_LOGPS_CHUNK_SIZE; the team hit the same error at chunk 2048 on it.
• Base branch: atwigg/mlperf.

Full evidence and notes (the previous description)

A bare #5701 below refers to maxtext#5701.

What

With context parallelism, the GRPO log-prob head does not compile on TPU at compute_logps_chunk_size 2048 or 4096
(team: 2048 on #2841; ours: parambole-lp4096, bench #5701 cell, 2026-10-13 23:26Z):

jax.errors.JaxRuntimeError: FAILED_PRECONDITION: Reshape should have supported layout before reaching the emitter.
HLO: %reshape.18785 = s32[4,1,4,1024]{3,0,1,2:T(4,128)} reshape(%param.0)
     op_name="jit(fwd_bwd)/jvp(jit(compute_per_token_logps))/transpose"

This PR replicates the chunk loop's two 32-bit arrays (int32 targets in, f32 log-probs out) over the sequence axes
when chunk_size > 1024. 4096 then compiles. At or below 1024 the traced program is unchanged, so today's runs are
bit-identical.

Why it fails

• The sequence is sharded over context outside the lax.scan in compute_chunked_logps (rows over
  fsdp x expert, 1 row per device at the bench cell; 16,384 tokens per device at cp=4).
• The scan slices the chunk axis each iteration, so the partitioner moves the cp split from the chunk-index dim to the
  in-chunk dim with an all-to-all on every array that crosses the loop boundary.
• For the 32-bit arrays, whose local source is 1-D, layout assignment folds the transpose into the all-to-all's input
  reshape (s32[4,1,4,1024]{3,0,1,2:T(4,128)} at chunk 4096). The TPU emitter cannot lower that reshape.
• At chunk <= 1024 (cp=4) it keeps a legal row-major reshape plus a separate transposing copy, so it compiles.
• There is a second instance in the backward: the f32 log-prob cotangent, f32[4,1,4,1024] in
  transpose(jvp(...))/transpose. Fixing only the ids fails on it (shown below).
• Measured on libtpu 0.0.51.dev20261010, which reproduces the Pathways raiden_20261009183300 failure down to the
  instruction name. The older raiden_20261002162900 build compiles unfixed 4096 (our earlier run parambole-bs-l4ko).

The change

• _replicate_sequence_dims(x, seq_dims): with_sharding_constraint(x, P(UNCONSTRAINED, ..., None at seq dims))
  under the active jax.set_mesh mesh, so the batch dim keeps whatever sharding the partitioner chose. It is a no-op
  without a mesh.
• In compute_chunked_logps, when chunk_size > 1024 it is applied to the int32 targets [B, L] before the chunk
  reshape and to the f32 log-probs (and entropies) [chunks, B, chunk] right after the scan.
• Each all-to-all becomes a one-shot all-gather of 4 bytes per token (256 KiB per device per array at the bench cell)
  plus a local slice. The bf16 hidden states and the vocabulary matmuls are untouched.
• TUNIX_LOGPS_REPLICATE_SEQUENCE=auto|always|never (default auto = the rule above) overrides the rule. Any other
  value raises.
  • At cp=8 with 1 row per device, chunk 1024 already fails (8 chunks per device). Use always there, or a chunk above
    1024.

Evidence

All AOT rows use libtpu 0.0.51.dev20261010 (CPU AOT for tpu7x). A CPU AOT pass is not a TPU compile.

Check	Result	Tag
Head reproducer (the image's own compute_chunked_logps inside value_and_grad, bench shardings), cp=4, 128 chips	unfixed: 2048 FAIL s32[8,1,4,512], 4096 FAIL s32[4,1,4,1024] (the TPU signature); fixed: 2048 / 4096 / 8192 PASS; 1024 PASS with identical temp bytes (7,714,230,624)	MEASURED
Same, cp=8, 256 chips (fsdp32 x cp8 x ep2)	unfixed: 1024 FAIL s32[8,1,8,128], 2048 FAIL s32[4,1,8,256], 4096 FAIL s32[2,1,8,512], 512 PASS; fixed: 2048 / 4096 PASS, 1024 PASS with always	MEASURED
Negative controls	never at 4096: FAIL (same reshape). Fixed file with no mesh set: FAIL (same reshape). Ids-only variant: FAIL on f32[4,1,4,1024] (the backward instance). Constraint on the batch dim instead of the sequence: FAIL	MEASURED
Full bench program (#5701 aot_trainer.py, exact bench env), with this exact file	chunk 4096: rc 0 (the unpatched file FAILs here with the TPU's reshape.18785), device peak 88.81 GiB of 94.74. Chunk 1024: rc 0, and its optimized jit_fwd_bwd HLO has 0 non-metadata lines that differ from the unpatched program (127,300 lines; only the source-location tables change). Control: the 1024 and 4096 programs differ in 185,945 lines	MEASURED
Unit tests, new file (CPU, 8-device mesh shaped like the trainer's)	34/34 pass	MEASURED
Existing tests/rl/common_test.py	83/83 pass, with and without the change	MEASURED
TPU, bench #5701 cell (128 v7x, fsdp32 x cp4 x ep2, ragged_buffer_factor=2.75), Pathways raiden_20261009183300, client trellis-experimental@sha256:b4647c18…, 2026-10-14	chunk 4096 with this change: compiled in 628.0 s and ran 5/5 steps (rc 0): 123.076 (step 0, warm-up) / 121.750 / 121.763 / 121.644 / 121.648 s/step; regular micro-batches 12.44-12.51 s. Chunk 1024 without the change, same build, three other pods: 130.779 / 130.803 / 130.805 s/step (13.50 s per regular micro-batch). -9.1 s/step (-6.9%), -1.0 s per regular micro-batch. 0 dropped and 0 replayed MoE assignments on 45/45 micro-batches. Forward probe sample_mask/mult_prob_error_mean 333,638.1-333,643.1 against 333,644.1 at chunk 1024 (within 1.8e-5 relative). The same-pod chunk-1024 arm was cut for time, so this delta is cross-pod (the 1024 base varied by 0.026 s/step across the three pods). Not run yet: the newer raiden_20261013212200 build, and the loss/gradient numerics (the bench payload gives 0 loss by construction; numerics rest on the CPU tests above)	MEASURED

What the new tests check:

• Program:
  • unchanged at chunk 512 / 1024 and with no mesh;
  • changed at 2048 / 4096;
  • the constraints land on exactly [B, L] i32 [{?}, {}] and [chunks, B, chunk] f32 [{}, {?}, {}] (forward and its
    cotangent).
• Numerics:
  • bit-identical to the unreplicated program at the same chunk (log-probs, entropies, both gradients; 2048 and 4096,
    with and without entropy, packed/padded sequence);
  • always at 1024 is bit-identical;
  • 4096 vs 1024 is within f32 summation order (measured max |d| 1.9e-6 on log-probs, 6.5e-5 on a kernel gradient of
    max 83);
  • a float64 reference agrees.
• Mode parsing: every branch, plus unknown values raising.

Each check was shown to fail on a deliberately wrong input:

• Five broken copies of the change are each caught by exactly the intended tests:
  • gate at >=: 4 failures;
  • constraint on the batch dim: 2;
  • no-mesh guard removed: 3 errors;
  • ids only: 2;
  • lenient mode parsing: 5.
• The two variants that do not fix the TPU compile (batch dim, ids only) pass every numerics test, are caught only by
  the constraint test, and FAIL the AOT compile. So that test is the one that matters.

Expected speed (pre-registered before the TPU run)

• Same bench cell and pod on the Oct-2 Pathways build, where unfixed 4096 compiles: logps 4096 ran 119.979 s/step vs
  129.007 at 1024 (steps 1-4), i.e. -9.03 s/step = -1.00 s per micro-step. MEASURED, parambole-bs-l4ko.
• Prediction for this change on the current build: -9.0 s/step (band -11 to -6) against chunk 1024.
• Measured on TPU (row "TPU" above): -9.1 s/step, inside the band.

Limits

• Only a mesh set with jax.set_mesh (as the MaxText training engine does) engages it. A legacy with mesh: block
  does not set the abstract mesh (MEASURED, jax 0.11.0), so there it stays a no-op, as today.
• The 1024 threshold is empirical for these compiler builds. The trigger looks like the number of chunks per device
  (<= 8 fails), which the code cannot see at trace time under Auto sharding; hence the override.
• This does not remove the per-chunk vocabulary-table all-gather or the redundant per-cp-rank head compute. The
  opt-in sharded head in [DO NOT MERGE] Opt-in sharded GRPO log-prob head: gather the vocab table once, scan local tokens (-11.8% step time, v7x 397B bench) #2873 does that, and when it engages (packed rows, TUNIX_LOGPS_SHARDED_HEAD=local)
  compute_per_token_logps returns its result without calling compute_chunked_logps, so this change only affects
  the default path. Both PRs edit tunix/rl/common.py; the only textual conflict is one import line
  (import math / import os). With both applied, [DO NOT MERGE] Opt-in sharded GRPO log-prob head: gather the vocab table once, scan local tokens (-11.8% step time, v7x 397B bench) #2873's tests (19/19), this PR's tests (34/34) and
  tests/rl/common_test.py (83/83) pass on CPU.