Repository navigation

TPU sequence parallelism: Pallas collective-matmul kernels, LTX-2 / HunyuanVideo SP, Wan padding + split RoPE - #153

Closed

shiyang-opensource wants to merge 23 commits into

mainfrom

tpu-ltx2-sp

Closed

TPU sequence parallelism: Pallas collective-matmul kernels, LTX-2 / HunyuanVideo SP, Wan padding + split RoPE#153
shiyang-opensource wants to merge 23 commits into
mainfrom
tpu-ltx2-sp

Commits on Nov 9, 2026
- feat(tpu, flux): T5-XXL on the chips, layer-pipelined across the ranks
  
  authored and
  SHIYANG CHENG
  committed
  5423c7e
- perf(tpu, flux): relayout-free RoPE by permuting q/k head channels at load
  
  authored and
  SHIYANG CHENG
  committed
  1c3d703
Commits on Nov 10, 2026
- perf(tpu, flux): sequence parallelism over collective matmuls (DIFFLET_TPU_SP=1)
  
  authored and
  SHIYANG CHENG
  committed
  7db19c8
- perf(tpu, ltx-2): encode the prompt over its real tokens, on rank 0 only
  
  authored and
  SHIYANG CHENG
  committed
  591c0ad
- fix(tpu, ltx-2): contiguous DiT inputs, so torch.compile stops recompiling per process
  
  authored and
  SHIYANG CHENG
  committed
  c31a991
- bench(tpu, ltx-2): CFG check of the text trim (batch-2 conditioning on the chips)
  
  authored and
  SHIYANG CHENG
  committed
  5c2cbbf
- perf(tpu, ltx-2): torch.compile the DiT by default
  
  authored and
  SHIYANG CHENG
  committed
  1a205f0
- fix(tpu, flux): pipelined T5 segments under transformers 4.57 (CI pin)
  
  authored and
  SHIYANG CHENG
  committed
  fd3750c
- test(tpu, flux): SP tests initialize the mesh the code will import
  
  authored and
  SHIYANG CHENG
  committed
  440bbc3
- Merge remote-tracking branch 'origin/tpu-flux-t5-chip' into tpu-ltx2-sp
  
  SHIYANG CHENG
  authored and
  SHIYANG CHENG
  committed
  fa5a2e2
- perf(tpu, ltx-2): sequence parallelism over collective matmuls (DIFFLET_TPU_SP=1)
  
  authored and
  SHIYANG CHENG
  committed
  71389c4
- docs(tpu): plan for Pallas collective-matmul kernels behind the SP API
  
  authored and
  SHIYANG CHENG
  committed
  ad5bf2e
- docs(tpu): the confirmed mechanism behind the in-graph padding failure
  
  authored and
  SHIYANG CHENG
  committed
  7df8208
- docs(tpu): step-1 profile of the SP rings (Wan, LTX-2) and the revised plan
  
  authored and
  SHIYANG CHENG
  committed
  6f78e40
Commits on Nov 11, 2026
- perf(tpu): Pallas all-gather + matmul kernel for the SP rings (DIFFLET_TPU_SP_KERNEL)
  
  authored and
  SHIYANG CHENG
  committed
  01d027f
- perf(tpu, hunyuan-video): single-stream blocks sequence-parallel too
  
  authored and
  SHIYANG CHENG
  committed
  bf658ac
- perf(tpu): the Pallas AG kernel by default for Wan and FLUX
  
  authored and
  SHIYANG CHENG
  committed
  f098f4e
- docs(tpu): P0, the Pallas AG kernel, other models and HunyuanVideo single-stream SP
  
  authored and
  SHIYANG CHENG
  committed
  868401e
- refactor(tpu): the Pallas AG kernel as its one streamed version, with full names
  
  authored and
  SHIYANG CHENG
  committed
  7af1ecd
- feat(tpu): Pallas matmul + reduce-scatter kernel for the SP rings (opt-in)
  
  authored and
  SHIYANG CHENG
  committed
  84fe6b1
- perf(tpu, wan): pad the SP sequence to whole tiles and rotate RoPE as split halves
  
  authored and
  SHIYANG CHENG
  committed
  79f0976
- docs(tpu): the RS kernel, Wan padding and split RoPE, the GELU epilogue; trace_diff
  
  authored and
  SHIYANG CHENG
  committed
  ca6abd3
- docs(tpu): how to reproduce the afternoon's A/B numbers, and what is not verified
  
  authored and
  SHIYANG CHENG
  committed
  917801b

Commits

Commits on Nov 9, 2026

Commits on Nov 10, 2026

Commits on Nov 11, 2026