Repository navigation
TPU sequence parallelism: Pallas collective-matmul kernels, LTX-2 / HunyuanVideo SP, Wan padding + split RoPE - #153
Open
shiyang-opensource wants to merge 23 commits into
Open
shiyang-opensource wants to merge 23 commits into
shiyang-opensource wants to merge 23 commits into
Conversation
The FLUX request spent 3.75 s of ~9.2 s in T5-XXL, fp32 on the host (rank 0, then broadcast). bf16 T5 is ~9.5 GB, more than one chip has free next to its DiT shard, so TpuPipelinedT5Encoder splits the 24 blocks into one contiguous segment per rank (~2.4 GB a chip) and hands the hidden state on with broadcast -- the stock HF blocks, no sharded modeling, and the one collective every port already uses. Each hop carries a status flag so a failing rank never strands the others in a collective. On by default; DIFFLET_TPU_T5_DEVICE=0 restores the host encoder. v5e-4, main's bench (flux_1_dev, 1024x1024, 28 steps), 2026-11-09, same commit: text_encode 3.748 -> 0.197 s denoise 5.264 -> 5.244 s, VAE 0.185 -> 0.186 s image vs host-T5 run: 39.6 dB PSNR, same composition (one prompt/seed; not compared against a higher-precision reference) embeds vs host fp32 (t5_device_probe, 3 prompts): cos 0.996-0.999, rel 5-9% fp32 on the chips (~4.8 GB a chip) runs the encoder in 63 ms but leaves rank 0 too little HBM for the 1024x1024 VAE decode (OOM), so bf16 is the default. benchmark/tpu_profile/t5_device_probe.py measures latency and error per dtype. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… load diffusers' FLUX rotary pairs channels (2i, 2i+1) of each head, so every q and k of every block goes through reshape(..., D/2, 2).unbind / stack.flatten -- a minor-dimension relayout that the 2026-11-03 trace put at ~4-7% of a compiled denoise step. q . k is invariant under one permutation applied to both, so at load the TPU path permutes each head's channels to [0, 2, ..., 1, 3, ...] in to_q / to_k / add_q_proj / add_k_proj (rows and bias) and in the per-head qk RMSNorm weights, permutes the cos/sin tables the same way once per forward (PermutedFluxPosEmbed), and applies the split-halves rotary (use_real_unbind_dim=-2): a slice and a concat. V and everything after the attention are untouched. - backends/tpu/flux/rope.py: the permutation, applied after the checkpoint is loaded; DIFFLET_TPU_FLUX_SPLIT_ROPE=0 turns it off. - models/flux/tp_sharding.py (TPU-only user): FluxTPAttnProcessor takes rope_unbind_dim, default -1 = unchanged behaviour. - CPU test: the permuted module reproduces diffusers' stock forward (atol 1e-5, non-trivial norm weights). v5e-4, main's bench, flux_1_dev 1024x1024 28 steps, same commit, back to back (2026-11-09): step 185.0 / 182.5 -> 173.9 / 173.5 ms (synced runs; -6%) denoise 5.176 -> 4.964 s image on vs off: 37.8 dB PSNR, mean |diff| 0.95/255, 0.46% of pixels > 16 (one prompt/seed; the change is exact in math, the difference is bf16 rounding in a differently fused graph carried through 28 steps; not compared against a higher-precision reference) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…T_TPU_SP=1)
FLUX's 57 blocks end every attention and MLP in an all-reduce (~20% of a
compiled step). As for Wan / Qwen-Image / HunyuanVideo, each becomes a
reduce-scatter fused into the row-parallel matmul and an all-gather fused into
the next column-parallel one; FLUX has no shared Megatron-SP to reuse, so
backends/tpu/flux/sp.py swaps the model and block forwards on the instance
(diffusers' modeling and tp_sharding unchanged).
- joint blocks: image and text streams each sequence-sharded (1024 + 128 rows a
rank at 1024x1024); AG+[to_q,to_k,to_v] and AG+[add_q,add_k,add_v], the TP
processor unchanged on the full sequence, to_out / to_add_out + RS, ff and
ff_context AG/RS.
- single blocks: one AG feeds q, k, v and proj_mlp; one RS for the sum of the
two row-parallel proj_out halves (rs over two (x, w) pairs, no concatenation).
The gathered order is [txt_0, img_0, txt_1, ...]; the cos/sin rows are
permuted to it once per forward (no mask, per-token RoPE).
- inputs sliced before the embedders, the 64-channel output gathered at the end.
- collective_matmul: AG over 2 and 4 weights, RS over 2 pairs; prepare() takes
the (kind, arity) list a model needs. The 1-pair RS and 1/3-weight AG keep
their outputs: Qwen-Image and Wan SP outputs are bit-identical before/after.
- CPU test: with the rings at their tp=1 meaning the SP forward reproduces
diffusers' (split RoPE on and off); the single-block order; refusal when a
stream does not split.
v5e-4, main's bench, flux_1_dev 1024x1024 28 steps, same commit, back to back
(2026-11-10):
step (synced runs) 176.2 / 173.3 -> 173.6 / 156.3 ms
denoise 4.953 -> 4.473 s (-10%)
image SP vs no SP 41.5 dB PSNR, mean |diff| 0.68/255 (one prompt/seed;
summation order only; no higher-precision reference)
regression Qwen-Image SP 02c485d72c20, Wan 2.1 SP 13d36a081762:
sha256 unchanged from before the collective_matmul change
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
diffusers left-pads every prompt to 1024 tokens and runs Gemma-3 12B (host
fp32) and the connectors' text_proj_in (3840*49 -> 3840) over all of them.
Nothing computed at a padding row reaches the DiT: the connectors'
masked-mean norm zero-fills it, text_proj_in has no bias, and the register
step overwrites it. backends/tpu/ltx_2/text_encode.py runs both over the last
n columns (n = longest prompt of the batch), with the position_ids of the
padded call so every token keeps its rotary angle, and zero-pads back.
Rank 0 now runs Gemma *and* the connectors and broadcasts the connector
output (2 x (B, 1024, 3840)) instead of Gemma's packed (B, 1024, 188160);
the other ranks no longer run the connectors on the host. A failure on rank 0
is carried to the broadcast so every rank raises. DIFFLET_TPU_LTX2_TRIM_TEXT=0
restores the padded path.
Measured on a v5e-4 (torch 2.14.1+cpu, torch_tpu 0.2.0, libtpu 0.0.51,
jax/jaxlib 0.12.0.1, diffusers 0.38.0, transformers 5.20.2, LTX-2 @ 47da56e2),
main's bench (480x704x49, 20 steps, eager), off/on back to back on one tree:
TRIM=0 TRIM=1
text_encode 22.13 s 2.95 s
denoise 14.73 s 14.86 s
vae_decode 1.20 s 1.10 s
request (sum) 38.06 s 18.91 s (-50%)
e2e incl. spawn+load 65.1 s 42.9 s
Output: PSNR 41.31 dB, mean |d| 1.20/255 between the two (frames 0/24/48,
sha256 2566a6b7f6b8 vs d362cabe92d6; each variant bit-identical across its
three runs); side by side the frames are indistinguishable. One prompt/seed
end to end. At the conditioning itself, benchmark/tpu_profile/ltx2_text_probe.py
over 3 prompts (1/9/49 tokens): Gemma 18.8/16.1/12.8 s -> 1.6/1.6/3.0 s,
connectors 2.1-2.9 -> 1.2-1.6 s, connector outputs within 2.5e-6 relative
(fp32 rounding: the matmul and softmax row counts change), masks bit-equal.
Not checked: CFG (guidance > 1) end to end, batch > 1 on the chip, and a
higher-precision reference.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…iling per process
The orchestrator's packed latents reach forward_dit as a transposed view
(hidden_states stride (2310, 1, 2310) at 480x704x49); the warmup's example
inputs are contiguous. Under DIFFLET_TPU_COMPILE=1 Dynamo's stride guard
failed on the first real request ("tensor 'kwargs['hidden_states']' stride
mismatch at index 0. expected 295680, actual 2310", TORCH_LOGS=recompiles)
and recompiled the whole DiT: that request's denoise took 72 s for 20 steps
of 0.23 s. The inputs are now made contiguous before the copy to the chip.
v5e-4, main's bench, 480x704x49, 20 steps, compile on (text trim on), same
versions as the previous commit:
recompiles 3 -> 0 (one per process); denoise 72.19 s -> 4.39 s;
steady step 228 / 220 ms (unchanged); output bit-identical to the run
before the fix (sha256 e641a64c7f40 both).
Compile stays off by default for LTX-2. Against eager the compiled output is
35.44 dB (mean |d| 2.1/255, max 119, sha256 d362cabe92d6 vs e641a64c7f40;
same frames, differences in the fast-moving fox), lower than the 49 dB seen on
Qwen-Image, and there is no higher-precision reference to say which side is
closer. Request with compile: text 2.89 + denoise 4.39 + VAE 1.08 = 8.4 s,
against 18.9 s eager; the first-call compile costs 155 s at load (tier-3
cache cold).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…n the chips) main's bench runs LTX-2 at guidance 1.0, so the negative + positive prompt path was not exercised on the device. This loads the application once, runs a 4-step CFG request (guidance 4) twice with the trim and twice after reinstalling the padded path on the same host pipeline, same seed, threads per rank as main's bench worker (cpu_count / world). v5e-4, same versions as 591c0ad: conditioning (2, 1024, 3840); encode trim 5.48 / 4.95 s, padded 40.15 / 40.70 s (first / second call); latents trim vs padded: video rel L2 3.0% (cos 0.99952), audio 2.3% (cos 0.99974) after 4 steps. Without the per-rank thread cap both paths took ~48 s (four ranks x 56 threads on 112 logical CPUs). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
As for Qwen-Image, Wan and FLUX.1-dev, LTX-2's DiT is now compiled unless DIFFLET_TPU_COMPILE=0. Measured on a v5e-4 (main's bench, 480x704x49, 20 steps, text trim on, after c31a991 removed the per-process recompile): eager compiled step 740 ms 220 ms denoise 14.86 s 4.39 s request 18.9 s 8.4 s load 15.8 s 83.8 s (warm tier-3 cache; ~155 s first compile cold) Compiled vs eager output: 35.44 dB, mean |d| 2.1/255 (frames 0/24/48, sha256 d362cabe92d6 vs e641a64c7f40), same content, differences in the fast-moving region. One prompt/seed; no higher-precision reference, so which side is closer is not known. The extra load time is paid once per process and belongs to the resident server, as for the other compiled models; one-shot runs can set DIFFLET_TPU_COMPILE=0. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
TpuPipelinedT5Encoder._segment called each T5Block positionally with no cache_position. transformers 4.57.6, which CI pins, then indexes cache_position[-1] in the first block (position_bias is None there and no query_length is given): "TypeError: 'NoneType' object is not subscriptable" in test_pipelined_t5_segments_reproduce_the_hf_encoder on both Trainium CI jobs. The TPU venv runs transformers 5.20.2, where the argument is gone, which is why the device runs did not see it. The blocks are now called with keywords and an explicit cache_position = arange(seq), and only outputs [0] and [1] are read; both releases return (hidden, position_bias, ...). Same math, so the device measurements of 5423c7e stand (not re-run: the TPU is released). tests/unit/models/flux: 59 passed, 5 skipped under transformers 4.57.6, 5.20.2 and 5.22.2 (python -m pytest inside `uv run --with`; the pytest entry script would pick up the project venv's transformers instead). Before the fix, 4.57.6: 1 failed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
test_flux_tpu_sp.py initialized the TPU parallel mesh on the parallel_mesh module bound at collection. sp.install and the SP forward import it when they run, and test_torch_megatron.py::test_weights_and_checkpoint_import_without_tpu deletes every difflet.backends.tpu.* entry from sys.modules without restoring them. When both land on one xdist worker in that order, the code imports a fresh, never-initialized copy: ParallelMeshNotInitializedError in 3 tests (trainium-native CI job, worker gw1). Test-only; the library code is unchanged. The fixture now also resolves the module at run time and initializes both. Reproduced locally by running that megatron test before this file (-p no:xdist): 3 failed before, 5 passed after; tests/unit/models/flux: 59 passed, 5 skipped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ET_TPU_SP=1)
LTX-2's video stream goes sequence-parallel as Wan, Qwen-Image and FLUX do:
each video-facing row-parallel output (self-attention, text cross-attention,
a2v, the K/V side of v2a, FFN) becomes a reduce-scatter fused into the matmul
plus an all-gather fused into the next column-parallel one
(backends/tpu/ltx_2/sp.py, block forward only; diffusers' model forward is
unchanged). The audio stream stays TP: a few dozen rows, small collectives.
Video rows are padded to a multiple of 2*tp (480x704x49: 2310 -> 2312 at tp=4)
**eagerly, before the compiled graph**, by the TPU transformer application
(sp.pad_video_inputs / unpad_video_output); pad keys are dropped in attention.
Root cause of the first version's wrong output (PSNR 14.5 dB vs TP), hangs and
core halts: it padded inside the compiled graph and sliced each rank's shard at
a rank-constant offset. torch_tpu compiles one program per process and inlines
the rings (every collective-permute has channel_id=1, paired by order); XLA
folded the pad away on the ranks whose shard lies in the unpadded input, the
per-rank programs diverged, and a chip halted in the first ring hop's matmul
fusion while the others waited in a collective (per-rank HLO dumps and
timer-stopped per-rank jax.profiler traces). Evidence, minimal repro
benchmark/tpu_profile/cm_padded_repro.py (random weights, AG->RS blocks):
- unpadded lengths 1024..6144, incl. 2312: right (rel 0.0062);
- pad in graph, then slice / slice, then pad the last shard: hang, core halt
or inf; AG alone and RS alone both affected;
- plain JAX, same ring code, 2312 (cm_pure_jax.py): right;
- pad in graph with the offset as a runtime tensor (every rank compiles the
same program, fingerprint identical): right (0.0045);
- pad eagerly before the graph: right (0.0045 / 0.0067).
collective_matmul.py's docstring now states the rule: rank-dependent values
enter compiled graphs as runtime tensors.
v5e-4, main's bench, 480x704x49, 20 steps, compiled, one tree (md5
2bb56c9e67e2), back to back; torch 2.14.1+cpu, torch_tpu 0.2.0, libtpu 0.0.51,
jax/jaxlib 0.12.0.1, diffusers 0.38.0, LTX-2 @ 47da56e2:
SP off SP on
step (synced) 218.0 ms 192.1 ms (-11.9%)
step (natural) 218.8 ms 190.8 ms
denoise 4.38 s 3.80 s (-13%)
request 8.59 s 8.02 s
Output SP vs TP: 34.32 dB, mean |d| 2.48/255 (frames 0/24/48; sha256
e641a64c7f40 vs ef75047e8a34, each bit-identical over its 3 runs; the TP
output equals the 2026-11-09 compiled one). Same order as compiled-vs-eager
(35.4 dB); side by side the frames match, differences in the fast-moving fox.
One prompt/seed, no higher-precision reference. CPU: tp=1 parity with diffusers
(fake rings) at a padded and an unpadded length, pad-key and unpadded-input
guards (tests/unit/backends/test_tpu_ltx2_sp.py).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Proposal, gated on a profile of the communication still exposed after SP. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Per-rank XLA dumps of the minimal repro (cm_padded_repro.py rsonly, 2310 rows,
fresh compile via --salt), compared with the new benchmark/tpu_profile/
hlo_rank_diff.py:
pad in graph, rank-constant offset (hangs / halts):
ranks 0-2: 303 ENTRY ops, fingerprint ce7702a1cd
rank 3: 306 ENTRY ops, fingerprint be695c1230 (extra copy-start and a
pad_slice_fusion feeding the first collective-permute)
collective-permute receive buffers: ranks 0-2 at 29884416 / 27459584 /
32309248 / 34734080; rank 3 at 48824320 / 46399488 / 51249152 / 53673984
(+18939904 B each, the growth of rank 3's temp allocation); sync flags at
the same offsets on all ranks.
pad in graph, offset as a runtime tensor (right, rel 0.0045):
all four ranks 585 ops, fingerprint b0474a894f, receive buffers at the
same offsets.
A collective-permute writes into its peer at the receive offset of its own
program, so a rank whose buffer assignment differs receives into the wrong
place: wrong values / inf, or a core halt with the other chips waiting in a
collective. collective_matmul.py's docstring now states this. cm_padded_repro
gains --salt (a per-run constant, folded away by XLA) so every cache layer
misses and the dumps are produced.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…d plan The rings hide 80-90% of the all-reduce; the remaining cost is the AG ring's dynamic_update_slice write-back (Wan 11%, LTX-2 7% of a step) plus a 3-7% collective wait. Next: assemble the AG output once (P0) before any Pallas kernel. benchmark/tpu_profile/exposed_comm.py splits a bench trace into ICI wait / local DMA wait / compute / idle (xplane fallback when the JSON export is truncated). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…T_TPU_SP_KERNEL) collective_matmul_pallas.ag_matmul_body: the bidirectional AG ring as one Pallas TPU kernel. The two halves stay resident in VMEM and travel by VMEM-to-VMEM remote DMA; one double-buffered weight-tile stream runs over every (hop, weight, column tile), a runtime loop per hop; each product tile is written once at its final rows, so the ring's dynamic_update_slice and its copies disappear. Weights may differ in width (one common tile, a multiple of 128). Rank-dependent values are runtime inputs (identical programs per rank). collective_matmul: - DIFFLET_TPU_SP_KERNEL=jax|pallas, else the model's default (set_kernel_default); unsupported shapes fall back to the ring with a notice. - DIFFLET_TPU_SP_AG_ASSEMBLE=dus|concat: build the ring's output once. - Op names carry a hash of both sources: torch_tpu caches a jax_op by name and signature, and reused a stale program after a kernel change. - The op cache key includes both modes. v5e-4, main's bench, step time JAX ring -> kernel: Wan 292.4 -> 268.0 ms, FLUX 156.5 -> 140.6, HunyuanVideo (single-stream SP) 500.7 -> 478.0, Qwen-Image 178.0 -> 179.7, LTX-2 190.1 -> 190.5. Op parity vs plain gather + linear is 0.0 through torch_tpu; end-to-end outputs within the TP/SP PSNR spread (details in docs/plans/2026-11-10-tpu-collective-matmul-pallas.md). Tests: CPU (TPU interpret mode, race detector on) for the kernel, incl. mixed widths; the concat assembly vs numpy. Tools: cm_pallas_bench (op level), pallas_matmul_sweep (plain Pallas matmul vs XLA), sp_step_drift (k-step output drift between SP variants). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The 40 single-stream blocks ran dense TP with an all-reduce each (1.0 s of a 10.3 s, 20-step denoise at 320x512x61). They now run on a contiguous chunk of the joint [latent | text] sequence (10496 = 4 x 2624; a chunk may straddle the boundary): q|k|v|mlp share one all-gather matmul, the two output projections one reduce-scatter (bias once), the shared attention runs on the gathered q/k/v. The first block cuts, the last gathers. Only the TPU module changes; tpu_application passes the text length. The cut is an index_select by a per-rank row buffer: a Python-constant slice of the in-graph concatenation let XLA fold it on ranks 0-2 but not on rank 3, the programs diverged and the first forward hung in a collective. DIFFLET_TPU_SP_SINGLE (default on) and, for this model, the Pallas AG kernel by default. v5e-4: TP 549.9 -> dual-only SP 516.1 -> + single SP 500.7 (ring) / 478.0 ms (kernel); ICI wait 1042 -> 125 ms per 20 steps; PSNR vs TP 42.9 dB (TP vs dual-only SP 45.6). CPU test: three SP single blocks vs the dense ones. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Measured on v5e-4 with main's bench, same tree, back to back: Wan 292.4 -> 268.0 ms per step, FLUX 156.5 -> 140.6 ms. Qwen-Image and LTX-2 keep the JAX ring (the kernel is +1% / flat there). DIFFLET_TPU_SP_KERNEL=jax restores the ring. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ngle-stream SP Results of 2026-11-10/11 on v5e-4: the concat assembly (Wan -6%), the kernel's versions and dead ends, op and full-model A/Bs for all five models, the Wan seed spread (outputs differ by seed-dependent trajectory drift, per step at the bf16 noise level), the compile-cache pitfall, HunyuanVideo single-stream SP. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… full names collective_matmul_pallas keeps only the variant the models run (halves resident in VMEM, one double-buffered weight-tile stream over every hop, weight and column tile, a runtime loop per hop). The emit_pipeline variants and the pre-transposed-weights option are dropped; they are in 01d027f and the plan. Names are spelled out: all_gather_matmul_body(x, weights, owner_shards, neighbors, ring_size=...), unsupported_reason, ring_tables(ring_position, ...), RING_SLOTS, VMEM_BUDGET_BYTES, DEFAULT_COLUMN_TILE_LIMIT. cm_pallas_bench times the kernel as the models call it ((N, K) weights, column tiles up to 256; the plan's version table used (K, N) weights and up to 512) and no longer fails at exit (torch.tpu stub lacked is_initialized). Checked on v5e-4 against 868401e, same node, back to back (main's bench): outputs sha256-identical for Wan (70c6559258cd), FLUX (de53c10893bc) and HunyuanVideo (7259d2e87ea0); step 268.8 -> 267.9, 144.7 -> 141.9, 480.3 -> 478.9 ms; no fallback; the new op names compiled fresh. Op-level the kernel is bit-exact to the ring at every bench shape. CPU interpret tests: 4 passed (race detector on), plus unsupported_reason. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…t-in)
collective_matmul_pallas.matmul_reduce_scatter_body: the bidirectional RS ring
as one kernel, tile by tile -- each output column tile is multiplied, the
upstream accumulator's same tile (sent about a hop earlier) is added in the
epilogue and the sum sent on at once; accumulators travel VMEM to VMEM, the
last hop writes its sums straight to the output, no partial product or
concatenation goes to HBM. One or two (x, weight) pairs (FLUX / HunyuanVideo
single blocks), rounding as the JAX ring. Three receive slots per direction:
tile-wise accumulators let a neighbour run two hops ahead in one direction.
collective_matmul: DIFFLET_TPU_SP_RS_KERNEL=pallas selects it (default jax);
unsupported shapes fall back to the ring with a notice.
Off by default because it does not pay (v5e-4, 2026-11-11, op-level, ms per
call, benchmark/tpu_profile/cm_rs_pallas_bench.py):
Wan attention out 0.518 -> 0.504 (ICI-bound)
LTX-2 attention out 0.214 -> 0.212 (ICI-bound)
Wan FFN down 0.928 -> 1.040 (~75 us input relayout: 1170 rows are not
a multiple of 8; aligned 1168 rows: 0.917
-> 0.964, the matmul is slower than XLA's)
FLUX single pair 0.627 -> 0.663 (column tile 512: 0.649)
Wan end to end (padded): step 262.2 -> 262.7 ms. One pair is bit-identical to
the ring op-level; in the model 1 step 50.4 dB, 4 steps 48.9 dB vs the ring
(seed 42), as the AG kernel. A first version with a separate add pass after
each hop was slower everywhere.
Tests: CPU interpret, 1 and 2 pairs, batch 1 and 2, race detector on (it
catches two receive slots; a missing wait_send it cannot, a remote copy reads
its source at start in interpret mode), plus the unsupported reasons.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… split halves Two changes to Wan's TPU sequence parallelism (backends/tpu/wan/sp.py, TPU only; modeling_wan.py, shared with Trainium, is untouched), both default on: - DIFFLET_TPU_WAN_SP_PAD: the sequence is padded so each ring half is a multiple of 16 rows (480x832x9: 4680 -> 4736 at tp=4). The model's own scatter / gather are turned off; the first block pads the patch sequence and keeps this rank's rows (index_select on a per-rank buffer, a runtime value, so every rank compiles the same program), the last gathers and cuts the pad, self-attention drops the pad keys (exact) and gives the pad rows the identity rotation. The rings' input splits and output assembly stop being relayout copies (138 ms of a 20-step denoise); cutting the pad keys and the wider final gather cost ~50. - DIFFLET_TPU_WAN_SPLIT_ROPE: as FLUX, each head's channels of self-attention's to_q / to_k rows and q/k RMSNorm weights are permuted once to [0, 2, ..., 1, 3, ...], so the pairwise rotation becomes split halves (a slice and a concat). Removes the (..., 64, 2) interleave and the q/k/v output relayout (~320 ms per denoise). The permutation runs on the host: done on the device with copy_, the parameters kept a layout every forward copied back (copy bf16[1280,5120] x 80 per step, 50 ms, plus ~130 ms idle). v5e-4, main's bench, same tree, back to back: step 269.2 (before) -> 262.2 (pad) -> 249.2 ms (pad + split RoPE), -7.4%; TP is 360.7. Output vs TP 42.2 dB (seed 42; the interleaved vs split layouts differ by the RMSNorm's summation order, which seed 42 amplifies to 26 dB over 20 steps, as before). Tests: the plain model vs the fused one with padding and split RoPE on and off (float64, 1e-10; the q/k RMSNorm swapped for a float64 one, the stock one computes in float32); fail without the pad-key cut or the norm permutation. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ue; trace_diff The plan gets the afternoon's results: the AG kernel cleanup checked against 868401e, the RS kernel's numbers and why it stays off, what Wan's row padding and split-halves RoPE removed (per-shape trace diff), and the reverted GELU epilogue (a Pallas output lives in plain HBM, so its consumer loses XLA's S(1) placement). benchmark/tpu_profile/trace_diff.py: device time per op kind of two traced denoises side by side; reads the xplane when the JSON export has no device events, and with --by-shape keys ops by opcode and output shape. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…not verified benchmark/tpu_profile/bench_case.sh: the one-case runner the Wan / FLUX / HunyuanVideo A/B rows were measured with (main's bench once, outputs and log under $OUT/<tag>, one summary line), with the node paths as overridable defaults and the tree's md5 at the top of the log. The plan lists the environment (versions, model entry, seed), the commands per table row, the run-to-run noise of one configuration, and what was not verified: same-tree TP baseline, split RoPE's own 1-step check, other seeds, what S(1) is. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
shiyang-opensource
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR now carries the TPU sequence-parallel work as one change set: LTX-2 sequence parallelism (the original scope, below), Pallas kernels for the collective matmuls, HunyuanVideo single-stream SP, and Wan row padding plus split-halves RoPE. Every file touched is TPU-only (
backends/tpu/,models/*/tpu_application.py,benchmark/tpu_profile/,docs/plans/, tests). SP itself stays behindDIFFLET_TPU_SP=1.Still stacked on #150 (LTX-2 text encode) and #149 (FLUX): the diff shrinks once they merge.
Step time per model (v5e-4, main's
bench.py, compiled, ms per denoise step)Output vs TP (seed 42, frames PSNR): Wan 42.2 dB, FLUX 39.5 dB, HunyuanVideo 42.9 dB, LTX-2 34.3 dB; each within that model's TP/SP or compiled/eager spread. Every A/B pair ran back to back on one node and tree;
docs/plans/2026-11-10-tpu-collective-matmul-pallas.mdhas the per-step tables, traces, the reproduction commands (benchmark/tpu_profile/bench_case.sh), and what was not verified.Commits added since the LTX-2 SP commit
6f78e40c: step-1 profile of the rings: the cost is the AG ring's write-back, not the wait.01d027fe: Pallas all-gather + matmul kernel (DIFFLET_TPU_SP_KERNEL): x halves resident in VMEM and passed VMEM to VMEM, one double-buffered weight-tile stream across hops and weights, each product tile written once at its final rows. Op names carry a source hash (torch_tpu caches ajax_opby name).bf658ac2: HunyuanVideo's 40 single-stream blocks under SP too (DIFFLET_TPU_SP_SINGLE, default on).f098f4eb: the AG kernel by default for Wan and FLUX (and HunyuanVideo); Qwen-Image and LTX-2 keep the ring.868401eb: plan: P0 (concat vs dus), kernel results, other models.7af1ecdc: kernel cleanup: only the streamed version, full names; outputs sha256-identical to868401ebon Wan, FLUX, HunyuanVideo.84fe6b13: Pallas matmul + reduce-scatter kernel, opt-in (DIFFLET_TPU_SP_RS_KERNEL=pallas): faster only where ICI-bound (Wan attention out 0.518 → 0.504 ms), slower where compute-bound (Wan FFN down 0.928 → 1.040), so off by default.79f0976d: Wan: SP sequence padded to whole tiles (4680 → 4736 rows; 269.2 → 262.2 ms) and RoPE as split halves via a one-time q/k channel permutation done on the host (→ 249.2 ms).ca6abd35,917801b8: plan updates (including a reverted GELU-epilogue experiment),trace_diff.py,bench_case.sh.Tests added
test_tpu_collective_matmul_pallas.py/_rs.py: both kernels in Pallas TPU interpret mode on 4 CPU devices against numpy, race detector on (checked to catch one slot too few); needjax(skipped in CI today).test_tpu_hunyuan_video_sp_single.py,test_tpu_wan_sp_pad.py: the SP layouts against the plain models at tp=1 in float64.Original scope: LTX-2 sequence parallelism
This PR adds sequence parallelism for LTX-2 on the TPU backend over collective matmuls, behind
DIFFLET_TPU_SP=1(off by default, as for Wan, Qwen-Image and FLUX). On a v5e-4 at 480x704x49, a denoise step goes from 218 ms to 192 ms (−12%). The change touches TPU-only files.Stacked on #150 (LTX-2 text encode) and #149 (FLUX): this branch includes their commits, plus a merge of #149, which adds the 2-weight all-gather ring that v2a needs. The new work is the last two commits:
71389c43: LTX-2 SP;ad5bf2e7: plan document.The diff becomes small once #149 and #150 merge.
What changes
backends/tpu/ltx_2/sp.py(new): replaces the block forward only; diffusers' model forward is unchanged.backends/tpu/ltx_2/transformer.py: pads the video rows eagerly, before the compiled graph, to a multiple of 2·tp (2310 → 2312), then cuts the output back. Pad keys are dropped in attention.ops_impl/collective_matmul.py: the docstring now states the rule below.benchmark/tpu_profile/gains:cm_padded_repro.py: a minimal repro, ring blocks on random weights, with pad modes and variants;cm_pure_jax.py: the same ring code in plain JAX;cm_ops_parity.py: each ring op against plain collectives;trace_tail.py: the last device ops per chip, from per-rank traces.docs/plans/2026-11-10-tpu-collective-matmul-pallas.md: a proposal to replace the rings' internals with Pallas kernels, gated on a profile.Root cause of the first (broken) version
The first version padded the sequence inside the compiled graph and sliced each rank's shard at a rank-constant offset. The output was wrong (14.5 dB against TP), and runs hung or halted a core. The chain:
collective-permutecarrieschannel_id=1, so ranks pair them by order.concatenateaway on the ranks whose shard lies inside the unpadded input. The per-rank programs diverged.jax.profilertraces stopped by a timer.The minimal repro separates the cases:
Rule (now in
collective_matmul.py): a value that differs by rank enters a compiled graph as a runtime tensor, not as a Python constant.Confirmed mechanism (
7df82087): per-rank XLA dumps, compared withbenchmark/tpu_profile/hlo_rank_diff.py, show the following.copy-startand apad_slice_fusionthat feed the firstcollective-permute.collective-permutereceive buffers sit +18.9 MB further into its temp allocation (29,884,416 → 48,824,320, …). The sync flags sit at the same offsets on all ranks.Measurements
bench.py, 480x704x49, 20 steps, compiled.2bb56c9e67e2), A and B back to back.47da56e2.Output, SP vs TP: 34.32 dB, mean |d| 2.48/255 (frames 0/24/48).
e641a64c7f40(TP) vsef75047e8a34(SP). Each is bit-identical across its 3 runs.Not checked:
Test plan
DIFFLET_BACKEND=cpu pytest tests/unit/backends/test_tpu_ltx2_sp.py tests/unit/backends/test_tpu_ltx2_text_encode.py tests/unit/models/ltx_2 tests/unit/models/flux: 80 passed, 5 skipped. The new tests cover:🤖 Generated with Claude Code