From 5423c7eff68752a7fa59c6228820c389a8c9de2e Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Mon, 9 Nov 2026 15:35:11 -0800
Subject: [PATCH 01/22] feat(tpu, flux): T5-XXL on the chips, layer-pipelined
 across the ranks

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>
---
 benchmark/tpu_profile/t5_device_probe.py      | 124 +++++++++++++++
 difflet/models/flux/tpu_application.py        | 142 +++++++++++++++++-
 .../models/flux/test_flux_tpu_application.py  |  78 ++++++++++
 3 files changed, 339 insertions(+), 5 deletions(-)
 create mode 100644 benchmark/tpu_profile/t5_device_probe.py

diff --git a/benchmark/tpu_profile/t5_device_probe.py b/benchmark/tpu_profile/t5_device_probe.py
new file mode 100644
index 00000000..628988c6
--- /dev/null
+++ b/benchmark/tpu_profile/t5_device_probe.py
@@ -0,0 +1,124 @@
+"""FLUX's T5-XXL on the chips (layer pipeline) against the host fp32 encoder.
+
+    python -m benchmark.tpu_profile.t5_device_probe [--dtypes bf16,fp32]
+
+One process per chip, as in serving. Per dtype, every rank builds
+``TpuPipelinedT5Encoder`` (its segment of the 24 blocks), encodes the prompts
+twice (the first call compiles) and rank 0 compares each result against the
+stock host encoder in fp32 — what ``TpuBroadcastT5Encoder`` computes today.
+Reports the warm encode latency, the host latency, and per prompt the cosine
+and the max / relative error of the (1, 512, 4096) prompt embeds.
+"""
+
+from __future__ import annotations
+
+import argparse
+import multiprocessing as mp
+import os
+import sys
+import time
+from pathlib import Path
+
+MODEL = os.environ.get(
+    "DIFFLET_FLUX_PATH",
+    "/mnt/models/hf/hub/models--black-forest-labs--FLUX.1-dev/snapshots",
+)
+PROMPTS = [
+    "A cat holding a sign that says hello world",
+    "A photorealistic portrait of an old fisherman at dawn, golden light, 85mm lens, shallow depth of field",
+    "",
+]
+
+
+def _model_path() -> str:
+    path = Path(MODEL)
+    if (path / "text_encoder_2").is_dir():
+        return str(path)
+    return str(next(p for p in sorted(path.iterdir()) if (p / "text_encoder_2").is_dir()))
+
+
+def _worker(rank, world, dtypes, reply_q):
+    try:
+        from difflet.backends.tpu.ops_impl import device_runtime
+
+        device_runtime.initialize_replica(rank, world)
+        import torch
+
+        from difflet.models.flux.tpu_application import (
+            FLUX_TEXT_SEQ_LEN,
+            TpuBroadcastT5Encoder,
+            TpuPipelinedT5Encoder,
+            _normalize_dtype,
+        )
+
+        path = _model_path()
+        results: dict = {"rank": rank}
+        if rank == 0:
+            ref = TpuBroadcastT5Encoder(path, seq_len=FLUX_TEXT_SEQ_LEN, dtype=torch.float32)
+            ref_out, host_s = [], []
+            for prompt in PROMPTS:
+                mark = time.monotonic()
+                ref_out.append(ref.encode_local(prompt).float())
+                host_s.append(time.monotonic() - mark)
+            results["host_fp32_s"] = [round(s, 3) for s in host_s]
+            del ref
+        for name in dtypes:
+            enc = TpuPipelinedT5Encoder(path, seq_len=FLUX_TEXT_SEQ_LEN, dtype=torch.float32,
+                                        encoder_dtype=_normalize_dtype(name))
+            mem_gb = device_runtime.memory_info()["bytes_used"] / 2**30
+            first = []
+            for prompt in PROMPTS:
+                mark = time.monotonic()
+                enc(prompt)
+                first.append(time.monotonic() - mark)
+            warm, outs = [], []
+            for prompt in PROMPTS:
+                mark = time.monotonic()
+                outs.append(enc(prompt).float())
+                warm.append(time.monotonic() - mark)
+            row = {"first_s": [round(s, 3) for s in first], "warm_s": [round(s, 3) for s in warm],
+                   "hbm_used_gb": round(mem_gb, 2)}
+            if rank == 0:
+                errs = []
+                for got, want in zip(outs, ref_out):
+                    diff = (got - want).abs()
+                    cos = torch.nn.functional.cosine_similarity(got.flatten(), want.flatten(), dim=0)
+                    errs.append({"cos": round(float(cos), 6), "max_abs": round(float(diff.max()), 4),
+                                 "rel": round(float(diff.norm() / want.norm()), 5)})
+                row["vs_host_fp32"] = errs
+            results[name] = row
+            del enc
+            import gc
+
+            gc.collect()
+        reply_q.put(results)
+    except Exception as exc:  # noqa: BLE001 - report, the parent prints it
+        import traceback
+
+        traceback.print_exc()
+        reply_q.put({"rank": rank, "error": repr(exc)})
+
+
+def main() -> None:
+    from difflet.backends.tpu.ops_impl import device_runtime
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dtypes", default="bf16,fp32")
+    args = parser.parse_args()
+    dtypes = [d for d in args.dtypes.split(",") if d]
+    world = len(list(Path("/dev/vfio").glob("[0-9]*"))) or 1
+    device_runtime.prepare_launch(world)
+    ctx = mp.get_context("spawn")
+    reply_q = ctx.Queue()
+    procs = [ctx.Process(target=_worker, args=(r, world, dtypes, reply_q)) for r in range(world)]
+    for p in procs:
+        p.start()
+    for _ in range(world):
+        msg = reply_q.get(timeout=1500)
+        print(msg, flush=True, file=sys.stdout if msg.get("rank") == 0 else sys.stderr)
+    for p in procs:
+        p.join(timeout=60)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/difflet/models/flux/tpu_application.py b/difflet/models/flux/tpu_application.py
index 0c88912a..f54b9e7e 100644
--- a/difflet/models/flux/tpu_application.py
+++ b/difflet/models/flux/tpu_application.py
@@ -1,4 +1,4 @@
-"""FLUX.1-dev on the TPU backend: the DiT and the VAE on the chips, the text encoders on the host.
+"""FLUX.1-dev on the TPU backend: the DiT, T5 and the VAE on the chips, CLIP on the host.
 
 Fifth model on the TPU backend, option (b) of
 docs/plans/2026-09-11-tpu-port-hunyuan-ltx2-flux.md: diffusers'
@@ -8,9 +8,12 @@
 
 Placement follows the other ports:
 
-* T5-XXL (4.7B) runs on the host in fp32 on XLA ordinal 0 only and its hidden
-  state is broadcast (``TpuBroadcastT5Encoder``); one fp32 copy is ~19 GB, and
-  bf16 is emulated on the EPYC host (Qwen measured 4.5x slower);
+* T5-XXL (4.7B) runs on the chips in bf16, its blocks split into one segment
+  per rank (``TpuPipelinedT5Encoder``, ~2.4 GB a chip): 0.20 s per request
+  against 3.75 s for the fp32 host encoder (2026-11-09). fp32 on the chips
+  (~4.8 GB) leaves rank 0 too little HBM for the VAE decode.
+  ``DIFFLET_TPU_T5_DEVICE=0`` restores the host path (``TpuBroadcastT5Encoder``:
+  fp32 on rank 0, broadcast; bf16 is emulated on the EPYC host);
 * CLIP-L (123M) runs on the host on every rank — the pooled projection is
   cheaper to recompute than to broadcast in lockstep;
 * the 2D VAE (84M) is resident on the chip of the primary replica
@@ -217,6 +220,130 @@ def __call__(self, prompt: str) -> torch.Tensor:
         return payload[0].cpu()
 
 
+#: ``0`` keeps T5 on the host (``TpuBroadcastT5Encoder``); anything else runs it
+#: on the chips, layer-pipelined (``TpuPipelinedT5Encoder``).
+T5_DEVICE_ENV = "DIFFLET_TPU_T5_DEVICE"
+#: dtype of the on-chip T5: ``bf16`` (default) or ``fp32`` — fp32 runs the
+#: encoder but then the 1024x1024 VAE decode on rank 0 runs out of HBM (2026-11-09).
+T5_DEVICE_DTYPE_ENV = "DIFFLET_TPU_T5_DEVICE_DTYPE"
+
+
+def _t5_on_device() -> bool:
+    return os.environ.get(T5_DEVICE_ENV, "1").strip().lower() not in {"0", "false", "no"}
+
+
+def layer_segments(num_layers: int, parts: int) -> list[range]:
+    """Split ``num_layers`` into ``parts`` contiguous ranges, the remainder spread from the front."""
+    base, extra = divmod(int(num_layers), int(parts))
+    out, start = [], 0
+    for index in range(int(parts)):
+        size = base + (1 if index < extra else 0)
+        out.append(range(start, start + size))
+        start += size
+    return out
+
+
+class TpuPipelinedT5Encoder:
+    """T5-XXL on the chips, its 24 blocks split into one contiguous segment per rank.
+
+    The host encoder costs ~3.4 s of a ~8.7 s FLUX request (fp32 on the EPYC
+    host). bf16 T5 is ~9.5 GB: too big for the ~9.7 GB a DiT shard leaves on one
+    chip, ~2.4 GB per chip split four ways. Tensor-parallel T5 would need its own
+    sharded modeling; a layer pipeline needs nothing but the stock HF blocks and
+    ``broadcast``, the one TorchTPU collective every port already relies on
+    (send/recv hangs). Only one chip computes at a time, but the whole encoder is
+    ~5 TFLOP — tens of milliseconds — and each hop moves one (1, 512, 4096) state.
+
+    Rank r runs its segment and broadcasts the hidden state from r; rank 0 also
+    owns the embedding, the last rank the final norm. Every rank keeps the
+    relative-attention-bias table (64x32) so a segment that does not start at
+    block 0 rebuilds the shared position bias locally, as block 0 would. Each hop
+    carries a status flag, so a failing rank never leaves the others blocked in
+    a collective: later segments skip their compute and every rank raises.
+    """
+
+    def __init__(self, model_path: str, *, seq_len: int, dtype: torch.dtype,
+                 encoder_dtype: torch.dtype | None = None) -> None:
+        from difflet.backends.tpu.ops_impl import device_runtime
+        from transformers import AutoTokenizer, T5EncoderModel
+
+        if encoder_dtype is None:
+            encoder_dtype = _normalize_dtype(os.environ.get(T5_DEVICE_DTYPE_ENV, "bf16"))
+        self.seq_len = int(seq_len)
+        self.dtype = dtype
+        self.encoder_dtype = encoder_dtype
+        self.rank = device_runtime.rank()
+        self.world = device_runtime.world_size()
+        self._device = device_runtime.device()
+        self.tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer_2"))
+        model = T5EncoderModel.from_pretrained(
+            os.path.join(model_path, "text_encoder_2"), dtype=encoder_dtype
+        ).eval()
+        model.requires_grad_(False)
+        self.hidden_size = int(model.config.d_model)
+        blocks = model.encoder.block
+        self.segments = layer_segments(len(blocks), self.world)
+        mine = self.segments[self.rank]
+        dev = self._device
+        self.blocks = torch.nn.ModuleList([blocks[i] for i in mine]).to(dev)
+        # The bias table lives in block 0's self-attention; only it computes the bias.
+        self.bias_attention = blocks[0].layer[0].SelfAttention
+        if mine.start != 0:
+            self.bias_attention.relative_attention_bias.to(dev)
+        self.embed = model.encoder.embed_tokens.to(dev) if self.rank == 0 else None
+        self.final_norm = model.encoder.final_layer_norm.to(dev) if self.rank == self.world - 1 else None
+        # Drop the host copies of everything this rank does not run.
+        del model, blocks
+        device_runtime.wait()
+
+    @property
+    def model(self):  # parity with TpuBroadcastT5Encoder for the load log line
+        return self.blocks
+
+    def _segment(self, hidden: torch.Tensor | None, ids: torch.Tensor) -> torch.Tensor:
+        if self.rank == 0:
+            hidden = self.embed(ids)
+        position_bias = None
+        if self.segments[self.rank].start != 0:
+            position_bias = self.bias_attention.compute_bias(self.seq_len, self.seq_len, device=self._device)
+        for block in self.blocks:
+            hidden, position_bias, _ = block(hidden, None, position_bias)
+        if self.final_norm is not None:
+            hidden = self.final_norm(hidden)
+        return hidden
+
+    def __call__(self, prompt: str) -> torch.Tensor:
+        from difflet.backends.tpu.ops_impl import device_runtime
+
+        ids = self.tokenizer(
+            [prompt], padding="max_length", max_length=self.seq_len, truncation=True,
+            return_overflowing_tokens=False, return_length=False, return_tensors="pt",
+        ).input_ids.to(self._device)
+        hidden = None
+        failed = False
+        for owner in range(self.world):
+            out = None
+            status = _ENCODE_FAILED if failed else _ENCODE_OK  # only the owner's value is sent
+            if self.rank == owner and not failed:
+                try:
+                    with torch.no_grad():
+                        out = self._segment(hidden, ids).to(self.encoder_dtype)
+                except Exception:  # noqa: BLE001 - re-raised below on every rank
+                    logger.exception("flux.tpu_t5_segment_failed")
+                    status = _ENCODE_FAILED
+            if out is None:
+                out = torch.zeros(1, self.seq_len, self.hidden_size, dtype=self.encoder_dtype, device=self._device)
+            payload = [out, torch.tensor([status], dtype=torch.int32).to(self._device)]
+            device_runtime.broadcast(payload, src=owner)
+            # Read the flag on every rank: it also completes the collective
+            # (an unread collective result leaves the other ranks waiting).
+            failed = failed or int(payload[1].cpu()[0]) != _ENCODE_OK
+            hidden = payload[0]
+        if failed:
+            raise RuntimeError("Flux T5 prompt encoding failed on a pipeline rank")
+        return hidden.to(self.dtype).cpu()
+
+
 class HostClipPooler:
     """CLIP-L pooled projection on the host (``_get_clip_prompt_embeds``)."""
 
@@ -361,7 +488,10 @@ def load_eager(self) -> None:
 
         if bool(self.kwargs.get("enable_host_pipeline", True)):
             mark = time.monotonic()
-            self.t5 = TpuBroadcastT5Encoder(self.model_path, seq_len=self.text_seq_len, dtype=self.dtype)
+            if _t5_on_device():
+                self.t5 = TpuPipelinedT5Encoder(self.model_path, seq_len=self.text_seq_len, dtype=self.dtype)
+            else:
+                self.t5 = TpuBroadcastT5Encoder(self.model_path, seq_len=self.text_seq_len, dtype=self.dtype)
             self.clip = HostClipPooler(self.model_path, dtype=self.dtype)
             self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
                 os.path.join(self.model_path, "scheduler")
@@ -534,6 +664,8 @@ def __call__(self, prompt: str, num_inference_steps: int = 28, height: int | Non
     "TpuBroadcastT5Encoder",
     "TpuDeviceImageVae",
     "TpuFluxApplication",
+    "TpuPipelinedT5Encoder",
+    "layer_segments",
     "device_euler_loop",
     "flux_sigmas_and_mu",
     "image_seq_len",
diff --git a/tests/unit/models/flux/test_flux_tpu_application.py b/tests/unit/models/flux/test_flux_tpu_application.py
index 2af95d2b..6a7b3f81 100644
--- a/tests/unit/models/flux/test_flux_tpu_application.py
+++ b/tests/unit/models/flux/test_flux_tpu_application.py
@@ -205,3 +205,81 @@ def test_denoise_refuses_a_shape_other_than_the_loaded_one(snapshot):
     app.scheduler = object()
     with pytest.raises(ValueError, match="loaded for 1024x1024"):
         app.denoise({}, num_inference_steps=4, guidance_scale=3.5, generator=None, height=512, width=512)
+
+
+def test_layer_segments_cover_every_layer_once():
+    assert [list(r) for r in mod.layer_segments(24, 4)] == [list(range(6 * i, 6 * i + 6)) for i in range(4)]
+    segs = mod.layer_segments(7, 3)
+    assert [len(r) for r in segs] == [3, 2, 2]
+    assert [i for r in segs for i in r] == list(range(7))
+
+
+def _tiny_t5():
+    from transformers import T5Config, T5EncoderModel
+
+    torch.manual_seed(0)
+    cfg = T5Config(vocab_size=50, d_model=16, d_kv=4, d_ff=32, num_layers=5, num_heads=4,
+                   relative_attention_num_buckets=8, feed_forward_proj="gated-gelu",
+                   dropout_rate=0.0)
+    return T5EncoderModel(cfg).eval()
+
+
+def _pipelined_ranks(monkeypatch, world: int, seq_len: int):
+    """One TpuPipelinedT5Encoder per simulated rank, all built from the same tiny T5."""
+    import transformers
+
+    from difflet.backends.tpu.ops_impl import device_runtime
+
+    reference = _tiny_t5()
+    state = reference.state_dict()
+
+    def load(*args, **kwargs):
+        model = _tiny_t5()
+        model.load_state_dict(state)
+        return model
+
+    monkeypatch.setattr(transformers.T5EncoderModel, "from_pretrained", staticmethod(load))
+    monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", staticmethod(lambda *a, **k: None))
+    monkeypatch.setattr(device_runtime, "device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(device_runtime, "world_size", lambda: world)
+    monkeypatch.setattr(device_runtime, "wait", lambda: None)
+    encoders = []
+    for rank in range(world):
+        monkeypatch.setattr(device_runtime, "rank", lambda rank=rank: rank)
+        encoders.append(mod.TpuPipelinedT5Encoder("unused", seq_len=seq_len, dtype=torch.float32,
+                                                  encoder_dtype=torch.float32))
+    return reference, encoders
+
+
+def test_pipelined_t5_segments_reproduce_the_hf_encoder(monkeypatch):
+    seq = 9
+    reference, encoders = _pipelined_ranks(monkeypatch, world=3, seq_len=seq)
+    assert [list(e.segments[e.rank]) for e in encoders] == [[0, 1], [2, 3], [4]]
+    ids = torch.randint(0, 50, (1, seq))
+    hidden = None
+    with torch.no_grad():
+        for enc in encoders:  # each rank's segment, the hidden state handed on as the broadcast would
+            hidden = enc._segment(hidden, ids)
+        want = reference(ids)[0]
+    torch.testing.assert_close(hidden, want, rtol=1e-5, atol=1e-5)
+
+
+def test_pipelined_t5_failure_reaches_every_rank_without_skipping_a_collective(monkeypatch):
+    from difflet.backends.tpu.ops_impl import device_runtime
+
+    _, encoders = _pipelined_ranks(monkeypatch, world=2, seq_len=4)
+    enc = encoders[1]  # the last rank: it receives segment 0, then runs its own
+    enc.tokenizer = lambda *a, **k: SimpleNamespace(input_ids=torch.zeros(1, 4, dtype=torch.long))
+    calls = []
+
+    def broadcast(tensors, src=0):
+        calls.append(src)
+        if src == 0:  # rank 0 reports a failure
+            tensors[1].fill_(mod._ENCODE_FAILED)
+
+    monkeypatch.setattr(device_runtime, "broadcast", broadcast)
+    ran = []
+    monkeypatch.setattr(enc, "_segment", lambda *a: ran.append(1))
+    with pytest.raises(RuntimeError, match="pipeline rank"):
+        enc("a fox")
+    assert calls == [0, 1] and ran == []  # every hop still reached, own segment skipped

From 1c3d703343d437ad9a83dabf46d900210f0a173e Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Mon, 9 Nov 2026 15:59:57 -0800
Subject: [PATCH 02/22] perf(tpu, flux): relayout-free RoPE by permuting q/k
 head channels at 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>
---
 difflet/backends/tpu/flux/rope.py             | 112 ++++++++++++++++++
 difflet/backends/tpu/flux/transformer.py      |  14 +++
 difflet/models/flux/tp_sharding.py            |  12 +-
 .../unit/models/flux/test_flux_tp_sharding.py |  25 ++++
 4 files changed, 161 insertions(+), 2 deletions(-)
 create mode 100644 difflet/backends/tpu/flux/rope.py

diff --git a/difflet/backends/tpu/flux/rope.py b/difflet/backends/tpu/flux/rope.py
new file mode 100644
index 00000000..c2639e5b
--- /dev/null
+++ b/difflet/backends/tpu/flux/rope.py
@@ -0,0 +1,112 @@
+"""FLUX's RoPE in the split-halves layout, by permuting q/k head channels once at load.
+
+diffusers' FLUX rotary embedding pairs channels ``(2i, 2i+1)`` of every head:
+``x.reshape(..., D/2, 2).unbind(-1)`` then ``stack(...).flatten`` — a relayout
+of the minor dimension on every q and k of every block, ~4-7% of a compiled
+TPU denoise step (trace 2026-11-03: ``embeddings.py`` ``apply_rotary_emb`` and
+most of the "data formatting" time).
+
+Attention only sees q and k through ``q . k``, which is unchanged when both are
+permuted the same way along the head dim. Permuting each head's channels to
+``P = [0, 2, 4, ..., D-2, 1, 3, ..., D-1]`` — the output rows of ``to_q``/``to_k``
+(and the text stream's ``add_q_proj``/``add_k_proj``), and the weights of the
+per-head qk RMSNorms that follow them — turns each pair ``(2i, 2i+1)`` into
+``(i, D/2 + i)``. That is the split-halves rotary (``use_real_unbind_dim=-2``):
+a slice and a concat, no interleave. The cos/sin tables are permuted the same
+way once per forward (``PermutedFluxPosEmbed``). V, the attention output and
+everything after are untouched, so the result equals the stock layout's up to
+the RMSNorm's summation order.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+ROPE_UNBIND_DIM = -2
+
+
+def pair_split_permutation(head_dim: int) -> torch.Tensor:
+    """``[0, 2, ..., D-2, 1, 3, ..., D-1]``: interleaved pairs -> split halves."""
+    if int(head_dim) % 2:
+        raise ValueError(f"rotary head_dim must be even, got {head_dim}")
+    idx = torch.arange(int(head_dim))
+    return torch.cat([idx[0::2], idx[1::2]])
+
+
+class PermutedFluxPosEmbed(nn.Module):
+    """``FluxPosEmbed`` whose (cos, sin) come out in the split-halves channel order."""
+
+    def __init__(self, inner: nn.Module, head_dim: int) -> None:
+        super().__init__()
+        self.inner = inner
+        self.register_buffer("perm", pair_split_permutation(head_dim), persistent=False)
+
+    def forward(self, ids: torch.Tensor):
+        cos, sin = self.inner(ids)
+        perm = self.perm.to(cos.device)
+        return cos.index_select(-1, perm), sin.index_select(-1, perm)
+
+
+def _permute_rows_per_head_(linear: nn.Module, head_dim: int, perm: torch.Tensor) -> None:
+    weight = linear.weight
+    out = int(weight.shape[0])
+    if out % int(head_dim):
+        raise ValueError(f"{type(linear).__name__} rows {out} are not whole heads of {head_dim}")
+    heads = out // int(head_dim)
+    perm = perm.to(weight.device)
+    with torch.no_grad():
+        w = weight.view(heads, int(head_dim), -1).index_select(1, perm).reshape(out, -1)
+        weight.copy_(w)
+        bias = getattr(linear, "bias", None)
+        if bias is not None:
+            bias.copy_(bias.view(heads, int(head_dim)).index_select(1, perm).reshape(out))
+
+
+def _permute_norm_(norm: nn.Module, perm: torch.Tensor) -> None:
+    weight = getattr(norm, "weight", None)
+    if weight is None:
+        return
+    with torch.no_grad():
+        weight.copy_(weight.index_select(0, perm.to(weight.device)))
+
+
+def permute_flux_rope_layout_(transformer: nn.Module, processor_cls=None) -> int:
+    """Rewrite a loaded (TP-sharded or not) FLUX transformer in place for split-halves RoPE.
+
+    Returns the number of attention modules rewritten. ``processor_cls`` builds
+    each attention's processor with ``rope_unbind_dim=-2``; it defaults to the
+    TP processor of ``models/flux/tp_sharding.py``.
+    """
+    if processor_cls is None:
+        from difflet.models.flux.tp_sharding import FluxTPAttnProcessor as processor_cls
+    if isinstance(transformer.pos_embed, PermutedFluxPosEmbed):
+        raise RuntimeError("FLUX RoPE layout is already permuted")
+    head_dim = None
+    count = 0
+    processor = processor_cls(rope_unbind_dim=ROPE_UNBIND_DIM)
+    for block in list(transformer.transformer_blocks) + list(transformer.single_transformer_blocks):
+        attn = block.attn
+        dim = int(attn.to_q.weight.shape[0]) // int(attn.heads)
+        head_dim = dim if head_dim is None else head_dim
+        perm = pair_split_permutation(dim)
+        for name in ("to_q", "to_k", "add_q_proj", "add_k_proj"):
+            linear = getattr(attn, name, None)
+            if linear is not None:
+                _permute_rows_per_head_(linear, dim, perm)
+        for name in ("norm_q", "norm_k", "norm_added_q", "norm_added_k"):
+            norm = getattr(attn, name, None)
+            if norm is not None:
+                _permute_norm_(norm, perm)
+        attn.processor = processor
+        count += 1
+    transformer.pos_embed = PermutedFluxPosEmbed(transformer.pos_embed, head_dim)
+    return count
+
+
+__all__ = [
+    "PermutedFluxPosEmbed",
+    "ROPE_UNBIND_DIM",
+    "pair_split_permutation",
+    "permute_flux_rope_layout_",
+]
diff --git a/difflet/backends/tpu/flux/transformer.py b/difflet/backends/tpu/flux/transformer.py
index 595d5e5b..3ce52aac 100644
--- a/difflet/backends/tpu/flux/transformer.py
+++ b/difflet/backends/tpu/flux/transformer.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import os
 import re
 from pathlib import Path
 
@@ -16,6 +17,15 @@
 
 logger = logging.getLogger(__name__)
 
+#: ``0`` keeps diffusers' interleaved-pair RoPE; otherwise q/k head channels are
+#: permuted at load for the relayout-free split-halves form (``rope.py``).
+SPLIT_ROPE_ENV = "DIFFLET_TPU_FLUX_SPLIT_ROPE"
+
+
+def _split_rope_enabled() -> bool:
+    return os.environ.get(SPLIT_ROPE_ENV, "1").strip().lower() not in {"0", "false", "no"}
+
+
 _SINGLE_PROJ_OUT = re.compile(r"^(single_transformer_blocks\.\d+\.)proj_out_(attn|mlp)\.(weight|bias)$")
 
 
@@ -87,6 +97,10 @@ def _prepare_module(self) -> torch.nn.Module:
             module, self.model_path, tp_size=spec.tp, tp_rank=parallel_mesh.get_tp_rank(),
             dtype=dtype, rename=make_checkpoint_key(self.config.inner_dim), strict=True,
         )
+        if _split_rope_enabled():
+            from difflet.backends.tpu.flux.rope import permute_flux_rope_layout_
+
+            permute_flux_rope_layout_(module)
         module.requires_grad_(False)
         self.module = module.eval()
         return self.module
diff --git a/difflet/models/flux/tp_sharding.py b/difflet/models/flux/tp_sharding.py
index 1cbcdb1e..2d1fd036 100644
--- a/difflet/models/flux/tp_sharding.py
+++ b/difflet/models/flux/tp_sharding.py
@@ -85,8 +85,15 @@ class FluxTPAttnProcessor:
     the joint sequence, shared RoPE, no mask; only the attention call differs
     (heads folded into the batch axis, ``tp_q/tp_k`` so the op knows the heads
     are already sharded).
+
+    ``rope_unbind_dim`` is diffusers' ``use_real_unbind_dim``: ``-1`` is FLUX's
+    interleaved (re, im) pairs; ``-2`` the split-halves layout, for q/k weights
+    whose head channels were permuted to match (``backends/tpu/flux/rope.py``).
     """
 
+    def __init__(self, rope_unbind_dim: int = -1) -> None:
+        self.rope_unbind_dim = int(rope_unbind_dim)
+
     def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None,
                  image_rotary_emb=None):
         from diffusers.models.embeddings import apply_rotary_emb
@@ -108,8 +115,9 @@ def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_ma
             key = torch.cat([ek, key], dim=1)
             value = torch.cat([ev, value], dim=1)
         if image_rotary_emb is not None:
-            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
-            key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
+            unbind = self.rope_unbind_dim
+            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=unbind)
+            key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=unbind)
 
         bsz, seq, _, head_dim = query.shape
         out_dtype = query.dtype
diff --git a/tests/unit/models/flux/test_flux_tp_sharding.py b/tests/unit/models/flux/test_flux_tp_sharding.py
index fdf80105..6945f067 100644
--- a/tests/unit/models/flux/test_flux_tp_sharding.py
+++ b/tests/unit/models/flux/test_flux_tp_sharding.py
@@ -125,3 +125,28 @@ def test_flux_1_dev_geometry():
     assert cfg.inner_dim == 3072 and cfg.image_seq_len == 4096
     with pytest.raises(ValueError, match="divisible by 16"):
         TpuFluxConfig(FLUX_1_DEV_CONFIG, height=1000, tp_degree=4).validate()
+
+
+def test_split_rope_permutation_reproduces_the_stock_forward():
+    """q/k head channels permuted at load + split-halves RoPE == diffusers' interleaved RoPE."""
+    from difflet.backends.tpu.flux.rope import pair_split_permutation, permute_flux_rope_layout_
+    from difflet.models.flux import tp_sharding as S
+
+    assert pair_split_permutation(8).tolist() == [0, 2, 4, 6, 1, 3, 5, 7]
+    cfg = TpuFluxConfig(_TINY, height=32, width=48, tp_degree=1, torch_dtype=torch.float32)
+    ref = build_flux_transformer(cfg).eval()
+    for module in ref.modules():  # non-trivial norm weights, so a missed permutation shows
+        if type(module).__name__ == "RMSNorm" and getattr(module, "weight", None) is not None:
+            module.weight.data.uniform_(0.5, 1.5)
+    permuted = build_flux_transformer(cfg).eval()
+    permuted.load_state_dict(ref.state_dict())
+    for block in list(ref.transformer_blocks) + list(ref.single_transformer_blocks):
+        block.attn.processor = S.FluxTPAttnProcessor()
+    assert permute_flux_rope_layout_(permuted) == 2
+    with pytest.raises(RuntimeError, match="already permuted"):
+        permute_flux_rope_layout_(permuted)
+    inputs = _inputs(cfg)
+    with torch.no_grad():
+        expected = ref(*inputs, return_dict=False)[0]
+        actual = permuted(*inputs, return_dict=False)[0]
+    torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-4)

From 7db19c880fa178e34159a4ab39d2dd3ccc61c4a2 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Mon, 9 Nov 2026 17:14:47 -0800
Subject: [PATCH 03/22] perf(tpu, flux): sequence parallelism over collective
 matmuls (DIFFLET_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>
---
 difflet/backends/tpu/flux/sp.py               | 207 ++++++++++++++++++
 difflet/backends/tpu/flux/transformer.py      |  15 ++
 .../tpu/ops_impl/collective_matmul.py         |  91 ++++++--
 tests/unit/models/flux/test_flux_tpu_sp.py    |  99 +++++++++
 4 files changed, 391 insertions(+), 21 deletions(-)
 create mode 100644 difflet/backends/tpu/flux/sp.py
 create mode 100644 tests/unit/models/flux/test_flux_tpu_sp.py

diff --git a/difflet/backends/tpu/flux/sp.py b/difflet/backends/tpu/flux/sp.py
new file mode 100644
index 00000000..edd3c09a
--- /dev/null
+++ b/difflet/backends/tpu/flux/sp.py
@@ -0,0 +1,207 @@
+"""FLUX.1-dev on TPU with sequence parallelism over collective matmuls (``DIFFLET_TPU_SP=1``).
+
+Under TP every block ends its attention and its MLP in an all-reduce (20% of a
+compiled denoise step, trace 2026-11-03). As for Wan and Qwen-Image, each
+all-reduce becomes a reduce-scatter fused into the row-parallel matmul and an
+all-gather fused into the next column-parallel one
+(``ops_impl.collective_matmul``); the per-token work between them runs on this
+rank's 1/tp of the tokens. FLUX has no shared Megatron-SP to reuse (diffusers'
+modeling, TP-sharded by ``models/flux/tp_sharding.py``), so :func:`install`
+swaps, on this module instance only, the model forward and every block forward.
+
+Layout. Each stream is split by sequence on its own: TP rank r holds text rows
+``[r*T/tp, (r+1)*T/tp)`` and image rows ``[r*I/tp, (r+1)*I/tp)`` (128 + 1024
+at 512 text tokens, 1024x1024). The inputs are sliced before the embedders, so
+the embedders already run on the shard; the 64-channel output is gathered at
+the end, after ``norm_out`` / ``proj_out``.
+
+    joint block   img -AG+[to_q,to_k,to_v]-> ; txt -AG+[add_q,add_k,add_v]->
+                  joint attention (the TP processor, unchanged, full sequence)
+                  img -to_out[0]+RS-> ; txt -to_add_out+RS->
+                  ff / ff_context: -AG+net[0].proj-> gelu -net[2]+RS->
+    single block  x = [txt_r | img_r]  -AG+[to_q,to_k,to_v,proj_mlp]->
+                  attention ; gelu(mlp)
+                  -[proj_out_attn | proj_out_mlp]+RS-> (one reduce-scatter
+                  for the sum of the two row-parallel halves)
+
+Token order. A single block gathers ``[txt_r | img_r]`` from every rank, so its
+full sequence is ``[txt_0, img_0, txt_1, img_1, ...]`` rather than ``[txt | img]``.
+FLUX's attention has no mask and RoPE is per token, so the attention result is
+the same for any order as long as every token keeps its own rotary angle: the
+cos/sin rows are permuted to the gathered order once per forward
+(``single_order``). The joint blocks gather each stream separately, so they see
+the natural order.
+"""
+
+from __future__ import annotations
+
+import types
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from difflet.backends.tpu.ops_impl import collective_matmul as cm
+
+#: The collective matmuls FLUX uses: RS over one and two pairs, AG over 1, 3, 4 weights.
+FLUX_OPS = (("rs", 1), ("rs", 2), ("ag", 1), ("ag", 3), ("ag", 4))
+
+
+def single_order(text_len: int, image_len: int, tp: int) -> torch.Tensor:
+    """Indices into the natural ``[txt | img]`` order, in the gathered single-block
+    order ``[txt_0, img_0, txt_1, img_1, ...]``."""
+    t, i = int(text_len) // int(tp), int(image_len) // int(tp)
+    parts = []
+    for rank in range(int(tp)):
+        parts.append(torch.arange(rank * t, (rank + 1) * t))
+        parts.append(int(text_len) + torch.arange(rank * i, (rank + 1) * i))
+    return torch.cat(parts)
+
+
+class _Precomputed:
+    """The ``attn`` the TP processor sees: projections return tensors already
+    computed by the fused all-gathers, the output projections pass through."""
+
+    def __init__(self, attn, qkv, added_qkv=None):
+        self._attn = attn
+        self.to_q, self.to_k, self.to_v = (lambda _x, t=t: t for t in qkv)
+        if added_qkv is not None:
+            self.add_q_proj, self.add_k_proj, self.add_v_proj = (lambda _x, t=t: t for t in added_qkv)
+        self.to_out = [nn.Identity(), *list(attn.to_out)[1:]] if hasattr(attn, "to_out") else None
+        self.to_add_out = nn.Identity()
+
+    def __getattr__(self, name):  # heads, norm_q, norm_k, added_kv_proj_dim, ...
+        return getattr(self._attn, name)
+
+
+def _ff(ff, x):
+    act = ff.net[0]
+    (h,) = cm.all_gather_matmul(x, (act.proj,))
+    h = F.gelu(h, approximate=getattr(act, "approximate", "tanh"))
+    return cm.matmul_reduce_scatter(h, ff.net[2])
+
+
+def _joint_block_forward(self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=None,
+                         joint_attention_kwargs=None):
+    # img (B, I/tp, D), txt (B, T/tp, D)
+    norm_h, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+    norm_e, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+        encoder_hidden_states, emb=temb)
+    attn = self.attn
+    qkv = cm.all_gather_matmul(norm_h, (attn.to_q, attn.to_k, attn.to_v))
+    added = cm.all_gather_matmul(norm_e, (attn.add_q_proj, attn.add_k_proj, attn.add_v_proj))
+    img_ctx, txt_ctx = attn.processor(
+        _Precomputed(attn, qkv, added),
+        hidden_states=qkv[0],          # read for nothing but shapes
+        encoder_hidden_states=added[0],  # the processor splits at its length
+        image_rotary_emb=image_rotary_emb,
+        **(joint_attention_kwargs or {}),
+    )
+    hidden_states = hidden_states + gate_msa.unsqueeze(1) * cm.matmul_reduce_scatter(
+        img_ctx, attn.to_out[0])
+    norm_h = self.norm2(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+    hidden_states = hidden_states + gate_mlp.unsqueeze(1) * _ff(self.ff, norm_h)
+
+    encoder_hidden_states = encoder_hidden_states + c_gate_msa.unsqueeze(1) * cm.matmul_reduce_scatter(
+        txt_ctx, attn.to_add_out)
+    norm_e = self.norm2_context(encoder_hidden_states) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+    encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * _ff(self.ff_context, norm_e)
+    if encoder_hidden_states.dtype == torch.float16:
+        encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+    return encoder_hidden_states, hidden_states
+
+
+def _single_block_forward(self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=None,
+                          joint_attention_kwargs=None):
+    text_len = encoder_hidden_states.shape[1]
+    hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)  # [txt_r | img_r]
+    residual = hidden_states
+    norm, gate = self.norm(hidden_states, emb=temb)
+    attn = self.attn
+    q, k, v, mlp = cm.all_gather_matmul(norm, (attn.to_q, attn.to_k, attn.to_v, self.proj_mlp))
+    attn_out = attn.processor(
+        _Precomputed(attn, (q, k, v)), hidden_states=q, image_rotary_emb=image_rotary_emb,
+        **(joint_attention_kwargs or {}),
+    )
+    proj = cm.matmul2_reduce_scatter(attn_out, self.proj_out_attn, self.act_mlp(mlp), self.proj_out_mlp,
+                                     bias=self.proj_out_attn.bias)
+    hidden_states = residual + gate.unsqueeze(1) * proj
+    if hidden_states.dtype == torch.float16:
+        hidden_states = hidden_states.clip(-65504, 65504)
+    return hidden_states[:, :text_len], hidden_states[:, text_len:]
+
+
+def _rows(x: torch.Tensor, rank: int, tp: int) -> torch.Tensor:
+    n = x.shape[1] // tp
+    return x[:, rank * n:(rank + 1) * n]
+
+
+def _model_forward(self, hidden_states, encoder_hidden_states=None, pooled_projections=None,
+                   timestep=None, img_ids=None, txt_ids=None, guidance=None,
+                   joint_attention_kwargs=None, return_dict=False, **unsupported):
+    """``FluxTransformer2DModel.forward`` (diffusers 0.38) with sequence-sharded streams."""
+    from difflet.backends.tpu.ops_impl import parallel_mesh
+    from difflet.ops import gather_from_sequence_parallel_region
+
+    if any(v is not None for v in unsupported.values()):
+        raise NotImplementedError(f"FLUX SP forward does not take {sorted(unsupported)}")
+    tp, rank = parallel_mesh.get_tp_size(), parallel_mesh.get_tp_rank()
+    # Slice before the embedders: they are per token, so they run on the shard.
+    hidden_states = self.x_embedder(_rows(hidden_states, rank, tp))
+    timestep = timestep.to(hidden_states.dtype) * 1000
+    if guidance is not None:
+        guidance = guidance.to(hidden_states.dtype) * 1000
+    temb = (self.time_text_embed(timestep, pooled_projections) if guidance is None
+            else self.time_text_embed(timestep, guidance, pooled_projections))
+    encoder_hidden_states = self.context_embedder(_rows(encoder_hidden_states, rank, tp))
+
+    ids = torch.cat((txt_ids, img_ids), dim=0)
+    cos, sin = self.pos_embed(ids)
+    joint_rope = (cos, sin)
+    order = self._difflet_sp_single_order.to(cos.device)
+    single_rope = (cos.index_select(0, order), sin.index_select(0, order))
+
+    for block in self.transformer_blocks:
+        encoder_hidden_states, hidden_states = block(
+            hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb,
+            image_rotary_emb=joint_rope, joint_attention_kwargs=joint_attention_kwargs)
+    for block in self.single_transformer_blocks:
+        encoder_hidden_states, hidden_states = block(
+            hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb,
+            image_rotary_emb=single_rope, joint_attention_kwargs=joint_attention_kwargs)
+
+    hidden_states = self.norm_out(hidden_states, temb)
+    output = self.proj_out(hidden_states)  # (B, I/tp, out_channels)
+    output = gather_from_sequence_parallel_region(output, dim=1)
+    return (output,) if not return_dict else types.SimpleNamespace(sample=output)
+
+
+def install(module: nn.Module, *, image_seq_len: int, text_seq_len: int) -> bool:
+    """Swap the SP forwards into a TP-sharded ``FluxTransformer2DModel``. Returns
+    False (and changes nothing) when a stream does not split."""
+    from difflet.backends.tpu.ops_impl import parallel_mesh
+
+    tp = parallel_mesh.get_tp_size()
+    for name, seq in (("image", image_seq_len), ("text", text_seq_len),
+                      ("joint", int(image_seq_len) + int(text_seq_len))):
+        if not cm.supports(seq, tp) or int(seq) % tp:
+            print(f"[difflet tpu] {cm.SP_ENV}: FLUX {name} sequence {seq} does not split into "
+                  f"2 x tp={tp}; SP not installed", flush=True)
+            return False
+    if not all(hasattr(b, "proj_out_attn") for b in module.single_transformer_blocks):
+        raise ValueError("install() needs a FLUX transformer sharded by shard_flux_transformer")
+    cm.prepare(FLUX_OPS)
+    for block in module.transformer_blocks:
+        block.forward = types.MethodType(_joint_block_forward, block)
+    for block in module.single_transformer_blocks:
+        block.forward = types.MethodType(_single_block_forward, block)
+    module.register_buffer("_difflet_sp_single_order",
+                           single_order(text_seq_len, image_seq_len, tp), persistent=False)
+    module.forward = types.MethodType(_model_forward, module)
+    print(f"[difflet tpu] FLUX sequence-parallel with collective matmuls "
+          f"({len(module.transformer_blocks)} joint + {len(module.single_transformer_blocks)} single "
+          f"blocks, S_img={image_seq_len}, S_txt={text_seq_len}, tp={tp})", flush=True)
+    return True
+
+
+__all__ = ["FLUX_OPS", "install", "single_order"]
diff --git a/difflet/backends/tpu/flux/transformer.py b/difflet/backends/tpu/flux/transformer.py
index 3ce52aac..2b3f3422 100644
--- a/difflet/backends/tpu/flux/transformer.py
+++ b/difflet/backends/tpu/flux/transformer.py
@@ -101,10 +101,25 @@ def _prepare_module(self) -> torch.nn.Module:
             from difflet.backends.tpu.flux.rope import permute_flux_rope_layout_
 
             permute_flux_rope_layout_(module)
+        self._maybe_install_sp(module)
         module.requires_grad_(False)
         self.module = module.eval()
         return self.module
 
+    def _maybe_install_sp(self, module) -> bool:
+        """``DIFFLET_TPU_SP=1``: sequence parallelism with the collectives fused into
+        the matmuls (``flux/sp.py``). Compiled path only: through jax_op, eager
+        re-exports every ring on every call."""
+        from difflet.backends.tpu.ops_impl import collective_matmul
+
+        if not collective_matmul.requested() or int(self._mesh().tp) <= 1:
+            return False
+        from difflet.backends.tpu.flux import sp
+
+        collective_matmul.require_compiled(default=True)  # FLUX compiles by default
+        return sp.install(module, image_seq_len=int(self.config.image_seq_len),
+                          text_seq_len=int(self.config.text_seq_len))
+
     def get_example_inputs(self) -> tuple:
         cfg = self.config
         dtype = getattr(cfg, "torch_dtype", torch.bfloat16)
diff --git a/difflet/backends/tpu/ops_impl/collective_matmul.py b/difflet/backends/tpu/ops_impl/collective_matmul.py
index 4a9d9c3b..70d570f0 100644
--- a/difflet/backends/tpu/ops_impl/collective_matmul.py
+++ b/difflet/backends/tpu/ops_impl/collective_matmul.py
@@ -73,12 +73,19 @@ def key(d):
     return sorted(devices, key=key)
 
 
-def prepare() -> None:
+#: The ops Wan / Qwen-Image / HunyuanVideo use: (kind, arity).
+DEFAULT_OPS = (("rs", 1), ("ag", 1), ("ag", 3))
+
+
+def prepare(ops=DEFAULT_OPS) -> None:
     """Build the mesh, the shard table and the jax_op ops, outside any compiled
     region (traced lazily, Dynamo would graph-break on the registration).
-    Collective: every TP rank must call it, in step."""
+    Collective: every TP rank must call it, in step. ``ops`` lists the
+    ``(kind, arity)`` a model needs; calling again with more builds the rest."""
     global _STATE
     if _STATE is not None:
+        for kind, arity in ops:
+            _op(kind, arity)
         return
     import jax
     import numpy as np
@@ -111,14 +118,16 @@ def prepare() -> None:
         "ids": torch.tensor([pos], dtype=torch.int32, device=dev),
         "ops": {},
     }
-    # Every op the Wan forwards use, built now: built lazily they would be
+    # Every op the model's forwards use, built now: built lazily they would be
     # built inside Dynamo's trace, which cannot run JAX.
-    for kind, n_weights in (("rs", 1), ("ag", 1), ("ag", 3)):
-        _op(kind, n_weights)
+    for kind, arity in ops:
+        _op(kind, arity)
 
 
 def _op(kind: str, n_weights: int):
-    """The jax_op for ``kind`` ('rs' or 'ag') over ``n_weights`` weights, built once."""
+    """The jax_op for ``kind`` built once: 'rs' over ``n_weights`` (input, weight)
+    pairs whose products are summed, 'ag' over ``n_weights`` weights sharing one
+    gathered input."""
     key = (kind, n_weights)
     ops = _STATE["ops"]
     if key in ops:
@@ -138,21 +147,33 @@ def _op(kind: str, n_weights: int):
         __import__("jax.experimental.shard_map", fromlist=["shard_map"]).shard_map
 
     if kind == "rs":
-        # x (B, S, K/tp) @ w (N, K/tp)^T, summed over the group, this rank's rows.
-        def body(x, w, ids):
+        # sum_i x_i (B, S, K_i/tp) @ w_i (N, K_i/tp)^T, summed over the group,
+        # this rank's rows. One pair is the plain row-parallel linear; two are a
+        # row-parallel linear over a concatenation (FLUX's single-block proj_out
+        # over [attn, mlp]) without materialising the concatenation.
+        def body(*args):
+            *pairs, ids = args
+            xs, ws = pairs[0::2], pairs[1::2]
             shard_of = jnp.asarray(shard_table, dtype=jnp.int32)
             j = ids[0]
-            rows = x.shape[1] // n
-            half = w.shape[0] // 2
-            wl, wr = w[:half], w[half:]
+            rows = xs[0].shape[1] // n
+            half = ws[0].shape[0] // 2
 
-            def chunk(pos):
+            def chunk(x, pos):
                 return jax.lax.dynamic_slice_in_dim(x, shard_of[pos % n] * rows, rows, axis=1)
 
+            def partial(pos, sl):
+                out = None
+                for x, w in zip(xs, ws):
+                    term = jnp.einsum("bsk,nk->bsn", chunk(x, pos), w[sl])
+                    out = term if out is None else out + term
+                return out
+
+            lo, hi = slice(None, half), slice(half, None)
             left = right = None
             for step in range(n):
-                pl = jnp.einsum("bsk,nk->bsn", chunk(j - 1 - step), wl)
-                pr = jnp.einsum("bsk,nk->bsn", chunk(j + 1 + step), wr)
+                pl = partial(j - 1 - step, lo)
+                pr = partial(j + 1 + step, hi)
                 # Fence each partial from the add that consumes the received
                 # accumulator: fused, the matmul would wait for the hop instead
                 # of running under it (88 us idle per hop on attn_out).
@@ -164,15 +185,25 @@ def chunk(pos):
                     right = jax.lax.ppermute(right, "tp", ccw) + pr
             return jnp.concatenate([left, right], axis=2)
 
-        in_specs = (P(None, None, "tp"), P(None, "tp"), P("tp"))
+        in_specs = (P(None, None, "tp"), P(None, "tp")) * n_weights + (P("tp"),)
         out_specs = (P(None, "tp", None),)
 
-        def fn(x, w, ids):
+        def reduce(*args):
             return shard_map(body, mesh=mesh, in_specs=in_specs, out_specs=out_specs[0],
-                             check_vma=False)(x, w, ids)
+                             check_vma=False)(*args)
 
-        fn.__annotations__ = {"x": jax.Array, "w": jax.Array, "ids": jax.Array,
-                              "return": jax.Array}
+        # Arities spelled out for jax_op, as for 'ag' below.
+        if n_weights == 1:
+            def fn(x, w, ids):
+                return reduce(x, w, ids)
+            names = ["x", "w", "ids"]
+        elif n_weights == 2:
+            def fn(x0, w0, x1, w1, ids):
+                return reduce(x0, w0, x1, w1, ids)
+            names = ["x0", "w0", "x1", "w1", "ids"]
+        else:
+            raise ValueError(f"matmul_reduce_scatter supports 1 or 2 pairs, got {n_weights}")
+        fn.__annotations__ = {**{k: jax.Array for k in names}, "return": jax.Array}
     else:
         # x (B, S/tp, K) gathered, @ each w_i (N_i/tp, K)^T -> (B, S, N_i/tp).
         def body(x, ids, *ws):
@@ -217,11 +248,17 @@ def ring(*args):
         if n_weights == 1:
             def fn(x, ids, w0):
                 return ring(x, ids, w0)
+        elif n_weights == 2:
+            def fn(x, ids, w0, w1):
+                return ring(x, ids, w0, w1)
         elif n_weights == 3:
             def fn(x, ids, w0, w1, w2):
                 return ring(x, ids, w0, w1, w2)
+        elif n_weights == 4:
+            def fn(x, ids, w0, w1, w2, w3):
+                return ring(x, ids, w0, w1, w2, w3)
         else:
-            raise ValueError(f"all_gather_matmul supports 1 or 3 weights, got {n_weights}")
+            raise ValueError(f"all_gather_matmul supports 1 to 4 weights, got {n_weights}")
         names = ["x", "ids"] + [f"w{i}" for i in range(n_weights)]
         fn.__annotations__ = {**{k: jax.Array for k in names},
                               "return": (tuple[(jax.Array,) * n_weights] if n_weights > 1
@@ -247,6 +284,17 @@ def matmul_reduce_scatter(x: torch.Tensor, linear) -> torch.Tensor:
     return out
 
 
+def matmul2_reduce_scatter(x0: torch.Tensor, linear0, x1: torch.Tensor, linear1,
+                           bias: torch.Tensor | None = None) -> torch.Tensor:
+    """``linear0(x0) + linear1(x1)`` for two row-parallel linears with the same
+    output, reduced and scattered as one: (B, S, K_i/tp) -> (B, S/tp, N). The
+    linears' own biases are ignored; ``bias`` is added once after the reduction."""
+    out = _op("rs", 2)(x0, linear0.weight, x1, linear1.weight, _STATE["ids"])
+    if bias is not None:
+        out = out + bias.to(out.dtype)
+    return out
+
+
 def all_gather_matmul(x: torch.Tensor, linears) -> list[torch.Tensor]:
     """``[l(gather(x)) for l in linears]`` for column-parallel ``linears`` sharing
     one input: (B, S/tp, K) -> (B, S, N_i/tp) each, the gather overlapped with
@@ -263,5 +311,6 @@ def supports(seq_len: int, tp: int) -> bool:
     return seq_len % (2 * tp) == 0
 
 
-__all__ = ["SP_ENV", "all_gather_matmul", "matmul_reduce_scatter", "prepare",
+__all__ = ["DEFAULT_OPS", "SP_ENV", "all_gather_matmul", "matmul2_reduce_scatter",
+           "matmul_reduce_scatter", "prepare",
            "require_compiled", "requested", "supports"]
diff --git a/tests/unit/models/flux/test_flux_tpu_sp.py b/tests/unit/models/flux/test_flux_tpu_sp.py
new file mode 100644
index 00000000..1cecadd8
--- /dev/null
+++ b/tests/unit/models/flux/test_flux_tpu_sp.py
@@ -0,0 +1,99 @@
+"""CPU checks of the FLUX sequence-parallel forwards (backends/tpu/flux/sp.py).
+
+The collective matmuls are JAX rings that only run on the chips; here they are
+replaced by their tp=1 meaning (all-gather = identity, reduce-scatter = sum of the
+matmuls), so the rewritten model / block forwards must reproduce diffusers' own
+forward exactly. The single-block token order for tp>1 is checked on its own.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from difflet.backends.tpu.flux import sp
+from difflet.backends.tpu.flux.config import TpuFluxConfig
+from difflet.backends.tpu.ops_impl import collective_matmul as cm
+from difflet.backends.tpu.ops_impl import parallel_mesh as pm
+from difflet.models.flux.tp_sharding import build_flux_transformer
+from difflet.pipeline.parallel_mesh import MeshSpec
+
+from .test_flux_tp_sharding import _TINY, _inputs
+
+
+@pytest.fixture(autouse=True)
+def _mesh():
+    pm.destroy_parallel_mesh()
+    pm.init_parallel_mesh(MeshSpec(tp=1))
+    yield
+    pm.destroy_parallel_mesh()
+
+
+def test_single_order_interleaves_each_ranks_text_then_image():
+    assert sp.single_order(4, 8, 2).tolist() == [0, 1, 4, 5, 6, 7, 2, 3, 8, 9, 10, 11]
+    assert sp.single_order(4, 8, 1).tolist() == list(range(12))
+
+
+def _fake_rings(monkeypatch):
+    def ag(x, linears):
+        return [F.linear(x, l.weight, l.bias) for l in linears]
+
+    def rs(x, linear):
+        return F.linear(x, linear.weight, linear.bias)
+
+    def rs2(x0, l0, x1, l1, bias=None):
+        out = F.linear(x0, l0.weight) + F.linear(x1, l1.weight)
+        return out if bias is None else out + bias
+
+    monkeypatch.setattr(cm, "all_gather_matmul", ag)
+    monkeypatch.setattr(cm, "matmul_reduce_scatter", rs)
+    monkeypatch.setattr(cm, "matmul2_reduce_scatter", rs2)
+    monkeypatch.setattr(cm, "prepare", lambda ops=None: None)
+    monkeypatch.setattr(cm, "supports", lambda seq, tp: True)
+
+
+@pytest.mark.parametrize("split_rope", [False, True])
+def test_sp_forward_reproduces_diffusers_at_tp1(monkeypatch, split_rope):
+    from difflet.backends.tpu.flux.rope import permute_flux_rope_layout_
+    from difflet.models.flux import tp_sharding as S
+
+    _fake_rings(monkeypatch)
+    cfg = TpuFluxConfig(_TINY, height=32, width=48, tp_degree=1, torch_dtype=torch.float32)
+    ref = build_flux_transformer(cfg).eval()
+    model = build_flux_transformer(cfg).eval()
+    model.load_state_dict(ref.state_dict())
+    # The TP surgery shard_flux_transformer does at tp>1 (it is a no-op at tp=1).
+    for block in model.single_transformer_blocks:
+        dim = block.proj_out.out_features
+        w, b = block.proj_out.weight.data.clone(), block.proj_out.bias.data.clone()
+        block.proj_out_attn = S.RowParallelLinear(dim, dim, bias=True, input_is_parallel=True,
+                                                  reduce_output=False, skip_bias_add=True)
+        block.proj_out_mlp = S.RowParallelLinear(int(block.mlp_hidden_dim), dim, bias=False,
+                                                 input_is_parallel=True, reduce_output=False)
+        block.proj_out_attn.weight.data.copy_(w[:, :dim])
+        block.proj_out_attn.bias.data.copy_(b)
+        block.proj_out_mlp.weight.data.copy_(w[:, dim:])
+        del block.proj_out
+        block.attn.processor = S.FluxTPAttnProcessor()
+    for block in model.transformer_blocks:
+        block.attn.processor = S.FluxTPAttnProcessor()
+    if split_rope:
+        permute_flux_rope_layout_(model)
+    hidden, text, pooled, t, img_ids, txt_ids, guidance = _inputs(cfg)
+    assert sp.install(model, image_seq_len=hidden.shape[1], text_seq_len=text.shape[1])
+    with torch.no_grad():
+        expected = ref(hidden, text, pooled, t, img_ids, txt_ids, guidance, return_dict=False)[0]
+        actual = model(hidden_states=hidden, encoder_hidden_states=text, pooled_projections=pooled,
+                       timestep=t, img_ids=img_ids, txt_ids=txt_ids, guidance=guidance,
+                       return_dict=False)[0]
+    torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-4)
+
+
+def test_install_refuses_a_sequence_that_does_not_split(monkeypatch):
+    monkeypatch.setattr(cm, "prepare", lambda ops=None: pytest.fail("must not build rings"))
+    cfg = TpuFluxConfig(_TINY, height=32, width=48, tp_degree=1, torch_dtype=torch.float32)
+    model = build_flux_transformer(cfg)
+    before = model.forward
+    assert not sp.install(model, image_seq_len=7, text_seq_len=5)
+    assert model.forward == before

From 591c0ad35094d9530124660dc74448b714e129c1 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Mon, 9 Nov 2026 19:29:52 -0800
Subject: [PATCH 04/22] perf(tpu, ltx-2): encode the prompt over its real
 tokens, on rank 0 only

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>
---
 benchmark/tpu_profile/image_psnr.py           |  55 +++++++
 benchmark/tpu_profile/ltx2_text_probe.py      | 119 ++++++++++++++
 difflet/backends/tpu/ltx_2/text_encode.py     | 131 +++++++++++++++
 difflet/models/ltx_2/tpu_application.py       |  87 ++++++++++
 .../backends/test_tpu_ltx2_text_encode.py     | 150 ++++++++++++++++++
 5 files changed, 542 insertions(+)
 create mode 100644 benchmark/tpu_profile/image_psnr.py
 create mode 100644 benchmark/tpu_profile/ltx2_text_probe.py
 create mode 100644 difflet/backends/tpu/ltx_2/text_encode.py
 create mode 100644 tests/unit/backends/test_tpu_ltx2_text_encode.py

diff --git a/benchmark/tpu_profile/image_psnr.py b/benchmark/tpu_profile/image_psnr.py
new file mode 100644
index 00000000..ae01ae8b
--- /dev/null
+++ b/benchmark/tpu_profile/image_psnr.py
@@ -0,0 +1,55 @@
+"""PSNR / mean abs difference between saved bench images, with their sha256.
+
+    python -m benchmark.tpu_profile.image_psnr A.png B.png [C.png D.png ...]
+
+Pairs are taken in order (A vs B, C vs D, ...). Images are compared as 8-bit
+RGB; identical files report ``inf``. Prints each file's sha256 so a reported
+number can be tied to the exact images it came from.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import math
+import sys
+
+import numpy as np
+from PIL import Image
+
+
+def _load(path: str) -> np.ndarray:
+    return np.asarray(Image.open(path).convert("RGB"), dtype=np.float64)
+
+
+def _sha(path: str) -> str:
+    with open(path, "rb") as handle:
+        return hashlib.sha256(handle.read()).hexdigest()[:12]
+
+
+def compare(a_path: str, b_path: str) -> dict:
+    a, b = _load(a_path), _load(b_path)
+    if a.shape != b.shape:
+        raise ValueError(f"shape mismatch {a.shape} vs {b.shape}")
+    mse = float(((a - b) ** 2).mean())
+    return {
+        "psnr_db": math.inf if mse == 0 else 10 * math.log10(255.0**2 / mse),
+        "mean_abs": float(np.abs(a - b).mean()),
+        "max_abs": float(np.abs(a - b).max()),
+        "sha_a": _sha(a_path),
+        "sha_b": _sha(b_path),
+    }
+
+
+def main(argv: list[str]) -> int:
+    if len(argv) < 2 or len(argv) % 2:
+        print(__doc__)
+        return 2
+    for a_path, b_path in zip(argv[0::2], argv[1::2]):
+        r = compare(a_path, b_path)
+        print(f"{a_path} vs {b_path}: PSNR {r['psnr_db']:.2f} dB, mean |d| {r['mean_abs']:.3f}/255, "
+              f"max {r['max_abs']:.0f}  (sha256 {r['sha_a']} / {r['sha_b']})")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))
diff --git a/benchmark/tpu_profile/ltx2_text_probe.py b/benchmark/tpu_profile/ltx2_text_probe.py
new file mode 100644
index 00000000..f2fc4157
--- /dev/null
+++ b/benchmark/tpu_profile/ltx2_text_probe.py
@@ -0,0 +1,119 @@
+"""LTX-2 prompt encoding: stock (1024 padded tokens) vs real tokens only, on the host.
+
+    python -m benchmark.tpu_profile.ltx2_text_probe [--threads N] [--max-len 1024]
+
+No TPU needed. Loads Gemma-3 12B (fp32) and the connectors from the LTX-2
+snapshot, then per prompt runs
+
+* stock: ``pipe._get_gemma_prompt_embeds`` + ``pipe.connectors`` (what the TPU
+  application runs today), and
+* trimmed: ``backends/tpu/ltx_2/text_encode.py``,
+
+each twice (the second timed), and reports the latency of both halves and the
+max abs / relative difference of the connector outputs (video, audio) and the
+masks. Packed Gemma states are compared on the valid rows only (the trimmed
+path zero-fills padding, which nothing reads). Prints package versions and the
+torch thread count so the numbers can be reproduced.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+MODEL = os.environ.get(
+    "DIFFLET_LTX2_PATH", "/mnt/models/hf/hub/models--Lightricks--LTX-2/snapshots"
+)
+PROMPTS = [
+    "",
+    "A cat walks across a sunny kitchen floor",
+    (
+        "A slow cinematic dolly shot through a rain-soaked neon street market at night, vendors "
+        "calling out, steam rising from food stalls, reflections shimmering in puddles, a woman in "
+        "a yellow raincoat turns toward the camera and smiles while distant thunder rolls"
+    ),
+]
+
+
+def _model_path() -> str:
+    path = Path(MODEL)
+    if (path / "text_encoder").is_dir():
+        return str(path)
+    return str(next(p for p in sorted(path.iterdir()) if (p / "text_encoder").is_dir()))
+
+
+def _timed(fn):
+    fn()
+    start = time.perf_counter()
+    out = fn()
+    return out, time.perf_counter() - start
+
+
+def _diff(a: torch.Tensor, b: torch.Tensor) -> str:
+    a, b = a.double(), b.double()
+    err = (a - b).abs().max().item()
+    rel = err / max(b.abs().max().item(), 1e-30)
+    return "bit-equal" if err == 0 else f"max {err:.3e} rel {rel:.3e}"
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--threads", type=int, default=0)
+    parser.add_argument("--max-len", type=int, default=1024)
+    args = parser.parse_args()
+    if args.threads:
+        torch.set_num_threads(args.threads)
+
+    import diffusers
+    import transformers
+    from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
+    from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
+
+    from difflet.backends.tpu.ltx_2 import text_encode as te
+
+    path = _model_path()
+    print(f"torch {torch.__version__} diffusers {diffusers.__version__} transformers "
+          f"{transformers.__version__} threads {torch.get_num_threads()} snapshot {Path(path).name}")
+    tokenizer = AutoTokenizer.from_pretrained(os.path.join(path, "tokenizer"))
+    encoder = Gemma3ForConditionalGeneration.from_pretrained(
+        os.path.join(path, "text_encoder"), dtype=torch.float32).eval()
+    connectors = LTX2TextConnectors.from_pretrained(
+        os.path.join(path, "connectors"), torch_dtype=torch.float32).eval()
+
+    def stock(prompt):
+        ids, mask = te.tokenize_like_diffusers(tokenizer, prompt, args.max_len)
+        with torch.no_grad():
+            out = encoder(input_ids=ids, attention_mask=mask, output_hidden_states=True)
+        packed = torch.stack(out.hidden_states, dim=-1).flatten(2, 3)
+        return packed, mask
+
+    def trimmed(prompt):
+        ids, mask = te.tokenize_like_diffusers(tokenizer, prompt, args.max_len)
+        return te.gemma_prompt_embeds(encoder, ids, mask), mask
+
+    ok = True
+    for prompt in PROMPTS:
+        (sp, smask), t_sg = _timed(lambda: stock(prompt))
+        (tp, tmask), t_tg = _timed(lambda: trimmed(prompt))
+        with torch.no_grad():
+            sc, t_sc = _timed(lambda: connectors(sp, smask, padding_side="left"))
+            tc, t_tc = _timed(lambda: te.connectors_forward(connectors, tp, tmask))
+        n = int(smask.sum())
+        valid = smask.bool()
+        print(f"\nprompt tokens {n}/{args.max_len}: {prompt[:48]!r}")
+        print(f"  gemma       stock {t_sg:7.3f} s   trimmed {t_tg:7.3f} s   "
+              f"valid rows: {_diff(tp[valid], sp[valid])}")
+        print(f"  connectors  stock {t_sc:7.3f} s   trimmed {t_tc:7.3f} s")
+        for name, a, b in zip(("video", "audio", "mask"), tc, sc):
+            print(f"    {name:6s} {_diff(a, b)}")
+            ok &= bool(torch.isfinite(a.double()).all())
+    return 0 if ok else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/difflet/backends/tpu/ltx_2/text_encode.py b/difflet/backends/tpu/ltx_2/text_encode.py
new file mode 100644
index 00000000..528a149e
--- /dev/null
+++ b/difflet/backends/tpu/ltx_2/text_encode.py
@@ -0,0 +1,131 @@
+"""LTX-2 prompt encoding over the prompt's real tokens only.
+
+diffusers' ``LTX2Pipeline._get_gemma_prompt_embeds`` left-pads every prompt to
+``max_sequence_length`` (1024) and runs Gemma-3 12B over all of it; the
+connectors then run ``text_proj_in`` -- a (3840*49 -> 3840) linear, 0.72 B
+weights -- over the same 1024 rows. A typical prompt is a few dozen tokens, and
+nothing computed at a padding position reaches the DiT:
+
+* ``per_layer_masked_mean_norm`` takes its mean / min / max over the valid
+  positions only and zero-fills the padding rows;
+* ``text_proj_in`` has no bias, so a zero row stays zero;
+* ``LTX2ConnectorTransformer1d`` then moves the valid rows to the front and
+  overwrites every padding row with its learnable registers.
+
+So both are computed here over the last ``n`` columns (``n`` = the longest
+prompt of the batch, left padding) and the result is zero-padded back to
+``max_sequence_length``. Gemma is given the ``position_ids`` the padded call
+would have used (``arange(L)``, transformers' default, which ignores the
+mask) so every real token keeps its rotary angle; its sliding window (1024)
+covers the whole prompt either way. The connector blocks still run on the
+full 1024 rows (registers attend too).
+
+Only LTX-2.0's connectors (``per_modality_projections=False``) are handled;
+anything else falls back to the stock call.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def tokenize_like_diffusers(tokenizer, prompts, max_sequence_length: int):
+    """``_get_gemma_prompt_embeds``' tokenization: left padding to max_length."""
+    tokenizer.padding_side = "left"
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    prompts = [prompts] if isinstance(prompts, str) else list(prompts)
+    inputs = tokenizer(
+        [p.strip() for p in prompts],
+        padding="max_length",
+        max_length=int(max_sequence_length),
+        truncation=True,
+        add_special_tokens=True,
+        return_tensors="pt",
+    )
+    return inputs.input_ids, inputs.attention_mask
+
+
+def real_token_window(attention_mask: torch.Tensor) -> int:
+    """Index of the first column any prompt of the batch uses (left padding)."""
+    length = int(attention_mask.shape[1])
+    n = int(attention_mask.sum(dim=1).max())
+    start = length - n
+    if start > 0 and bool(attention_mask[:, :start].any()):
+        raise ValueError("expected left-padded prompts (a valid token precedes the window)")
+    return start
+
+
+@torch.no_grad()
+def gemma_prompt_embeds(text_encoder, input_ids, attention_mask, *, dtype=None):
+    """Packed (B, L, hidden * (layers + 1)) Gemma states, computed on the real tokens.
+
+    Padding rows are zero; the stock call leaves whatever Gemma produced there,
+    which nothing downstream reads.
+    """
+    batch, length = int(input_ids.shape[0]), int(input_ids.shape[1])
+    start = real_token_window(attention_mask)
+    position_ids = torch.arange(start, length, device=input_ids.device).unsqueeze(0).expand(batch, -1)
+    out = text_encoder(
+        input_ids=input_ids[:, start:],
+        attention_mask=attention_mask[:, start:],
+        position_ids=position_ids,
+        output_hidden_states=True,
+    )
+    packed = torch.stack(out.hidden_states, dim=-1).flatten(2, 3)
+    if dtype is not None:
+        packed = packed.to(dtype)
+    full = packed.new_zeros(batch, length, packed.shape[-1])
+    full[:, start:] = packed
+    return full
+
+
+@torch.no_grad()
+def connectors_forward(connectors, text_encoder_hidden_states, attention_mask,
+                       padding_side: str = "left", scale_factor: int = 8):
+    """``LTX2TextConnectors.forward`` (diffusers 0.38) with ``text_proj_in`` on the real rows.
+
+    Same return value: ``(video_text_embedding, audio_text_embedding, binary_mask)``.
+    """
+    from diffusers.pipelines.ltx2.connectors import per_layer_masked_mean_norm
+
+    cfg = connectors.config
+    if cfg.per_modality_projections or padding_side != "left":
+        return connectors(text_encoder_hidden_states, attention_mask,
+                          padding_side=padding_side, scale_factor=scale_factor)
+    hidden = text_encoder_hidden_states
+    if hidden.ndim == 3:
+        hidden = hidden.unflatten(2, (cfg.caption_channels, -1))
+    batch, length = int(hidden.shape[0]), int(hidden.shape[1])
+    start = real_token_window(attention_mask)
+    window_mask = attention_mask[:, start:]
+    norm = per_layer_masked_mean_norm(
+        text_hidden_states=hidden[:, start:],
+        sequence_lengths=window_mask.sum(dim=-1),
+        device=hidden.device,
+        padding_side="left",
+        scale_factor=scale_factor,
+    )
+    proj = connectors.text_proj_in(norm)
+    text_emb_proj = proj.new_zeros(batch, length, proj.shape[-1])
+    text_emb_proj[:, start:] = proj
+
+    # From here on, the stock forward unchanged.
+    text_dtype = text_emb_proj.dtype
+    mask = (attention_mask.to(torch.int64) - 1).to(text_dtype)
+    mask = mask.reshape(mask.shape[0], 1, -1, mask.shape[-1])
+    add_attn_mask = mask * torch.finfo(text_dtype).max
+    video, video_attn_mask = connectors.video_connector(text_emb_proj, add_attn_mask)
+    binary = (video_attn_mask < 1e-6).to(torch.int64)
+    binary = binary.reshape(video.shape[0], video.shape[1], 1)
+    video = video * binary
+    audio, _ = connectors.audio_connector(text_emb_proj, add_attn_mask)
+    return video, audio, binary.squeeze(-1)
+
+
+__all__ = [
+    "connectors_forward",
+    "gemma_prompt_embeds",
+    "real_token_window",
+    "tokenize_like_diffusers",
+]
diff --git a/difflet/models/ltx_2/tpu_application.py b/difflet/models/ltx_2/tpu_application.py
index c6c51925..b588cf56 100644
--- a/difflet/models/ltx_2/tpu_application.py
+++ b/difflet/models/ltx_2/tpu_application.py
@@ -11,6 +11,10 @@
 host) and its packed hidden states are broadcast to the other ranks, the
 same trick Wan uses for umT5 and HunyuanVideo for Llama. Four copies would
 not fit beside four DiT ranks on a 188 GB box.
+
+By default (``DIFFLET_TPU_LTX2_TRIM_TEXT``, on) rank 0 runs Gemma and the
+connectors over the prompt's real tokens only and broadcasts the connectors'
+output instead (``backends/tpu/ltx_2/text_encode.py``).
 """
 
 from __future__ import annotations
@@ -85,6 +89,9 @@ def load_tpu_host_pipeline(model_path: str, *, dtype: torch.dtype, encoder_dtype
     )
     pipe.to("cpu")
     pipe.set_progress_bar_config(disable=True)
+    if _trim_text_enabled():
+        _install_rank0_text_conditioning(pipe, is_encoder=is_encoder)
+        return pipe
     # The wire/return dtype follows the host pipeline (fp32), not the DiT.
     _install_broadcast_prompt_encoder(
         pipe, is_encoder=is_encoder, dtype=torch.float32, packed_width=packed_text_dim(model_path)
@@ -92,6 +99,86 @@ def load_tpu_host_pipeline(model_path: str, *, dtype: torch.dtype, encoder_dtype
     return pipe
 
 
+#: ``0`` restores the padded path: Gemma over all ``max_sequence_length`` tokens
+#: on rank 0, its packed states broadcast, the connectors on every rank.
+TRIM_TEXT_ENV = "DIFFLET_TPU_LTX2_TRIM_TEXT"
+
+
+def _trim_text_enabled() -> bool:
+    return os.environ.get(TRIM_TEXT_ENV, "1") != "0"
+
+
+def _install_rank0_text_conditioning(pipe, *, is_encoder: bool) -> None:
+    """Gemma and the connectors on rank 0 over the real tokens; the connector output broadcast.
+
+    ``backends/tpu/ltx_2/text_encode.py`` explains why only the real tokens are
+    computed. Rank 0 runs both halves; the other ranks skip Gemma (their
+    ``_get_gemma_prompt_embeds`` returns a zero-width placeholder of the right
+    batch and length) and receive the connectors' output -- two
+    (B, L, 3840) tensors and the mask, instead of the (B, L, 3840*49) packed
+    states the padded path broadcasts.
+
+    A failure on rank 0 (Gemma or the connectors) is held until the broadcast,
+    where every rank learns of it and raises: a rank that raised alone would
+    leave the others waiting in the collective.
+    """
+    from difflet.backends.tpu.ltx_2 import text_encode as te
+
+    failed: list[BaseException] = []
+    connectors = pipe.connectors
+    caption = int(connectors.config.caption_channels)
+
+    def _encode(prompt, num_videos_per_prompt=1, max_sequence_length=1024, scale_factor=8,
+                device=None, dtype=None):
+        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+        repeat = int(num_videos_per_prompt)
+        batch, length = len(prompts) * repeat, int(max_sequence_length)
+        if is_encoder and not failed:
+            try:
+                ids, mask = te.tokenize_like_diffusers(pipe.tokenizer, prompts, length)
+                embeds = te.gemma_prompt_embeds(pipe.text_encoder, ids, mask, dtype=torch.float32)
+                embeds = embeds.repeat_interleave(repeat, dim=0)
+                return embeds, mask.repeat_interleave(repeat, dim=0)
+            except Exception as exc:  # noqa: BLE001 - raised on every rank at the broadcast
+                logger.exception("ltx_2.tpu_prompt_encode_failed")
+                failed.append(exc)
+        del device, dtype, scale_factor
+        return torch.zeros(batch, length, 0), torch.zeros(batch, length, dtype=torch.int64)
+
+    def _connectors(text_encoder_hidden_states, attention_mask, padding_side="left", scale_factor=8):
+        from difflet.backends.tpu.ops_impl import device_runtime
+
+        batch, length = int(attention_mask.shape[0]), int(attention_mask.shape[1])
+        video = torch.zeros(batch, length, caption)
+        audio = torch.zeros(batch, length, caption)
+        mask = torch.zeros(batch, length, dtype=torch.int32)
+        status = _ENCODE_FAILED if failed else _ENCODE_OK
+        if is_encoder and not failed:
+            try:
+                video, audio, mask = te.connectors_forward(
+                    connectors, text_encoder_hidden_states, attention_mask,
+                    padding_side=padding_side, scale_factor=scale_factor)
+                video, audio, mask = video.float(), audio.float(), mask.to(torch.int32)
+            except Exception as exc:  # noqa: BLE001 - raised on every rank below
+                logger.exception("ltx_2.tpu_connectors_failed")
+                failed.append(exc)
+                status = _ENCODE_FAILED
+        xla = device_runtime.device()
+        payload = [video.to(xla), audio.to(xla), mask.to(xla),
+                   torch.tensor([status], dtype=torch.int32).to(xla)]
+        device_runtime.broadcast(payload)
+        ok = int(payload[3].cpu()[0]) == _ENCODE_OK
+        error = failed[0] if failed else None
+        failed.clear()
+        if not ok:
+            raise RuntimeError("LTX-2 prompt encoding failed on the encoder rank") from error
+        return payload[0].cpu(), payload[1].cpu(), payload[2].cpu().to(torch.int64)
+
+    pipe._get_gemma_prompt_embeds = _encode
+    # An instance attribute: nn.Module.__call__ looks ``forward`` up on the instance.
+    connectors.forward = _connectors
+
+
 def _install_broadcast_prompt_encoder(
     pipe, *, is_encoder: bool, dtype: torch.dtype, packed_width: int
 ) -> None:
diff --git a/tests/unit/backends/test_tpu_ltx2_text_encode.py b/tests/unit/backends/test_tpu_ltx2_text_encode.py
new file mode 100644
index 00000000..578ef84f
--- /dev/null
+++ b/tests/unit/backends/test_tpu_ltx2_text_encode.py
@@ -0,0 +1,150 @@
+"""LTX-2 prompt encoding over the real tokens (backends/tpu/ltx_2/text_encode.py)
+against the stock padded computation, on tiny random models."""
+
+from __future__ import annotations
+
+import types
+
+import pytest
+import torch
+
+from difflet.backends.tpu.ltx_2 import text_encode as te
+
+L = 16
+
+
+def _gemma():
+    from transformers import Gemma3ForCausalLM, Gemma3TextConfig
+
+    torch.manual_seed(0)
+    cfg = Gemma3TextConfig(
+        vocab_size=64, hidden_size=32, intermediate_size=64, num_hidden_layers=3,
+        num_attention_heads=2, num_key_value_heads=1, head_dim=16, sliding_window=L,
+        max_position_embeddings=64, layer_types=["sliding_attention", "full_attention", "sliding_attention"],
+    )
+    return Gemma3ForCausalLM(cfg).eval().double()
+
+
+def _connectors():
+    from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
+
+    torch.manual_seed(1)
+    return LTX2TextConnectors(
+        caption_channels=32, text_proj_in_factor=4, video_connector_num_attention_heads=2,
+        video_connector_attention_head_dim=16, video_connector_num_layers=1,
+        video_connector_num_learnable_registers=4, audio_connector_num_attention_heads=2,
+        audio_connector_attention_head_dim=16, audio_connector_num_layers=1,
+        audio_connector_num_learnable_registers=4, rope_type="split",
+    ).eval().double()
+
+
+def _left_padded(lengths):
+    ids = torch.zeros(len(lengths), L, dtype=torch.int64)
+    mask = torch.zeros(len(lengths), L, dtype=torch.int64)
+    g = torch.Generator().manual_seed(2)
+    for row, n in enumerate(lengths):
+        ids[row, L - n:] = torch.randint(3, 64, (n,), generator=g)
+        mask[row, L - n:] = 1
+    return ids, mask
+
+
+def test_gemma_on_real_tokens_matches_the_padded_call():
+    model = _gemma()
+    ids, mask = _left_padded([5, 9])
+    with torch.no_grad():
+        out = model(input_ids=ids, attention_mask=mask, output_hidden_states=True)
+    stock = torch.stack(out.hidden_states, dim=-1).flatten(2, 3)
+    trimmed = te.gemma_prompt_embeds(model, ids, mask)
+    valid = mask.bool()
+    assert trimmed.shape == stock.shape
+    torch.testing.assert_close(trimmed[valid], stock[valid], rtol=1e-10, atol=1e-10)
+    assert not trimmed[~valid].any()
+
+
+def test_connectors_on_real_rows_match_the_stock_forward():
+    conn = _connectors()
+    ids, mask = _left_padded([5, 9])
+    hidden = torch.randn(2, L, 32 * 4, dtype=torch.float64)
+    garbage = hidden.clone()
+    garbage[~mask.bool()] = 1e3 * torch.randn(int((~mask.bool()).sum()), 32 * 4, dtype=torch.float64)
+    with torch.no_grad():
+        stock = conn(garbage, mask, padding_side="left")
+    zeroed = hidden.clone()
+    zeroed[~mask.bool()] = 0
+    ours = te.connectors_forward(conn, zeroed, mask)
+    for a, b in zip(ours, stock):
+        torch.testing.assert_close(a, b, rtol=1e-10, atol=1e-10)
+
+
+def test_right_padding_is_refused():
+    mask = torch.tensor([[1, 1, 0, 0], [1, 1, 1, 0]])
+    with pytest.raises(ValueError, match="left-padded"):
+        te.real_token_window(mask)
+
+
+class _Ring:
+    """Two simulated ranks: rank 0's broadcast payload is what rank 1 receives."""
+
+    def __init__(self):
+        self.sent = None
+
+    def broadcast(self, tensors, src=0):
+        if self.sent is None:
+            self.sent = [t.clone() for t in tensors]
+        else:
+            for t, s in zip(tensors, self.sent):
+                t.copy_(s)
+
+
+def _pipe(is_encoder, monkeypatch, ring, gemma=None):
+    from difflet.backends.tpu.ops_impl import device_runtime
+    from difflet.models.ltx_2 import tpu_application as mod
+
+    monkeypatch.setattr(device_runtime, "device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(device_runtime, "broadcast", ring.broadcast)
+
+    class Tok:
+        padding_side, pad_token, eos_token = "left", "<pad>", "<eos>"
+
+        def __call__(self, prompts, max_length, **_):
+            lengths = [len(p.split()) + 1 for p in prompts]
+            ids, mask = _left_padded(lengths)
+            return types.SimpleNamespace(input_ids=ids[:, :max_length], attention_mask=mask[:, :max_length])
+
+    conn = _connectors().float()
+    pipe = types.SimpleNamespace(tokenizer=Tok(), text_encoder=gemma, connectors=conn,
+                                 _get_gemma_prompt_embeds=None)
+    mod._install_rank0_text_conditioning(pipe, is_encoder=is_encoder)
+    return pipe
+
+
+def _condition(pipe, prompts):
+    embeds, mask = pipe._get_gemma_prompt_embeds(prompts, max_sequence_length=L, device="cpu",
+                                                 dtype=torch.bfloat16)
+    return pipe.connectors(embeds, mask, padding_side="left")
+
+
+def test_rank0_conditioning_reaches_every_rank(monkeypatch):
+    ring = _Ring()
+    gemma = _gemma().float()
+    rank0 = _condition(_pipe(True, monkeypatch, ring, gemma), ["a red fox", "snow"])
+    rank1_pipe = _pipe(False, monkeypatch, ring)
+    embeds, _ = rank1_pipe._get_gemma_prompt_embeds(["a red fox", "snow"], max_sequence_length=L)
+    assert embeds.shape == (2, L, 0)  # no Gemma off rank 0
+    rank1 = rank1_pipe.connectors(embeds, torch.zeros(2, L, dtype=torch.int64))
+    for a, b in zip(rank0, rank1):
+        assert torch.equal(a, b)
+    assert rank0[0].shape == (2, L, 32) and rank0[2].dtype is torch.int64
+
+
+def test_rank0_failure_raises_on_every_rank(monkeypatch):
+    ring = _Ring()
+
+    def broken(**_):
+        raise OSError("disk")
+
+    with pytest.raises(RuntimeError, match="encoder rank") as info:
+        _condition(_pipe(True, monkeypatch, ring, gemma=broken), ["a fox"])
+    assert isinstance(info.value.__cause__, OSError)
+    with pytest.raises(RuntimeError, match="encoder rank"):
+        _condition(_pipe(False, monkeypatch, ring), ["a fox"])

From c31a991651a0e8140a8a20757b3d3b647121f023 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Mon, 9 Nov 2026 19:44:44 -0800
Subject: [PATCH 05/22] fix(tpu, ltx-2): contiguous DiT inputs, so
 torch.compile stops recompiling 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>
---
 difflet/models/ltx_2/tpu_application.py       |  9 +++++-
 .../models/ltx_2/test_ltx2_tpu_application.py | 30 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/difflet/models/ltx_2/tpu_application.py b/difflet/models/ltx_2/tpu_application.py
index b588cf56..0d8afe7b 100644
--- a/difflet/models/ltx_2/tpu_application.py
+++ b/difflet/models/ltx_2/tpu_application.py
@@ -468,8 +468,15 @@ def forward_dit(self, bundle: LTX2DiTInputBundle):
         transformer = self._require_transformer("forward")
         validate_ltx_2_dit_inputs(bundle, config=transformer.config, dtype=self.dtype)
         device = device_runtime.device()
+        # Contiguous before the copy: the orchestrator's packed latents arrive as a
+        # transposed view (hidden_states stride (2310, 1, 2310) at 480x704x49), the
+        # warmup's example inputs are contiguous, and under DIFFLET_TPU_COMPILE=1
+        # Dynamo's stride guard then recompiled the whole DiT on the first real
+        # request of every process (~60 s, 2026-11-10).
         with torch.no_grad():
-            video, audio = transformer.forward(*[t.to(device) for t in bundle.as_model_inputs()])
+            video, audio = transformer.forward(
+                *[t.contiguous().to(device) for t in bundle.as_model_inputs()]
+            )
         return video.cpu(), audio.cpu()
 
     def __call__(self, *args: Any, **kwargs: Any):
diff --git a/tests/unit/models/ltx_2/test_ltx2_tpu_application.py b/tests/unit/models/ltx_2/test_ltx2_tpu_application.py
index 83bb52a3..cc43d3cb 100644
--- a/tests/unit/models/ltx_2/test_ltx2_tpu_application.py
+++ b/tests/unit/models/ltx_2/test_ltx2_tpu_application.py
@@ -109,3 +109,33 @@ def __init__(self):
     z = torch.ones(1, 2, 1, 1, 1)
     out = _denormalize_ltx_2_video_latents(z, wrapper)
     assert out[0, :, 0, 0, 0].tolist() == [4.0, 6.0]  # z * std + mean, not z
+
+
+def test_forward_dit_hands_the_dit_contiguous_inputs(snapshot, monkeypatch):
+    """The packed latents arrive as a transposed view; under torch.compile a stride
+    guard on them recompiled the whole DiT on every process's first request."""
+    from difflet.backends.tpu.ops_impl import device_runtime
+    from difflet.models.ltx_2 import tpu_application as mod
+    from difflet.models.ltx_2.contract import LTX2DiTInputBundle
+
+    monkeypatch.setattr(device_runtime, "device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(mod, "validate_ltx_2_dit_inputs", lambda *a, **k: None)
+    app = create_ltx_2_application(
+        model_path=str(snapshot), parallel=DiffletParallelConfig(tp_degree=4), dtype="bf16",
+        shape={"height": 512, "width": 768, "num_frames": 121}, backend="tpu",
+    )
+    seen = []
+
+    def forward(*tensors):
+        seen.extend(t.is_contiguous() for t in tensors)
+        return tensors[0], tensors[1]
+
+    monkeypatch.setattr(app.transformer, "forward", forward)
+    inputs = [torch.zeros(1, 4, 8) for _ in range(10)]
+    inputs[0] = torch.zeros(1, 8, 4).transpose(1, 2)  # a view, like the packed latents
+    assert not inputs[0].is_contiguous()
+    keys = ("hidden_states", "audio_hidden_states", "encoder_hidden_states",
+            "audio_encoder_hidden_states", "timestep", "sigma", "encoder_attention_mask",
+            "audio_encoder_attention_mask", "video_coords", "audio_coords")
+    app.forward_dit(LTX2DiTInputBundle(**dict(zip(keys, inputs))))
+    assert seen and all(seen)

From 5c2cbbfb021dcf2d44d0af5b1adf0ebd0359cb95 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Mon, 9 Nov 2026 19:58:50 -0800
Subject: [PATCH 06/22] bench(tpu, ltx-2): CFG check of the text trim (batch-2
 conditioning on 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 591c0ad3: 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>
---
 benchmark/tpu_profile/ltx2_cfg_trim_check.py | 125 +++++++++++++++++++
 1 file changed, 125 insertions(+)
 create mode 100644 benchmark/tpu_profile/ltx2_cfg_trim_check.py

diff --git a/benchmark/tpu_profile/ltx2_cfg_trim_check.py b/benchmark/tpu_profile/ltx2_cfg_trim_check.py
new file mode 100644
index 00000000..0754a0d7
--- /dev/null
+++ b/benchmark/tpu_profile/ltx2_cfg_trim_check.py
@@ -0,0 +1,125 @@
+"""LTX-2 with CFG (guidance > 1): text trim on vs off, one process per chip.
+
+    python -m benchmark.tpu_profile.ltx2_cfg_trim_check [--steps 4 --guidance 4.0]
+
+main's bench runs LTX-2 at guidance 1.0, so the batch-2 conditioning path
+(negative + positive prompt through Gemma and the connectors, then a CFG
+denoise) is not covered there. Each rank loads the application once with
+``DIFFLET_TPU_LTX2_TRIM_TEXT=1``, runs a short CFG request to latents, then
+reinstalls the padded path on the same host pipeline and repeats with the same
+seed. Rank 0 prints the conditioning time of each and the relative L2 / cosine
+of the video and audio latents (trim vs padded).
+"""
+
+from __future__ import annotations
+
+import argparse
+import multiprocessing as mp
+import os
+import sys
+import time
+from pathlib import Path
+
+MODEL = os.environ.get("DIFFLET_LTX2_PATH", "/mnt/models/hf/hub/models--Lightricks--LTX-2/snapshots")
+PROMPT = "a cinematic shot of a red fox running through a snowy forest"
+NEGATIVE = "blurry, low quality, distorted"
+
+
+def _model_path() -> str:
+    path = Path(MODEL)
+    if (path / "transformer").is_dir():
+        return str(path)
+    return str(next(p for p in sorted(path.iterdir()) if (p / "transformer").is_dir()))
+
+
+def _worker(rank, world, args, reply_q):
+    try:
+        os.environ["DIFFLET_TPU_LTX2_TRIM_TEXT"] = "1"
+        from difflet.backends.tpu.ops_impl import device_runtime
+
+        device_runtime.initialize_replica(rank, world)
+        import torch
+
+        # As main's bench worker does (benchmark/adapters/tpu.py): share the host.
+        torch.set_num_threads(max(1, (os.cpu_count() or world) // world))
+
+        from difflet.models.ltx_2 import tpu_application as mod
+        from difflet.models.ltx_2.entry import create_ltx_2_application
+        from difflet.pipeline.parallel_config import DiffletParallelConfig
+
+        app = create_ltx_2_application(
+            model_path=_model_path(), parallel=DiffletParallelConfig(tp_degree=world),
+            dtype=torch.bfloat16, shape={"height": 480, "width": 704, "num_frames": 49},
+            backend="tpu", enable_host_pipeline=True,
+        )
+        app.load_eager()
+        orch, pipe = app.pipeline, app.host_pipeline
+
+        def request():
+            mark = time.monotonic()
+            cond = orch.prepare_conditioning(prompt=PROMPT, negative_prompt=NEGATIVE,
+                                             guidance_scale=args.guidance,
+                                             max_sequence_length=orch.text_seq_len)
+            encode = time.monotonic() - mark
+            states, audio_states, mask, audio_mask = cond
+            out = orch(encoder_hidden_states=states, audio_encoder_hidden_states=audio_states,
+                       encoder_attention_mask=mask, audio_encoder_attention_mask=audio_mask,
+                       num_inference_steps=args.steps, guidance_scale=args.guidance,
+                       generator=torch.Generator().manual_seed(42), output_type="latent")
+            return encode, out.latents.float(), out.audio_latents.float(), tuple(states.shape)
+
+        trim_first = request()
+        trim = request()
+        # Back to the padded path on the same pipeline.
+        pipe._get_gemma_prompt_embeds = type(pipe)._get_gemma_prompt_embeds.__get__(pipe)
+        del pipe.connectors.forward
+        mod._install_broadcast_prompt_encoder(
+            pipe, is_encoder=rank == 0, dtype=torch.float32,
+            packed_width=mod.packed_text_dim(_model_path()))
+        padded_first = request()
+        padded = request()
+
+        def err(a, b):
+            return {"rel_l2": float((a - b).norm() / b.norm()),
+                    "cos": float(torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0))}
+
+        reply_q.put({"rank": rank, "cond_shape": trim[3], "threads": torch.get_num_threads(),
+                     "encode_s_trim": [round(trim_first[0], 2), round(trim[0], 2)],
+                     "encode_s_padded": [round(padded_first[0], 2), round(padded[0], 2)], "video": err(trim[1], padded[1]),
+                     "audio": err(trim[2], padded[2])})
+    except BaseException as exc:  # noqa: BLE001
+        import traceback
+
+        reply_q.put({"rank": rank, "error": f"{type(exc).__name__}: {exc}",
+                     "tb": traceback.format_exc()[-2500:]})
+        raise
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--steps", type=int, default=4)
+    parser.add_argument("--guidance", type=float, default=4.0)
+    parser.add_argument("--world", type=int, default=4)
+    args = parser.parse_args()
+    from difflet.backends.tpu.ops_impl import device_runtime
+
+    device_runtime.prepare_launch(args.world)
+    ctx = mp.get_context("spawn")
+    reply_q = ctx.Queue()
+    procs = [ctx.Process(target=_worker, args=(r, args.world, args, reply_q)) for r in range(args.world)]
+    for p in procs:
+        p.start()
+    results = sorted((reply_q.get(timeout=1800) for _ in procs), key=lambda r: r["rank"])
+    for p in procs:
+        p.join(timeout=60)
+    ok = all("error" not in r for r in results)
+    for r in results:
+        if "error" in r:
+            print(f"rank {r['rank']} FAILED: {r['error']}\n{r['tb']}", flush=True)
+    if ok:
+        print(results[0], flush=True)
+    return 0 if ok else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

From 1a205f0ee4f829d27d6449147d6d3cd6f960aff9 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Mon, 9 Nov 2026 20:27:40 -0800
Subject: [PATCH 07/22] perf(tpu, ltx-2): torch.compile the DiT by default

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 c31a9916 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>
---
 difflet/backends/tpu/ops_impl/device_runtime.py | 2 +-
 difflet/models/ltx_2/tpu_application.py         | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/difflet/backends/tpu/ops_impl/device_runtime.py b/difflet/backends/tpu/ops_impl/device_runtime.py
index c8222c89..ce529cb2 100644
--- a/difflet/backends/tpu/ops_impl/device_runtime.py
+++ b/difflet/backends/tpu/ops_impl/device_runtime.py
@@ -450,7 +450,7 @@ def maybe_compile(module, *, default: bool = False):
 
     ``default`` is the call site's choice for its model; ``DIFFLET_TPU_COMPILE``
     overrides it either way (``1`` forces on, ``0`` forces off). On by default
-    only where it has been measured — Qwen-Image, Wan and FLUX.1-dev — and off
+    only where it has been measured — Qwen-Image, Wan, FLUX.1-dev and LTX-2 — and off
     for anything that has not run under it.
 
     Why compile: DeferAndFuse fuses heuristically and left torch_tpu 21% behind
diff --git a/difflet/models/ltx_2/tpu_application.py b/difflet/models/ltx_2/tpu_application.py
index 0d8afe7b..66723084 100644
--- a/difflet/models/ltx_2/tpu_application.py
+++ b/difflet/models/ltx_2/tpu_application.py
@@ -394,8 +394,11 @@ def load_eager(self) -> None:
         module = transformer._prepare_module().to(device)
         device_runtime.wait()
         _hbm_probe("DiT weights on the chip", next(module.parameters()))
-        # Not yet measured under torch.compile, so off unless DIFFLET_TPU_COMPILE=1.
-        module = device_runtime.maybe_compile(module, default=False)
+        # Compiled by default (DIFFLET_TPU_COMPILE=0 turns it off): measured on a
+        # v5e-4 at 480x704x49, 20 steps, a step is 220 ms against 740 ms eager
+        # (request 8.4 s vs 18.9 s); the first call compiles for ~155 s with a
+        # cold tier-3 cache. Output vs eager: 35.4 dB, differences in fast motion.
+        module = device_runtime.maybe_compile(module, default=True)
         transformer.module = module
         seconds = transformer.warmup_eager(module, device)
         print(f"[ltx_2] tpu warmup (first compile) in {seconds:.1f}s", flush=True)

From fd3750cf301e65dae696f8601eb2d015c6e63846 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Mon, 9 Nov 2026 20:52:20 -0800
Subject: [PATCH 08/22] fix(tpu, flux): pipelined T5 segments under
 transformers 4.57 (CI pin)

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 5423c7ef 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>
---
 difflet/models/flux/tpu_application.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/difflet/models/flux/tpu_application.py b/difflet/models/flux/tpu_application.py
index f54b9e7e..8abdc798 100644
--- a/difflet/models/flux/tpu_application.py
+++ b/difflet/models/flux/tpu_application.py
@@ -306,8 +306,16 @@ def _segment(self, hidden: torch.Tensor | None, ids: torch.Tensor) -> torch.Tens
         position_bias = None
         if self.segments[self.rank].start != 0:
             position_bias = self.bias_attention.compute_bias(self.seq_len, self.seq_len, device=self._device)
+        # Keywords, an explicit cache_position and only outputs [0] and [1]: the
+        # call has to hold across transformers releases. 4.57 (the CI pin)
+        # indexes cache_position when it is None and no query_length is given;
+        # 5.x has dropped the argument (it lands in **kwargs). Both return
+        # (hidden, position_bias, ...).
+        cache_position = torch.arange(int(hidden.shape[1]), device=hidden.device)
         for block in self.blocks:
-            hidden, position_bias, _ = block(hidden, None, position_bias)
+            out = block(hidden, attention_mask=None, position_bias=position_bias,
+                        cache_position=cache_position)
+            hidden, position_bias = out[0], out[1]
         if self.final_norm is not None:
             hidden = self.final_norm(hidden)
         return hidden

From 440bbc3352906297a330d7503ad90b2c73bb326c Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Mon, 9 Nov 2026 21:08:43 -0800
Subject: [PATCH 09/22] test(tpu, flux): SP tests initialize the mesh the code
 will import

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>
---
 tests/unit/models/flux/test_flux_tpu_sp.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/tests/unit/models/flux/test_flux_tpu_sp.py b/tests/unit/models/flux/test_flux_tpu_sp.py
index 1cecadd8..c788e542 100644
--- a/tests/unit/models/flux/test_flux_tpu_sp.py
+++ b/tests/unit/models/flux/test_flux_tpu_sp.py
@@ -8,6 +8,8 @@
 
 from __future__ import annotations
 
+import importlib
+
 import pytest
 import torch
 import torch.nn.functional as F
@@ -24,10 +26,18 @@
 
 @pytest.fixture(autouse=True)
 def _mesh():
-    pm.destroy_parallel_mesh()
-    pm.init_parallel_mesh(MeshSpec(tp=1))
+    # Resolved now, not at collection: sp.install / the SP forward import
+    # parallel_mesh when they run, and another test in the same worker may have
+    # dropped difflet.backends.tpu.* from sys.modules in between
+    # (test_torch_megatron.py::test_weights_and_checkpoint_import_without_tpu),
+    # leaving the module bound above a stale, never-initialized copy.
+    mesh = importlib.import_module("difflet.backends.tpu.ops_impl.parallel_mesh")
+    for module in {mesh, pm}:
+        module.destroy_parallel_mesh()
+        module.init_parallel_mesh(MeshSpec(tp=1))
     yield
-    pm.destroy_parallel_mesh()
+    for module in {mesh, pm}:
+        module.destroy_parallel_mesh()
 
 
 def test_single_order_interleaves_each_ranks_text_then_image():

From 71389c430982f8cf67542f2f02a40d29286eac7a Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Tue, 10 Nov 2026 14:13:02 -0800
Subject: [PATCH 10/22] perf(tpu, ltx-2): sequence parallelism over collective
 matmuls (DIFFLET_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>
---
 benchmark/tpu_profile/cm_ops_parity.py        | 155 ++++++++
 benchmark/tpu_profile/cm_padded_repro.py      | 269 ++++++++++++++
 benchmark/tpu_profile/cm_pure_jax.py          |  91 +++++
 benchmark/tpu_profile/trace_tail.py           |  44 +++
 difflet/backends/tpu/ltx_2/sp.py              | 336 ++++++++++++++++++
 difflet/backends/tpu/ltx_2/transformer.py     |  38 ++
 .../tpu/ops_impl/collective_matmul.py         |  11 +
 tests/unit/backends/test_tpu_ltx2_sp.py       | 162 +++++++++
 8 files changed, 1106 insertions(+)
 create mode 100644 benchmark/tpu_profile/cm_ops_parity.py
 create mode 100644 benchmark/tpu_profile/cm_padded_repro.py
 create mode 100644 benchmark/tpu_profile/cm_pure_jax.py
 create mode 100644 benchmark/tpu_profile/trace_tail.py
 create mode 100644 difflet/backends/tpu/ltx_2/sp.py
 create mode 100644 tests/unit/backends/test_tpu_ltx2_sp.py

diff --git a/benchmark/tpu_profile/cm_ops_parity.py b/benchmark/tpu_profile/cm_ops_parity.py
new file mode 100644
index 00000000..462eaa70
--- /dev/null
+++ b/benchmark/tpu_profile/cm_ops_parity.py
@@ -0,0 +1,155 @@
+"""Collective-matmul rings vs plain collectives, compiled, one process per chip.
+
+    python -m benchmark.tpu_profile.cm_ops_parity [--rows 578 --hidden 4096]
+
+For each (kind, arity) in the list below, runs the ring op and its meaning with
+plain collectives (gather + linear, linear + reduce-scatter) inside
+``torch.compile``, on the same random bf16 input, and prints the relative L2 of
+ring against plain on rank 0. Defaults are LTX-2's video shapes at 480x704x49
+(2312 padded rows / 4 = 578 per rank, hidden 4096, 1024 heads-dims per rank).
+"""
+
+from __future__ import annotations
+
+import argparse
+import multiprocessing as mp
+import sys
+
+OPS = (("ag", 1), ("ag", 2), ("ag", 3), ("rs", 1))
+
+
+def _worker(rank, world, args, reply_q):
+    try:
+        from difflet.backends.tpu.ops_impl import device_runtime
+
+        device_runtime.initialize_replica(rank, world)
+        import torch
+        import torch.nn.functional as F
+
+        from difflet.backends.tpu.ops_impl import collective_matmul as cm
+        from difflet.backends.tpu.ops_impl import parallel_mesh
+        from difflet.ops import (
+            gather_from_sequence_parallel_region,
+            reduce_scatter_to_sequence_parallel_region,
+        )
+        from difflet.pipeline.parallel_mesh import MeshSpec
+
+        parallel_mesh.init_parallel_mesh(MeshSpec(tp=world))
+        dev = device_runtime.device()
+        cm.prepare(OPS)
+        g = torch.Generator().manual_seed(100 + rank)  # every rank its own shard / weights
+        bf = torch.bfloat16
+        hidden, local = args.hidden, args.hidden // world
+
+        def lin(n_out, n_in, bias=True):
+            m = torch.nn.Linear(n_in, n_out, bias=bias)
+            with torch.no_grad():
+                m.weight.copy_(torch.randn(n_out, n_in, generator=g) / n_in ** 0.5)
+                if bias:
+                    m.bias.copy_(torch.randn(n_out, generator=g) * 0.1)
+            return m.to(dev, bf)
+
+        x_shard = torch.randn(1, args.rows, hidden, generator=g).to(dev, bf)
+        x_full_local = torch.randn(1, args.rows * world, local, generator=g).to(dev, bf)
+        cols = [lin(local, hidden) for _ in range(3)]
+        row = lin(hidden, local)
+        # A row-parallel bias is replicated: the same on every rank.
+        bias = torch.randn(hidden, generator=torch.Generator().manual_seed(5)).to(dev, bf)
+        with torch.no_grad():
+            row.bias.copy_(bias)
+
+        def ring_ag(n):
+            return lambda x: cm.all_gather_matmul(x, cols[:n])
+
+        def plain_ag(n):
+            def f(x):
+                full = gather_from_sequence_parallel_region(x, dim=1)
+                return [F.linear(full, c.weight, c.bias) for c in cols[:n]]
+            return f
+
+        def ring_rs(x):
+            return [cm.matmul_reduce_scatter(x, row)]
+
+        def plain_rs(x):
+            out = reduce_scatter_to_sequence_parallel_region(F.linear(x, row.weight), dim=1)
+            return [out + row.bias]
+
+        res = {}
+        for kind, n in OPS:
+            if kind == "ag":
+                ring, plain, x = ring_ag(n), plain_ag(n), x_shard
+            else:
+                ring, plain, x = ring_rs, plain_rs, x_full_local
+            with torch.no_grad():
+                a = torch.compile(ring, backend="tpu")(x)
+                b = torch.compile(plain, backend="tpu")(x)
+            errs = []
+            for u, v in zip(a, b):
+                u, v = u.float().cpu(), v.float().cpu()
+                errs.append(round(float((u - v).norm() / v.norm()), 5))
+            res[f"{kind}{n}"] = errs
+        # Twice in one compiled graph, different weights (as consecutive blocks do).
+        cols2 = [lin(local, hidden) for _ in range(3)]
+        row2 = lin(hidden, local)
+        with torch.no_grad():
+            row2.bias.copy_(bias)
+        for kind, n in OPS:
+            if kind == "ag":
+                def ring2(x, n=n):
+                    a = cm.all_gather_matmul(x, cols[:n])
+                    b = cm.all_gather_matmul(x * 0.5, cols2[:n])
+                    return list(a) + list(b)
+
+                def plain2(x, n=n):
+                    full = gather_from_sequence_parallel_region(x, dim=1)
+                    full2 = gather_from_sequence_parallel_region(x * 0.5, dim=1)
+                    return ([F.linear(full, c.weight, c.bias) for c in cols[:n]]
+                            + [F.linear(full2, c.weight, c.bias) for c in cols2[:n]])
+                x = x_shard
+            else:
+                def ring2(x):
+                    return [cm.matmul_reduce_scatter(x, row), cm.matmul_reduce_scatter(x * 0.5, row2)]
+
+                def plain2(x):
+                    a = reduce_scatter_to_sequence_parallel_region(F.linear(x, row.weight), dim=1)
+                    b = reduce_scatter_to_sequence_parallel_region(F.linear(x * 0.5, row2.weight), dim=1)
+                    return [a + row.bias, b + row2.bias]
+                x = x_full_local
+            with torch.no_grad():
+                a = torch.compile(ring2, backend="tpu")(x)
+                b = torch.compile(plain2, backend="tpu")(x)
+            res[f"{kind}{n}_twice"] = [round(float((u.float().cpu() - v.float().cpu()).norm()
+                                                  / v.float().cpu().norm()), 5) for u, v in zip(a, b)]
+        reply_q.put({"rank": rank, **res})
+    except BaseException as exc:  # noqa: BLE001
+        import traceback
+
+        reply_q.put({"rank": rank, "error": f"{type(exc).__name__}: {exc}",
+                     "tb": traceback.format_exc()[-2500:]})
+        raise
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--rows", type=int, default=578)
+    parser.add_argument("--hidden", type=int, default=4096)
+    args = parser.parse_args()
+    from difflet.backends.tpu.ops_impl import device_runtime
+
+    world = 4
+    device_runtime.prepare_launch(world)
+    ctx = mp.get_context("spawn")
+    reply_q = ctx.Queue()
+    procs = [ctx.Process(target=_worker, args=(r, world, args, reply_q)) for r in range(world)]
+    for p in procs:
+        p.start()
+    results = sorted((reply_q.get(timeout=1200) for _ in procs), key=lambda r: r["rank"])
+    for p in procs:
+        p.join(timeout=60)
+    for r in results:
+        print(r, flush=True)
+    return 0 if all("error" not in r for r in results) else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/benchmark/tpu_profile/cm_padded_repro.py b/benchmark/tpu_profile/cm_padded_repro.py
new file mode 100644
index 00000000..3290a484
--- /dev/null
+++ b/benchmark/tpu_profile/cm_padded_repro.py
@@ -0,0 +1,269 @@
+"""Minimal repro: collective-matmul rings across two transformer-like blocks, padded sequence.
+
+    python -m benchmark.tpu_profile.cm_padded_repro [--seq 2310] [--variant full]
+
+LTX-2's SP forward is right at 6144 video rows (no padding) and wrong -- or the
+core halts -- at 2310 rows padded to 2312 once two blocks share one compiled
+graph (2026-11-10). This strips the model away: random weights, a block of
+
+    n = rms(x); q,k,v = AG(n); attn(q, k[:S], v[:S]); x += RS(o)
+    n = rms(x); h = AG(n); x += RS(gelu(h))
+
+run ``--blocks`` times inside one ``torch.compile`` graph on the chips, against
+an fp32 CPU reference that rank 0 rebuilds from the same seeds (no TPU
+collective in the reference). ``--variant`` drops pieces to bisect:
+``full`` | ``noattn`` (o = q) | ``attnonly`` (no FFN) | ``ffonly`` | ``rsonly``
+(x += RS(AG(x)), no norm or activation).
+"""
+
+from __future__ import annotations
+
+import argparse
+import multiprocessing as mp
+import sys
+
+HID, HEADS, HD, FF = 4096, 32, 128, 8192
+
+
+def _weights(tp, rank, block, device=None, dtype=None):
+    import torch
+
+    g = torch.Generator().manual_seed(1000 * block + rank)
+    loc = HID // tp
+    ffl = FF // tp
+
+    def w(o, i):
+        t = torch.randn(o, i, generator=g) / i ** 0.5
+        return t if device is None else t.to(device, dtype)
+
+    return {"q": w(loc, HID), "k": w(loc, HID), "v": w(loc, HID), "o": w(HID, loc),
+            "f1": w(ffl, HID), "f2": w(HID, ffl)}
+
+
+def _worker(rank, world, args, reply_q):
+    try:
+        import os
+
+        base = os.environ.get("REPRO_DUMP_BASE")
+        if base:  # per-rank compiler dumps, set before the runtime starts
+            os.environ["XLA_FLAGS"] = f"--xla_dump_to={base}/rank{rank}/xla --xla_dump_hlo_as_text"
+            os.environ["TORCH_TPU_INTERNAL_DUMP_MLIR_PASSES_TO"] = f"{base}/rank{rank}/mlir"
+            os.environ["TORCH_TPU_TIER3_COMPILATION_CACHE_ROOT"] = f"{base}/rank{rank}/t3"
+        from difflet.backends.tpu.ops_impl import device_runtime
+
+        device_runtime.initialize_replica(rank, world)
+        import torch
+        import torch.nn.functional as F
+
+        from difflet.backends.tpu.ltx_2 import sp
+        from difflet.backends.tpu.ops_impl import collective_matmul as cm
+        from difflet.backends.tpu.ops_impl import parallel_mesh
+        from difflet.ops import attention, gather_from_sequence_parallel_region
+        from difflet.pipeline.parallel_mesh import MeshSpec
+
+        parallel_mesh.init_parallel_mesh(MeshSpec(tp=world))
+        cm.prepare(sp.LTX2_OPS)
+        tp_rank = parallel_mesh.get_tp_rank()
+        dev, bf = device_runtime.device(), torch.bfloat16
+        seq = args.seq
+        padded = sp.padded_length(seq, world) if args.pad else seq
+        n = padded // world
+        heads = HEADS // world
+
+        class Lin:  # what the rings read: .weight / .bias
+            def __init__(self, w):
+                self.weight, self.bias = w, None
+
+        blocks = [{k: Lin(v) for k, v in _weights(world, tp_rank, b, dev, bf).items()}
+                  for b in range(args.blocks)]
+        g = torch.Generator().manual_seed(7)
+        x_full = torch.randn(1, seq, HID, generator=g)
+
+        def rms(t):
+            return t * torch.rsqrt(t.float().pow(2).mean(-1, keepdim=True) + 1e-6).to(t.dtype)
+
+        def block_fn(x, w):
+            if args.variant in ("full", "noattn", "attnonly"):
+                q, k, v = cm.all_gather_matmul(rms(x), (w["q"], w["k"], w["v"]))
+                if args.variant in ("full", "attnonly"):
+                    b, s, inner = q.shape
+                    k, v = k[:, :seq], v[:, :seq]
+
+                    def bh(t):
+                        return t.reshape(b, t.shape[1], heads, HD).permute(0, 2, 1, 3).reshape(
+                            b * heads, t.shape[1], HD)
+
+                    o = attention(bh(q), bh(k), bh(v), scale=HD ** -0.5, causal=False)
+                    o = o.reshape(b, heads, s, HD).permute(0, 2, 1, 3).reshape(b, s, inner)
+                else:
+                    o = q
+                x = x + cm.matmul_reduce_scatter(o, w["o"])
+            if args.variant in ("full", "noattn", "ffonly"):
+                (h,) = cm.all_gather_matmul(rms(x), (w["f1"],))
+                x = x + cm.matmul_reduce_scatter(F.gelu(h, approximate="tanh"), w["f2"])
+            if args.variant == "rsonly":
+                (h,) = cm.all_gather_matmul(x, (w["f1"],))
+                x = x + cm.matmul_reduce_scatter(h, w["f2"])
+            if args.variant == "agonly":
+                # AG ring only: this rank's rows of the gathered product, projected
+                # back with a local (B, n, FF/tp) x (FF/tp, HID) -- no collective.
+                (h,) = cm.all_gather_matmul(x, (w["f1"],))
+                own = h[:, tp_rank * n:(tp_rank + 1) * n]
+                x = x + own @ w["f2"].weight.T
+            if args.variant == "rsonly2":
+                # RS ring only: a full-sequence input built locally (the shard
+                # tiled tp times), reduced and scattered back to this rank's rows.
+                full = x.repeat(1, world, 1)
+                part = full @ w["f1"].weight.T  # (B, S, FF/tp), local
+                x = x + cm.matmul_reduce_scatter(part, w["f2"])
+            return x
+
+        def model(xf, off):
+            if args.pad_mode == "dynamic":
+                # Padded in the graph as in "full", but the shard offset is a
+                # runtime input, so every rank compiles the same program.
+                xp = sp._pad_seq(xf, 1, padded, 0.0)
+                x = torch.index_select(xp, 1, off + torch.arange(n, device=xf.device))
+            elif args.pad_mode == "outside":
+                # xf is already padded (eagerly, outside the graph): no cat in here.
+                x = xf[:, tp_rank * n:(tp_rank + 1) * n]
+            elif args.pad_mode == "shard":
+                # Slice the shard off the graph input first, then pad only that
+                # shard (only the last rank gets any rows).
+                x = sp._pad_seq(xf[:, tp_rank * n:min((tp_rank + 1) * n, seq)], 1, n, 0.0)
+            else:  # "full": pad the whole sequence, then slice (the original)
+                x = sp._pad_seq(xf, 1, padded, 0.0)[:, tp_rank * n:(tp_rank + 1) * n]
+            for w in blocks:
+                x = block_fn(x, w)
+            out = gather_from_sequence_parallel_region(x, dim=1)
+            return out if args.pad_mode == "outside" else out[:, :seq]
+
+        stopper = None
+        if args.trace:
+            # Started before the run; a timer thread stops it, so a hang still
+            # leaves a trace of what each chip completed before it stopped.
+            import threading
+
+            import jax
+
+            jax.profiler.start_trace(f"{args.trace}/rank{rank}")
+
+            def _stop():
+                print(f"[trace] rank{rank}: stopping trace after {args.trace_after}s", flush=True)
+                jax.profiler.stop_trace()
+                print(f"[trace] rank{rank}: trace written", flush=True)
+
+            stopper = threading.Timer(args.trace_after, _stop)
+            stopper.daemon = True
+            stopper.start()
+        with torch.no_grad():
+            xin = x_full.to(dev, bf)
+            if args.pad_mode == "outside":
+                xin = sp._pad_seq(xin, 1, padded, 0.0)  # eager, before the compiled graph
+            off = torch.tensor([tp_rank * n], dtype=torch.int64).to(dev)
+            got = torch.compile(model, backend="tpu")(xin, off)
+            got = got[:, :seq].float().cpu()  # unpad eagerly too
+        if stopper is not None and stopper.is_alive():
+            stopper.cancel()
+            import jax
+
+            jax.profiler.stop_trace()
+        res = {"rank": rank, "finite": bool(torch.isfinite(got).all())}
+        if rank == 0:
+            # fp32 CPU reference: every rank's weights rebuilt from the seeds.
+            ws = [[_weights(world, r, b) for r in range(world)] for b in range(args.blocks)]
+            x = x_full.clone()
+
+            def cat_rows(key, b):
+                return torch.cat([ws[b][r][key] for r in range(world)], dim=0)
+
+            def cat_cols(key, b):
+                return torch.cat([ws[b][r][key] for r in range(world)], dim=1)
+
+            def rms32(t):
+                return t * torch.rsqrt(t.pow(2).mean(-1, keepdim=True) + 1e-6)
+
+            xp_state = torch.cat([x, x.new_zeros(1, padded - seq, HID)], dim=1)
+            for b in range(args.blocks):
+                if args.variant in ("full", "noattn", "attnonly"):
+                    nrm = rms32(x)
+                    q, k, v = (nrm @ cat_rows(key, b).T for key in ("q", "k", "v"))
+                    if args.variant in ("full", "attnonly"):
+                        def heads_of(t):
+                            return t.reshape(1, seq, HEADS, HD).transpose(1, 2)
+
+                        o = F.scaled_dot_product_attention(heads_of(q), heads_of(k), heads_of(v),
+                                                           scale=HD ** -0.5)
+                        o = o.transpose(1, 2).reshape(1, seq, HID)
+                    else:
+                        o = q
+                    x = x + o @ cat_cols("o", b).T
+                if args.variant in ("full", "noattn", "ffonly"):
+                    h = rms32(x) @ cat_rows("f1", b).T
+                    x = x + F.gelu(h, approximate="tanh") @ cat_cols("f2", b).T
+                if args.variant == "rsonly":
+                    x = x + (x @ cat_rows("f1", b).T) @ cat_cols("f2", b).T
+                if args.variant == "agonly":
+                    # Each rank r keeps its own rows and only its f1 / f2 slice.
+                    xp = xp_state
+                    outs = []
+                    for r in range(world):
+                        rows = xp[:, r * n:(r + 1) * n]
+                        outs.append(rows + (rows @ ws[b][r]["f1"].T) @ ws[b][r]["f2"].T)
+                    xp_state = torch.cat(outs, dim=1)
+                    x = xp_state[:, :seq]
+                if args.variant == "rsonly2":
+                    # The pad rows carry state here (they feed the tiled input).
+                    xp = xp_state
+                    shards = [xp[:, r * n:(r + 1) * n] for r in range(world)]
+                    total = sum((shards[r].repeat(1, world, 1) @ ws[b][r]["f1"].T) @ ws[b][r]["f2"].T
+                                for r in range(world))  # (1, padded, HID): the reduce
+                    xp_state = torch.cat([shards[r] + total[:, r * n:(r + 1) * n]
+                                          for r in range(world)], dim=1)
+                    x = xp_state[:, :seq]
+            err = (got - x).norm() / x.norm()
+            rows = ((got - x).norm(dim=-1) / x.norm(dim=-1).clamp_min(1e-6))[0]
+            res.update(rel=round(float(err), 5),
+                       worst_rows=[int(i) for i in rows.topk(6).indices],
+                       worst_vals=[round(float(v), 3) for v in rows.topk(6).values])
+        reply_q.put(res)
+    except BaseException as exc:  # noqa: BLE001
+        import traceback
+
+        reply_q.put({"rank": rank, "error": f"{type(exc).__name__}: {exc}",
+                     "tb": traceback.format_exc()[-2000:]})
+        raise
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--seq", type=int, default=2310)
+    parser.add_argument("--blocks", type=int, default=2)
+    parser.add_argument("--pad", type=int, default=1)
+    parser.add_argument("--pad-mode", default="full", choices=("full", "shard", "outside", "dynamic"))
+    parser.add_argument("--trace", default="", help="jax.profiler trace dir (one subdir per rank)")
+    parser.add_argument("--trace-after", type=float, default=90.0,
+                        help="seconds after which a timer thread stops the trace (hang case)")
+    parser.add_argument("--variant", default="full", choices=("full", "noattn", "ffonly", "rsonly", "attnonly", "agonly", "rsonly2"))
+    args = parser.parse_args()
+    from difflet.backends.tpu.ops_impl import device_runtime
+
+    world = 4
+    device_runtime.prepare_launch(world)
+    ctx = mp.get_context("spawn")
+    reply_q = ctx.Queue()
+    procs = [ctx.Process(target=_worker, args=(r, world, args, reply_q)) for r in range(world)]
+    for p in procs:
+        p.start()
+    results = sorted((reply_q.get(timeout=900) for _ in procs), key=lambda r: r["rank"])
+    for p in procs:
+        p.join(timeout=60)
+    print(vars(args), results[0], [r.get("finite") for r in results], flush=True)
+    for r in results:
+        if "error" in r:
+            print(f"rank {r['rank']} FAILED: {r['error']}\n{r['tb']}", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/benchmark/tpu_profile/cm_pure_jax.py b/benchmark/tpu_profile/cm_pure_jax.py
new file mode 100644
index 00000000..aac07012
--- /dev/null
+++ b/benchmark/tpu_profile/cm_pure_jax.py
@@ -0,0 +1,91 @@
+"""The collective-matmul rings in plain JAX, without torch_tpu: AG -> RS blocks at a padded length.
+
+    python -m benchmark.tpu_profile.cm_pure_jax [--seq 2312] [--blocks 2]
+
+The SP rings misbehave on torch_tpu at 2312 rows (hang / core halt / wrong
+values; 6144 is fine; 2026-11-10). This runs the *same* ring functions --
+``collective_matmul._op`` builds them; ``torch.tpu.pallas.jax_op`` is stubbed
+to hand back the JAX function itself -- under ``jax.jit`` in one process that
+owns all four chips, no torch_tpu, and compares against an fp32 numpy
+reference. A block is ``x += RS(AG(x) @ w1^T, w2)``, the ``rsonly`` variant of
+cm_padded_repro. Right here and wrong there puts the bug in the torch_tpu
+jax_op path; wrong here puts it in the ring code or XLA.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import types
+
+HID, FF = 4096, 8192
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--seq", type=int, default=2312)
+    parser.add_argument("--blocks", type=int, default=2)
+    args = parser.parse_args()
+
+    import jax
+    import jax.numpy as jnp
+    import numpy as np
+    import torch
+    from jax.sharding import NamedSharding
+    from jax.sharding import PartitionSpec as P
+
+    # Stub the torch_tpu entry point: _op() then returns the raw JAX function.
+    torch.tpu = types.SimpleNamespace(pallas=types.SimpleNamespace(
+        jax_op=lambda name, fn, **kw: fn))
+    from difflet.backends.tpu.ops_impl import collective_matmul as cm
+
+    devices = jax.devices()
+    tp = len(devices)
+    mesh = jax.sharding.Mesh(np.array(cm._snake(devices)), ("tp",))
+    # One process: the shard of ring position p is simply shard p.
+    cm._STATE = {"mesh": mesh, "tp": tp, "shard_of_pos": tuple(range(tp)), "ids": None, "ops": {}}
+    ag = cm._op("ag", 1)
+    rs = cm._op("rs", 1)
+
+    seq = args.seq
+    if seq % (2 * tp):
+        raise SystemExit(f"--seq must divide by 2*tp={2 * tp}")
+    rng = np.random.default_rng(0)
+    x0 = rng.standard_normal((1, seq, HID), dtype=np.float32)
+    w1s = [rng.standard_normal((FF, HID), dtype=np.float32) / HID ** 0.5 for _ in range(args.blocks)]
+    w2s = [rng.standard_normal((HID, FF), dtype=np.float32) / FF ** 0.5 for _ in range(args.blocks)]
+
+    def put(a, spec, dtype=jnp.bfloat16):
+        return jax.device_put(jnp.asarray(a, dtype), NamedSharding(mesh, spec))
+
+    ids = put(np.arange(tp, dtype=np.int32), P("tp"), jnp.int32)
+    x = put(x0, P(None, "tp", None))
+    w1 = [put(w, P("tp", None)) for w in w1s]
+    w2 = [put(w, P(None, "tp")) for w in w2s]
+
+    @jax.jit
+    def model(x, ids, w1, w2):
+        for a, b in zip(w1, w2):
+            h = ag(x, ids, a)  # (1, S, FF) sharded on FF
+            x = x + rs(h, b, ids)  # (1, S, HID) sharded on S
+        return x
+
+    print(f"jax {jax.__version__} devices {[d.id for d in mesh.devices.flat]} seq {seq} "
+          f"rows/rank {seq // tp} half {seq // tp // 2} blocks {args.blocks}", flush=True)
+    got = np.asarray(model(x, ids, w1, w2).astype(jnp.float32))
+
+    ref = x0.astype(np.float64)
+    for a, b in zip(w1s, w2s):
+        # The device rounds the operands to bf16; so does the reference.
+        bf = lambda t: np.asarray(jnp.asarray(t, jnp.bfloat16).astype(jnp.float32), np.float64)  # noqa: E731
+        ref = ref + (bf(ref) @ bf(a).T) @ bf(b).T
+    err = np.linalg.norm(got - ref) / np.linalg.norm(ref)
+    rows = np.linalg.norm(got - ref, axis=-1)[0] / np.maximum(np.linalg.norm(ref, axis=-1)[0], 1e-6)
+    worst = np.argsort(rows)[-6:][::-1]
+    print(f"RESULT seq {seq} blocks {args.blocks}: rel {err:.5f} finite {bool(np.isfinite(got).all())} "
+          f"worst rows {worst.tolist()} {np.round(rows[worst], 3).tolist()}", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/benchmark/tpu_profile/trace_tail.py b/benchmark/tpu_profile/trace_tail.py
new file mode 100644
index 00000000..a770142f
--- /dev/null
+++ b/benchmark/tpu_profile/trace_tail.py
@@ -0,0 +1,44 @@
+"""The last device ops each chip completed, from per-rank jax.profiler traces.
+
+    python -m benchmark.tpu_profile.trace_tail <trace_dir> [N]
+
+For a hung run (cm_padded_repro --trace): every rank wrote a trace under
+``<trace_dir>/rank<r>``. Prints, per TPU device in each trace, the last N
+"XLA Ops" events (name, end time relative to that device's first op) and the
+last module, so the chips' stopping points can be compared side by side.
+"""
+import glob
+import gzip
+import json
+import sys
+
+
+def main() -> int:
+    root, n = sys.argv[1], int(sys.argv[2]) if len(sys.argv) > 2 else 8
+    for path in sorted(glob.glob(root + "/rank*/plugins/profile/*/*.trace.json.gz")):
+        ev = json.load(gzip.open(path))["traceEvents"]
+        procs = {e["pid"]: e["args"]["name"] for e in ev
+                 if e.get("ph") == "M" and e.get("name") == "process_name"}
+        threads = {(e["pid"], e["tid"]): e["args"]["name"] for e in ev
+                   if e.get("ph") == "M" and e.get("name") == "thread_name"}
+        print(f"== rank{path.split('/rank')[1].split('/')[0]}")
+        for pid, name in sorted(procs.items(), key=lambda kv: kv[1]):
+            if not name.startswith("/device:TPU"):
+                continue
+            ops = sorted((e for e in ev if e.get("ph") == "X" and e["pid"] == pid
+                          and threads.get((pid, e["tid"])) == "XLA Ops"), key=lambda e: e["ts"] + e["dur"])
+            mods = sorted((e for e in ev if e.get("ph") == "X" and e["pid"] == pid
+                           and threads.get((pid, e["tid"])) == "XLA Modules"), key=lambda e: e["ts"])
+            if not ops:
+                print(f"  {name}: no ops")
+                continue
+            t0 = ops[0]["ts"]
+            print(f"  {name}: {len(ops)} ops, {len(mods)} modules; last module "
+                  f"{mods[-1]['name'][:60] if mods else '-'}")
+            for e in ops[-n:]:
+                print(f"    +{(e['ts'] + e['dur'] - t0) / 1e3:9.2f} ms  {e['name'][:90]}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/difflet/backends/tpu/ltx_2/sp.py b/difflet/backends/tpu/ltx_2/sp.py
new file mode 100644
index 00000000..18b9e0ae
--- /dev/null
+++ b/difflet/backends/tpu/ltx_2/sp.py
@@ -0,0 +1,336 @@
+"""LTX-2 on TPU with sequence parallelism over collective matmuls (``DIFFLET_TPU_SP=1``).
+
+Under TP every LTX-2 block ends eight row-parallel linears in an all-reduce: the
+video and audio self-attention, text cross-attention and a2v/v2a outputs, and
+both FFNs. As for Wan, Qwen-Image and FLUX, the **video** ones become a
+reduce-scatter fused into the row-parallel matmul plus an all-gather fused into
+the next column-parallel one (``ops_impl.collective_matmul``), and the per-token
+work between them runs on this rank's 1/tp of the video rows. The audio stream
+(a few dozen to a few hundred rows) stays TP and replicated: its collectives are
+small, and sharding it would not pay for the rings.
+
+Layout. TP rank r holds video rows ``[r*Sp/tp, (r+1)*Sp/tp)`` of the sequence
+padded to ``Sp``, a multiple of ``2*tp`` (the rings split each shard in two):
+480x704x49 has 2310 rows, padded to 2312 at tp=4. The pad rows are dropped as
+keys (video self-attention, and v2a where video is the key side), so the real
+rows see exactly the unpadded attention; as queries they compute values that
+are cut off after the model.
+
+Padding happens **outside the compiled graph** (:func:`pad_video_inputs` /
+:func:`unpad_video_output`, called eagerly by the TPU transformer application):
+the latents and the video coordinates come in already padded, and the output
+is cut back after. Padding inside the graph -- a ``cat`` of the sequence, then
+the shard slice -- broke the rings on torch_tpu 0.2.0 (hangs, core halts, inf or
+wrong values; v5e-4, 2026-11-10), whether the whole sequence or only the last
+shard was padded, while the same rings were right at every unpadded length
+(1024..6144, incl. 2312 itself), in plain JAX at 2312, and with the padding done
+eagerly before the graph (benchmark/tpu_profile/cm_padded_repro.py).
+
+    attn1     norm -AG+[to_q,to_k,to_v]-> attention (keys < S) -to_out+RS->
+    attn2     norm -AG+[to_q]-> ; text -to_k,to_v (TP)-> attention -to_out+RS->
+    a2v       video norm -AG+[to_q]-> ; audio -to_k,to_v (TP)-> -to_out+RS->
+    v2a       audio -to_q (TP)-> ; video norm -AG+[to_k,to_v]-> (keys < S)
+              -to_out (TP all-reduce, audio rows)->
+    ff        -AG+net[0].proj-> gelu -net[2]+RS->
+    audio_*   unchanged (TP)
+
+The blocks receive the full (padded) video stream; the first block keeps this
+rank's rows, the last one gathers them back, so the model forward (embedders,
+AdaLN tables, RoPE, ``proj_out``) is diffusers' own.
+"""
+
+from __future__ import annotations
+
+import math
+import types
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from difflet.backends.tpu.ops_impl import collective_matmul as cm
+
+#: The collective matmuls LTX-2 uses: RS over one pair, AG over 1, 2, 3 weights.
+LTX2_OPS = (("rs", 1), ("ag", 1), ("ag", 2), ("ag", 3))
+
+def padded_length(seq_len: int, tp: int) -> int:
+    """``seq_len`` rounded up to a multiple of ``2*tp``."""
+    step = 2 * int(tp)
+    return (int(seq_len) + step - 1) // step * step
+
+
+def _pad_seq(x: torch.Tensor, dim: int, length: int, value: float) -> torch.Tensor:
+    pad = int(length) - int(x.shape[dim])
+    if pad <= 0:
+        return x
+    shape = list(x.shape)
+    shape[dim] = pad
+    return torch.cat([x, x.new_full(shape, value)], dim=dim)
+
+
+def pad_video_inputs(hidden_states: torch.Tensor, video_coords: torch.Tensor,
+                     padded: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pad the packed latents (B, S, C) with zeros and the video coordinates
+    (B, 3, S, 2) by repeating the last row, to ``padded`` rows. Eager only:
+    see the module docstring."""
+    pad = int(padded) - int(hidden_states.shape[1])
+    if pad <= 0:
+        return hidden_states, video_coords
+    hidden_states = torch.cat([hidden_states, hidden_states.new_zeros(
+        hidden_states.shape[0], pad, hidden_states.shape[2])], dim=1)
+    video_coords = torch.cat([video_coords, video_coords[:, :, -1:].expand(
+        -1, -1, pad, -1)], dim=2)
+    return hidden_states, video_coords
+
+
+def unpad_video_output(video: torch.Tensor, seq_len: int) -> torch.Tensor:
+    """The real rows of the model's video output (eager)."""
+    return video[:, :int(seq_len)].contiguous()
+
+
+def _to_bh(x: torch.Tensor, heads: int) -> torch.Tensor:
+    b, s, inner = x.shape
+    return x.reshape(b, s, heads, inner // heads).permute(0, 2, 1, 3).reshape(b * heads, s, -1)
+
+
+def _attend(attn, q, k, v, *, q_rope=None, k_rope=None, text_mask=None, valid_keys=None):
+    """The TP processor's attention on precomputed projections: (global) qk norm,
+    RoPE, attention; returns the pre-``to_out`` context (B, Sq, inner_local)."""
+    import diffusers.models.transformers.transformer_ltx2 as ltx2
+
+    from difflet.ops import attention
+
+    norm = getattr(attn.processor, "_global_rms_norm", None)
+    if norm is not None:  # TP: the RMS spans every rank's heads
+        q, k = norm(attn.norm_q, q), norm(attn.norm_k, k)
+    else:  # tp == 1: the stock across-heads norm is already global
+        q, k = attn.norm_q(q), attn.norm_k(k)
+    if q_rope is not None:
+        rotate = (ltx2.apply_interleaved_rotary_emb if attn.rope_type == "interleaved"
+                  else ltx2.apply_split_rotary_emb)
+        q = rotate(q, q_rope)
+        k = rotate(k, k_rope if k_rope is not None else q_rope)
+
+    heads = int(attn.heads)
+    b, sq, inner = q.shape
+    sk = int(k.shape[1])
+    kwargs = {}
+    if text_mask is not None:
+        from difflet.backends.tpu.ltx_2.sharding import _key_window
+
+        mask = attn.prepare_attention_mask(text_mask, sk, b)
+        mask = mask.view(b, heads, -1, mask.shape[-1])
+        kwargs = _key_window(mask, b * heads, sq, sk)
+    elif valid_keys is not None and int(valid_keys) < sk:
+        # The pad keys are dropped: exact, and the plain flash path TP uses.
+        k, v = k[:, :int(valid_keys)], v[:, :int(valid_keys)]
+    out = attention(_to_bh(q, heads), _to_bh(k, heads), _to_bh(v, heads),
+                    scale=1.0 / math.sqrt(inner // heads), causal=False,
+                    tp_q=True, tp_k=True, tp_out=False, **kwargs)
+    return out.reshape(b, heads, sq, -1).permute(0, 2, 1, 3).reshape(b, sq, inner).to(q.dtype)
+
+
+def _ag(x, linears):
+    return cm.all_gather_matmul(x, linears)
+
+
+def _rs(x, linear):
+    return cm.matmul_reduce_scatter(x, linear)
+
+
+def _out_rs(attn, ctx):
+    """``to_out`` of a video-query attention: row-parallel matmul + reduce-scatter."""
+    return attn.to_out[1](_rs(ctx, attn.to_out[0]))
+
+
+def _ff(ff, x):
+    """diffusers ``FeedForward``: [GELU(proj), Dropout, Linear, (Dropout)]."""
+    act = ff.net[0]
+    (h,) = _ag(x, (act.proj,))
+    h = ff.net[1](F.gelu(h, approximate=getattr(act, "approximate", "none")))
+    out = _rs(h, ff.net[2])
+    return ff.net[3](out) if len(ff.net) > 3 else out
+
+
+def _block_forward(
+    self,
+    hidden_states,
+    audio_hidden_states,
+    encoder_hidden_states,
+    audio_encoder_hidden_states,
+    temb,
+    temb_audio,
+    temb_ca_scale_shift,
+    temb_ca_audio_scale_shift,
+    temb_ca_gate,
+    temb_ca_audio_gate,
+    temb_prompt=None,
+    temb_prompt_audio=None,
+    video_rotary_emb=None,
+    audio_rotary_emb=None,
+    ca_video_rotary_emb=None,
+    ca_audio_rotary_emb=None,
+    encoder_attention_mask=None,
+    audio_encoder_attention_mask=None,
+    self_attention_mask=None,
+    audio_self_attention_mask=None,
+    a2v_cross_attention_mask=None,
+    v2a_cross_attention_mask=None,
+    use_a2v_cross_attention=True,
+    use_v2a_cross_attention=True,
+    perturbation_mask=None,
+    all_perturbed=None,
+):
+    """``LTX2VideoTransformerBlock.forward`` (diffusers 0.38) with the video rows sharded."""
+    sp = self._difflet_sp
+    if (self_attention_mask is not None or a2v_cross_attention_mask is not None
+            or v2a_cross_attention_mask is not None or perturbation_mask is not None):
+        raise NotImplementedError("LTX-2 SP: self/a2v/v2a masks and STG are not supported")
+    if sp.first:
+        if int(hidden_states.shape[1]) != sp.padded:
+            raise ValueError(f"LTX-2 SP expects the video stream padded to {sp.padded} rows "
+                             f"(pad_video_inputs), got {hidden_states.shape[1]}")
+        n = sp.padded // sp.tp
+        hidden_states = hidden_states[:, sp.rank * n:(sp.rank + 1) * n]
+
+    batch_size = hidden_states.size(0)
+
+    # 1.1 video self-attention
+    video_ada = self.get_mod_params(self.scale_shift_table, temb, batch_size)
+    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = video_ada[:6]
+    if self.video_cross_attn_adaln:
+        shift_text_q, scale_text_q, gate_text_q = video_ada[6:9]
+    norm = self.norm1(hidden_states) * (1 + scale_msa) + shift_msa
+    attn = self.attn1
+    q, k, v = _ag(norm, (attn.to_q, attn.to_k, attn.to_v))
+    ctx = _attend(attn, q, k, v, q_rope=video_rotary_emb, valid_keys=sp.seq)
+    hidden_states = hidden_states + (_out_rs(attn, ctx)) * gate_msa
+
+    # 1.2 audio self-attention (TP, unchanged)
+    audio_ada = self.get_mod_params(self.audio_scale_shift_table, temb_audio, batch_size)
+    (audio_shift_msa, audio_scale_msa, audio_gate_msa,
+     audio_shift_mlp, audio_scale_mlp, audio_gate_mlp) = audio_ada[:6]
+    if self.audio_cross_attn_adaln:
+        audio_shift_text_q, audio_scale_text_q, audio_gate_text_q = audio_ada[6:9]
+    norm_audio = self.audio_norm1(audio_hidden_states) * (1 + audio_scale_msa) + audio_shift_msa
+    audio_hidden_states = audio_hidden_states + (self.audio_attn1(
+        hidden_states=(norm_audio), encoder_hidden_states=None,
+        query_rotary_emb=audio_rotary_emb, attention_mask=audio_self_attention_mask,
+    )) * audio_gate_msa
+
+    # 2. text cross-attention
+    if self.cross_attn_adaln:
+        shift_text_kv, scale_text_kv = self.get_mod_params(
+            self.prompt_scale_shift_table, temb_prompt, batch_size)
+        audio_shift_text_kv, audio_scale_text_kv = self.get_mod_params(
+            self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size)
+    norm = self.norm2(hidden_states)
+    if self.video_cross_attn_adaln:
+        norm = norm * (1 + scale_text_q) + shift_text_q
+    if self.cross_attn_adaln:
+        encoder_hidden_states = encoder_hidden_states * (1 + scale_text_kv) + shift_text_kv
+    attn = self.attn2
+    (q,) = _ag((norm), (attn.to_q,))
+    ctx = _attend(attn, q, attn.to_k(encoder_hidden_states), attn.to_v(encoder_hidden_states),
+                  text_mask=encoder_attention_mask)
+    out = (_out_rs(attn, ctx))
+    if self.video_cross_attn_adaln:
+        out = out * gate_text_q
+    hidden_states = hidden_states + out
+
+    norm_audio = self.audio_norm2(audio_hidden_states)
+    if self.audio_cross_attn_adaln:
+        norm_audio = norm_audio * (1 + audio_scale_text_q) + audio_shift_text_q
+    if self.cross_attn_adaln:
+        audio_encoder_hidden_states = (
+            audio_encoder_hidden_states * (1 + audio_scale_text_kv) + audio_shift_text_kv)
+    out_audio = (self.audio_attn2(
+        (norm_audio), encoder_hidden_states=audio_encoder_hidden_states,
+        query_rotary_emb=None, attention_mask=audio_encoder_attention_mask,
+    ))
+    if self.audio_cross_attn_adaln:
+        out_audio = out_audio * audio_gate_text_q
+    audio_hidden_states = audio_hidden_states + out_audio
+
+    # 3. audio-to-video and video-to-audio cross-attention
+    if use_a2v_cross_attention or use_v2a_cross_attention:
+        norm = self.audio_to_video_norm(hidden_states)
+        norm_audio = self.video_to_audio_norm(audio_hidden_states)
+        video_ca = self.get_mod_params(
+            self.video_a2v_cross_attn_scale_shift_table[:4, :], temb_ca_scale_shift, batch_size)
+        video_gate = self.get_mod_params(
+            self.video_a2v_cross_attn_scale_shift_table[4:, :], temb_ca_gate, batch_size)
+        video_a2v_scale, video_a2v_shift, video_v2a_scale, video_v2a_shift = video_ca
+        a2v_gate = video_gate[0].squeeze(2)
+        audio_ca = self.get_mod_params(
+            self.audio_a2v_cross_attn_scale_shift_table[:4, :], temb_ca_audio_scale_shift, batch_size)
+        audio_gate = self.get_mod_params(
+            self.audio_a2v_cross_attn_scale_shift_table[4:, :], temb_ca_audio_gate, batch_size)
+        audio_a2v_scale, audio_a2v_shift, audio_v2a_scale, audio_v2a_shift = audio_ca
+        v2a_gate = audio_gate[0].squeeze(2)
+
+        if use_a2v_cross_attention:  # Q video (sharded), K/V audio (replicated)
+            mod = norm * (1 + video_a2v_scale.squeeze(2)) + video_a2v_shift.squeeze(2)
+            mod_audio = norm_audio * (1 + audio_a2v_scale.squeeze(2)) + audio_a2v_shift.squeeze(2)
+            attn = self.audio_to_video_attn
+            (q,) = _ag((mod), (attn.to_q,))
+            ctx = _attend(attn, q, attn.to_k(mod_audio), attn.to_v(mod_audio),
+                          q_rope=ca_video_rotary_emb, k_rope=ca_audio_rotary_emb)
+            hidden_states = hidden_states + a2v_gate * (_out_rs(attn, ctx))
+
+        if use_v2a_cross_attention:  # Q audio (replicated), K/V video (sharded)
+            mod = norm * (1 + video_v2a_scale.squeeze(2)) + video_v2a_shift.squeeze(2)
+            mod_audio = norm_audio * (1 + audio_v2a_scale.squeeze(2)) + audio_v2a_shift.squeeze(2)
+            attn = self.video_to_audio_attn
+            k, v = _ag((mod), (attn.to_k, attn.to_v))
+            ctx = _attend(attn, attn.to_q(mod_audio), k, v,
+                          q_rope=ca_audio_rotary_emb, k_rope=ca_video_rotary_emb, valid_keys=sp.seq)
+            out_audio = (attn.to_out[1](attn.to_out[0](ctx)))  # TP all-reduce, audio rows
+            audio_hidden_states = audio_hidden_states + v2a_gate * out_audio
+
+    # 4. feed-forward
+    norm = self.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp
+    hidden_states = hidden_states + (_ff(self.ff, (norm))) * gate_mlp
+    norm_audio = self.audio_norm3(audio_hidden_states) * (1 + audio_scale_mlp) + audio_shift_mlp
+    audio_hidden_states = audio_hidden_states + (
+        self.audio_ff((norm_audio))) * audio_gate_mlp
+
+    if sp.last:
+        from difflet.ops import gather_from_sequence_parallel_region
+
+        hidden_states = gather_from_sequence_parallel_region(hidden_states, dim=1)
+    return hidden_states, audio_hidden_states
+
+
+def install(module: nn.Module, *, video_seq_len: int) -> bool:
+    """Swap the SP block forward into a TP-sharded ``LTX2VideoTransformer3DModel``.
+
+    Returns False (and changes nothing) when there is nothing to shard.
+    """
+    from difflet.backends.tpu.ops_impl import parallel_mesh
+
+    tp, rank = parallel_mesh.get_tp_size(), parallel_mesh.get_tp_rank()
+    blocks = list(module.transformer_blocks)
+    if not blocks:
+        return False
+    for block in blocks:
+        if getattr(block, "perturbed_attn", False):
+            raise NotImplementedError("LTX-2 SP does not support perturbed attention (STG)")
+        if getattr(block.attn1, "to_gate_logits", None) is not None:
+            raise NotImplementedError("LTX-2 SP does not support gated attention")
+    seq = int(video_seq_len)
+    padded = padded_length(seq, tp)
+    cm.prepare(LTX2_OPS)
+    for index, block in enumerate(blocks):
+        block._difflet_sp = types.SimpleNamespace(
+            seq=seq, padded=padded, tp=tp, rank=rank,
+            first=index == 0, last=index == len(blocks) - 1,
+        )
+        block.forward = types.MethodType(_block_forward, block)
+    module._difflet_sp_rows = (seq, padded)  # read by the application to pad / unpad
+    print(f"[difflet tpu] LTX-2 sequence-parallel with collective matmuls ({len(blocks)} blocks, "
+          f"video S={seq} padded to {padded}, tp={tp}; audio stays TP)", flush=True)
+    return True
+
+
+__all__ = ["LTX2_OPS", "install", "pad_video_inputs", "padded_length", "unpad_video_output"]
diff --git a/difflet/backends/tpu/ltx_2/transformer.py b/difflet/backends/tpu/ltx_2/transformer.py
index 555b669e..29d8fa5a 100644
--- a/difflet/backends/tpu/ltx_2/transformer.py
+++ b/difflet/backends/tpu/ltx_2/transformer.py
@@ -134,10 +134,27 @@ def _prepare_module(self) -> torch.nn.Module:
             dtype=dtype,
             strict=True,
         )
+        self._maybe_install_sp(module)
         module.requires_grad_(False)
         self.module = module.eval()
         return self.module
 
+    def _maybe_install_sp(self, module) -> bool:
+        """``DIFFLET_TPU_SP=1``: sequence parallelism with the collectives fused into
+        the matmuls (``ltx_2/sp.py``). Compiled path only: through jax_op, eager
+        re-exports every ring on every call."""
+        from difflet.backends.tpu.ops_impl import collective_matmul
+
+        if not collective_matmul.requested() or int(self._mesh().tp) <= 1:
+            return False
+        from difflet.backends.tpu.ltx_2 import sp
+
+        collective_matmul.require_compiled(default=True)  # LTX-2 compiles by default
+        installed = sp.install(module, video_seq_len=int(self.config.video_seq_len))
+        if installed:
+            self._sp_rows = module._difflet_sp_rows  # (real, padded) video rows
+        return installed
+
     def get_example_inputs(self) -> tuple:
         """Fixed-shape inputs matching ``LTX2DiTInputBundle`` (10 tensors)."""
         cfg = self.config
@@ -178,6 +195,27 @@ def forward(
             self._init_runtime()
             self._prepare_module()
         cfg = self.config
+        rows = getattr(self, "_sp_rows", None)
+        if rows is not None and rows[1] != rows[0]:
+            # SP pads the video stream here, eagerly, never inside the compiled
+            # graph (see ltx_2/sp.py), and cuts the output back below.
+            from difflet.backends.tpu.ltx_2 import sp
+
+            hidden_states, video_coords = sp.pad_video_inputs(hidden_states, video_coords, rows[1])
+            video, audio = self._call_module(
+                hidden_states, audio_hidden_states, encoder_hidden_states,
+                audio_encoder_hidden_states, timestep, sigma, encoder_attention_mask,
+                audio_encoder_attention_mask, video_coords, audio_coords)
+            return sp.unpad_video_output(video, rows[0]), audio
+        return self._call_module(
+            hidden_states, audio_hidden_states, encoder_hidden_states,
+            audio_encoder_hidden_states, timestep, sigma, encoder_attention_mask,
+            audio_encoder_attention_mask, video_coords, audio_coords)
+
+    def _call_module(self, hidden_states, audio_hidden_states, encoder_hidden_states,
+                     audio_encoder_hidden_states, timestep, sigma, encoder_attention_mask,
+                     audio_encoder_attention_mask, video_coords, audio_coords):
+        cfg = self.config
         return self.module(
             hidden_states=hidden_states,
             audio_hidden_states=audio_hidden_states,
diff --git a/difflet/backends/tpu/ops_impl/collective_matmul.py b/difflet/backends/tpu/ops_impl/collective_matmul.py
index 70d570f0..a8847187 100644
--- a/difflet/backends/tpu/ops_impl/collective_matmul.py
+++ b/difflet/backends/tpu/ops_impl/collective_matmul.py
@@ -25,6 +25,17 @@
 compiled region must not receive a view at a nonzero storage offset (the
 J5ag ring then hangs before compiling; see sp_pair_probe).
 
+**Every rank must compile the same program around the rings.** torch_tpu
+compiles one program per process and inlines these rings into it (their
+``collective-permute``s all carry ``channel_id=1`` and pair by order). Values
+that differ by rank therefore enter the compiled graph as runtime tensors (the
+ring position is ``_STATE["ids"]`` for this reason), never as Python constants
+the compiler can specialize on. LTX-2 broke this by padding its sequence inside
+the graph and slicing each rank's shard at a constant offset: XLA folded the pad
+away on some ranks only, the programs diverged, and a chip halted in the first
+ring hop (v5e-4, torch_tpu 0.2.0, 2026-11-10; ``benchmark/tpu_profile/
+cm_padded_repro.py``). Pad before the graph, or pass the offset as a tensor.
+
 **Shard order.** The model's sequence shard for TP rank ``r`` is rows
 ``[r*S/tp, (r+1)*S/tp)`` (``scatter_tp_dim``). The ring runs in physical
 order, which is not torch rank order (rank 0 may hold TPU_2), so every chunk
diff --git a/tests/unit/backends/test_tpu_ltx2_sp.py b/tests/unit/backends/test_tpu_ltx2_sp.py
new file mode 100644
index 00000000..c1697132
--- /dev/null
+++ b/tests/unit/backends/test_tpu_ltx2_sp.py
@@ -0,0 +1,162 @@
+"""CPU checks of LTX-2's sequence-parallel block forward (backends/tpu/ltx_2/sp.py).
+
+The collective matmuls are JAX rings that only run on the chips; here they are
+replaced by their tp=1 meaning (all-gather = identity, reduce-scatter = the
+matmul), so the SP model must reproduce diffusers' own forward. The video
+sequence is odd, so it is padded (to a multiple of 2*tp) and the pad keys must
+be masked for the result to match.
+"""
+
+from __future__ import annotations
+
+import importlib
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from difflet.pipeline.parallel_mesh import MeshSpec
+
+
+@pytest.fixture(autouse=True)
+def _mesh():
+    # Resolved at run time: another test may drop difflet.backends.tpu.* from
+    # sys.modules, and sp.install imports parallel_mesh when it runs.
+    pm = importlib.import_module("difflet.backends.tpu.ops_impl.parallel_mesh")
+    pm.destroy_parallel_mesh()
+    pm.init_parallel_mesh(MeshSpec(tp=1))
+    yield
+    pm.destroy_parallel_mesh()
+
+
+def _tiny():
+    from diffusers.models.transformers.transformer_ltx2 import LTX2VideoTransformer3DModel
+
+    torch.manual_seed(0)
+    return LTX2VideoTransformer3DModel(
+        in_channels=8, out_channels=8, patch_size=1, patch_size_t=1,
+        num_attention_heads=2, attention_head_dim=8, cross_attention_dim=16,
+        vae_scale_factors=(8, 32, 32), pos_embed_max_pos=20, base_height=2048,
+        base_width=2048, audio_in_channels=8, audio_out_channels=8,
+        audio_patch_size=16, audio_patch_size_t=1, audio_num_attention_heads=2,
+        audio_attention_head_dim=8, audio_cross_attention_dim=16, audio_scale_factor=4,
+        audio_pos_embed_max_pos=20, audio_sampling_rate=16000, audio_hop_length=160,
+        num_layers=2, caption_channels=16, rope_double_precision=False,
+        use_prompt_embeddings=True, rope_type="split",
+    ).eval().double()
+
+
+def _fake_rings(monkeypatch, cm):
+    def ag(x, linears):
+        return [F.linear(x, lin.weight, lin.bias) for lin in linears]
+
+    def rs(x, linear):
+        return F.linear(x, linear.weight, linear.bias)
+
+    monkeypatch.setattr(cm, "all_gather_matmul", ag)
+    monkeypatch.setattr(cm, "matmul_reduce_scatter", rs)
+    monkeypatch.setattr(cm, "prepare", lambda ops=None: None)
+
+
+def _inputs(frames=1, height=3, width=3, audio_frames=5, text=6, batch=2):
+    g = torch.Generator().manual_seed(1)
+    seq = frames * height * width
+    mask = torch.zeros(batch, text, dtype=torch.int64)
+    mask[0, 2:] = 1  # left padding, a different length per row
+    mask[1, 4:] = 1
+    return dict(
+        hidden_states=torch.randn(batch, seq, 8, generator=g, dtype=torch.float64),
+        audio_hidden_states=torch.randn(batch, audio_frames, 8, generator=g, dtype=torch.float64),
+        encoder_hidden_states=torch.randn(batch, text, 16, generator=g, dtype=torch.float64),
+        audio_encoder_hidden_states=torch.randn(batch, text, 16, generator=g, dtype=torch.float64),
+        timestep=torch.tensor([300.0, 700.0], dtype=torch.float64),
+        sigma=torch.tensor([0.3, 0.7], dtype=torch.float64),
+        encoder_attention_mask=mask, audio_encoder_attention_mask=mask,
+        num_frames=frames, height=height, width=width, audio_num_frames=audio_frames,
+        return_dict=False,
+    )
+
+
+def _coords(model, inputs):
+    return model.rope.prepare_video_coords(
+        inputs["hidden_states"].shape[0], inputs["num_frames"], inputs["height"],
+        inputs["width"], inputs["hidden_states"].device, fps=24.0)
+
+
+def _sp_forward(sp, model, inputs):
+    """What the TPU application does around the compiled model: pad eagerly,
+    run, cut the video output back."""
+    seq, padded = model._difflet_sp_rows
+    hidden, coords = sp.pad_video_inputs(inputs["hidden_states"], _coords(model, inputs), padded)
+    video, audio = model(**dict(inputs, hidden_states=hidden, video_coords=coords))
+    return sp.unpad_video_output(video, seq), audio
+
+
+@pytest.mark.parametrize("dims", [(1, 3, 3), (2, 2, 2)])  # 9 rows (padded to 10), 8 rows (no pad)
+def test_sp_forward_reproduces_diffusers_at_tp1(monkeypatch, dims):
+    sp = importlib.import_module("difflet.backends.tpu.ltx_2.sp")
+    cm = importlib.import_module("difflet.backends.tpu.ops_impl.collective_matmul")
+    _fake_rings(monkeypatch, cm)
+    frames, height, width = dims
+    ref_model = _tiny()
+    inputs = _inputs(frames, height, width)
+    with torch.no_grad():
+        want = ref_model(**inputs)
+    model = _tiny()
+    seq = frames * height * width
+    assert sp.install(model, video_seq_len=seq)
+    assert model._difflet_sp_rows == (seq, sp.padded_length(seq, 1))
+    with torch.no_grad():
+        got = _sp_forward(sp, model, inputs)
+    for g, w in zip(got, want):
+        assert g.shape == w.shape
+        # 1e-6, not float64 rounding: difflet's CPU attention op takes its
+        # softmax in fp32 (diffusers' SDPA stays in float64), ~5e-8 here.
+        torch.testing.assert_close(g, w, rtol=1e-6, atol=1e-6)
+
+
+def test_padding_without_dropping_the_pad_keys_would_differ(monkeypatch):
+    """Guards the pad-key handling: attending over the pad rows changes the result."""
+    sp = importlib.import_module("difflet.backends.tpu.ltx_2.sp")
+    cm = importlib.import_module("difflet.backends.tpu.ops_impl.collective_matmul")
+    _fake_rings(monkeypatch, cm)
+    inputs = _inputs(1, 3, 3)
+    with torch.no_grad():
+        want = _tiny()(**inputs)
+    model = _tiny()
+    sp.install(model, video_seq_len=9)
+    real = sp._attend
+    monkeypatch.setattr(sp, "_attend", lambda *a, valid_keys=None, **k: real(*a, **k))
+    with torch.no_grad():
+        got = _sp_forward(sp, model, inputs)
+    assert (got[0] - want[0]).abs().max() > 1e-3
+
+
+def test_unpadded_stream_is_refused(monkeypatch):
+    """Padding inside the compiled graph broke the rings on torch_tpu; the first
+    block insists on an already padded stream."""
+    sp = importlib.import_module("difflet.backends.tpu.ltx_2.sp")
+    cm = importlib.import_module("difflet.backends.tpu.ops_impl.collective_matmul")
+    _fake_rings(monkeypatch, cm)
+    model = _tiny()
+    sp.install(model, video_seq_len=9)
+    with pytest.raises(ValueError, match="padded to 10"), torch.no_grad():
+        model(**_inputs(1, 3, 3))
+
+
+def test_pad_and_unpad_video_rows():
+    sp = importlib.import_module("difflet.backends.tpu.ltx_2.sp")
+    h = torch.arange(2 * 3 * 2, dtype=torch.float32).reshape(2, 3, 2)
+    c = torch.arange(2 * 3 * 3 * 2, dtype=torch.float32).reshape(2, 3, 3, 2)
+    hp, cp = sp.pad_video_inputs(h, c, 4)
+    assert hp.shape == (2, 4, 2) and not hp[:, 3].any()
+    assert cp.shape == (2, 3, 4, 2) and torch.equal(cp[:, :, 3], c[:, :, 2])
+    assert torch.equal(sp.unpad_video_output(hp, 3), h)
+    assert sp.pad_video_inputs(h, c, 3)[0] is h
+
+
+def test_padded_length():
+    sp = importlib.import_module("difflet.backends.tpu.ltx_2.sp")
+    assert sp.padded_length(2310, 4) == 2312
+    assert sp.padded_length(6144, 4) == 6144
+    assert sp.padded_length(9, 1) == 10

From ad5bf2e79e3266b36836689767a3af848c8cd4a8 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Tue, 10 Nov 2026 14:13:03 -0800
Subject: [PATCH 11/22] docs(tpu): plan for Pallas collective-matmul kernels
 behind the SP API

Proposal, gated on a profile of the communication still exposed after SP.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 ...2026-11-10-tpu-collective-matmul-pallas.md | 149 ++++++++++++++++++
 1 file changed, 149 insertions(+)
 create mode 100644 docs/plans/2026-11-10-tpu-collective-matmul-pallas.md

diff --git a/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md b/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
new file mode 100644
index 00000000..2938e3b9
--- /dev/null
+++ b/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
@@ -0,0 +1,149 @@
+# Pallas collective-matmul kernels for TPU sequence parallelism
+
+Date: 2026-11-10 · Branch: `tpu-ltx2-sp` · Status: proposal, gated on step 1
+
+## Goal
+
+Replace the internals of `difflet/backends/tpu/ops_impl/collective_matmul.py`
+(today a JAX `shard_map` ring run through `torch.tpu.pallas.jax_op`) with
+hand-written Pallas TPU kernels, behind the same three functions every SP model
+already calls:
+
+| function | used by |
+|---|---|
+| `all_gather_matmul(x, linears)` — AG + column-parallel matmul, 1–4 weights | Wan, Qwen-Image, HunyuanVideo, FLUX, LTX-2 |
+| `matmul_reduce_scatter(x, linear)` — row-parallel matmul + RS | same |
+| `matmul2_reduce_scatter(x0, l0, x1, l1)` — two pairs, one RS | FLUX single blocks |
+
+No model file changes. Trainium and CUDA paths are untouched (TPU-only module).
+
+## Why
+
+The JAX ring hides about 55% of the communication it replaces
+(`sp_pair_probe`, Wan block, 2026-11-03: 6.33 → 4.95 ms per block of 8.8 ms;
+the all-reduce was 2.5 ms). XLA decides the overlap of each hop's
+`collective-permute` with its einsum, and every hop's partial result goes
+through HBM. A kernel controls both: DMA start/wait around the matmul with
+double buffering, and accumulators kept in VMEM.
+
+Expected gain, Wan as the example: ~1.1 ms/block still exposed after SP; hiding
+85–90% of it is 0.6–0.9 ms/block, **~7–10% of a step** on top of today's SP.
+LTX-2 (all-reduce ~33% of a TP step) should gain at least as much; FLUX and
+HunyuanVideo less.
+
+Secondary: fusing work the JAX ring cannot (bias, gelu, an RMSNorm prologue,
+group-wise int8 dequantization — the MiniMax-H3 probe OOMed at group 128 because
+the compiled graph materialized every dequantized weight).
+
+## Lessons from 2026-11-10 that shape the design
+
+LTX-2 SP produced wrong output, hangs and core halts at 480x704x49 (2310 video
+rows padded to 2312). Root cause, from `benchmark/tpu_profile/cm_padded_repro.py`,
+per-rank HLO dumps and per-rank `jax.profiler` traces:
+
+- torch_tpu compiles **one program per process**, and the ring is **inlined**
+  into each (`collective-permute`s, all `channel_id=1`, paired by order only).
+- Padding inside the compiled graph with **rank-constant** slice offsets let XLA
+  specialize each rank differently (ranks whose shard lies inside the unpadded
+  input folded the `concatenate` away); the per-rank programs diverged, a chip
+  halted in the ring's first matmul fusion, the others waited in a collective.
+- Every unpadded length 1024..6144 (incl. 2312) was right; plain JAX at 2312 was
+  right; padding eagerly before the graph was right; padding in the graph with
+  the offset as a **runtime tensor** (identical programs on every rank) was right.
+
+Design rules that follow, for the kernels and for their callers:
+
+1. **Rank-dependent values enter compiled graphs as runtime tensors**, never as
+   Python constants; the program around an SP collective must be structurally
+   identical on every rank. Document this in `collective_matmul.py`.
+2. Sequence padding stays **outside** the compiled graph (as LTX-2 now does).
+3. A Pallas kernel is an opaque custom call, so XLA cannot fold neighbouring ops
+   into it — this removes the failure class above, but not rule 1 (the kernel's
+   call order must still match across ranks).
+
+## Design
+
+### AG + matmul
+
+```
+in : x_local (B, S/tp, K); W_i (N_i/tp, K), i = 1..n
+out: (B, S, N_i/tp)
+split x_local into two halves; ring over the 2x2 snake, both directions
+for hop in 0..tp-1:
+    start remote copy: top half -> cw neighbour, bottom half -> ccw neighbour
+    matmul the halves held now with every W_i, tiles [tm x tk] x [tk x tn] in VMEM
+    store each result block at its owner's rows (shard_of_pos table, as today)
+    wait for this hop's DMA (semaphores)
+```
+
+Double-buffered send/compute buffers. n weights share one DMA (ag2/3/4).
+
+### matmul + RS
+
+```
+in : x (B, S, K/tp); W (N, K/tp); out: (B, S/tp, N)
+for hop in 0..tp-1:
+    partial = x[chunk for the rank `hop` steps away] @ W^T
+    acc = received accumulator + partial      (VMEM if it fits, else HBM double buffer)
+    start remote copy acc -> next rank
+final acc = this rank's rows; add bias once
+```
+
+`rs2` sums two (x, W) pairs into the same accumulator.
+
+### Shapes
+
+| dimension | range seen today |
+|---|---|
+| rows per rank | 578 (LTX-2), 1170 (Wan), 1536, 1024 + 128 (FLUX streams) — often not a tile multiple |
+| K, N | 1024 – 16384, multiples of 128 |
+| weights per AG | 1 – 4 |
+| dtype | bf16 (int8 later) |
+| topology | v5e-4, 2x2, tp = 4 |
+
+Ragged rows: mask the last tile in-kernel, or require callers to pad (eagerly)
+to a multiple of `16 * 2 * tp`. Start with the latter; it keeps the kernel
+simpler and matches rule 2.
+
+### Switch
+
+`DIFFLET_TPU_SP_KERNEL=pallas|jax` (default `jax` until step 5 passes).
+Anything the kernel does not support (tp != 4, unsupported arity or shape)
+falls back to the JAX ring with a one-line notice.
+
+## Steps
+
+| # | Work | Validation | Estimate |
+|---|---|---|---|
+| 1 | **Decision gate.** Profile every model's compiled SP step; measure the communication still exposed (`trace_breakdown.py`, `sp_pair_probe.py`). | Go if exposed comm ≥ 8% of a step on LTX-2 or Wan; stop if < 5%. | 0.5 day |
+| 2 | Pallas AG (1 weight) and RS (1 pair) in **plain JAX**, one process owning 4 chips (harness like `cm_pure_jax.py`). | vs numpy fp32 reference; per-op time vs the JAX ring at Wan / LTX-2 / FLUX FFN shapes; overlap ratio. | 2–3 days |
+| 3 | Through `jax_op` into torch_tpu behind the switch. | `cm_ops_parity.py`; `cm_padded_repro.py` length sweep and all pad modes; per-rank program fingerprints identical. | 1–2 days |
+| 4 | AG with 2/3/4 weights, RS with 2 pairs, bias. | As step 3, plus `sp_pair_probe.py` per block. | 1–2 days |
+| 5 | Per-model A/B with main's `bench.py`: Wan, Qwen-Image, HunyuanVideo, FLUX, LTX-2. Same commit, back to back, kernel vs JAX ring. | PSNR within the noise floor of compile-vs-eager; image sha256, versions and commit recorded (reproducibility rule). | 1–2 days |
+| 6 (optional) | Fusions: gelu after the up-projection, RMSNorm prologue, group-wise int8 dequant. | FFN micro-benchmark; MiniMax-H3 int8 probe at group 128. | as needed |
+
+Total for steps 1–5: about 1–1.5 weeks.
+
+## Success criteria
+
+- Per-op overlap ≥ 85% (today 55%).
+- LTX-2 and Wan step time ≥ 5% below today's SP.
+- Output of every model within its compile-vs-eager noise floor; no hang or core
+  halt across the length sweep and pad modes.
+- One environment variable reverts to the JAX ring.
+
+## Risks
+
+| Risk | Mitigation |
+|---|---|
+| Pallas TPU remote DMA / semaphore API on jax 0.12 | Step 2 in plain JAX first; start from the collective-matmul example in JAX's "Distributed computing in Pallas for TPUs" guide. |
+| Deadlocks and out-of-bounds access are hard to debug | The 2026-11-10 tooling: minimal repro, per-rank HLO fingerprints, timer-stopped per-rank traces. |
+| torch_tpu support for Pallas custom calls | The TPU flash attention is already a Pallas kernel on the same path. |
+| Tied to v5e-4 (2x2, tp = 4) | Topology and tile sizes as parameters; other configs fall back explicitly. |
+| Gain smaller than estimated | Step 1 measures it before any kernel work. |
+
+## Cheaper alternative to try first
+
+Split each shard into 4 or 8 chunks instead of 2 in the existing JAX ring, so
+XLA has more hops to overlap. No kernel; a day of work; tells how much of the
+gap is granularity rather than control.

From 7df82087cb7636107c6af833e09d46625a678323 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Tue, 10 Nov 2026 14:26:44 -0800
Subject: [PATCH 12/22] docs(tpu): the confirmed mechanism behind the in-graph
 padding failure

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>
---
 benchmark/tpu_profile/cm_padded_repro.py      |  7 ++
 benchmark/tpu_profile/hlo_rank_diff.py        | 75 +++++++++++++++++++
 .../tpu/ops_impl/collective_matmul.py         | 16 ++--
 3 files changed, 93 insertions(+), 5 deletions(-)
 create mode 100644 benchmark/tpu_profile/hlo_rank_diff.py

diff --git a/benchmark/tpu_profile/cm_padded_repro.py b/benchmark/tpu_profile/cm_padded_repro.py
index 3290a484..6cf46277 100644
--- a/benchmark/tpu_profile/cm_padded_repro.py
+++ b/benchmark/tpu_profile/cm_padded_repro.py
@@ -119,6 +119,11 @@ def bh(t):
             return x
 
         def model(xf, off):
+            if args.salt:
+                # A distinct constant per run: a new graph hash, so every cache
+                # layer misses and the compiler really runs (and dumps). XLA
+                # folds the zero away; the optimized program is unchanged.
+                xf = xf + xf.new_tensor(args.salt) * 0
             if args.pad_mode == "dynamic":
                 # Padded in the graph as in "full", but the shard offset is a
                 # runtime input, so every rank compiles the same program.
@@ -241,6 +246,8 @@ def main() -> int:
     parser.add_argument("--blocks", type=int, default=2)
     parser.add_argument("--pad", type=int, default=1)
     parser.add_argument("--pad-mode", default="full", choices=("full", "shard", "outside", "dynamic"))
+    parser.add_argument("--salt", type=float, default=0.0,
+                        help="nonzero: force a fresh compile (see model())")
     parser.add_argument("--trace", default="", help="jax.profiler trace dir (one subdir per rank)")
     parser.add_argument("--trace-after", type=float, default=90.0,
                         help="seconds after which a timer thread stops the trace (hang case)")
diff --git a/benchmark/tpu_profile/hlo_rank_diff.py b/benchmark/tpu_profile/hlo_rank_diff.py
new file mode 100644
index 00000000..91b908f6
--- /dev/null
+++ b/benchmark/tpu_profile/hlo_rank_diff.py
@@ -0,0 +1,75 @@
+"""Compare the per-rank compiled programs of a torch_tpu run (structure and collective buffers).
+
+    python -m benchmark.tpu_profile.hlo_rank_diff <dump_base>
+
+``<dump_base>/rank<r>/xla`` are XLA dumps (``--xla_dump_to``, one directory per
+rank, e.g. ``cm_padded_repro.py`` with ``REPRO_DUMP_BASE``; pass ``--salt`` so a
+cached program does not skip the compile). For the compiled graph's module it
+prints, per rank: the ENTRY op count and a fingerprint with numbers stripped,
+the first differing regions against rank 0, and the offsets of the
+``collective-permute`` receive buffers from the buffer assignment.
+
+On torch_tpu every rank compiles its own program and the SP rings are inlined
+into it; a collective-permute writes into its peer at the receive buffer's
+offset, so the programs -- and those offsets -- must agree across ranks
+(2026-11-10: LTX-2's in-graph padding moved rank 3's by +18.9 MB).
+"""
+
+import difflib
+import glob
+import hashlib
+import re
+import sys
+
+MODULE = "*fx_to_mlir*"
+
+
+def _read(base, rank, suffix):
+    files = sorted(glob.glob(f"{base}/rank{rank}/xla/{MODULE}{suffix}"))
+    return open(files[0]).read() if files else None
+
+
+def _entry(text):
+    m = re.search(r"\nENTRY [^\n]*\{\n(.*?)\n\}", text, re.S)
+    return m.group(1).splitlines() if m else []
+
+
+def _norm(line):
+    line = re.sub(r", metadata=\{[^}]*\}", "", line)
+    line = re.sub(r", backend_config=.*", "", line)
+    return re.sub(r"\d+", "#", line).strip()
+
+
+def main() -> int:
+    base = sys.argv[1]
+    entries = {}
+    print("== ENTRY computation per rank (numbers stripped for the fingerprint)")
+    for r in range(4):
+        text = _read(base, r, ".after_optimizations.txt")
+        if text is None:
+            print(f"rank{r}: no dump")
+            continue
+        entries[r] = [_norm(line) for line in _entry(text)]
+        digest = hashlib.md5("\n".join(entries[r]).encode()).hexdigest()[:10]
+        print(f"rank{r}: {len(entries[r])} ops, fingerprint {digest}")
+    for r in (1, 2, 3):
+        if 0 not in entries or r not in entries:
+            continue
+        sm = difflib.SequenceMatcher(None, entries[0], entries[r], autojunk=False)
+        diffs = [op for op in sm.get_opcodes() if op[0] != "equal"]
+        print(f"rank0 vs rank{r}: {len(diffs)} differing regions")
+        for tag, i1, i2, j1, j2 in diffs[:3]:
+            print(f"  {tag} rank0[{i1}:{i2}] {entries[0][i1:i2][:1]}")
+            print(f"  {' ' * len(tag)} rank{r}[{j1}:{j2}] {entries[r][j1:j2][:1]}")
+    print("== collective-permute receive buffers (buffer assignment)")
+    for r in range(4):
+        text = _read(base, r, ".after_optimizations-buffer-assignment.txt")
+        if text is None:
+            continue
+        found = re.findall(r"<\d+ (collective-permute-start[.\d]*)\{1\} @1> \(size=(\d+),offset=(\d+)\)", text)
+        print(f"rank{r}: " + " ".join(f"{n}@{o}" for n, _s, o in found[:4]))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/difflet/backends/tpu/ops_impl/collective_matmul.py b/difflet/backends/tpu/ops_impl/collective_matmul.py
index a8847187..4dccb29b 100644
--- a/difflet/backends/tpu/ops_impl/collective_matmul.py
+++ b/difflet/backends/tpu/ops_impl/collective_matmul.py
@@ -30,11 +30,17 @@
 ``collective-permute``s all carry ``channel_id=1`` and pair by order). Values
 that differ by rank therefore enter the compiled graph as runtime tensors (the
 ring position is ``_STATE["ids"]`` for this reason), never as Python constants
-the compiler can specialize on. LTX-2 broke this by padding its sequence inside
-the graph and slicing each rank's shard at a constant offset: XLA folded the pad
-away on some ranks only, the programs diverged, and a chip halted in the first
-ring hop (v5e-4, torch_tpu 0.2.0, 2026-11-10; ``benchmark/tpu_profile/
-cm_padded_repro.py``). Pad before the graph, or pass the offset as a tensor.
+the compiler can specialize on. A ``collective-permute`` writes into its peer at the
+receive buffer's offset in *its own* program's buffer assignment, so those
+offsets must agree too. LTX-2 broke this by padding its sequence inside the
+graph and slicing each rank's shard at a constant offset: XLA kept the pad only
+on rank 3 (whose shard holds the pad rows), rank 3's program grew a copy and a
+fusion, and its receive buffers moved by +18.9 MB; the peers' writes landed on
+other live buffers, and the run hung, halted a core or produced wrong values
+(v5e-4, torch_tpu 0.2.0, 2026-11-10; ``benchmark/tpu_profile/
+cm_padded_repro.py`` and ``hlo_rank_diff.py``). With the offset passed as a
+tensor all four programs, and their receive-buffer offsets, were identical and
+the result right. Pad before the graph, or pass rank-dependent values as tensors.
 
 **Shard order.** The model's sequence shard for TP rank ``r`` is rows
 ``[r*S/tp, (r+1)*S/tp)`` (``scatter_tp_dim``). The ring runs in physical

From 6f78e40ca65c26a9cfb1eedd940f90a547412a08 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Tue, 10 Nov 2026 14:55:35 -0800
Subject: [PATCH 13/22] docs(tpu): step-1 profile of the SP rings (Wan, LTX-2)
 and the revised 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>
---
 benchmark/tpu_profile/exposed_comm.py         | 133 ++++++++++++++++++
 ...2026-11-10-tpu-collective-matmul-pallas.md |  31 ++++
 2 files changed, 164 insertions(+)
 create mode 100644 benchmark/tpu_profile/exposed_comm.py

diff --git a/benchmark/tpu_profile/exposed_comm.py b/benchmark/tpu_profile/exposed_comm.py
new file mode 100644
index 00000000..484680ed
--- /dev/null
+++ b/benchmark/tpu_profile/exposed_comm.py
@@ -0,0 +1,133 @@
+"""How much communication a traced TPU denoise still waits on.
+
+    python -m benchmark.tpu_profile.exposed_comm <trace_dir> [<trace_dir> ...]
+
+Reads the rank-0 ``jax.profiler`` trace that main's bench writes with
+``DIFFLET_BENCH_TRACE_DIR`` (the natural, warm, compiled request's denoise) and
+splits device time three ways:
+
+* **ICI wait** -- ops whose time is the core waiting on the interconnect:
+  synchronous collectives (``all-reduce``, ``all-gather``, ``reduce-scatter``,
+  ``all-to-all``) and the ``-done`` half of asynchronous ones
+  (``collective-permute-done``, ...). The ``-start`` halves only issue the
+  transfer and are counted as compute;
+* **local DMA wait** -- ``copy-done`` / ``slice-done``: waits on on-chip async
+  copies (HBM <-> VMEM), not on other chips;
+* **compute** -- every other device op;
+* **idle** -- the window minus the union of all op intervals.
+
+This is the decision input of docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
+(step 1): the share of a step a better-overlapped collective could still save.
+"""
+
+from __future__ import annotations
+
+import collections
+import glob
+import gzip
+import json
+import re
+import sys
+
+#: Ops that wait on the interconnect: synchronous collectives and the -done half
+#: of asynchronous ones. copy-done / slice-done are local async DMAs, not ICI.
+WAIT = re.compile(r"^(all-reduce|all-gather|reduce-scatter|all-to-all|collective-permute"
+                  r"|async-collective)(-done)?(\.\d+)?$")
+LOCAL = re.compile(r"^(copy|slice)-done(\.\d+)?$")
+ASYNC_START = re.compile(r"-start(\.\d+)?$")
+#: "%name = <shape> <opcode>(" -- the opcode of an HLO instruction line.
+HLO_LINE = re.compile(r"^%?[\w.-]+ = .*?\s([a-z][a-z0-9-]*)\(")
+
+
+def _union(intervals):
+    total, cur_s, cur_e = 0.0, None, None
+    for s, e in sorted(intervals):
+        if cur_e is None or s > cur_e:
+            if cur_e is not None:
+                total += cur_e - cur_s
+            cur_s, cur_e = s, e
+        else:
+            cur_e = max(cur_e, e)
+    return total + ((cur_e - cur_s) if cur_e is not None else 0.0)
+
+
+def _from_json(trace_dir):
+    path = sorted(glob.glob(trace_dir + "/plugins/profile/*/*.trace.json.gz"))[-1]
+    ev = json.load(gzip.open(path))["traceEvents"]
+    procs = {e["pid"]: e["args"]["name"] for e in ev
+             if e.get("ph") == "M" and e.get("name") == "process_name"}
+    threads = {(e["pid"], e["tid"]): e["args"]["name"] for e in ev
+               if e.get("ph") == "M" and e.get("name") == "thread_name"}
+    dev = [e for e in ev if e.get("ph") == "X" and procs.get(e["pid"]) == "/device:TPU:0"]
+    ops = [e for e in dev if threads.get((e["pid"], e["tid"])) == "XLA Ops"]
+    mods = [e for e in dev if threads.get((e["pid"], e["tid"])) == "XLA Modules"]
+    return ops, mods
+
+
+def _from_xplane(trace_dir):
+    """The raw xplane, for traces whose JSON export dropped the device events
+    (it stops at 5M events; a host-heavy run fills them with /host:CPU)."""
+    import jax.profiler
+
+    path = sorted(glob.glob(trace_dir + "/plugins/profile/*/*.xplane.pb"))[-1]
+    data = jax.profiler.ProfileData.from_file(path)
+    plane = next(p for p in data.planes if p.name == "/device:TPU:0")
+    lines = {line.name: line for line in plane.lines}
+
+    def short(text):
+        # xplane names are full HLO lines: "%x.3 = bf16[..] collective-permute-done(..)";
+        # keep the opcode so the patterns below see what the JSON export calls them.
+        m = HLO_LINE.match(text)
+        return m.group(1) if m else text
+
+    def events(name):
+        return [{"name": short(e.name), "ts": e.start_ns / 1e3, "dur": e.duration_ns / 1e3}
+                for e in lines[name].events] if name in lines else []
+
+    return events("XLA Ops"), events("XLA Modules")
+
+
+def analyze(trace_dir: str) -> dict:
+    ops, mods = _from_json(trace_dir)
+    if not mods:
+        ops, mods = _from_xplane(trace_dir)
+    t0 = min(e["ts"] for e in mods)
+    t1 = max(e["ts"] + e["dur"] for e in mods)
+    wall = t1 - t0
+    busy = _union((e["ts"], e["ts"] + e["dur"]) for e in ops)
+    waits = [e for e in ops if WAIT.search(e["name"]) and not ASYNC_START.search(e["name"])]
+    wait = _union((e["ts"], e["ts"] + e["dur"]) for e in waits)
+    local = _union((e["ts"], e["ts"] + e["dur"]) for e in ops if LOCAL.search(e["name"]))
+    by_kind = collections.Counter()
+    for e in waits:
+        by_kind[re.sub(r"(\.\d+)$", "", e["name"])] += e["dur"]
+    by_op = collections.Counter()
+    for e in ops:
+        by_op[re.sub(r"(\.\d+)+$", "", e["name"])] += e["dur"]
+    return {
+        "top_ops": by_op.most_common(8),
+        "dir": trace_dir, "window_ms": wall / 1e3, "busy_ms": busy / 1e3,
+        "wait_ms": wait / 1e3, "local_ms": local / 1e3, "idle_ms": (wall - busy) / 1e3,
+        "modules": len(mods), "top_wait": by_kind.most_common(5),
+    }
+
+
+def main() -> int:
+    print(f"{'trace':24s} {'window':>9s} {'ICI wait':>16s} {'local DMA wait':>16s} "
+          f"{'compute':>16s} {'idle':>14s}")
+    for d in sys.argv[1:]:
+        r = analyze(d)
+        w = r["window_ms"]
+        compute = r["busy_ms"] - r["wait_ms"] - r["local_ms"]
+        print(f"{d.rstrip('/').split('/')[-1]:24s} {w:8.1f}ms "
+              f"{r['wait_ms']:8.1f}ms {100 * r['wait_ms'] / w:5.1f}% "
+              f"{r['local_ms']:8.1f}ms {100 * r['local_ms'] / w:5.1f}% "
+              f"{compute:8.1f}ms {100 * compute / w:5.1f}% "
+              f"{r['idle_ms']:7.1f}ms {100 * r['idle_ms'] / w:4.1f}%")
+        print("    top waits (ms): " + ", ".join(f"{k} {v / 1e3:.1f}" for k, v in r["top_wait"]))
+        print("    top ops (ms):   " + ", ".join(f"{k} {v / 1e3:.1f}" for k, v in r["top_ops"]))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md b/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
index 2938e3b9..844df024 100644
--- a/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
+++ b/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
@@ -111,6 +111,37 @@ simpler and matches rule 2.
 Anything the kernel does not support (tp != 4, unsupported arity or shape)
 falls back to the JAX ring with a one-line notice.
 
+## Step 1 result (2026-11-10): the cost is the AG ring's write-back, not the wait
+
+Rank-0 traces of the natural, warm, compiled denoise (main's bench with
+`DIFFLET_BENCH_TRACE_DIR`), split by `benchmark/tpu_profile/exposed_comm.py`
+(reads the xplane directly when the JSON export is truncated):
+
+| 20 steps | Wan TP | Wan SP | LTX-2 TP | LTX-2 SP |
+|---|---|---|---|---|
+| window | 7200.7 ms | 5837.2 ms (-19%) | 4339.8 ms | 3793.5 ms (-12.6%) |
+| ICI wait | 1978.6 (27.5%) | 192.9 (3.3%) | 1275.5 (29.4%) | 258.5 (6.8%) |
+| compute | 4987.8 | 5398.9 (+411) | 2585.4 | 2993.5 (+408) |
+| of which `dynamic-update-slice` | -- | 643.0 (11.0%) | -- | 264.7 (7.0%) |
+| idle | 3.1% | 3.6% | 10.5% | 13.2% |
+
+- The rings already hide 80-90% of the all-reduce (the op-level probe's 55%
+  understated it); what is left waiting is 3-7% of a step.
+- SP adds ~410 ms of compute per 20 steps, mostly the AG ring writing each
+  hop's product into the full output with `dynamic_update_slice` (tp=4: eight
+  writes per AG). The matmuls themselves get *faster* under SP (LTX-2: 82% ->
+  90% of peak).
+- LTX-2 also spends 176 ms in small all-reduces (global qk RMS norm, audio TP)
+  and idles 13% between steps -- separate, cheaper targets.
+
+Revised order: **P0** -- assemble the AG output once (collect the hops,
+concatenate in ring-relative order, one reorder by the runtime ring position)
+instead of eight `dynamic_update_slice`s; same math, so outputs must be
+bit-identical; measure with `sp_pair_probe.py` and a Wan / LTX-2 bench A/B. If
+P0 removes most of the write-back, the kernel's remaining target is the 3-7%
+wait and the kernel is likely not worth it; otherwise continue with the Pallas
+AG kernel below, which writes each tile of the product in place.
+
 ## Steps
 
 | # | Work | Validation | Estimate |

From 01d027fe076d93a6a241c82de04430447aa3176a Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Tue, 10 Nov 2026 22:24:02 -0800
Subject: [PATCH 14/22] perf(tpu): Pallas all-gather + matmul kernel for the SP
 rings (DIFFLET_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>
---
 benchmark/tpu_profile/cm_pallas_bench.py      | 172 +++++++
 benchmark/tpu_profile/pallas_matmul_sweep.py  | 101 ++++
 benchmark/tpu_profile/sp_step_drift.py        |  54 +++
 .../tpu/ops_impl/collective_matmul.py         | 115 ++++-
 .../tpu/ops_impl/collective_matmul_pallas.py  | 436 ++++++++++++++++++
 .../test_tpu_collective_matmul_ag_assemble.py |  65 +++
 .../test_tpu_collective_matmul_pallas.py      |  70 +++
 7 files changed, 1009 insertions(+), 4 deletions(-)
 create mode 100644 benchmark/tpu_profile/cm_pallas_bench.py
 create mode 100644 benchmark/tpu_profile/pallas_matmul_sweep.py
 create mode 100644 benchmark/tpu_profile/sp_step_drift.py
 create mode 100644 difflet/backends/tpu/ops_impl/collective_matmul_pallas.py
 create mode 100644 tests/unit/backends/test_tpu_collective_matmul_ag_assemble.py
 create mode 100644 tests/unit/backends/test_tpu_collective_matmul_pallas.py

diff --git a/benchmark/tpu_profile/cm_pallas_bench.py b/benchmark/tpu_profile/cm_pallas_bench.py
new file mode 100644
index 00000000..f3ba9121
--- /dev/null
+++ b/benchmark/tpu_profile/cm_pallas_bench.py
@@ -0,0 +1,172 @@
+"""All-gather + matmul on the chips: JAX ring (dus / concat) vs the Pallas kernel.
+
+    python -m benchmark.tpu_profile.cm_pallas_bench [--iters 50] [--shapes wan_qkv,ltx_ffn]
+
+Plain JAX, one process owning the four chips (no torch_tpu): the JAX rings are
+``collective_matmul._op`` with ``torch.tpu.pallas.jax_op`` stubbed out (as in
+cm_pure_jax), the kernel is ``collective_matmul_pallas.ag_matmul_body``. For
+each shape, prints the time per call and the kernel's difference from the dus
+ring (bf16 inputs, fp32 accumulation in both). Step 2 of
+docs/plans/2026-11-10-tpu-collective-matmul-pallas.md.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+import types
+
+#: name -> (rows per device, K, output columns per device of each weight)
+SHAPES = {
+    "wan_qkv": (1170, 5120, (1280, 1280, 1280)),
+    "wan_ffn": (1170, 5120, (3456,)),
+    "ltx_qkv": (578, 4096, (1024, 1024, 1024)),
+    "ltx_ffn": (578, 4096, (4096,)),
+}
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--iters", type=int, default=50)
+    parser.add_argument("--shapes", default=",".join(SHAPES))
+    parser.add_argument("--tn", type=int, default=512)
+    parser.add_argument("--ablate", action="store_true",
+                        help="also time the kernel without comm / without compute, and one-chip XLA einsums")
+    parser.add_argument("--trace", default="", help="jax.profiler trace dir for the kernel calls")
+    args = parser.parse_args()
+
+    import jax
+    import jax.numpy as jnp
+    import numpy as np
+    import torch
+    from jax.sharding import NamedSharding
+    from jax.sharding import PartitionSpec as P
+
+    torch.tpu = types.SimpleNamespace(pallas=types.SimpleNamespace(
+        jax_op=lambda name, fn, **kw: fn))
+    from difflet.backends.tpu.ops_impl import collective_matmul as cm
+    from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kp
+
+    devices = jax.devices()
+    n = len(devices)
+    mesh = jax.sharding.Mesh(np.array(cm._snake(devices)), ("tp",))
+    shard_of_pos = tuple(range(n))  # one process: ring position p holds shard p
+    cm._STATE = {"mesh": mesh, "tp": n, "shard_of_pos": shard_of_pos, "ids": None, "ops": {}}
+
+    def put(a, spec, dtype=jnp.bfloat16):
+        return jax.device_put(jnp.asarray(a, dtype), NamedSharding(mesh, spec))
+
+    def timed(f, *xs):
+        out = f(*xs)
+        jax.block_until_ready(out)
+        for _ in range(3):
+            jax.block_until_ready(f(*xs))
+        t = time.perf_counter()
+        for _ in range(args.iters):
+            out = f(*xs)
+        jax.block_until_ready(out)
+        return out, (time.perf_counter() - t) / args.iters * 1e3
+
+    print(f"jax {jax.__version__} {devices[0].device_kind} x{n}, iters {args.iters}, tn<={args.tn}",
+          flush=True)
+    ids = put(np.arange(n), P("tp"), jnp.int32)
+    rng = np.random.default_rng(0)
+    for name in args.shapes.split(","):
+        rows, k, cols = SHAPES[name]
+        x = put(rng.standard_normal((1, rows * n, k), dtype=np.float32), P(None, "tp", None))
+        ws = [put(rng.standard_normal((c * n, k), dtype=np.float32) / k ** 0.5, P("tp", None))
+              for c in cols]
+        res = {}
+        for mode in ("dus", "concat"):
+            os.environ[cm.AG_ASSEMBLE_ENV] = mode
+            ring = jax.jit(cm._op("ag", len(ws)))
+            out, ms = timed(ring, x, ids, *ws)
+            res[mode] = (out if isinstance(out, tuple) else (out,), ms)
+
+        def body(x, ids, *ws):
+            owners, nbrs = kp.ring_tables(ids[0], shard_of_pos, n)
+            return kp.ag_matmul_body(x, ws, owners, nbrs, n=n, tn=args.tn)
+
+        kern = jax.jit(jax.shard_map(
+            body, mesh=mesh, in_specs=(P(None, "tp", None), P("tp")) + (P("tp", None),) * len(ws),
+            out_specs=(P(None, None, "tp"),) * len(ws), check_vma=False))
+        try:
+            res["pallas"] = timed(kern, x, ids, *ws)
+            if args.trace:
+                with jax.profiler.trace(f"{args.trace}/{name}"):
+                    for _ in range(5):
+                        jax.block_until_ready(kern(x, ids, *ws))
+        except Exception as exc:  # noqa: BLE001 -- report and go on to the next shape
+            print(f"{name}: pallas failed: {type(exc).__name__}: {str(exc)[:600]}", flush=True)
+            continue
+        if args.ablate:
+            extra = []
+            for ab in ("comm", "compute"):
+                def body_ab(x, ids, *ws, ab=ab):
+                    owners, nbrs = kp.ring_tables(ids[0], shard_of_pos, n)
+                    return kp.ag_matmul_body(x, ws, owners, nbrs, n=n, tn=args.tn, ablate=ab)
+                f = jax.jit(jax.shard_map(
+                    body_ab, mesh=mesh,
+                    in_specs=(P(None, "tp", None), P("tp")) + (P("tp", None),) * len(ws),
+                    out_specs=(P(None, None, "tp"),) * len(ws), check_vma=False))
+                extra.append(f"no-{ab} {timed(f, x, ids, *ws)[1]:.3f}")
+            # The same FLOPs on one chip, plain XLA: the full sequence times each local weight.
+            d0 = devices[0]
+            xf = jax.device_put(jnp.zeros((1, rows * n, k), jnp.bfloat16), d0)
+            wl = [jax.device_put(jnp.zeros((c, k), jnp.bfloat16), d0) for c in cols]
+            ein = jax.jit(lambda x, *ws: tuple(jnp.einsum("bsk,nk->bsn", x, w) for w in ws))
+            extra.append(f"xla-1chip {timed(ein, xf, *wl)[1]:.3f}")
+            print(f"{name:8s} ablation (ms): " + ", ".join(extra), flush=True)
+        # v2: weights pre-transposed to (K, N) once, outside the timed call.
+        wts = [jax.device_put(w.T, NamedSharding(mesh, P(None, "tp"))) for w in ws]
+
+        def body_kn(x, ids, *ws):
+            owners, nbrs = kp.ring_tables(ids[0], shard_of_pos, n)
+            return kp.ag_matmul_body(x, ws, owners, nbrs, n=n, tn=args.tn, weights_kn=True)
+
+        kern_kn = jax.jit(jax.shard_map(
+            body_kn, mesh=mesh, in_specs=(P(None, "tp", None), P("tp")) + (P(None, "tp"),) * len(ws),
+            out_specs=(P(None, None, "tp"),) * len(ws), check_vma=False))
+        res["pallas-kn"] = timed(kern_kn, x, ids, *wts)
+
+        # v3: halves resident in VMEM, remote DMA VMEM -> VMEM, pipelines stream weights only.
+        def body_res(x, ids, *ws, ab="", manual=False):
+            owners, nbrs = kp.ring_tables(ids[0], shard_of_pos, n)
+            return kp.ag_matmul_body(x, ws, owners, nbrs, n=n, tn=args.tn, weights_kn=True,
+                                     resident=True, manual=manual, ablate=ab)
+
+        def mk(ab="", manual=False):
+            return jax.jit(jax.shard_map(
+                lambda x, ids, *ws: body_res(x, ids, *ws, ab=ab, manual=manual), mesh=mesh,
+                in_specs=(P(None, "tp", None), P("tp")) + (P(None, "tp"),) * len(ws),
+                out_specs=(P(None, None, "tp"),) * len(ws), check_vma=False))
+        try:
+            res["pallas-res"] = timed(mk(), x, ids, *wts)
+            if args.ablate:
+                print(f"{name:8s} resident ablation (ms): no-comm {timed(mk('comm'), x, ids, *wts)[1]:.3f}",
+                      flush=True)
+        except Exception as exc:  # noqa: BLE001
+            print(f"{name}: pallas-res failed: {type(exc).__name__}: {str(exc)[:800]}", flush=True)
+        # v4: one hand-written weight-tile stream over every (hop, weight, tile).
+        try:
+            res["pallas-man"] = timed(mk(manual=True), x, ids, *wts)
+            if args.ablate:
+                print(f"{name:8s} manual ablation (ms): no-comm "
+                      f"{timed(mk('comm', manual=True), x, ids, *wts)[1]:.3f}", flush=True)
+        except Exception as exc:  # noqa: BLE001
+            print(f"{name}: pallas-man failed: {type(exc).__name__}: {str(exc)[:800]}", flush=True)
+        ref = [np.asarray(o.astype(jnp.float32)) for o in res["dus"][0]]
+        line = [f"{name:8s} rows {rows} K {k} N {cols}:"]
+        for mode in [m for m in ("dus", "concat", "pallas", "pallas-kn", "pallas-res", "pallas-man") if m in res]:
+            outs, ms = res[mode]
+            err = max(float(np.abs(np.asarray(o.astype(jnp.float32)) - r).max() / np.abs(r).max())
+                      for o, r in zip(outs, ref))
+            line.append(f"{mode} {ms:.3f} ms (max|d|/max {err:.1e})")
+        print("  ".join(line), flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/benchmark/tpu_profile/pallas_matmul_sweep.py b/benchmark/tpu_profile/pallas_matmul_sweep.py
new file mode 100644
index 00000000..cfd0303b
--- /dev/null
+++ b/benchmark/tpu_profile/pallas_matmul_sweep.py
@@ -0,0 +1,101 @@
+"""How fast a plain Pallas matmul gets on one v5e chip at the SP ring's shapes.
+
+    python -m benchmark.tpu_profile.pallas_matmul_sweep [--iters 50]
+
+The all-gather + matmul kernel (collective_matmul_pallas) hides its
+communication fully, but its per-tile matmuls run 20-45% slower than XLA's
+einsum at the same FLOPs. This isolates the matmul: out (M, N) = x (M, K) @
+w (K, N), x held in VMEM, w streamed in (tk, tn) tiles, fp32 accumulation over
+K, R = 16 distinct weights per call (one call is ~100 us of dispatch on its own),
+for M in {the ring half, a multiple of 8 above it, both halves together}
+and a few tile shapes, next to XLA. No communication, one chip.
+"""
+
+from __future__ import annotations
+
+import argparse
+import itertools
+import sys
+import time
+
+#: name -> (M candidates, K, N)
+SHAPES = {
+    "wan": ((585, 592, 1170, 1184), 5120, 1280),
+    "ltx": ((289, 296, 578, 592), 4096, 1024),
+}
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--iters", type=int, default=20)
+    parser.add_argument("--reps", type=int, default=16, help="matmuls (distinct weights) per call")
+    args = parser.parse_args()
+
+    import jax
+    import jax.numpy as jnp
+    from jax.experimental import pallas as pl
+    from jax.experimental.pallas import tpu as pltpu
+
+    def timed(f, *xs):
+        jax.block_until_ready(f(*xs))
+        for _ in range(3):
+            jax.block_until_ready(f(*xs))
+        t = time.perf_counter()
+        for _ in range(args.iters):
+            out = f(*xs)
+        jax.block_until_ready(out)
+        return (time.perf_counter() - t) / args.iters / args.reps * 1e3  # per matmul
+
+    def pallas_mm(m, k, n, tn, tk):
+        nk = k // tk
+
+        def kernel(x_ref, w_ref, o_ref, acc_ref):
+            kk = pl.program_id(2)
+
+            @pl.when(kk == 0)
+            def _():
+                acc_ref[...] = jnp.zeros_like(acc_ref)
+
+            acc_ref[...] += jnp.dot(x_ref[:, pl.ds(pl.multiple_of(kk * tk, tk), tk)], w_ref[...],
+                                    preferred_element_type=jnp.float32)
+
+            @pl.when(kk == nk - 1)
+            def _():
+                o_ref[...] = acc_ref[...].astype(o_ref.dtype)
+
+        return pl.pallas_call(
+            kernel,
+            out_shape=jax.ShapeDtypeStruct((args.reps, m, n), jnp.bfloat16),
+            grid=(args.reps, n // tn, nk),
+            in_specs=[pl.BlockSpec((m, k), lambda r, j, kk: (0, 0)),  # x resident (index constant)
+                      pl.BlockSpec((pl.Squeezed(), tk, tn), lambda r, j, kk: (r, kk, j))],
+            out_specs=pl.BlockSpec((pl.Squeezed(), m, tn), lambda r, j, kk: (r, 0, j)),
+            scratch_shapes=[pltpu.VMEM((m, tn), jnp.float32)],
+            compiler_params=pltpu.CompilerParams(
+                dimension_semantics=("arbitrary", "parallel", "arbitrary"), vmem_limit_bytes=100 * 2**20),
+        )
+
+    dev = jax.devices()[0]
+    print(f"jax {jax.__version__} {dev.device_kind}, iters {args.iters}", flush=True)
+    for name, (ms, k, n) in SHAPES.items():
+        w = jax.device_put(jnp.ones((args.reps, k, n), jnp.bfloat16) * 0.01, dev)
+        for m in ms:
+            x = jax.device_put(jnp.ones((m, k), jnp.bfloat16), dev)
+            flops = 2 * m * k * n
+            xla = timed(jax.jit(lambda x, w: jnp.einsum("mk,rkn->rmn", x, w)), x, w)
+            line = [f"{name} M {m:4d} K {k} N {n}: xla {xla * 1e3:7.1f} us "
+                    f"({flops / xla / 1e9:5.1f} TF/s)"]
+            for tn, tk in itertools.product((256, 512), (k, 1024, 512)):
+                if n % tn or k % tk:
+                    continue
+                try:
+                    t = timed(jax.jit(pallas_mm(m, k, n, tn, tk)), x, w)
+                    line.append(f"tn{tn}/tk{tk} {t * 1e3:7.1f}")
+                except Exception as exc:  # noqa: BLE001
+                    line.append(f"tn{tn}/tk{tk} FAIL({type(exc).__name__})")
+            print("  ".join(line), flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/benchmark/tpu_profile/sp_step_drift.py b/benchmark/tpu_profile/sp_step_drift.py
new file mode 100644
index 00000000..5c6a5c6f
--- /dev/null
+++ b/benchmark/tpu_profile/sp_step_drift.py
@@ -0,0 +1,54 @@
+"""How far an SP variant's output drifts after k denoising steps (TPU).
+
+    DIFFLET_TPU_SP=1 DIFFLET_TPU_SP_KERNEL=pallas \\
+        python -m benchmark.tpu_profile.sp_step_drift --model wan_2_1 --steps 1,4 --out DIR --tag pallas
+
+Runs main's TPU bench adapter on the matrix entry with only the step count
+changed (a local copy of the config; ``benchmark/models.py`` is untouched) and
+saves the first synced request's frames to ``DIR/<tag>-s<k>``. Run it once per
+variant (TP, SP + JAX ring, SP + Pallas kernel; the variant comes from the
+environment, as for the bench), then compare the frames with ``image_psnr``.
+
+The question it answers (2026-11-11): after 20 steps the Pallas AG kernel's
+video differs from the JAX ring's by 26 dB (Wan) / 31 dB (LTX-2) although the
+op is bit-identical in isolation. If the 1-step outputs differ by about as much
+as TP vs the ring, the 20-step gap is trajectory amplification of rounding
+differences in the surrounding fusions; if they already differ by much more,
+something in the model computes differently and needs bisecting.
+"""
+
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import sys
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", default="wan_2_1")
+    parser.add_argument("--steps", default="1,4")
+    parser.add_argument("--out", required=True)
+    parser.add_argument("--tag", required=True)
+    parser.add_argument("--seed", type=int, default=None, help="override the matrix seed")
+    args = parser.parse_args()
+
+    from benchmark.adapters.tpu import TpuAdapter
+    from benchmark.models import MATRIX
+
+    for k in (int(s) for s in args.steps.split(",")):
+        cfg = dataclasses.replace(MATRIX[args.model], steps=k)
+        if args.seed is not None:
+            cfg = dataclasses.replace(cfg, seed=args.seed)
+        adapter = TpuAdapter(save_dir=f"{args.out}/{args.tag}-s{k}")
+        try:
+            adapter.prepare(cfg)
+            g = adapter.run_generate(cfg)
+            print(f"[drift] {args.tag} steps={k}: output {g.get('output')}", flush=True)
+        finally:
+            adapter.shutdown()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/difflet/backends/tpu/ops_impl/collective_matmul.py b/difflet/backends/tpu/ops_impl/collective_matmul.py
index 4dccb29b..091970d3 100644
--- a/difflet/backends/tpu/ops_impl/collective_matmul.py
+++ b/difflet/backends/tpu/ops_impl/collective_matmul.py
@@ -56,6 +56,69 @@
 
 SP_ENV = "DIFFLET_TPU_SP"
 
+#: How the all-gather ring builds its output. ``dus`` (default): each hop's
+#: products are written into the full output with ``dynamic_update_slice``.
+#: ``concat``: every hop's products are kept and the output is built once at
+#: the end (a gather by the runtime ring position). Same values; the profile of
+#: 2026-11-10 put the per-hop writes at 7-11% of an SP step (Wan, LTX-2).
+AG_ASSEMBLE_ENV = "DIFFLET_TPU_SP_AG_ASSEMBLE"
+
+
+def ag_assemble_mode() -> str:
+    mode = os.environ.get(AG_ASSEMBLE_ENV, "dus").strip().lower() or "dus"
+    if mode not in ("dus", "concat"):
+        raise ValueError(f"{AG_ASSEMBLE_ENV} must be 'dus' or 'concat', got {mode!r}")
+    return mode
+
+
+#: What runs the all-gather + matmul (overrides the model's default, see
+#: :func:`set_kernel_default`). ``jax``: the ring below, XLA schedules it.
+#: ``pallas``: ``collective_matmul_pallas``'s kernel (VMEM-resident halves, remote
+#: DMA overlapped with a streamed matmul, each product tile written once at its
+#: final rows); shapes it does not take fall back to the JAX ring with a notice. Op-level, v5e-4, 2026-11-11: Wan q/k/v 1.259 -> 1.185
+#: ms, Wan FFN up 1.050 -> 1.056, LTX-2 q/k/v 0.445 -> 0.406 ms, outputs equal.
+SP_KERNEL_ENV = "DIFFLET_TPU_SP_KERNEL"
+
+
+#: What runs the all-gather + matmul when ``SP_KERNEL_ENV`` is unset; a model's SP
+#: ``install`` sets it (:func:`set_kernel_default`) before :func:`prepare`. The kernel
+#: pays where the ring's ``dynamic_update_slice`` is costly -- Wan -8.3%, FLUX -10.2%,
+#: HunyuanVideo (with single-stream SP) -4.5% per step; Qwen-Image +1%, LTX-2 0%
+#: (v5e-4, 2026-11-11) -- so it is per model.
+_KERNEL_DEFAULT = "jax"
+
+
+def set_kernel_default(mode: str) -> None:
+    """The model's default for :func:`sp_kernel_mode` (``SP_KERNEL_ENV`` still wins)."""
+    global _KERNEL_DEFAULT
+    if mode not in ("jax", "pallas"):
+        raise ValueError(f"kernel default must be 'jax' or 'pallas', got {mode!r}")
+    _KERNEL_DEFAULT = mode
+
+
+def sp_kernel_mode() -> str:
+    mode = os.environ.get(SP_KERNEL_ENV, "").strip().lower() or _KERNEL_DEFAULT
+    if mode not in ("jax", "pallas"):
+        raise ValueError(f"{SP_KERNEL_ENV} must be 'jax' or 'pallas', got {mode!r}")
+    return mode
+
+
+def _source_tag() -> str:
+    """A short hash of this module's and the Pallas kernel's source, put in every
+    op name. torch_tpu's compilation cache keys a ``jax_op`` by its name and
+    signature, not by what it lowers to: after the kernel learned weights of
+    different widths, FLUX's single-block ag4 kept running the cached JAX-ring
+    program under the unchanged name (2026-11-11, num_compilations=0)."""
+    import hashlib
+    from pathlib import Path
+
+    here = Path(__file__).resolve().parent
+    h = hashlib.sha256()
+    for name in ("collective_matmul.py", "collective_matmul_pallas.py"):
+        h.update((here / name).read_bytes())
+    return h.hexdigest()[:8]
+
+
 _STATE: dict | None = None
 
 
@@ -145,7 +208,8 @@ def _op(kind: str, n_weights: int):
     """The jax_op for ``kind`` built once: 'rs' over ``n_weights`` (input, weight)
     pairs whose products are summed, 'ag' over ``n_weights`` weights sharing one
     gathered input."""
-    key = (kind, n_weights)
+    key = (kind, n_weights,
+           (ag_assemble_mode(), sp_kernel_mode()) if kind == "ag" else None)
     ops = _STATE["ops"]
     if key in ops:
         return ops[key]
@@ -158,6 +222,8 @@ def _op(kind: str, n_weights: int):
 
     mesh, n = _STATE["mesh"], _STATE["tp"]
     shard_table = list(_STATE["shard_of_pos"])
+    assemble = ag_assemble_mode()
+    use_kernel = kind == "ag" and sp_kernel_mode() == "pallas"
     cw = [(i, (i + 1) % n) for i in range(n)]
     ccw = [(i, (i - 1) % n) for i in range(n)]
     shard_map = jax.shard_map if hasattr(jax, "shard_map") else \
@@ -226,11 +292,47 @@ def fn(x0, w0, x1, w1, ids):
         def body(x, ids, *ws):
             shard_of = jnp.asarray(shard_table, dtype=jnp.int32)
             j = ids[0]
+            if use_kernel:
+                from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kp
+
+                why = kp.unsupported(x.shape, [w.shape for w in ws], x.dtype.itemsize)
+                if why is None:
+                    # The ring position stays a runtime value (rule above): the
+                    # kernel's owners / neighbours are computed from it.
+                    owners, nbrs = kp.ring_tables(j, shard_table, n)
+                    outs = kp.ag_matmul_body(x, ws, owners, nbrs, n=n, resident=True,
+                                             manual=True)
+                    return outs if len(outs) > 1 else outs[0]
+                print(f"[difflet tpu] {SP_KERNEL_ENV}=pallas: ag x{len(ws)} {tuple(x.shape)} "
+                      f"falls back to the JAX ring ({why})", flush=True)
             rows = x.shape[1]
             hr = rows // 2
             seq = rows * n
-            outs = [jnp.zeros((x.shape[0], seq, w.shape[0]), x.dtype) for w in ws]
             top, bot = x[:, :hr], x[:, hr:]
+            if assemble == "concat":
+                # Keep every hop's products and build each output once at the end
+                # (see AG_ASSEMBLE_ENV): rank r's top half was computed at hop
+                # (j - pos_of[r]) % n, its bottom half at hop (pos_of[r] - j) % n.
+                tops = [[] for _ in ws]
+                bots = [[] for _ in ws]
+                for step in range(n):
+                    if step:
+                        top = jax.lax.ppermute(top, "tp", cw)
+                        bot = jax.lax.ppermute(bot, "tp", ccw)
+                    for i, w in enumerate(ws):
+                        tops[i].append(jnp.einsum("bsk,nk->bsn", top, w))
+                        bots[i].append(jnp.einsum("bsk,nk->bsn", bot, w))
+                pos_of = jnp.asarray([shard_table.index(r) for r in range(n)], dtype=jnp.int32)
+                idx_top, idx_bot = (j - pos_of) % n, (pos_of - j) % n
+                outs = []
+                for i in range(len(ws)):
+                    t = jnp.take(jnp.stack(tops[i]), idx_top, axis=0)  # (n, B, hr, N)
+                    b = jnp.take(jnp.stack(bots[i]), idx_bot, axis=0)
+                    o = jnp.stack([t, b], axis=1)                      # (n, 2, B, hr, N)
+                    outs.append(jnp.transpose(o, (2, 0, 1, 3, 4)).reshape(
+                        x.shape[0], seq, o.shape[-1]))
+                return tuple(outs) if len(outs) > 1 else outs[0]
+            outs = [jnp.zeros((x.shape[0], seq, w.shape[0]), x.dtype) for w in ws]
             for step in range(n):
                 if step:
                     top = jax.lax.ppermute(top, "tp", cw)
@@ -282,7 +384,11 @@ def fn(x, ids, w0, w1, w2, w3):
                                          else jax.Array)}
 
     op = torch.tpu.pallas.jax_op(
-        f"difflet::sp_{kind}_matmul_{n_weights}", fn, mesh=mesh,
+        f"difflet::sp_{kind}_matmul_{n_weights}"
+        + ("_cat" if kind == "ag" and assemble == "concat" else "")
+        + ("_pk" if use_kernel else "")
+        + "_" + _source_tag(),
+        fn, mesh=mesh,
         input_partition_specs=in_specs,
         # A tuple of specs, one per output: see the module docstring.
         output_partition_specs=tuple(out_specs),
@@ -328,6 +434,7 @@ def supports(seq_len: int, tp: int) -> bool:
     return seq_len % (2 * tp) == 0
 
 
-__all__ = ["DEFAULT_OPS", "SP_ENV", "all_gather_matmul", "matmul2_reduce_scatter",
+__all__ = ["DEFAULT_OPS", "SP_ENV", "SP_KERNEL_ENV", "all_gather_matmul", "matmul2_reduce_scatter",
+           "set_kernel_default",
            "matmul_reduce_scatter", "prepare",
            "require_compiled", "requested", "supports"]
diff --git a/difflet/backends/tpu/ops_impl/collective_matmul_pallas.py b/difflet/backends/tpu/ops_impl/collective_matmul_pallas.py
new file mode 100644
index 00000000..a0f57a07
--- /dev/null
+++ b/difflet/backends/tpu/ops_impl/collective_matmul_pallas.py
@@ -0,0 +1,436 @@
+"""Pallas TPU kernel for the SP all-gather + column-parallel matmul ring.
+
+Step 2 of docs/plans/2026-11-10-tpu-collective-matmul-pallas.md: the same
+bidirectional ring as ``collective_matmul``'s JAX ``ag`` op, but each hop's
+product tile goes from VMEM straight to its final rows of the output (one HBM
+write), with no temporary product and no ``dynamic_update_slice`` / gather
+assembly afterwards. The remote copy of the next half overlaps the matmul of
+the current one.
+
+``ag_matmul_body`` is the per-device body, called inside ``shard_map`` over a
+1-D mesh axis (default ``"tp"``) whose order is the ring order. Rank-dependent
+values (the ring position, owners, neighbours) are runtime int32 inputs, never
+Python constants -- rule 1 of the plan; every device runs the same program.
+
+Layout per device (x is this device's sequence shard, split in two halves that
+travel in opposite directions):
+
+    x     (B, R, K)            R = rows per device, even
+    w_i   (N_i, K)             this device's output columns
+    out_i (B, n * R, N_i)      viewed as (B, n, 2, R/2, N_i): [owner shard, half]
+
+The ``pallas_call``s are not marked ``has_side_effects``: through torch_tpu's
+``jax_op`` the program goes through XLA's SPMD partitioner, which rejects a
+side-effecting custom call without a sharding ("Side-effect HLO must have
+sharding", 2026-11-11). The outputs are always consumed, so nothing is dropped.
+
+Communication buffers are an extra output (2 directions x 3 slots x the half
+shape) in HBM. Three slots make the reuse race-free: a neighbour is at most one
+hop behind (each device waits for both neighbours' previous hop before going
+on), so the slot written at hop s+1 is never the one a neighbour still reads.
+"""
+
+from __future__ import annotations
+
+SLOTS = 3
+
+
+def ag_matmul_body(x, ws, owners, nbrs, *, n: int, axis: str = "tp", tn: int = 256,
+                   vmem_limit_bytes: int = 96 * 2**20, interpret=False, ablate: str = "",
+                   weights_kn: bool = False, resident: bool = False,
+                   manual: bool = False):
+    """Per-device all-gather + matmul with every weight in ``ws``.
+
+    owners: int32 (2, n) -- owners[0, s] / owners[1, s] is the shard (output row
+        block) whose top / bottom half this device holds after s hops.
+    nbrs:   int32 (2,)   -- mesh index along ``axis`` of the clockwise
+        neighbour (receives the top halves) and the counter-clockwise one.
+    weights_kn: the weights are given transposed, (K, N_i): the kernel then runs a
+        plain (M, K) x (K, N) matmul per tile instead of contracting both
+        operands on their last axis.
+    resident: keep the halves in VMEM for the whole call -- the local ones are
+        copied in once, the remote ones arrive by VMEM-to-VMEM DMA (they never
+        touch HBM), and each matmul pipeline streams only weight tiles. Needs
+        2 x 3 x B x R/2 x K elements of VMEM (Wan: 36 MB).
+    manual: with ``resident``, replace the per-(hop, weight) ``emit_pipeline``s by
+        one hand-written double-buffered stream over every (hop, weight, column
+        tile): the next weight tile is always in flight, across weights and hops,
+        so only the very first tile load is exposed. Output tiles are written back
+        asynchronously, two in flight. The tiles of a hop are a runtime loop, so
+        compile time does not grow with the tile count (a fully unrolled stream
+        took 225 s to compile at Wan's FFN shape). Weights may differ in width;
+        one tile width (a multiple of 128) must divide them all.
+    ablate: profiling only -- ``"comm"`` skips the remote copies (every hop
+        multiplies the local halves; wrong values), ``"compute"`` skips the matmuls.
+    Returns a tuple of (B, n * R, N_i), one per weight.
+    """
+    import jax
+    import jax.numpy as jnp
+    from jax.experimental import pallas as pl
+    from jax.experimental.pallas import tpu as pltpu
+
+    b, rows, k = x.shape
+    if rows % 2:
+        raise ValueError(f"rows per device must be even (two ring halves), got {rows}")
+    hr = rows // 2
+    tns = []
+    kd, nd = (0, 1) if weights_kn else (1, 0)  # weight axis of K / of the columns
+    cols = []
+    for w in ws:
+        if w.shape[kd] != k:
+            raise ValueError(f"weight {w.shape} does not match K={k}")
+        cols.append(w.shape[nd])
+        tns.append(_col_tile(w.shape[nd], tn))
+    nw = len(ws)
+    tm0 = _common_tile(cols, tn)  # the manual stream's one tile width
+
+    def kernel(owners_ref, nbrs_ref, x_ref, *refs):
+        w_refs = refs[:nw]
+        out_refs = refs[nw:2 * nw]
+        buf_ref, send_sem, recv_sem = refs[2 * nw:]
+
+        # Every neighbour must be inside this kernel (its buffers live) before
+        # anything is written into them.
+        barrier = pltpu.get_barrier_semaphore()
+        for d in range(2):
+            pl.semaphore_signal(barrier, 1, device_id={axis: nbrs_ref[d]})
+        pl.semaphore_wait(barrier, 2)
+
+        def held(d, s):
+            # The half travelling in direction d after s hops.
+            return x_ref.at[:, d] if s == 0 or ablate == "comm" else buf_ref.at[d, s % SLOTS]
+
+        def matmul(s, w_ref, out_ref, tn_i):
+            # Both halves held after s hops against one weight tile: each tile
+            # of W is read from HBM once per hop, not once per half.
+            own_top, own_bot = owners_ref[0, s], owners_ref[1, s]
+
+            def body(xt_blk, xb_blk, w_blk, ot_blk, ob_blk):
+                for x_blk, o_blk in ((xt_blk, ot_blk), (xb_blk, ob_blk)):
+                    o_blk[...] = jax.lax.dot_general(
+                        x_blk[...], w_blk[...], (((1,), (kd,)), ((), ())),
+                        preferred_element_type=jnp.float32).astype(o_blk.dtype)
+
+            x_spec = pl.BlockSpec((pl.Squeezed(), hr, k), lambda bi, j: (bi, 0, 0))
+            o_shape = (pl.Squeezed(), pl.Squeezed(), pl.Squeezed(), hr, tn_i)
+            w_spec = (pl.BlockSpec((k, tn_i), lambda bi, j: (0, j)) if weights_kn
+                      else pl.BlockSpec((tn_i, k), lambda bi, j: (j, 0)))
+            pltpu.emit_pipeline(
+                body,
+                grid=(b, w_ref.shape[nd] // tn_i),
+                in_specs=[x_spec, x_spec, w_spec],
+                out_specs=[pl.BlockSpec(o_shape, lambda bi, j: (bi, own_top, 0, 0, j)),
+                           pl.BlockSpec(o_shape, lambda bi, j: (bi, own_bot, 1, 0, j))],
+            )(held(0, s), held(1, s), w_ref, out_ref, out_ref)
+
+        for s in range(n):
+            copies = []
+            if s + 1 < n and ablate != "comm":
+                slot = (s + 1) % SLOTS
+                for d in range(2):
+                    c = pltpu.make_async_remote_copy(
+                        held(d, s), buf_ref.at[d, slot], send_sem.at[d, slot],
+                        recv_sem.at[d, slot], device_id={axis: nbrs_ref[d]})
+                    c.start()
+                    copies.append(c)
+            for w_ref, out_ref, tn_i in zip(w_refs, out_refs, tns):
+                if ablate != "compute":
+                    matmul(s, w_ref, out_ref, tn_i)
+            for c in copies:
+                c.wait_send()
+                c.wait_recv()
+
+    def kernel_resident(owners_ref, nbrs_ref, x_ref, *refs):
+        w_refs = refs[:nw]
+        out_refs = refs[nw:2 * nw]
+        xbuf, load_sem, send_sem, recv_sem = refs[2 * nw:]
+
+        barrier = pltpu.get_barrier_semaphore()
+        for d in range(2):
+            pl.semaphore_signal(barrier, 1, device_id={axis: nbrs_ref[d]})
+        pl.semaphore_wait(barrier, 2)
+
+        loads = [pltpu.make_async_copy(x_ref.at[:, d], xbuf.at[d, 0], load_sem.at[d])
+                 for d in range(2)]
+        for c in loads:
+            c.start()
+        for c in loads:
+            c.wait()
+
+        def held(d, s):
+            return xbuf.at[d, 0 if ablate == "comm" else s % SLOTS]
+
+        def matmul(s, w_ref, out_ref, tn_i):
+            own_top, own_bot = owners_ref[0, s], owners_ref[1, s]
+            xt, xb = held(0, s), held(1, s)
+
+            def body(w_blk, ot_blk, ob_blk):
+                w = w_blk[...]
+                for x_r, o_blk in ((xt, ot_blk), (xb, ob_blk)):
+                    for bi in range(b):
+                        o_blk[bi] = jax.lax.dot_general(
+                            x_r[bi], w, (((1,), (kd,)), ((), ())),
+                            preferred_element_type=jnp.float32).astype(o_blk.dtype)
+
+            o_shape = (b, pl.Squeezed(), pl.Squeezed(), hr, tn_i)
+            w_spec = (pl.BlockSpec((k, tn_i), lambda j: (0, j)) if weights_kn
+                      else pl.BlockSpec((tn_i, k), lambda j: (j, 0)))
+            pltpu.emit_pipeline(
+                body,
+                grid=(w_ref.shape[nd] // tn_i,),
+                in_specs=[w_spec],
+                out_specs=[pl.BlockSpec(o_shape, lambda j: (0, own_top, 0, 0, j)),
+                           pl.BlockSpec(o_shape, lambda j: (0, own_bot, 1, 0, j))],
+            )(w_ref, out_ref, out_ref)
+
+        for s in range(n):
+            copies = []
+            if s + 1 < n and ablate != "comm":
+                slot = (s + 1) % SLOTS
+                for d in range(2):
+                    c = pltpu.make_async_remote_copy(
+                        held(d, s), xbuf.at[d, slot], send_sem.at[d, slot],
+                        recv_sem.at[d, slot], device_id={axis: nbrs_ref[d]})
+                    c.start()
+                    copies.append(c)
+            for w_ref, out_ref, tn_i in zip(w_refs, out_refs, tns):
+                if ablate != "compute":
+                    matmul(s, w_ref, out_ref, tn_i)
+            for c in copies:
+                c.wait_send()
+                c.wait_recv()
+
+    def kernel_manual(owners_ref, nbrs_ref, x_ref, *refs):
+        w_refs = refs[:nw]
+        out_refs = refs[nw:2 * nw]
+        xbuf, wbuf, obuf, load_sem, send_sem, recv_sem, w_sem, o_sem = refs[2 * nw:]
+        tn0 = tm0                    # one column-tile width for every weight
+        cum = [0]                    # first tile of each weight within a hop
+        for c in cols:
+            cum.append(cum[-1] + c // tn0)
+        tph = cum[-1]                # tiles per hop
+        total = n * tph
+
+        barrier = pltpu.get_barrier_semaphore()
+        for d in range(2):
+            pl.semaphore_signal(barrier, 1, device_id={axis: nbrs_ref[d]})
+        pl.semaphore_wait(barrier, 2)
+
+        def cols_of(j):
+            return pl.ds(pl.multiple_of(j * tn0, tn0), tn0)
+
+        def w_src(ii, j):
+            return w_refs[ii].at[:, cols_of(j)] if weights_kn else w_refs[ii].at[cols_of(j), :]
+
+        def which(t):
+            # Global tile t -> (weight, column tile of that weight). Weights may
+            # differ in width (HunyuanVideo / FLUX single blocks: q|k|v|mlp), so
+            # each has its own run of tiles within a hop.
+            u = jnp.asarray(t, jnp.int32) % tph
+            i = sum(((u >= cum[kk]).astype(jnp.int32) for kk in range(1, nw)), jnp.int32(0))
+            base = sum((jnp.where(i == kk, cum[kk], 0) for kk in range(1, nw)), jnp.int32(0))
+            return i, u - base
+
+        def w_start(t, slot):
+            # The weights are separate refs, so one guarded DMA start per weight.
+            i, j = which(t)
+            for ii in range(nw):
+                @pl.when(i == ii)
+                def _():
+                    pltpu.make_async_copy(w_src(ii, j), wbuf.at[slot], w_sem.at[slot]).start()
+
+        def w_wait(slot):
+            pltpu.make_async_copy(w_src(0, 0), wbuf.at[slot], w_sem.at[slot]).wait()
+
+        def o_wait(slot):
+            for d in range(2):
+                pltpu.make_async_copy(obuf.at[slot, d], out_refs[0].at[:, 0, d, :, cols_of(0)],
+                                      o_sem.at[slot, d]).wait()
+
+        w_start(0, 0)  # the first weight tile overlaps the local loads
+        loads = [pltpu.make_async_copy(x_ref.at[:, d], xbuf.at[d, 0], load_sem.at[d])
+                 for d in range(2)]
+        for c in loads:
+            c.start()
+        for c in loads:
+            c.wait()
+
+        def slot_of(s):
+            return 0 if ablate == "comm" else s % SLOTS
+
+        copies = []
+        for s in range(n):  # hops: static (n = 4), tiles within a hop: a runtime loop
+            for c in copies:  # hop s's halves have arrived
+                c.wait_send()
+                c.wait_recv()
+            copies = []
+            if s + 1 < n and ablate != "comm":
+                for d in range(2):
+                    c = pltpu.make_async_remote_copy(
+                        xbuf.at[d, slot_of(s)], xbuf.at[d, (s + 1) % SLOTS],
+                        send_sem.at[d, (s + 1) % SLOTS], recv_sem.at[d, (s + 1) % SLOTS],
+                        device_id={axis: nbrs_ref[d]})
+                    c.start()
+                    copies.append(c)
+
+            def tile(u, carry, s=s):
+                t = s * tph + u
+                slot = t % 2
+
+                @pl.when(t + 1 < total)
+                def _():
+                    w_start(t + 1, 1 - slot)  # always one weight tile in flight
+
+                w_wait(slot)
+
+                @pl.when(t >= 2)
+                def _():
+                    o_wait(slot)  # this output slot's previous write-back is done
+
+                if ablate != "compute":
+                    w = wbuf[slot]
+                    for d in range(2):
+                        for bi in range(b):
+                            obuf[slot, d, bi] = jax.lax.dot_general(
+                                xbuf[d, slot_of(s), bi], w, (((1,), (kd,)), ((), ())),
+                                preferred_element_type=jnp.float32).astype(obuf.dtype)
+                i, j = which(t)
+                for ii in range(nw):
+                    @pl.when(i == ii)
+                    def _():
+                        for d in range(2):
+                            pltpu.make_async_copy(
+                                obuf.at[slot, d],
+                                out_refs[ii].at[:, owners_ref[d, s], d, :, cols_of(j)],
+                                o_sem.at[slot, d]).start()
+                return carry
+
+            jax.lax.fori_loop(0, tph, tile, 0)
+        for c in copies:
+            c.wait_send()
+            c.wait_recv()
+        for t in range(max(0, total - 2), total):  # the last write-backs
+            o_wait(t % 2)
+
+    any_spec = pl.BlockSpec(memory_space=pl.ANY)
+    if resident and manual:
+        if tm0 is None:
+            raise ValueError(f"no column tile (a multiple of 128) divides every width {cols}")
+        wshape = (k, tm0) if weights_kn else (tm0, k)
+        outs = pl.pallas_call(
+            kernel_manual,
+            out_shape=[jax.ShapeDtypeStruct((b, n, 2, hr, c), x.dtype) for c in cols],
+            grid_spec=pltpu.PrefetchScalarGridSpec(
+                num_scalar_prefetch=2,
+                in_specs=[any_spec] * (1 + nw),
+                out_specs=[any_spec] * nw,
+                scratch_shapes=[pltpu.VMEM((2, SLOTS, b, hr, k), x.dtype),
+                                pltpu.VMEM((2,) + wshape, x.dtype),
+                                pltpu.VMEM((2, 2, b, hr, tm0), x.dtype),
+                                pltpu.SemaphoreType.DMA((2,)),
+                                pltpu.SemaphoreType.DMA((2, SLOTS)),
+                                pltpu.SemaphoreType.DMA((2, SLOTS)),
+                                pltpu.SemaphoreType.DMA((2,)),
+                                pltpu.SemaphoreType.DMA((2, 2))],
+            ),
+            compiler_params=pltpu.CompilerParams(collective_id=0,
+                                                 vmem_limit_bytes=vmem_limit_bytes),
+            interpret=interpret,
+        )(owners, nbrs, x.reshape(b, 2, hr, k), *ws)
+        return tuple(o.reshape(b, n * rows, o.shape[-1]) for o in outs)
+    if resident:
+        outs = pl.pallas_call(
+            kernel_resident,
+            out_shape=[jax.ShapeDtypeStruct((b, n, 2, hr, c), x.dtype) for c in cols],
+            grid_spec=pltpu.PrefetchScalarGridSpec(
+                num_scalar_prefetch=2,
+                in_specs=[any_spec] * (1 + nw),
+                out_specs=[any_spec] * nw,
+                scratch_shapes=[pltpu.VMEM((2, SLOTS, b, hr, k), x.dtype),
+                                pltpu.SemaphoreType.DMA((2,)),
+                                pltpu.SemaphoreType.DMA((2, SLOTS)),
+                                pltpu.SemaphoreType.DMA((2, SLOTS))],
+            ),
+            compiler_params=pltpu.CompilerParams(collective_id=0,
+                                                 vmem_limit_bytes=vmem_limit_bytes),
+            interpret=interpret,
+        )(owners, nbrs, x.reshape(b, 2, hr, k), *ws)
+        return tuple(o.reshape(b, n * rows, o.shape[-1]) for o in outs)
+
+    out_shape = [jax.ShapeDtypeStruct((b, n, 2, hr, c), x.dtype) for c in cols]
+    out_shape.append(jax.ShapeDtypeStruct((2, SLOTS, b, hr, k), x.dtype))
+    outs = pl.pallas_call(
+        kernel,
+        out_shape=out_shape,
+        grid_spec=pltpu.PrefetchScalarGridSpec(
+            num_scalar_prefetch=2,
+            in_specs=[any_spec] * (1 + nw),
+            out_specs=[any_spec] * (nw + 1),
+            scratch_shapes=[pltpu.SemaphoreType.DMA((2, SLOTS)),
+                            pltpu.SemaphoreType.DMA((2, SLOTS))],
+        ),
+        compiler_params=pltpu.CompilerParams(collective_id=0,
+                                             vmem_limit_bytes=vmem_limit_bytes),
+        interpret=interpret,
+    )(owners, nbrs, x.reshape(b, 2, hr, k), *ws)
+    return tuple(o.reshape(b, n * rows, o.shape[-1]) for o in outs[:nw])
+
+
+#: VMEM the kernel may plan for (v5e has 128 MiB per core; leave room for Mosaic).
+VMEM_BUDGET = 96 * 2**20
+
+
+def unsupported(x_shape, w_shapes, itemsize: int, tn: int = 512) -> str | None:
+    """Why the streamed kernel (``resident=True, manual=True``, weights (N, K))
+    cannot take these shapes, or None. Callers fall back to the JAX ring."""
+    b, rows, k = x_shape
+    if rows % 2:
+        return f"odd rows per device ({rows})"
+    if any(ws[1] != k for ws in w_shapes):
+        return "weight K does not match the input"
+    widths = [ws[0] for ws in w_shapes]
+    t = _common_tile(widths, tn)
+    if t is None:
+        return f"no column tile (a multiple of 128) divides every width {sorted(set(widths))}"
+    hr = rows // 2
+    need = itemsize * (2 * SLOTS * b * hr * k   # resident halves
+                       + 2 * t * k              # weight tiles, double-buffered
+                       + 2 * 2 * b * hr * t)    # output tiles, double-buffered
+    need += 4 * hr * t                          # one fp32 product
+    if need > VMEM_BUDGET:
+        return f"needs {need / 2**20:.0f} MiB of VMEM (budget {VMEM_BUDGET / 2**20:.0f})"
+    return None
+
+
+def _common_tile(widths, limit: int) -> int | None:
+    """The largest multiple of 128 up to ``limit`` that divides every width, or
+    None (q|k|v 768 + mlp 3072 per device: 384)."""
+    for t in range(limit - limit % 128, 0, -128):
+        if all(w % t == 0 for w in widths):
+            return t
+    return None
+
+
+def _col_tile(cols: int, limit: int) -> int:
+    """The largest multiple of 128 up to ``limit`` that divides ``cols``, else ``cols``
+    (Wan's FFN has 3456 columns per device: 384)."""
+    for t in range(limit - limit % 128, 0, -128):
+        if cols % t == 0:
+            return t
+    return cols
+
+
+def ring_tables(j, shard_of_pos, n: int):
+    """Runtime owners / neighbours for ring position ``j`` (a traced int32).
+
+    After s hops a device at position j holds position (j - s)'s top half
+    (sent clockwise) and (j + s)'s bottom half; ``shard_of_pos`` maps a ring
+    position to the shard (output row block) it owns.
+    """
+    import jax.numpy as jnp
+
+    shard_of = jnp.asarray(shard_of_pos, dtype=jnp.int32)
+    s = jnp.arange(n, dtype=jnp.int32)
+    owners = jnp.stack([shard_of[(j - s) % n], shard_of[(j + s) % n]])
+    nbrs = jnp.stack([(j + 1) % n, (j - 1) % n]).astype(jnp.int32)
+    return owners, nbrs
diff --git a/tests/unit/backends/test_tpu_collective_matmul_ag_assemble.py b/tests/unit/backends/test_tpu_collective_matmul_ag_assemble.py
new file mode 100644
index 00000000..3e5af89c
--- /dev/null
+++ b/tests/unit/backends/test_tpu_collective_matmul_ag_assemble.py
@@ -0,0 +1,65 @@
+"""The all-gather ring's two output assemblies give the same values (CPU, 4 fake devices).
+
+``collective_matmul._op`` builds the JAX ring; ``torch.tpu.pallas.jax_op`` is
+stubbed to return it, and the mesh is four XLA host devices. The ring position
+-> TP rank table is a non-identity permutation, so the runtime reordering of
+``DIFFLET_TPU_SP_AG_ASSEMBLE=concat`` is exercised. Needs ``jax``; skipped
+without it (it is not a test dependency).
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+import textwrap
+
+import pytest
+
+pytest.importorskip("jax")
+
+SCRIPT = textwrap.dedent('''
+    import os, sys, types
+    os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4"
+    os.environ["JAX_PLATFORMS"] = "cpu"
+    os.environ["DIFFLET_TPU_SP_AG_ASSEMBLE"] = sys.argv[1]
+    import jax, jax.numpy as jnp, numpy as np, torch
+    from jax.sharding import NamedSharding, PartitionSpec as P
+    torch.tpu = types.SimpleNamespace(pallas=types.SimpleNamespace(jax_op=lambda name, fn, **kw: fn))
+    from difflet.backends.tpu.ops_impl import collective_matmul as cm
+
+    n = 4
+    mesh = jax.sharding.Mesh(np.array(jax.devices()[:n]), ("tp",))
+    shard_of_pos = (2, 0, 3, 1)  # ring position -> TP rank, not the identity
+    cm._STATE = {"mesh": mesh, "tp": n, "shard_of_pos": shard_of_pos, "ids": None, "ops": {}}
+    ag = cm._op("ag", 2)
+
+    rows, K, N = 6, 8, 4  # rows per rank (even: two ring halves)
+    rng = np.random.default_rng(0)
+    # The model's sequence shard of TP rank r sits at ring position pos_of[r].
+    x_by_rank = [rng.standard_normal((1, rows, K)).astype(np.float32) for _ in range(n)]
+    pos_of = [shard_of_pos.index(r) for r in range(n)]
+    x_by_pos = [x_by_rank[shard_of_pos[p]] for p in range(n)]
+    w0 = rng.standard_normal((N * n, K)).astype(np.float32)
+    w1 = rng.standard_normal((N * n, K)).astype(np.float32)
+
+    put = lambda a, spec: jax.device_put(jnp.asarray(a), NamedSharding(mesh, spec))
+    x = put(np.concatenate(x_by_pos, axis=1), P(None, "tp", None))
+    ids = put(np.arange(n, dtype=np.int32), P("tp"))
+    o0, o1 = jax.jit(ag)(x, ids, put(w0, P("tp", None)), put(w1, P("tp", None)))
+
+    full = np.concatenate(x_by_rank, axis=1)  # the sequence in TP-rank order
+    for o, w in ((o0, w0), (o1, w1)):
+        got = np.asarray(o)
+        want = full @ w.T  # each device holds its N columns; the global array is (1, S, N*n)
+        assert got.shape == want.shape, (got.shape, want.shape)
+        np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-5)
+    print("OK", sys.argv[1])
+''')
+
+
+@pytest.mark.parametrize("mode", ["dus", "concat"])
+def test_ag_ring_assembly_matches_numpy(mode):
+    proc = subprocess.run([sys.executable, "-c", SCRIPT, mode], capture_output=True, text=True,
+                          timeout=300)
+    assert proc.returncode == 0, proc.stdout[-2000:] + proc.stderr[-4000:]
+    assert f"OK {mode}" in proc.stdout
diff --git a/tests/unit/backends/test_tpu_collective_matmul_pallas.py b/tests/unit/backends/test_tpu_collective_matmul_pallas.py
new file mode 100644
index 00000000..2bcf3693
--- /dev/null
+++ b/tests/unit/backends/test_tpu_collective_matmul_pallas.py
@@ -0,0 +1,70 @@
+"""The Pallas all-gather + matmul ring kernel, in TPU interpret mode on 4 CPU devices.
+
+Checks values against numpy with a non-identity ring position -> shard table,
+two weights of different widths, and batch 1 and 2, with Pallas' race detector
+on (two communication slots instead of three is caught by it). Needs ``jax``;
+skipped without it (it is not a test dependency).
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+import textwrap
+
+import pytest
+
+pytest.importorskip("jax")
+
+SCRIPT = textwrap.dedent('''
+    import os, sys
+    os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4"
+    os.environ["JAX_PLATFORMS"] = "cpu"
+    import jax, jax.numpy as jnp, numpy as np
+    import jax._src.tpu_info as ti
+    from jax.sharding import NamedSharding, PartitionSpec as P
+    from jax.experimental.pallas import tpu as pltpu
+    # The pipeline's tiling inference needs a TPU generation: treat the host as a v5e.
+    ti.registry["cpu"] = lambda: ti.get_tpu_info_for_chip(ti.ChipVersion.TPU_V5E, 1)
+    from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kp
+
+    n, B, rows, K = 4, int(sys.argv[1]), 16, 256
+    kn = "kn" in sys.argv[2]
+    resident = "res" in sys.argv[2]
+    manual = "man" in sys.argv[2]
+    mesh = jax.sharding.Mesh(np.array(jax.devices()[:n]), ("tp",))
+    shard_of_pos = (2, 0, 3, 1)  # ring position -> shard, not the identity
+    rng = np.random.default_rng(0)
+    x_by_shard = [rng.standard_normal((B, rows, K)).astype(np.float32) for _ in range(n)]
+    ws = [rng.standard_normal((N * n, K)).astype(np.float32) / K ** 0.5 for N in ((256, 384) if "mixed" in sys.argv[2] else (256, 256) if manual else (256, 384))]
+    put = lambda a, spec: jax.device_put(jnp.asarray(a), NamedSharding(mesh, spec))
+    x = put(np.concatenate([x_by_shard[shard_of_pos[p]] for p in range(n)], axis=1),
+            P(None, "tp", None))
+    ids = put(np.arange(n, dtype=np.int32), P("tp"))
+
+    def body(x, ids, *ws):
+        owners, nbrs = kp.ring_tables(ids[0], shard_of_pos, n)
+        return kp.ag_matmul_body(x, ws, owners, nbrs, n=n, tn=128 if manual else 256,
+                                 weights_kn=kn, resident=resident or manual, manual=manual,
+                                 interpret=pltpu.InterpretParams(detect_races=True))
+
+    f = jax.jit(jax.shard_map(body, mesh=mesh,
+                              in_specs=(P(None, "tp", None), P("tp"))
+                              + ((P(None, "tp"),) if kn else (P("tp", None),)) * 2,
+                              out_specs=(P(None, None, "tp"),) * 2, check_vma=False))
+    outs = f(x, ids, *[put(w.T, P(None, "tp")) if kn else put(w, P("tp", None)) for w in ws])
+    full = np.concatenate(x_by_shard, axis=1)  # the sequence in shard order
+    for o, w in zip(outs, ws):
+        np.testing.assert_allclose(np.asarray(o), full @ w.T, rtol=2e-5, atol=2e-5)
+    print("OK", B, sys.argv[2])
+''')
+
+
+@pytest.mark.parametrize("batch,layout", [(1, "nk"), (2, "nk"), (2, "kn"), (1, "kn-res"), (2, "nk-res"),
+                                         (1, "kn-man"), (2, "nk-man"),
+                                         (2, "nk-man-mixed")])
+def test_pallas_ag_matmul_matches_numpy(batch, layout):
+    proc = subprocess.run([sys.executable, "-c", SCRIPT, str(batch), layout], capture_output=True,
+                          text=True, timeout=600)
+    assert proc.returncode == 0, proc.stdout[-2000:] + proc.stderr[-4000:]
+    assert f"OK {batch} {layout}" in proc.stdout

From bf658ac20eb27035fd56af6a21ad42307d081508 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Tue, 10 Nov 2026 22:24:02 -0800
Subject: [PATCH 15/22] perf(tpu, hunyuan-video): single-stream blocks
 sequence-parallel too

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>
---
 difflet/backends/tpu/hunyuan_video/sp.py      | 110 +++++++++++++++++-
 .../models/hunyuan_video/tpu_application.py   |   3 +-
 .../test_tpu_hunyuan_video_sp_single.py       |  92 +++++++++++++++
 3 files changed, 202 insertions(+), 3 deletions(-)
 create mode 100644 tests/unit/backends/test_tpu_hunyuan_video_sp_single.py

diff --git a/difflet/backends/tpu/hunyuan_video/sp.py b/difflet/backends/tpu/hunyuan_video/sp.py
index 9b19e5b6..df9fb806 100644
--- a/difflet/backends/tpu/hunyuan_video/sp.py
+++ b/difflet/backends/tpu/hunyuan_video/sp.py
@@ -77,7 +77,107 @@ def __getattr__(self, name):  # approximate, ...
         return getattr(self._act, name)
 
 
-def install(module: torch.nn.Module, *, latent_seq_len: int) -> bool:
+#: ``DIFFLET_TPU_SP_SINGLE`` (with ``DIFFLET_TPU_SP=1``; default on, ``0`` turns it
+#: off): the 40 single-stream blocks run sequence-parallel too, instead of dense TP
+#: with an all-reduce per block (1.0 s of a 10.3 s, 20-step denoise on v5e-4 at
+#: 320x512x61, 2026-11-11). Step 516.1 -> 500.7 ms (ring) / 478.0 ms (kernel).
+SINGLE_ENV = "DIFFLET_TPU_SP_SINGLE"
+
+#: The single-stream blocks' ops: q|k|v|mlp share one all-gather, the two row-parallel
+#: projections one reduce-scatter (as FLUX's single blocks).
+SINGLE_OPS = (("ag", 4), ("rs", 2))
+
+
+def _single_requested() -> bool:
+    import os
+
+    return os.environ.get(SINGLE_ENV, "1").strip().lower() in ("1", "true", "yes", "on")
+
+
+class _PrecomputedSingle(_Precomputed):
+    """The ``self`` the shared attention sees in a single-stream block: the joint
+    q/k/v are precomputed and the attention has no output projection (pre_only)."""
+
+    def __init__(self, attn, qkv):
+        self._attn = attn
+        self.to_q, self.to_k, self.to_v = (lambda _x, t=t: t for t in qkv)
+        self.to_out = None
+        self.sp_enabled = False
+
+
+def _single_block_forward(self, hidden_states, encoder_hidden_states, temb,
+                          attention_mask=None, image_rotary_emb=None, *args, **kwargs):
+    """A single-stream block on this rank's contiguous chunk of the joint
+    [latent | text] sequence (B, (S_lat + S_txt)/tp, H).
+
+    The shared model forward hands the first single block the full latent and text
+    (it gathers the latent at the dual -> single transition) and expects full
+    streams back from the last; in between, the blocks pass the chunk along as
+    ``hidden_states``. A single-stream block treats every joint token alike outside
+    attention, so a chunk may straddle the latent / text boundary; attention runs on
+    the gathered q/k/v in the natural joint order, with the shared mask and rotary.
+    """
+    del args, kwargs
+    from difflet.models.hunyuan_video.modeling_hunyuan_video import HunyuanVideoAttention
+    from difflet.ops import gather_from_sequence_parallel_region
+
+    first, last, text_len = self._difflet_sp_single
+    if first:
+        joint = torch.cat([hidden_states, encoder_hidden_states], dim=1)
+        # This rank's rows come from a buffer (runtime data), not a Python-constant
+        # slice: with a constant, XLA folded the concatenation away on ranks 0-2
+        # (their chunk lies inside the latent) but not on rank 3 (its chunk holds
+        # the text), the per-rank programs diverged and the first forward hung in a
+        # collective (v5e-4, 2026-11-11) -- the LTX-2 failure, see collective_matmul.
+        hidden_states = joint.index_select(1, self._difflet_sp_rows)
+    residual = hidden_states
+    norm, gate = self.norm(hidden_states, emb=temb)
+    attn = self.attn
+    q, k, v, mlp = cm.all_gather_matmul(norm, (attn.to_q, attn.to_k, attn.to_v, self.proj_mlp))
+    s_lat = q.shape[1] - text_len
+    latent_out, text_out = HunyuanVideoAttention.forward(
+        _PrecomputedSingle(attn, (q, k, v)),
+        hidden_states=q[:, :s_lat],  # read for nothing but the latent length
+        encoder_hidden_states=q[:, s_lat:],
+        attention_mask=attention_mask,
+        image_rotary_emb=image_rotary_emb,
+    )
+    attn_out = torch.cat([latent_out, text_out], dim=1)
+    # proj_out_attn has the bias (skip_bias_add, added once after the reduction);
+    # proj_out_mlp has none.
+    proj = cm.matmul2_reduce_scatter(attn_out, self.proj_out_attn, self.act_mlp(mlp),
+                                     self.proj_out_mlp, bias=self.proj_out_attn.bias)
+    hidden_states = residual + gate.unsqueeze(1) * proj
+    if last:
+        full = gather_from_sequence_parallel_region(hidden_states, dim=1)
+        return full[:, :-text_len], full[:, -text_len:]
+    return hidden_states, encoder_hidden_states  # the text is not read again until the last
+
+
+def _install_single(module, *, latent_seq_len: int, text_seq_len: int, tp: int) -> bool:
+    from difflet.backends.tpu.ops_impl import parallel_mesh
+
+    joint = int(latent_seq_len) + int(text_seq_len)
+    if not cm.supports(joint, tp):
+        print(f"[difflet tpu] {SINGLE_ENV}: joint sequence {joint} does not split into "
+              f"2 x tp={tp}; single-stream blocks stay dense", flush=True)
+        return False
+    cm.prepare(SINGLE_OPS)
+    blocks = list(module.single_transformer_blocks)
+    rank, n = parallel_mesh.get_tp_rank(), joint // tp
+    device = next(module.parameters()).device
+    for i, block in enumerate(blocks):
+        block._difflet_sp_single = (i == 0, i == len(blocks) - 1, int(text_seq_len))
+        block.forward = types.MethodType(_single_block_forward, block)
+    blocks[0].register_buffer(
+        "_difflet_sp_rows", torch.arange(rank * n, (rank + 1) * n, device=device), persistent=False)
+    print(f"[difflet tpu] HunyuanVideo single-stream blocks sequence-parallel too "
+          f"({len(blocks)} blocks, joint S={joint}, tp={tp})", flush=True)
+    return True
+
+
+def install(module: torch.nn.Module, *, latent_seq_len: int,
+            text_seq_len: int | None = None) -> bool:
     """Swap the fused forwards into the dual-stream blocks of a model built with
     ``sp_enabled``. Returns False (and changes nothing) when the latent sequence does
     not split, so the caller keeps the shared Megatron-SP path."""
@@ -92,13 +192,19 @@ def install(module: torch.nn.Module, *, latent_seq_len: int) -> bool:
         print(f"[difflet tpu] {cm.SP_ENV}: latent sequence {latent_seq_len} does not split "
               f"into 2 x tp={tp}; using the unfused Megatron-SP path", flush=True)
         return False
+    # The Pallas AG kernel: with single-stream SP, -4.5% per step at 320x512x61
+    # (v5e-4, 2026-11-11); DIFFLET_TPU_SP_KERNEL=jax keeps the ring.
+    cm.set_kernel_default("pallas")
     cm.prepare()
     for block in blocks:
         block.attn.forward = types.MethodType(_attention_forward, block.attn)
         block.ff.forward = types.MethodType(_ffn_forward, block.ff)
+    single = (_single_requested() and text_seq_len is not None
+              and _install_single(module, latent_seq_len=latent_seq_len,
+                                  text_seq_len=text_seq_len, tp=tp))
     print(f"[difflet tpu] HunyuanVideo sequence-parallel with collective matmuls "
           f"({len(blocks)} dual-stream blocks, S_latent={latent_seq_len}, tp={tp}; "
-          f"single-stream blocks stay dense)", flush=True)
+          f"single-stream blocks {'SP' if single else 'stay dense'})", flush=True)
     return True
 
 
diff --git a/difflet/models/hunyuan_video/tpu_application.py b/difflet/models/hunyuan_video/tpu_application.py
index a9596f5a..54504c37 100644
--- a/difflet/models/hunyuan_video/tpu_application.py
+++ b/difflet/models/hunyuan_video/tpu_application.py
@@ -243,7 +243,8 @@ def load_eager(self) -> None:
             from difflet.backends.tpu.ops_impl import collective_matmul
 
             collective_matmul.require_compiled(default=False)  # as maybe_compile below
-            sp.install(module, latent_seq_len=int(self.config.image_seq_len))
+            sp.install(module, latent_seq_len=int(self.config.image_seq_len),
+                       text_seq_len=int(self.config.text_seq_len))
         # Not yet measured under torch.compile, so off unless DIFFLET_TPU_COMPILE=1.
         module = device_runtime.maybe_compile(module, default=False)
         # Compile the request-shaped graph now, a few ranks at a time (see
diff --git a/tests/unit/backends/test_tpu_hunyuan_video_sp_single.py b/tests/unit/backends/test_tpu_hunyuan_video_sp_single.py
new file mode 100644
index 00000000..2e8b6441
--- /dev/null
+++ b/tests/unit/backends/test_tpu_hunyuan_video_sp_single.py
@@ -0,0 +1,92 @@
+"""CPU check of HunyuanVideo's sequence-parallel single-stream blocks (backends/tpu/hunyuan_video/sp.py).
+
+The collective matmuls are JAX rings that only run on the chips; here they are
+replaced by their tp=1 meaning (all-gather = identity, reduce-scatter = the
+matmul(s) plus the bias once). A chain of three SP single blocks -- the first
+cuts the joint [latent | text] sequence into its chunk, the last gathers it back
+-- must reproduce the shared dense blocks, including a padded text mask.
+"""
+
+from __future__ import annotations
+
+import importlib
+import types
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from difflet.pipeline.parallel_mesh import MeshSpec
+
+
+@pytest.fixture(autouse=True)
+def _mesh(monkeypatch):
+    # The shared HunyuanVideo layers resolve their implementation per backend; the
+    # TPU one needs no Neuron SDK.
+    monkeypatch.setenv("DIFFLET_BACKEND", "tpu")
+    pm = importlib.import_module("difflet.backends.tpu.ops_impl.parallel_mesh")
+    pm.destroy_parallel_mesh()
+    pm.init_parallel_mesh(MeshSpec(tp=1))
+    yield
+    pm.destroy_parallel_mesh()
+
+
+def _fake_rings(monkeypatch, cm):
+    def ag(x, linears):
+        return [F.linear(x, lin.weight, lin.bias) for lin in linears]
+
+    def rs2(x0, l0, x1, l1, bias=None):
+        out = F.linear(x0, l0.weight) + F.linear(x1, l1.weight)
+        return out + bias if bias is not None else out
+
+    monkeypatch.setattr(cm, "all_gather_matmul", ag)
+    monkeypatch.setattr(cm, "matmul2_reduce_scatter", rs2)
+    monkeypatch.setattr(cm, "prepare", lambda ops=None: None)
+
+
+def _blocks(n=3, heads=2, head_dim=8):
+    from difflet.models.hunyuan_video.modeling_hunyuan_video import (
+        HunyuanVideoSingleTransformerBlock,
+    )
+
+    torch.manual_seed(0)
+    blocks = [HunyuanVideoSingleTransformerBlock(heads, head_dim, mlp_ratio=2.0,
+                                                 reduce_dtype=torch.float64)
+              for _ in range(n)]
+    for b in blocks:
+        with torch.no_grad():
+            for p in b.parameters():  # zero-initialised biases / gates would hide bugs
+                p.copy_(torch.randn_like(p) * 0.2)
+    return [b.double().eval() for b in blocks]
+
+
+def test_sp_single_blocks_match_dense(monkeypatch):
+    sp = importlib.import_module("difflet.backends.tpu.hunyuan_video.sp")
+    _fake_rings(monkeypatch, sp.cm)
+    heads, head_dim, s_lat, s_txt, batch = 2, 8, 12, 4, 2
+    hidden = heads * head_dim
+    g = torch.Generator().manual_seed(1)
+    latent = torch.randn(batch, s_lat, hidden, generator=g, dtype=torch.float64)
+    text = torch.randn(batch, s_txt, hidden, generator=g, dtype=torch.float64)
+    temb = torch.randn(batch, hidden, generator=g, dtype=torch.float64)
+    mask = torch.ones(batch, 1, 1, s_lat + s_txt, dtype=torch.bool)
+    mask[1, ..., -2:] = False  # row 1: the last two text tokens are padding
+    cos = torch.randn(s_lat, head_dim, generator=g, dtype=torch.float64)
+    sin = torch.randn(s_lat, head_dim, generator=g, dtype=torch.float64)
+
+    dense = _blocks()
+    h, e = latent, text
+    for b in dense:
+        h, e = b(h, e, temb, mask, (cos, sin))
+
+    sp_blocks = _blocks()  # same seed: same weights
+    for i, b in enumerate(sp_blocks):
+        b._difflet_sp_single = (i == 0, i == len(sp_blocks) - 1, s_txt)
+        b.forward = types.MethodType(sp._single_block_forward, b)
+    sp_blocks[0].register_buffer("_difflet_sp_rows", torch.arange(s_lat + s_txt), persistent=False)
+    h2, e2 = latent, text
+    for b in sp_blocks:
+        h2, e2 = b(h2, e2, temb, mask, (cos, sin))
+
+    torch.testing.assert_close(h2, h, rtol=1e-10, atol=1e-10)
+    torch.testing.assert_close(e2, e, rtol=1e-10, atol=1e-10)

From f098f4eb1df033fff5ce817c00ac3543e40b8171 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Tue, 10 Nov 2026 22:24:02 -0800
Subject: [PATCH 16/22] perf(tpu): the Pallas AG kernel by default for Wan and
 FLUX

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>
---
 difflet/backends/tpu/flux/sp.py | 3 +++
 difflet/backends/tpu/wan/sp.py  | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/difflet/backends/tpu/flux/sp.py b/difflet/backends/tpu/flux/sp.py
index edd3c09a..8f7e72bc 100644
--- a/difflet/backends/tpu/flux/sp.py
+++ b/difflet/backends/tpu/flux/sp.py
@@ -190,6 +190,9 @@ def install(module: nn.Module, *, image_seq_len: int, text_seq_len: int) -> bool
             return False
     if not all(hasattr(b, "proj_out_attn") for b in module.single_transformer_blocks):
         raise ValueError("install() needs a FLUX transformer sharded by shard_flux_transformer")
+    # The Pallas AG kernel: -10.2% per step at 1024x1024 (v5e-4, 2026-11-11);
+    # DIFFLET_TPU_SP_KERNEL=jax keeps the ring.
+    cm.set_kernel_default("pallas")
     cm.prepare(FLUX_OPS)
     for block in module.transformer_blocks:
         block.forward = types.MethodType(_joint_block_forward, block)
diff --git a/difflet/backends/tpu/wan/sp.py b/difflet/backends/tpu/wan/sp.py
index 18e3e8db..e4f41a00 100644
--- a/difflet/backends/tpu/wan/sp.py
+++ b/difflet/backends/tpu/wan/sp.py
@@ -82,6 +82,9 @@ def install(module: torch.nn.Module, *, seq_len: int) -> bool:
         print(f"[difflet tpu] {cm.SP_ENV}: sequence {seq_len} does not split into "
               f"2 x tp={tp}; using the unfused Megatron-SP path", flush=True)
         return False
+    # The Pallas AG kernel: -8.3% per step at 480x832x9 (v5e-4, 2026-11-11);
+    # DIFFLET_TPU_SP_KERNEL=jax keeps the ring.
+    cm.set_kernel_default("pallas")
     cm.prepare()
     for block in module.blocks:
         for attn in (block.attn1, block.attn2):

From 868401ebaedfba09cfbbea64c0fef0848d01b2e1 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Tue, 10 Nov 2026 22:24:02 -0800
Subject: [PATCH 17/22] docs(tpu): P0, the Pallas AG kernel, other models and
 HunyuanVideo single-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>
---
 ...2026-11-10-tpu-collective-matmul-pallas.md | 162 ++++++++++++++++++
 1 file changed, 162 insertions(+)

diff --git a/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md b/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
index 844df024..ee7dde63 100644
--- a/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
+++ b/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
@@ -142,6 +142,168 @@ P0 removes most of the write-back, the kernel's remaining target is the 3-7%
 wait and the kernel is likely not worth it; otherwise continue with the Pallas
 AG kernel below, which writes each tile of the product in place.
 
+## P0 result (2026-11-10): concat helps Wan, not LTX-2, and is not output-equal
+
+`DIFFLET_TPU_SP_AG_ASSEMBLE=concat` (default `dus`). Same tree (md5 77c0c1ad4ed5),
+back to back, main's bench, rank-0 traces:
+
+| | Wan dus | Wan concat | LTX-2 dus | LTX-2 concat |
+|---|---|---|---|---|
+| step | 292.6 ms | 275.1 ms (-6.0%) | 189.7 ms | 194.2 ms (+2.4%) |
+| window (20 steps) | 5848.5 | 5468.8 | 3814.4 | 3806.6 |
+| compute | 5399.0 | 5022.4 | 2993.6 | 2870.8 |
+| ICI wait | 192.9 | 152.5 | 258.6 | 390.1 |
+
+- `dynamic-update-slice` is gone; the assembly costs instead ~480 ms (Wan:
+  copy +130, reshape +105, a `while` from the runtime-indexed `take` +106,
+  fusion +110) -- net -380 ms.
+- LTX-2 (289-row halves) is ICI-bound: the compute saved (-123 ms) became wait
+  (+131 ms). Less HBM traffic alone does not help it.
+- Op parity at Wan shapes (1170 rows/rank, hidden 5120): ag1/2/3 vs plain
+  gather + linear are 0.0 in **both** modes.
+- End to end the outputs differ: Wan dus vs concat 29.06 dB, while TP vs dus is
+  43.87 dB and TP vs concat 29.43 dB (sha256 TP 8e646ad9b22e, dus e41edd1e23ab,
+  concat 495847095139, each deterministic over 3 runs). LTX-2 dus vs concat
+  43.75 dB (ef75047e8a34 / ec137f0d0fec). The Wan difference follows the
+  content (largest on the subject), with no band at shard boundaries -- a
+  trajectory drift ~5x the TP/SP noise, cause not found (a one-step forward
+  comparison is the next check). Not checked: Qwen, FLUX, HunyuanVideo.
+
+Decision: keep `dus` as default; continue with the Pallas AG kernel, which
+writes each product tile once at its final rows (removing the remaining
+assembly cost on Wan). For LTX-2 the kernel's value is overlap control, not
+HBM traffic.
+
+## Kernel result (2026-11-11): Wan -8.3% per step, LTX-2 flat, outputs within the TP/SP spread
+
+`difflet/backends/tpu/ops_impl/collective_matmul_pallas.py`, behind
+`DIFFLET_TPU_SP_KERNEL=pallas` (default `jax`), AG only. What got it there,
+op-level (plain JAX, v5e-4, `benchmark/tpu_profile/cm_pallas_bench.py`, ms per call):
+
+| version | Wan q/k/v | Wan FFN up | LTX-2 q/k/v | LTX-2 FFN up |
+|---|---|---|---|---|
+| JAX ring (dus) | 1.259 | 1.050 | 0.445 | 0.500 |
+| v1: emit_pipeline per (hop, weight, half) | 1.541 | 1.181 | 0.606 | 0.630 |
+| v3: halves resident in VMEM, VMEM->VMEM remote DMA | 1.241 | 1.074 | 0.464 | 0.540 |
+| v4: one weight-tile stream over every (hop, weight, tile), runtime loop | **1.185** | **1.056** | **0.406** | **0.517** |
+
+- Ablation: communication is fully hidden from v3 on (kernel == kernel without
+  remote copies, within 0.001 ms); what is left is the matmul.
+- Dead ends: pre-transposed weights (no change); wider column tiles (worse:
+  each pipeline's first tile load is exposed); padding the halves to 16 rows and
+  one (2*hrp)-row dot per tile ("v5": +2% on Wan q/k/v, nothing elsewhere, and
+  Mosaic only slices tiled dims in whole tiles, so it needs an XLA pad first).
+- A plain Pallas matmul equals XLA at the same shape
+  (`pallas_matmul_sweep.py`); the kernel's gap to the one-chip XLA einsum is the
+  small M (585 / 289 rows per dot), not the kernel.
+- Fully unrolled, the stream took 225 s to compile at Wan's FFN shape; a
+  `fori_loop` over a hop's tiles: 5-10 s.
+- `has_side_effects=True` breaks torch_tpu (SPMD partitioner: "Side-effect HLO
+  must have sharding"); the kernel does not set it.
+
+Through torch_tpu (`cm_ops_parity.py`, `DIFFLET_TPU_SP_KERNEL=pallas`, Wan and
+LTX-2 shapes, all four ranks): ag1/ag2/ag3 vs plain gather + linear 0.0,
+including two calls in one graph. No fallback.
+
+Full models, main's bench, same tree (md5 011bdf616420), back to back:
+
+| | Wan jax | Wan pallas | LTX-2 jax | LTX-2 pallas |
+|---|---|---|---|---|
+| step | 292.4 ms | **268.0 ms (-8.3%)** | 190.1 ms | 190.5 ms |
+| window (20 steps) | 5832.4 | 5350.9 | 3798.4 | 3735.3 |
+| `dynamic-update-slice` | 643.1 | 0 | 264.7 | 0 |
+| ICI wait | 192.9 | 85.5 | 258.5 | 111.8 |
+
+- Wan pays ~250 ms of relayout copies/reshapes instead: 585-row halves are not
+  a multiple of the (8, 128) HBM tile, so `(B, n, 2, 585, N) -> (B, 4680, N)` and
+  the input split are real copies. Removing them needs rows per rank padded to a
+  multiple of 16 outside the graph (as LTX-2 pads), i.e. a model change.
+- LTX-2 is ICI-bound: the compute saved becomes wait, as with concat.
+
+Output (Wan, `sp_step_drift.py`, PSNR of the decoded frames):
+
+| | TP vs jax | jax vs pallas | TP vs pallas |
+|---|---|---|---|
+| 1 step, seed 42 | 49.88 | 50.47 | 49.88 |
+| 4 steps, seed 42 | 49.04 | 49.47 | 49.40 |
+| 20 steps, seed 42 | 43.87 | 26.09 | 26.23 |
+| 20 steps, seed 7 | 38.87 | 45.85 | 38.38 |
+| 20 steps, seed 1234 | 43.46 | 41.14 | 44.75 |
+
+Per step the three agree at the bf16 noise level; after 20 steps which pair
+drifts apart depends on the seed (each configuration is deterministic over
+three runs). Read: the kernel is another valid rounding, like TP vs SP. Not
+checked: more seeds, a seed-independent quality metric (fp32 reference,
+perceptual), LTX-2's seed spread, a visual review of the seed-42 output, Qwen /
+FLUX / HunyuanVideo. Open risk: consecutive kernels share the barrier semaphore
+(`collective_id=0`); no failure seen, not proven safe.
+
+### The other models (2026-11-11, same tree md5 011bdf616420, main's bench)
+
+| model | step: JAX ring -> kernel | PSNR jax vs kernel | PSNR TP vs jax (floor) | note |
+|---|---|---|---|---|
+| Qwen-Image | 178.0 -> 179.7 ms | inf (bit-identical) | 47.03 | no `dynamic-update-slice` in the jax trace; ICI -140 ms, compute +189 ms |
+| FLUX | 156.1 -> 157.3 ms | 43.90 | 41.52 | single blocks' ag4 (768x3 + 3072) fell back to the ring; dual blocks: dus -253 ms, kernel +545 vs fusion -186 |
+| HunyuanVideo | 517.8 -> 518.1 ms | 42.52 | 45.55 | the step is bound by the dense single blocks' all-reduce (1.0 s / 20 steps) |
+
+The kernel pays only where the ring's `dynamic-update-slice` is expensive (Wan);
+where XLA already fuses the ring's writes (Qwen-Image) or the dus is small,
+its matmul is a little slower than XLA's einsum and it loses 1-2%.
+
+### HunyuanVideo single-stream blocks under SP (`DIFFLET_TPU_SP_SINGLE=1`)
+
+`backends/tpu/hunyuan_video/sp.py`: the 40 single-stream blocks run on a
+contiguous chunk of the joint [latent | text] sequence (10496 = 4 x 2624; a chunk
+may straddle the boundary), ag4 (q|k|v|mlp) + rs2 as FLUX's single blocks; the
+first block cuts, the last gathers. Same tree (md5 3d32d564e9bd), back to back:
+
+| | dual-only SP | dual + single SP |
+|---|---|---|
+| step | 516.1 ms | **502.2 ms (-2.7%; -8.7% vs TP 549.9)** |
+| window (20 steps) | 10322.1 | 10003.7 |
+| ICI wait | 1042.5 (all-reduce 1025.7) | **141.6** |
+| `dynamic-update-slice` | -- | 646.3 (the ag4 rings) |
+| PSNR vs TP | 45.55 | 44.35 (dual-only vs single: 45.68) |
+
+First attempt hung in the first forward: the cut was a Python-constant slice of
+the in-graph concatenation; ranks 0-2's chunks lie inside the latent, rank 3's
+holds the text, so XLA folded the concatenation on 0-2 only and the programs
+diverged (rule 1 again). The cut is now `index_select` with a per-rank row
+buffer.
+
+The ag4's dus, next (all single-stream SP; tree d9e05f83d1aa for jax / pallas,
+the kernel now takes weights of different widths through one common tile,
+q|k|v 768 + mlp 3072 -> 384):
+
+| HunyuanVideo | step | window | `dynamic-update-slice` | PSNR vs TP |
+|---|---|---|---|---|
+| TP | 549.9 ms | | | |
+| dual-only SP | 516.1 | 10322.1 | -- | 45.55 |
+| + single SP, JAX ring | 500.7 | 10023.6 | 646.3 | 44.35 |
+| + single SP, concat assembly | 493.9 | 9846.1 | 0 (copy +142, while +224) | 42.86 |
+| + single SP, **Pallas kernel** | **478.0 (-13.1% vs TP)** | **9525.4** | 0 | 42.86 |
+
+The kernel's video is bit-identical to the concat assembly's (sha256
+f13f00fb7aba both); the JAX ring's is reproducible across trees (03e10d9b7119).
+Each configuration is deterministic over three runs.
+
+### FLUX with the mixed-width kernel (single blocks' ag4 on the kernel)
+
+A first rerun looked identical to the old one -- torch_tpu's compilation cache
+keys a `jax_op` by name and signature, so the unchanged op name reused the
+ring program (num_compilations=0, 18.8 s warm-up, same sha256). Op names now
+carry a hash of the two source files. Fresh compile (78 s, num_compilations=1,
+no fallback), tree 077ab253f18d:
+
+| FLUX (28 steps) | TP | SP, JAX ring | SP, Pallas kernel |
+|---|---|---|---|
+| step | 175.2 ms | 156.5 | **140.6 (-10.2%; -19.7% vs TP)** |
+| window | | 4329.5 | 3890.5 |
+| `dynamic-update-slice` | | 1171.1 | 0 |
+| PSNR vs TP | | 41.52 | 39.52 (jax vs kernel 44.91) |
+
+Not checked: other seeds, a visual review (TP vs kernel max |d| 199/255).
+
 ## Steps
 
 | # | Work | Validation | Estimate |

From 7af1ecdc29418f523a7275061aa12ab339a87b9a Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Wed, 11 Nov 2026 12:20:00 -0800
Subject: [PATCH 18/22] refactor(tpu): the Pallas AG kernel as its one streamed
 version, with 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 01d027fe 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 868401eb, 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>
---
 benchmark/tpu_profile/cm_pallas_bench.py      |  93 +--
 .../tpu/ops_impl/collective_matmul.py         |  19 +-
 .../tpu/ops_impl/collective_matmul_pallas.py  | 648 +++++++-----------
 .../test_tpu_collective_matmul_pallas.py      |  63 +-
 4 files changed, 340 insertions(+), 483 deletions(-)

diff --git a/benchmark/tpu_profile/cm_pallas_bench.py b/benchmark/tpu_profile/cm_pallas_bench.py
index f3ba9121..b5d66af9 100644
--- a/benchmark/tpu_profile/cm_pallas_bench.py
+++ b/benchmark/tpu_profile/cm_pallas_bench.py
@@ -4,10 +4,14 @@
 
 Plain JAX, one process owning the four chips (no torch_tpu): the JAX rings are
 ``collective_matmul._op`` with ``torch.tpu.pallas.jax_op`` stubbed out (as in
-cm_pure_jax), the kernel is ``collective_matmul_pallas.ag_matmul_body``. For
+cm_pure_jax), the kernel is ``collective_matmul_pallas.all_gather_matmul_body``. For
 each shape, prints the time per call and the kernel's difference from the dus
 ring (bf16 inputs, fp32 accumulation in both). Step 2 of
 docs/plans/2026-11-10-tpu-collective-matmul-pallas.md.
+
+The kernel is timed as the models run it: weights (N, K), column tiles up to
+``DEFAULT_COLUMN_TILE_LIMIT``. The plan's op-level version table was taken at
+commit 01d027fe with pre-transposed (K, N) weights and tiles up to 512.
 """
 
 from __future__ import annotations
@@ -31,7 +35,8 @@ def main() -> int:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--iters", type=int, default=50)
     parser.add_argument("--shapes", default=",".join(SHAPES))
-    parser.add_argument("--tn", type=int, default=512)
+    parser.add_argument("--column-tile-limit", type=int, default=None,
+                        help="the kernel's column tile limit (default: the production value)")
     parser.add_argument("--ablate", action="store_true",
                         help="also time the kernel without comm / without compute, and one-chip XLA einsums")
     parser.add_argument("--trace", default="", help="jax.profiler trace dir for the kernel calls")
@@ -44,10 +49,15 @@ def main() -> int:
     from jax.sharding import NamedSharding
     from jax.sharding import PartitionSpec as P
 
-    torch.tpu = types.SimpleNamespace(pallas=types.SimpleNamespace(
-        jax_op=lambda name, fn, **kw: fn))
+    # No torch_tpu here: jax_op returns the JAX function itself. is_initialized is
+    # asked at interpreter exit.
+    torch.tpu = types.SimpleNamespace(is_initialized=lambda: False,
+                                      pallas=types.SimpleNamespace(jax_op=lambda name, fn, **kw: fn))
     from difflet.backends.tpu.ops_impl import collective_matmul as cm
-    from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kp
+    from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kernel
+
+    if args.column_tile_limit is None:
+        args.column_tile_limit = kernel.DEFAULT_COLUMN_TILE_LIMIT
 
     devices = jax.devices()
     n = len(devices)
@@ -69,7 +79,7 @@ def timed(f, *xs):
         jax.block_until_ready(out)
         return out, (time.perf_counter() - t) / args.iters * 1e3
 
-    print(f"jax {jax.__version__} {devices[0].device_kind} x{n}, iters {args.iters}, tn<={args.tn}",
+    print(f"jax {jax.__version__} {devices[0].device_kind} x{n}, iters {args.iters}, column tile <= {args.column_tile_limit}",
           flush=True)
     ids = put(np.arange(n), P("tp"), jnp.int32)
     rng = np.random.default_rng(0)
@@ -85,33 +95,28 @@ def timed(f, *xs):
             out, ms = timed(ring, x, ids, *ws)
             res[mode] = (out if isinstance(out, tuple) else (out,), ms)
 
-        def body(x, ids, *ws):
-            owners, nbrs = kp.ring_tables(ids[0], shard_of_pos, n)
-            return kp.ag_matmul_body(x, ws, owners, nbrs, n=n, tn=args.tn)
+        def kernel_call(ablation=""):
+            def body(x, ids, *ws):
+                owner_shards, neighbors = kernel.ring_tables(ids[0], shard_of_pos, n)
+                return kernel.all_gather_matmul_body(
+                    x, ws, owner_shards, neighbors, ring_size=n,
+                    column_tile_limit=args.column_tile_limit, profile_ablation=ablation)
+            return jax.jit(jax.shard_map(
+                body, mesh=mesh, in_specs=(P(None, "tp", None), P("tp")) + (P("tp", None),) * len(ws),
+                out_specs=(P(None, None, "tp"),) * len(ws), check_vma=False))
 
-        kern = jax.jit(jax.shard_map(
-            body, mesh=mesh, in_specs=(P(None, "tp", None), P("tp")) + (P("tp", None),) * len(ws),
-            out_specs=(P(None, None, "tp"),) * len(ws), check_vma=False))
         try:
-            res["pallas"] = timed(kern, x, ids, *ws)
+            res["pallas"] = timed(kernel_call(), x, ids, *ws)
             if args.trace:
                 with jax.profiler.trace(f"{args.trace}/{name}"):
                     for _ in range(5):
-                        jax.block_until_ready(kern(x, ids, *ws))
+                        jax.block_until_ready(kernel_call()(x, ids, *ws))
         except Exception as exc:  # noqa: BLE001 -- report and go on to the next shape
             print(f"{name}: pallas failed: {type(exc).__name__}: {str(exc)[:600]}", flush=True)
             continue
         if args.ablate:
-            extra = []
-            for ab in ("comm", "compute"):
-                def body_ab(x, ids, *ws, ab=ab):
-                    owners, nbrs = kp.ring_tables(ids[0], shard_of_pos, n)
-                    return kp.ag_matmul_body(x, ws, owners, nbrs, n=n, tn=args.tn, ablate=ab)
-                f = jax.jit(jax.shard_map(
-                    body_ab, mesh=mesh,
-                    in_specs=(P(None, "tp", None), P("tp")) + (P("tp", None),) * len(ws),
-                    out_specs=(P(None, None, "tp"),) * len(ws), check_vma=False))
-                extra.append(f"no-{ab} {timed(f, x, ids, *ws)[1]:.3f}")
+            extra = [f"no-{ablation} {timed(kernel_call(ablation), x, ids, *ws)[1]:.3f}"
+                     for ablation in ("communication", "compute")]
             # The same FLOPs on one chip, plain XLA: the full sequence times each local weight.
             d0 = devices[0]
             xf = jax.device_put(jnp.zeros((1, rows * n, k), jnp.bfloat16), d0)
@@ -119,47 +124,9 @@ def body_ab(x, ids, *ws, ab=ab):
             ein = jax.jit(lambda x, *ws: tuple(jnp.einsum("bsk,nk->bsn", x, w) for w in ws))
             extra.append(f"xla-1chip {timed(ein, xf, *wl)[1]:.3f}")
             print(f"{name:8s} ablation (ms): " + ", ".join(extra), flush=True)
-        # v2: weights pre-transposed to (K, N) once, outside the timed call.
-        wts = [jax.device_put(w.T, NamedSharding(mesh, P(None, "tp"))) for w in ws]
-
-        def body_kn(x, ids, *ws):
-            owners, nbrs = kp.ring_tables(ids[0], shard_of_pos, n)
-            return kp.ag_matmul_body(x, ws, owners, nbrs, n=n, tn=args.tn, weights_kn=True)
-
-        kern_kn = jax.jit(jax.shard_map(
-            body_kn, mesh=mesh, in_specs=(P(None, "tp", None), P("tp")) + (P(None, "tp"),) * len(ws),
-            out_specs=(P(None, None, "tp"),) * len(ws), check_vma=False))
-        res["pallas-kn"] = timed(kern_kn, x, ids, *wts)
-
-        # v3: halves resident in VMEM, remote DMA VMEM -> VMEM, pipelines stream weights only.
-        def body_res(x, ids, *ws, ab="", manual=False):
-            owners, nbrs = kp.ring_tables(ids[0], shard_of_pos, n)
-            return kp.ag_matmul_body(x, ws, owners, nbrs, n=n, tn=args.tn, weights_kn=True,
-                                     resident=True, manual=manual, ablate=ab)
-
-        def mk(ab="", manual=False):
-            return jax.jit(jax.shard_map(
-                lambda x, ids, *ws: body_res(x, ids, *ws, ab=ab, manual=manual), mesh=mesh,
-                in_specs=(P(None, "tp", None), P("tp")) + (P(None, "tp"),) * len(ws),
-                out_specs=(P(None, None, "tp"),) * len(ws), check_vma=False))
-        try:
-            res["pallas-res"] = timed(mk(), x, ids, *wts)
-            if args.ablate:
-                print(f"{name:8s} resident ablation (ms): no-comm {timed(mk('comm'), x, ids, *wts)[1]:.3f}",
-                      flush=True)
-        except Exception as exc:  # noqa: BLE001
-            print(f"{name}: pallas-res failed: {type(exc).__name__}: {str(exc)[:800]}", flush=True)
-        # v4: one hand-written weight-tile stream over every (hop, weight, tile).
-        try:
-            res["pallas-man"] = timed(mk(manual=True), x, ids, *wts)
-            if args.ablate:
-                print(f"{name:8s} manual ablation (ms): no-comm "
-                      f"{timed(mk('comm', manual=True), x, ids, *wts)[1]:.3f}", flush=True)
-        except Exception as exc:  # noqa: BLE001
-            print(f"{name}: pallas-man failed: {type(exc).__name__}: {str(exc)[:800]}", flush=True)
         ref = [np.asarray(o.astype(jnp.float32)) for o in res["dus"][0]]
         line = [f"{name:8s} rows {rows} K {k} N {cols}:"]
-        for mode in [m for m in ("dus", "concat", "pallas", "pallas-kn", "pallas-res", "pallas-man") if m in res]:
+        for mode in [m for m in ("dus", "concat", "pallas") if m in res]:
             outs, ms = res[mode]
             err = max(float(np.abs(np.asarray(o.astype(jnp.float32)) - r).max() / np.abs(r).max())
                       for o, r in zip(outs, ref))
diff --git a/difflet/backends/tpu/ops_impl/collective_matmul.py b/difflet/backends/tpu/ops_impl/collective_matmul.py
index 091970d3..891030c0 100644
--- a/difflet/backends/tpu/ops_impl/collective_matmul.py
+++ b/difflet/backends/tpu/ops_impl/collective_matmul.py
@@ -293,18 +293,19 @@ def body(x, ids, *ws):
             shard_of = jnp.asarray(shard_table, dtype=jnp.int32)
             j = ids[0]
             if use_kernel:
-                from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kp
+                from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kernel
 
-                why = kp.unsupported(x.shape, [w.shape for w in ws], x.dtype.itemsize)
-                if why is None:
+                reason = kernel.unsupported_reason(x.shape, [w.shape for w in ws],
+                                                   x.dtype.itemsize)
+                if reason is None:
                     # The ring position stays a runtime value (rule above): the
-                    # kernel's owners / neighbours are computed from it.
-                    owners, nbrs = kp.ring_tables(j, shard_table, n)
-                    outs = kp.ag_matmul_body(x, ws, owners, nbrs, n=n, resident=True,
-                                             manual=True)
-                    return outs if len(outs) > 1 else outs[0]
+                    # kernel's owner shards / neighbours are computed from it.
+                    owner_shards, neighbors = kernel.ring_tables(j, shard_table, n)
+                    outputs = kernel.all_gather_matmul_body(x, ws, owner_shards, neighbors,
+                                                            ring_size=n)
+                    return outputs if len(outputs) > 1 else outputs[0]
                 print(f"[difflet tpu] {SP_KERNEL_ENV}=pallas: ag x{len(ws)} {tuple(x.shape)} "
-                      f"falls back to the JAX ring ({why})", flush=True)
+                      f"falls back to the JAX ring ({reason})", flush=True)
             rows = x.shape[1]
             hr = rows // 2
             seq = rows * n
diff --git a/difflet/backends/tpu/ops_impl/collective_matmul_pallas.py b/difflet/backends/tpu/ops_impl/collective_matmul_pallas.py
index a0f57a07..c8fff541 100644
--- a/difflet/backends/tpu/ops_impl/collective_matmul_pallas.py
+++ b/difflet/backends/tpu/ops_impl/collective_matmul_pallas.py
@@ -1,436 +1,316 @@
 """Pallas TPU kernel for the SP all-gather + column-parallel matmul ring.
 
-Step 2 of docs/plans/2026-11-10-tpu-collective-matmul-pallas.md: the same
-bidirectional ring as ``collective_matmul``'s JAX ``ag`` op, but each hop's
-product tile goes from VMEM straight to its final rows of the output (one HBM
-write), with no temporary product and no ``dynamic_update_slice`` / gather
-assembly afterwards. The remote copy of the next half overlaps the matmul of
-the current one.
-
-``ag_matmul_body`` is the per-device body, called inside ``shard_map`` over a
-1-D mesh axis (default ``"tp"``) whose order is the ring order. Rank-dependent
-values (the ring position, owners, neighbours) are runtime int32 inputs, never
-Python constants -- rule 1 of the plan; every device runs the same program.
-
-Layout per device (x is this device's sequence shard, split in two halves that
-travel in opposite directions):
-
-    x     (B, R, K)            R = rows per device, even
-    w_i   (N_i, K)             this device's output columns
-    out_i (B, n * R, N_i)      viewed as (B, n, 2, R/2, N_i): [owner shard, half]
-
-The ``pallas_call``s are not marked ``has_side_effects``: through torch_tpu's
+The same bidirectional ring as ``collective_matmul``'s JAX ``ag`` op
+(docs/plans/2026-11-10-tpu-collective-matmul-pallas.md), as one kernel:
+
+- Each device's sequence shard is split into a top and a bottom half. The top
+  halves travel clockwise, the bottom halves counter-clockwise; after
+  ``ring_size - 1`` hops every device has seen every half.
+- The halves stay in VMEM for the whole call: the local ones are copied in once,
+  the remote ones arrive by VMEM-to-VMEM DMA and never touch HBM. The copy of the
+  next hop's halves overlaps the current hop's matmuls.
+- The weights are streamed through VMEM one column tile at a time, as a single
+  double-buffered stream over every (hop, weight, column tile): the next tile is
+  always in flight, across weights and hops, so only the first load is exposed.
+- Each product tile is written once, asynchronously, at its final rows of the
+  output: no temporary product, no ``dynamic_update_slice`` assembly afterwards.
+
+:func:`all_gather_matmul_body` is the per-device body, called inside ``shard_map``
+over a 1-D mesh axis (default ``"tp"``) whose order is the ring order.
+Rank-dependent values (the ring position, the owner shards, the neighbours) are
+runtime int32 inputs, never Python constants -- rule 1 of the plan: every device
+must run the same program.
+
+Layout per device:
+
+    x          (batch, rows, in_features)       rows per device, even
+    weight_i   (out_features_i, in_features)    this device's output columns
+    output_i   (batch, ring_size * rows, out_features_i)
+               written as (batch, ring_size, 2, rows / 2, out_features_i):
+               [owner shard, half]
+
+Communication slots: the halves have three VMEM slots per direction. A
+neighbour is at most one hop behind (each device waits for the previous hop's
+copies before starting the next), so the slot written at hop h + 1 is never the
+one a neighbour still reads; Pallas' race detector catches two slots.
+
+The ``pallas_call`` is not marked ``has_side_effects``: through torch_tpu's
 ``jax_op`` the program goes through XLA's SPMD partitioner, which rejects a
 side-effecting custom call without a sharding ("Side-effect HLO must have
 sharding", 2026-11-11). The outputs are always consumed, so nothing is dropped.
 
-Communication buffers are an extra output (2 directions x 3 slots x the half
-shape) in HBM. Three slots make the reuse race-free: a neighbour is at most one
-hop behind (each device waits for both neighbours' previous hop before going
-on), so the slot written at hop s+1 is never the one a neighbour still reads.
+Earlier variants (one ``emit_pipeline`` per hop and weight, with the halves in
+HBM or in VMEM; pre-transposed weights) were measured and dropped; they are in
+commit 01d027fe and the numbers in the plan.
 """
 
 from __future__ import annotations
 
-SLOTS = 3
-
-
-def ag_matmul_body(x, ws, owners, nbrs, *, n: int, axis: str = "tp", tn: int = 256,
-                   vmem_limit_bytes: int = 96 * 2**20, interpret=False, ablate: str = "",
-                   weights_kn: bool = False, resident: bool = False,
-                   manual: bool = False):
-    """Per-device all-gather + matmul with every weight in ``ws``.
-
-    owners: int32 (2, n) -- owners[0, s] / owners[1, s] is the shard (output row
-        block) whose top / bottom half this device holds after s hops.
-    nbrs:   int32 (2,)   -- mesh index along ``axis`` of the clockwise
-        neighbour (receives the top halves) and the counter-clockwise one.
-    weights_kn: the weights are given transposed, (K, N_i): the kernel then runs a
-        plain (M, K) x (K, N) matmul per tile instead of contracting both
-        operands on their last axis.
-    resident: keep the halves in VMEM for the whole call -- the local ones are
-        copied in once, the remote ones arrive by VMEM-to-VMEM DMA (they never
-        touch HBM), and each matmul pipeline streams only weight tiles. Needs
-        2 x 3 x B x R/2 x K elements of VMEM (Wan: 36 MB).
-    manual: with ``resident``, replace the per-(hop, weight) ``emit_pipeline``s by
-        one hand-written double-buffered stream over every (hop, weight, column
-        tile): the next weight tile is always in flight, across weights and hops,
-        so only the very first tile load is exposed. Output tiles are written back
-        asynchronously, two in flight. The tiles of a hop are a runtime loop, so
-        compile time does not grow with the tile count (a fully unrolled stream
-        took 225 s to compile at Wan's FFN shape). Weights may differ in width;
-        one tile width (a multiple of 128) must divide them all.
-    ablate: profiling only -- ``"comm"`` skips the remote copies (every hop
-        multiplies the local halves; wrong values), ``"compute"`` skips the matmuls.
-    Returns a tuple of (B, n * R, N_i), one per weight.
+#: Ring slots per direction for the travelling halves (see the module docstring).
+RING_SLOTS = 3
+
+#: VMEM the kernel may plan for (v5e has 128 MiB per core; leave room for Mosaic).
+VMEM_BUDGET_BYTES = 96 * 2**20
+
+#: Upper bound on the column tile width. The end-to-end results in the plan
+#: (Wan, FLUX, HunyuanVideo) were measured with this value.
+DEFAULT_COLUMN_TILE_LIMIT = 256
+
+
+def all_gather_matmul_body(x, weights, owner_shards, neighbors, *, ring_size: int,
+                           axis: str = "tp",
+                           column_tile_limit: int = DEFAULT_COLUMN_TILE_LIMIT,
+                           vmem_limit_bytes: int = VMEM_BUDGET_BYTES,
+                           interpret=False, profile_ablation: str = ""):
+    """Per-device all-gather of ``x`` along the ring + matmul with every weight.
+
+    x:            (batch, rows, in_features), this device's sequence shard.
+    weights:      sequence of (out_features_i, in_features); widths may differ, but
+                  one column tile width (a multiple of 128, at most
+                  ``column_tile_limit``) must divide them all.
+    owner_shards: int32 (2, ring_size) -- owner_shards[0, hop] / owner_shards[1, hop]
+                  is the shard (output row block) whose top / bottom half this
+                  device holds after ``hop`` hops. From :func:`ring_tables`.
+    neighbors:    int32 (2,) -- mesh index along ``axis`` of the clockwise neighbour
+                  (receives the top halves) and of the counter-clockwise one.
+    interpret:    False on the chips; ``pltpu.InterpretParams(...)`` for CPU tests.
+    profile_ablation: profiling only, wrong values -- ``"communication"`` skips
+                  the remote copies (every hop multiplies the local halves),
+                  ``"compute"`` skips the matmuls.
+
+    Returns a tuple of (batch, ring_size * rows, out_features_i), one per weight.
     """
     import jax
     import jax.numpy as jnp
     from jax.experimental import pallas as pl
     from jax.experimental.pallas import tpu as pltpu
 
-    b, rows, k = x.shape
+    if profile_ablation not in ("", "communication", "compute"):
+        raise ValueError(f"profile_ablation must be '', 'communication' or 'compute', "
+                         f"got {profile_ablation!r}")
+    batch, rows, in_features = x.shape
     if rows % 2:
         raise ValueError(f"rows per device must be even (two ring halves), got {rows}")
-    hr = rows // 2
-    tns = []
-    kd, nd = (0, 1) if weights_kn else (1, 0)  # weight axis of K / of the columns
-    cols = []
-    for w in ws:
-        if w.shape[kd] != k:
-            raise ValueError(f"weight {w.shape} does not match K={k}")
-        cols.append(w.shape[nd])
-        tns.append(_col_tile(w.shape[nd], tn))
-    nw = len(ws)
-    tm0 = _common_tile(cols, tn)  # the manual stream's one tile width
-
-    def kernel(owners_ref, nbrs_ref, x_ref, *refs):
-        w_refs = refs[:nw]
-        out_refs = refs[nw:2 * nw]
-        buf_ref, send_sem, recv_sem = refs[2 * nw:]
-
-        # Every neighbour must be inside this kernel (its buffers live) before
-        # anything is written into them.
-        barrier = pltpu.get_barrier_semaphore()
-        for d in range(2):
-            pl.semaphore_signal(barrier, 1, device_id={axis: nbrs_ref[d]})
-        pl.semaphore_wait(barrier, 2)
-
-        def held(d, s):
-            # The half travelling in direction d after s hops.
-            return x_ref.at[:, d] if s == 0 or ablate == "comm" else buf_ref.at[d, s % SLOTS]
-
-        def matmul(s, w_ref, out_ref, tn_i):
-            # Both halves held after s hops against one weight tile: each tile
-            # of W is read from HBM once per hop, not once per half.
-            own_top, own_bot = owners_ref[0, s], owners_ref[1, s]
-
-            def body(xt_blk, xb_blk, w_blk, ot_blk, ob_blk):
-                for x_blk, o_blk in ((xt_blk, ot_blk), (xb_blk, ob_blk)):
-                    o_blk[...] = jax.lax.dot_general(
-                        x_blk[...], w_blk[...], (((1,), (kd,)), ((), ())),
-                        preferred_element_type=jnp.float32).astype(o_blk.dtype)
-
-            x_spec = pl.BlockSpec((pl.Squeezed(), hr, k), lambda bi, j: (bi, 0, 0))
-            o_shape = (pl.Squeezed(), pl.Squeezed(), pl.Squeezed(), hr, tn_i)
-            w_spec = (pl.BlockSpec((k, tn_i), lambda bi, j: (0, j)) if weights_kn
-                      else pl.BlockSpec((tn_i, k), lambda bi, j: (j, 0)))
-            pltpu.emit_pipeline(
-                body,
-                grid=(b, w_ref.shape[nd] // tn_i),
-                in_specs=[x_spec, x_spec, w_spec],
-                out_specs=[pl.BlockSpec(o_shape, lambda bi, j: (bi, own_top, 0, 0, j)),
-                           pl.BlockSpec(o_shape, lambda bi, j: (bi, own_bot, 1, 0, j))],
-            )(held(0, s), held(1, s), w_ref, out_ref, out_ref)
-
-        for s in range(n):
-            copies = []
-            if s + 1 < n and ablate != "comm":
-                slot = (s + 1) % SLOTS
-                for d in range(2):
-                    c = pltpu.make_async_remote_copy(
-                        held(d, s), buf_ref.at[d, slot], send_sem.at[d, slot],
-                        recv_sem.at[d, slot], device_id={axis: nbrs_ref[d]})
-                    c.start()
-                    copies.append(c)
-            for w_ref, out_ref, tn_i in zip(w_refs, out_refs, tns):
-                if ablate != "compute":
-                    matmul(s, w_ref, out_ref, tn_i)
-            for c in copies:
-                c.wait_send()
-                c.wait_recv()
-
-    def kernel_resident(owners_ref, nbrs_ref, x_ref, *refs):
-        w_refs = refs[:nw]
-        out_refs = refs[nw:2 * nw]
-        xbuf, load_sem, send_sem, recv_sem = refs[2 * nw:]
-
-        barrier = pltpu.get_barrier_semaphore()
-        for d in range(2):
-            pl.semaphore_signal(barrier, 1, device_id={axis: nbrs_ref[d]})
-        pl.semaphore_wait(barrier, 2)
-
-        loads = [pltpu.make_async_copy(x_ref.at[:, d], xbuf.at[d, 0], load_sem.at[d])
-                 for d in range(2)]
-        for c in loads:
-            c.start()
-        for c in loads:
-            c.wait()
-
-        def held(d, s):
-            return xbuf.at[d, 0 if ablate == "comm" else s % SLOTS]
-
-        def matmul(s, w_ref, out_ref, tn_i):
-            own_top, own_bot = owners_ref[0, s], owners_ref[1, s]
-            xt, xb = held(0, s), held(1, s)
-
-            def body(w_blk, ot_blk, ob_blk):
-                w = w_blk[...]
-                for x_r, o_blk in ((xt, ot_blk), (xb, ob_blk)):
-                    for bi in range(b):
-                        o_blk[bi] = jax.lax.dot_general(
-                            x_r[bi], w, (((1,), (kd,)), ((), ())),
-                            preferred_element_type=jnp.float32).astype(o_blk.dtype)
-
-            o_shape = (b, pl.Squeezed(), pl.Squeezed(), hr, tn_i)
-            w_spec = (pl.BlockSpec((k, tn_i), lambda j: (0, j)) if weights_kn
-                      else pl.BlockSpec((tn_i, k), lambda j: (j, 0)))
-            pltpu.emit_pipeline(
-                body,
-                grid=(w_ref.shape[nd] // tn_i,),
-                in_specs=[w_spec],
-                out_specs=[pl.BlockSpec(o_shape, lambda j: (0, own_top, 0, 0, j)),
-                           pl.BlockSpec(o_shape, lambda j: (0, own_bot, 1, 0, j))],
-            )(w_ref, out_ref, out_ref)
-
-        for s in range(n):
-            copies = []
-            if s + 1 < n and ablate != "comm":
-                slot = (s + 1) % SLOTS
-                for d in range(2):
-                    c = pltpu.make_async_remote_copy(
-                        held(d, s), xbuf.at[d, slot], send_sem.at[d, slot],
-                        recv_sem.at[d, slot], device_id={axis: nbrs_ref[d]})
-                    c.start()
-                    copies.append(c)
-            for w_ref, out_ref, tn_i in zip(w_refs, out_refs, tns):
-                if ablate != "compute":
-                    matmul(s, w_ref, out_ref, tn_i)
-            for c in copies:
-                c.wait_send()
-                c.wait_recv()
-
-    def kernel_manual(owners_ref, nbrs_ref, x_ref, *refs):
-        w_refs = refs[:nw]
-        out_refs = refs[nw:2 * nw]
-        xbuf, wbuf, obuf, load_sem, send_sem, recv_sem, w_sem, o_sem = refs[2 * nw:]
-        tn0 = tm0                    # one column-tile width for every weight
-        cum = [0]                    # first tile of each weight within a hop
-        for c in cols:
-            cum.append(cum[-1] + c // tn0)
-        tph = cum[-1]                # tiles per hop
-        total = n * tph
-
+    half_rows = rows // 2
+    output_widths = []
+    for weight in weights:
+        if weight.shape[1] != in_features:
+            raise ValueError(f"weight {weight.shape} does not match in_features={in_features}")
+        output_widths.append(weight.shape[0])
+    num_weights = len(weights)
+    tile_width = _common_column_tile(output_widths, column_tile_limit)
+    if tile_width is None:
+        raise ValueError(f"no column tile (a multiple of 128) divides every width {output_widths}")
+
+    # The global tile index runs over every (hop, weight, column tile) in that
+    # order; within a hop, weight i owns tiles [first_tile_of_weight[i],
+    # first_tile_of_weight[i + 1]).
+    first_tile_of_weight = [0]
+    for width in output_widths:
+        first_tile_of_weight.append(first_tile_of_weight[-1] + width // tile_width)
+    tiles_per_hop = first_tile_of_weight[-1]
+    total_tiles = ring_size * tiles_per_hop
+    skip_communication = profile_ablation == "communication"
+    skip_compute = profile_ablation == "compute"
+
+    def kernel(owner_shards_ref, neighbors_ref, x_ref, *refs):
+        weight_refs = refs[:num_weights]
+        output_refs = refs[num_weights:2 * num_weights]
+        (halves, weight_tiles, output_tiles, load_semaphores, send_semaphores,
+         receive_semaphores, weight_semaphores, output_semaphores) = refs[2 * num_weights:]
+
+        # Both neighbours must be inside this kernel (their VMEM slots belong to
+        # it) before anything is copied into them.
         barrier = pltpu.get_barrier_semaphore()
-        for d in range(2):
-            pl.semaphore_signal(barrier, 1, device_id={axis: nbrs_ref[d]})
+        for direction in range(2):
+            pl.semaphore_signal(barrier, 1, device_id={axis: neighbors_ref[direction]})
         pl.semaphore_wait(barrier, 2)
 
-        def cols_of(j):
-            return pl.ds(pl.multiple_of(j * tn0, tn0), tn0)
-
-        def w_src(ii, j):
-            return w_refs[ii].at[:, cols_of(j)] if weights_kn else w_refs[ii].at[cols_of(j), :]
-
-        def which(t):
-            # Global tile t -> (weight, column tile of that weight). Weights may
-            # differ in width (HunyuanVideo / FLUX single blocks: q|k|v|mlp), so
-            # each has its own run of tiles within a hop.
-            u = jnp.asarray(t, jnp.int32) % tph
-            i = sum(((u >= cum[kk]).astype(jnp.int32) for kk in range(1, nw)), jnp.int32(0))
-            base = sum((jnp.where(i == kk, cum[kk], 0) for kk in range(1, nw)), jnp.int32(0))
-            return i, u - base
-
-        def w_start(t, slot):
-            # The weights are separate refs, so one guarded DMA start per weight.
-            i, j = which(t)
-            for ii in range(nw):
-                @pl.when(i == ii)
+        def columns_of(column_tile):
+            return pl.ds(pl.multiple_of(column_tile * tile_width, tile_width), tile_width)
+
+        def locate_tile(tile):
+            """Global tile index -> (weight index, column tile of that weight)."""
+            tile_in_hop = jnp.asarray(tile, jnp.int32) % tiles_per_hop
+            weight_index = sum(
+                ((tile_in_hop >= first_tile_of_weight[i]).astype(jnp.int32)
+                 for i in range(1, num_weights)), jnp.int32(0))
+            weight_first_tile = sum(
+                (jnp.where(weight_index == i, first_tile_of_weight[i], 0)
+                 for i in range(1, num_weights)), jnp.int32(0))
+            return weight_index, tile_in_hop - weight_first_tile
+
+        def start_weight_tile_load(tile, buffer_slot):
+            # The weights are separate refs: one guarded DMA start per weight.
+            weight_index, column_tile = locate_tile(tile)
+            for i in range(num_weights):
+                @pl.when(weight_index == i)
                 def _():
-                    pltpu.make_async_copy(w_src(ii, j), wbuf.at[slot], w_sem.at[slot]).start()
-
-        def w_wait(slot):
-            pltpu.make_async_copy(w_src(0, 0), wbuf.at[slot], w_sem.at[slot]).wait()
-
-        def o_wait(slot):
-            for d in range(2):
-                pltpu.make_async_copy(obuf.at[slot, d], out_refs[0].at[:, 0, d, :, cols_of(0)],
-                                      o_sem.at[slot, d]).wait()
-
-        w_start(0, 0)  # the first weight tile overlaps the local loads
-        loads = [pltpu.make_async_copy(x_ref.at[:, d], xbuf.at[d, 0], load_sem.at[d])
-                 for d in range(2)]
-        for c in loads:
-            c.start()
-        for c in loads:
-            c.wait()
-
-        def slot_of(s):
-            return 0 if ablate == "comm" else s % SLOTS
-
-        copies = []
-        for s in range(n):  # hops: static (n = 4), tiles within a hop: a runtime loop
-            for c in copies:  # hop s's halves have arrived
-                c.wait_send()
-                c.wait_recv()
-            copies = []
-            if s + 1 < n and ablate != "comm":
-                for d in range(2):
-                    c = pltpu.make_async_remote_copy(
-                        xbuf.at[d, slot_of(s)], xbuf.at[d, (s + 1) % SLOTS],
-                        send_sem.at[d, (s + 1) % SLOTS], recv_sem.at[d, (s + 1) % SLOTS],
-                        device_id={axis: nbrs_ref[d]})
-                    c.start()
-                    copies.append(c)
-
-            def tile(u, carry, s=s):
-                t = s * tph + u
-                slot = t % 2
-
-                @pl.when(t + 1 < total)
+                    pltpu.make_async_copy(
+                        weight_refs[i].at[columns_of(column_tile), :],
+                        weight_tiles.at[buffer_slot], weight_semaphores.at[buffer_slot]).start()
+
+        def wait_weight_tile_load(buffer_slot):
+            # A wait only needs a descriptor of the same size: any weight's tile 0.
+            pltpu.make_async_copy(weight_refs[0].at[columns_of(0), :],
+                                  weight_tiles.at[buffer_slot],
+                                  weight_semaphores.at[buffer_slot]).wait()
+
+        def wait_output_write_back(buffer_slot):
+            for direction in range(2):
+                pltpu.make_async_copy(output_tiles.at[buffer_slot, direction],
+                                      output_refs[0].at[:, 0, direction, :, columns_of(0)],
+                                      output_semaphores.at[buffer_slot, direction]).wait()
+
+        def ring_slot(hop):
+            return 0 if skip_communication else hop % RING_SLOTS
+
+        # The first weight tile load overlaps the local halves' load.
+        start_weight_tile_load(0, 0)
+        local_loads = [pltpu.make_async_copy(x_ref.at[:, direction], halves.at[direction, 0],
+                                             load_semaphores.at[direction])
+                       for direction in range(2)]
+        for copy in local_loads:
+            copy.start()
+        for copy in local_loads:
+            copy.wait()
+
+        remote_copies = []
+        # Hops: a static loop (ring_size = 4). Tiles within a hop: a runtime loop, so
+        # compile time does not grow with the tile count (fully unrolled, Wan's FFN
+        # took 225 s to compile).
+        for hop in range(ring_size):
+            for copy in remote_copies:  # this hop's halves have arrived
+                copy.wait_send()
+                copy.wait_recv()
+            remote_copies = []
+            if hop + 1 < ring_size and not skip_communication:
+                next_slot = (hop + 1) % RING_SLOTS
+                for direction in range(2):
+                    copy = pltpu.make_async_remote_copy(
+                        halves.at[direction, ring_slot(hop)], halves.at[direction, next_slot],
+                        send_semaphores.at[direction, next_slot],
+                        receive_semaphores.at[direction, next_slot],
+                        device_id={axis: neighbors_ref[direction]})
+                    copy.start()
+                    remote_copies.append(copy)
+
+            def process_tile(tile_in_hop, carry, hop=hop):
+                tile = hop * tiles_per_hop + tile_in_hop
+                buffer_slot = tile % 2
+
+                @pl.when(tile + 1 < total_tiles)
                 def _():
-                    w_start(t + 1, 1 - slot)  # always one weight tile in flight
+                    start_weight_tile_load(tile + 1, 1 - buffer_slot)  # always one in flight
 
-                w_wait(slot)
+                wait_weight_tile_load(buffer_slot)
 
-                @pl.when(t >= 2)
+                @pl.when(tile >= 2)
                 def _():
-                    o_wait(slot)  # this output slot's previous write-back is done
-
-                if ablate != "compute":
-                    w = wbuf[slot]
-                    for d in range(2):
-                        for bi in range(b):
-                            obuf[slot, d, bi] = jax.lax.dot_general(
-                                xbuf[d, slot_of(s), bi], w, (((1,), (kd,)), ((), ())),
-                                preferred_element_type=jnp.float32).astype(obuf.dtype)
-                i, j = which(t)
-                for ii in range(nw):
-                    @pl.when(i == ii)
+                    wait_output_write_back(buffer_slot)  # this slot's previous tile is out
+
+                if not skip_compute:
+                    weight_tile = weight_tiles[buffer_slot]
+                    for direction in range(2):
+                        for batch_index in range(batch):
+                            output_tiles[buffer_slot, direction, batch_index] = jax.lax.dot_general(
+                                halves[direction, ring_slot(hop), batch_index], weight_tile,
+                                (((1,), (1,)), ((), ())),
+                                preferred_element_type=jnp.float32).astype(output_tiles.dtype)
+                weight_index, column_tile = locate_tile(tile)
+                for i in range(num_weights):
+                    @pl.when(weight_index == i)
                     def _():
-                        for d in range(2):
+                        for direction in range(2):
+                            owner = owner_shards_ref[direction, hop]
                             pltpu.make_async_copy(
-                                obuf.at[slot, d],
-                                out_refs[ii].at[:, owners_ref[d, s], d, :, cols_of(j)],
-                                o_sem.at[slot, d]).start()
+                                output_tiles.at[buffer_slot, direction],
+                                output_refs[i].at[:, owner, direction, :, columns_of(column_tile)],
+                                output_semaphores.at[buffer_slot, direction]).start()
                 return carry
 
-            jax.lax.fori_loop(0, tph, tile, 0)
-        for c in copies:
-            c.wait_send()
-            c.wait_recv()
-        for t in range(max(0, total - 2), total):  # the last write-backs
-            o_wait(t % 2)
-
-    any_spec = pl.BlockSpec(memory_space=pl.ANY)
-    if resident and manual:
-        if tm0 is None:
-            raise ValueError(f"no column tile (a multiple of 128) divides every width {cols}")
-        wshape = (k, tm0) if weights_kn else (tm0, k)
-        outs = pl.pallas_call(
-            kernel_manual,
-            out_shape=[jax.ShapeDtypeStruct((b, n, 2, hr, c), x.dtype) for c in cols],
-            grid_spec=pltpu.PrefetchScalarGridSpec(
-                num_scalar_prefetch=2,
-                in_specs=[any_spec] * (1 + nw),
-                out_specs=[any_spec] * nw,
-                scratch_shapes=[pltpu.VMEM((2, SLOTS, b, hr, k), x.dtype),
-                                pltpu.VMEM((2,) + wshape, x.dtype),
-                                pltpu.VMEM((2, 2, b, hr, tm0), x.dtype),
-                                pltpu.SemaphoreType.DMA((2,)),
-                                pltpu.SemaphoreType.DMA((2, SLOTS)),
-                                pltpu.SemaphoreType.DMA((2, SLOTS)),
-                                pltpu.SemaphoreType.DMA((2,)),
-                                pltpu.SemaphoreType.DMA((2, 2))],
-            ),
-            compiler_params=pltpu.CompilerParams(collective_id=0,
-                                                 vmem_limit_bytes=vmem_limit_bytes),
-            interpret=interpret,
-        )(owners, nbrs, x.reshape(b, 2, hr, k), *ws)
-        return tuple(o.reshape(b, n * rows, o.shape[-1]) for o in outs)
-    if resident:
-        outs = pl.pallas_call(
-            kernel_resident,
-            out_shape=[jax.ShapeDtypeStruct((b, n, 2, hr, c), x.dtype) for c in cols],
-            grid_spec=pltpu.PrefetchScalarGridSpec(
-                num_scalar_prefetch=2,
-                in_specs=[any_spec] * (1 + nw),
-                out_specs=[any_spec] * nw,
-                scratch_shapes=[pltpu.VMEM((2, SLOTS, b, hr, k), x.dtype),
-                                pltpu.SemaphoreType.DMA((2,)),
-                                pltpu.SemaphoreType.DMA((2, SLOTS)),
-                                pltpu.SemaphoreType.DMA((2, SLOTS))],
-            ),
-            compiler_params=pltpu.CompilerParams(collective_id=0,
-                                                 vmem_limit_bytes=vmem_limit_bytes),
-            interpret=interpret,
-        )(owners, nbrs, x.reshape(b, 2, hr, k), *ws)
-        return tuple(o.reshape(b, n * rows, o.shape[-1]) for o in outs)
-
-    out_shape = [jax.ShapeDtypeStruct((b, n, 2, hr, c), x.dtype) for c in cols]
-    out_shape.append(jax.ShapeDtypeStruct((2, SLOTS, b, hr, k), x.dtype))
-    outs = pl.pallas_call(
+            jax.lax.fori_loop(0, tiles_per_hop, process_tile, 0)
+        for copy in remote_copies:
+            copy.wait_send()
+            copy.wait_recv()
+        for tile in range(max(0, total_tiles - 2), total_tiles):  # the last write-backs
+            wait_output_write_back(tile % 2)
+
+    any_memory = pl.BlockSpec(memory_space=pl.ANY)
+    outputs = pl.pallas_call(
         kernel,
-        out_shape=out_shape,
+        out_shape=[jax.ShapeDtypeStruct((batch, ring_size, 2, half_rows, width), x.dtype)
+                   for width in output_widths],
         grid_spec=pltpu.PrefetchScalarGridSpec(
             num_scalar_prefetch=2,
-            in_specs=[any_spec] * (1 + nw),
-            out_specs=[any_spec] * (nw + 1),
-            scratch_shapes=[pltpu.SemaphoreType.DMA((2, SLOTS)),
-                            pltpu.SemaphoreType.DMA((2, SLOTS))],
+            in_specs=[any_memory] * (1 + num_weights),
+            out_specs=[any_memory] * num_weights,
+            scratch_shapes=[
+                pltpu.VMEM((2, RING_SLOTS, batch, half_rows, in_features), x.dtype),  # halves
+                pltpu.VMEM((2, tile_width, in_features), x.dtype),             # weight tiles
+                pltpu.VMEM((2, 2, batch, half_rows, tile_width), x.dtype),     # output tiles
+                pltpu.SemaphoreType.DMA((2,)),               # local loads, per direction
+                pltpu.SemaphoreType.DMA((2, RING_SLOTS)),    # remote sends
+                pltpu.SemaphoreType.DMA((2, RING_SLOTS)),    # remote receives
+                pltpu.SemaphoreType.DMA((2,)),               # weight tile loads
+                pltpu.SemaphoreType.DMA((2, 2)),             # output write-backs
+            ],
         ),
         compiler_params=pltpu.CompilerParams(collective_id=0,
                                              vmem_limit_bytes=vmem_limit_bytes),
         interpret=interpret,
-    )(owners, nbrs, x.reshape(b, 2, hr, k), *ws)
-    return tuple(o.reshape(b, n * rows, o.shape[-1]) for o in outs[:nw])
-
-
-#: VMEM the kernel may plan for (v5e has 128 MiB per core; leave room for Mosaic).
-VMEM_BUDGET = 96 * 2**20
+    )(owner_shards, neighbors, x.reshape(batch, 2, half_rows, in_features), *weights)
+    return tuple(output.reshape(batch, ring_size * rows, output.shape[-1])
+                 for output in outputs)
 
 
-def unsupported(x_shape, w_shapes, itemsize: int, tn: int = 512) -> str | None:
-    """Why the streamed kernel (``resident=True, manual=True``, weights (N, K))
-    cannot take these shapes, or None. Callers fall back to the JAX ring."""
-    b, rows, k = x_shape
+def unsupported_reason(x_shape, weight_shapes, itemsize: int,
+                       column_tile_limit: int = DEFAULT_COLUMN_TILE_LIMIT) -> str | None:
+    """Why :func:`all_gather_matmul_body` cannot take these shapes, or None.
+    Callers fall back to the JAX ring."""
+    batch, rows, in_features = x_shape
     if rows % 2:
         return f"odd rows per device ({rows})"
-    if any(ws[1] != k for ws in w_shapes):
-        return "weight K does not match the input"
-    widths = [ws[0] for ws in w_shapes]
-    t = _common_tile(widths, tn)
-    if t is None:
+    if any(shape[1] != in_features for shape in weight_shapes):
+        return "weight in_features does not match the input"
+    widths = [shape[0] for shape in weight_shapes]
+    tile_width = _common_column_tile(widths, column_tile_limit)
+    if tile_width is None:
         return f"no column tile (a multiple of 128) divides every width {sorted(set(widths))}"
-    hr = rows // 2
-    need = itemsize * (2 * SLOTS * b * hr * k   # resident halves
-                       + 2 * t * k              # weight tiles, double-buffered
-                       + 2 * 2 * b * hr * t)    # output tiles, double-buffered
-    need += 4 * hr * t                          # one fp32 product
-    if need > VMEM_BUDGET:
-        return f"needs {need / 2**20:.0f} MiB of VMEM (budget {VMEM_BUDGET / 2**20:.0f})"
+    half_rows = rows // 2
+    needed = itemsize * (2 * RING_SLOTS * batch * half_rows * in_features  # resident halves
+                         + 2 * tile_width * in_features                    # weight tiles
+                         + 2 * 2 * batch * half_rows * tile_width)         # output tiles
+    needed += 4 * half_rows * tile_width                                   # one fp32 product
+    if needed > VMEM_BUDGET_BYTES:
+        return (f"needs {needed / 2**20:.0f} MiB of VMEM "
+                f"(budget {VMEM_BUDGET_BYTES / 2**20:.0f})")
     return None
 
 
-def _common_tile(widths, limit: int) -> int | None:
-    """The largest multiple of 128 up to ``limit`` that divides every width, or
-    None (q|k|v 768 + mlp 3072 per device: 384)."""
-    for t in range(limit - limit % 128, 0, -128):
-        if all(w % t == 0 for w in widths):
-            return t
+def _common_column_tile(widths, limit: int) -> int | None:
+    """The largest multiple of 128 up to ``limit`` that divides every width, or None
+    (q|k|v 768 + mlp 3072 per device, limit 256: 256)."""
+    for tile_width in range(limit - limit % 128, 0, -128):
+        if all(width % tile_width == 0 for width in widths):
+            return tile_width
     return None
 
 
-def _col_tile(cols: int, limit: int) -> int:
-    """The largest multiple of 128 up to ``limit`` that divides ``cols``, else ``cols``
-    (Wan's FFN has 3456 columns per device: 384)."""
-    for t in range(limit - limit % 128, 0, -128):
-        if cols % t == 0:
-            return t
-    return cols
-
-
-def ring_tables(j, shard_of_pos, n: int):
-    """Runtime owners / neighbours for ring position ``j`` (a traced int32).
+def ring_tables(ring_position, shard_of_position, ring_size: int):
+    """Runtime owner shards and neighbours for ``ring_position`` (a traced int32).
 
-    After s hops a device at position j holds position (j - s)'s top half
-    (sent clockwise) and (j + s)'s bottom half; ``shard_of_pos`` maps a ring
-    position to the shard (output row block) it owns.
+    After h hops the device at position p holds position (p - h)'s top half (sent
+    clockwise) and position (p + h)'s bottom half; ``shard_of_position`` maps a
+    ring position to the shard (output row block) it owns.
+    Returns (owner_shards int32 (2, ring_size), neighbors int32 (2,)).
     """
     import jax.numpy as jnp
 
-    shard_of = jnp.asarray(shard_of_pos, dtype=jnp.int32)
-    s = jnp.arange(n, dtype=jnp.int32)
-    owners = jnp.stack([shard_of[(j - s) % n], shard_of[(j + s) % n]])
-    nbrs = jnp.stack([(j + 1) % n, (j - 1) % n]).astype(jnp.int32)
-    return owners, nbrs
+    shard_of = jnp.asarray(shard_of_position, dtype=jnp.int32)
+    hops = jnp.arange(ring_size, dtype=jnp.int32)
+    owner_shards = jnp.stack([shard_of[(ring_position - hops) % ring_size],
+                              shard_of[(ring_position + hops) % ring_size]])
+    neighbors = jnp.stack([(ring_position + 1) % ring_size,
+                           (ring_position - 1) % ring_size]).astype(jnp.int32)
+    return owner_shards, neighbors
diff --git a/tests/unit/backends/test_tpu_collective_matmul_pallas.py b/tests/unit/backends/test_tpu_collective_matmul_pallas.py
index 2bcf3693..ce0bc468 100644
--- a/tests/unit/backends/test_tpu_collective_matmul_pallas.py
+++ b/tests/unit/backends/test_tpu_collective_matmul_pallas.py
@@ -1,8 +1,8 @@
 """The Pallas all-gather + matmul ring kernel, in TPU interpret mode on 4 CPU devices.
 
 Checks values against numpy with a non-identity ring position -> shard table,
-two weights of different widths, and batch 1 and 2, with Pallas' race detector
-on (two communication slots instead of three is caught by it). Needs ``jax``;
+two weights of equal or different widths, and batch 1 and 2, with Pallas' race detector
+on (two ring slots instead of three is caught by it). Needs ``jax``;
 skipped without it (it is not a test dependency).
 """
 
@@ -26,45 +26,54 @@
     from jax.experimental.pallas import tpu as pltpu
     # The pipeline's tiling inference needs a TPU generation: treat the host as a v5e.
     ti.registry["cpu"] = lambda: ti.get_tpu_info_for_chip(ti.ChipVersion.TPU_V5E, 1)
-    from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kp
+    from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kernel
 
-    n, B, rows, K = 4, int(sys.argv[1]), 16, 256
-    kn = "kn" in sys.argv[2]
-    resident = "res" in sys.argv[2]
-    manual = "man" in sys.argv[2]
-    mesh = jax.sharding.Mesh(np.array(jax.devices()[:n]), ("tp",))
-    shard_of_pos = (2, 0, 3, 1)  # ring position -> shard, not the identity
+    ring_size, batch, rows, in_features = 4, int(sys.argv[1]), 16, 256
+    widths = (256, 384) if sys.argv[2] == "mixed" else (256, 256)
+    mesh = jax.sharding.Mesh(np.array(jax.devices()[:ring_size]), ("tp",))
+    shard_of_position = (2, 0, 3, 1)  # ring position -> shard, not the identity
     rng = np.random.default_rng(0)
-    x_by_shard = [rng.standard_normal((B, rows, K)).astype(np.float32) for _ in range(n)]
-    ws = [rng.standard_normal((N * n, K)).astype(np.float32) / K ** 0.5 for N in ((256, 384) if "mixed" in sys.argv[2] else (256, 256) if manual else (256, 384))]
+    x_by_shard = [rng.standard_normal((batch, rows, in_features)).astype(np.float32)
+                  for _ in range(ring_size)]
+    weights = [rng.standard_normal((width * ring_size, in_features)).astype(np.float32)
+               / in_features ** 0.5 for width in widths]
     put = lambda a, spec: jax.device_put(jnp.asarray(a), NamedSharding(mesh, spec))
-    x = put(np.concatenate([x_by_shard[shard_of_pos[p]] for p in range(n)], axis=1),
+    x = put(np.concatenate([x_by_shard[shard_of_position[p]] for p in range(ring_size)], axis=1),
             P(None, "tp", None))
-    ids = put(np.arange(n, dtype=np.int32), P("tp"))
+    ids = put(np.arange(ring_size, dtype=np.int32), P("tp"))
 
-    def body(x, ids, *ws):
-        owners, nbrs = kp.ring_tables(ids[0], shard_of_pos, n)
-        return kp.ag_matmul_body(x, ws, owners, nbrs, n=n, tn=128 if manual else 256,
-                                 weights_kn=kn, resident=resident or manual, manual=manual,
-                                 interpret=pltpu.InterpretParams(detect_races=True))
+    def body(x, ids, *weights):
+        owner_shards, neighbors = kernel.ring_tables(ids[0], shard_of_position, ring_size)
+        return kernel.all_gather_matmul_body(
+            x, weights, owner_shards, neighbors, ring_size=ring_size, column_tile_limit=128,
+            interpret=pltpu.InterpretParams(detect_races=True))
 
     f = jax.jit(jax.shard_map(body, mesh=mesh,
-                              in_specs=(P(None, "tp", None), P("tp"))
-                              + ((P(None, "tp"),) if kn else (P("tp", None),)) * 2,
+                              in_specs=(P(None, "tp", None), P("tp")) + (P("tp", None),) * 2,
                               out_specs=(P(None, None, "tp"),) * 2, check_vma=False))
-    outs = f(x, ids, *[put(w.T, P(None, "tp")) if kn else put(w, P("tp", None)) for w in ws])
+    outputs = f(x, ids, *[put(w, P("tp", None)) for w in weights])
     full = np.concatenate(x_by_shard, axis=1)  # the sequence in shard order
-    for o, w in zip(outs, ws):
-        np.testing.assert_allclose(np.asarray(o), full @ w.T, rtol=2e-5, atol=2e-5)
-    print("OK", B, sys.argv[2])
+    for output, weight in zip(outputs, weights):
+        np.testing.assert_allclose(np.asarray(output), full @ weight.T, rtol=2e-5, atol=2e-5)
+    print("OK", batch, sys.argv[2])
 ''')
 
 
-@pytest.mark.parametrize("batch,layout", [(1, "nk"), (2, "nk"), (2, "kn"), (1, "kn-res"), (2, "nk-res"),
-                                         (1, "kn-man"), (2, "nk-man"),
-                                         (2, "nk-man-mixed")])
+@pytest.mark.parametrize("batch,layout", [(1, "equal"), (2, "equal"), (2, "mixed")])
 def test_pallas_ag_matmul_matches_numpy(batch, layout):
     proc = subprocess.run([sys.executable, "-c", SCRIPT, str(batch), layout], capture_output=True,
                           text=True, timeout=600)
     assert proc.returncode == 0, proc.stdout[-2000:] + proc.stderr[-4000:]
     assert f"OK {batch} {layout}" in proc.stdout
+
+
+def test_unsupported_reason():
+    from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kernel
+
+    # Wan q/k/v and FLUX / HunyuanVideo single blocks (q|k|v 768 + mlp 3072) fit.
+    assert kernel.unsupported_reason((1, 1170, 5120), [(1280, 5120)] * 3, 2) is None
+    assert kernel.unsupported_reason((1, 1024, 3072), [(768, 3072)] * 3 + [(3072, 3072)], 2) is None
+    assert "odd rows" in kernel.unsupported_reason((1, 1171, 5120), [(1280, 5120)], 2)
+    assert "in_features" in kernel.unsupported_reason((1, 1170, 5120), [(1280, 4096)], 2)
+    assert "no column tile" in kernel.unsupported_reason((1, 1170, 5120), [(1000, 5120)], 2)
+    assert "VMEM" in kernel.unsupported_reason((1, 8192, 8192), [(1024, 8192)], 2)

From 84fe6b1312b0e14a9f75b5fd45788b5d3fa9202f Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Wed, 11 Nov 2026 12:20:19 -0800
Subject: [PATCH 19/22] feat(tpu): Pallas matmul + reduce-scatter kernel for
 the SP rings (opt-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>
---
 benchmark/tpu_profile/cm_rs_pallas_bench.py   | 127 +++++++
 .../tpu/ops_impl/collective_matmul.py         |  37 ++-
 .../tpu/ops_impl/collective_matmul_pallas.py  | 310 ++++++++++++++++++
 .../test_tpu_collective_matmul_pallas_rs.py   |  87 +++++
 4 files changed, 558 insertions(+), 3 deletions(-)
 create mode 100644 benchmark/tpu_profile/cm_rs_pallas_bench.py
 create mode 100644 tests/unit/backends/test_tpu_collective_matmul_pallas_rs.py

diff --git a/benchmark/tpu_profile/cm_rs_pallas_bench.py b/benchmark/tpu_profile/cm_rs_pallas_bench.py
new file mode 100644
index 00000000..7f159443
--- /dev/null
+++ b/benchmark/tpu_profile/cm_rs_pallas_bench.py
@@ -0,0 +1,127 @@
+"""Matmul + reduce-scatter on the chips: the JAX ring vs the Pallas kernel.
+
+    python -m benchmark.tpu_profile.cm_rs_pallas_bench [--iters 50] [--shapes wan_attn_out,...]
+
+Plain JAX, one process owning the four chips (no torch_tpu), as cm_pallas_bench:
+both are ``collective_matmul._op("rs", pairs)`` with ``torch.tpu.pallas.jax_op``
+stubbed out, built once with ``DIFFLET_TPU_SP_RS_KERNEL=jax`` and once with
+``pallas``, so the kernel runs through the same entry as in the models. For each
+shape, prints the time per call and the kernel's largest difference from the
+ring relative to the ring's largest value (bf16 inputs).
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+import types
+
+#: name -> (rows per device, this device's in_features of each pair, out_features)
+SHAPES = {
+    "wan_attn_out": (1170, (1280,), 5120),
+    "wan_ffn_down": (1170, (3456,), 5120),
+    # Wan's FFN down with rows a multiple of 8: no relayout of the input before the kernel.
+    "wan_ffn_down_aligned": (1168, (3456,), 5120),
+    "ltx_attn_out": (578, (1024,), 4096),
+    "flux_single": (1152, (768, 3072), 3072),
+}
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--iters", type=int, default=50)
+    parser.add_argument("--shapes", default=",".join(SHAPES))
+    parser.add_argument("--ablate", action="store_true",
+                        help="also time the kernel without communication / without compute")
+    parser.add_argument("--column-tile-limits", default="",
+                        help="comma-separated: also time the kernel at these column tile limits")
+    args = parser.parse_args()
+
+    import jax
+    import jax.numpy as jnp
+    import numpy as np
+    import torch
+    from jax.sharding import NamedSharding
+    from jax.sharding import PartitionSpec as P
+
+    # No torch_tpu here: jax_op returns the JAX function itself. is_initialized is
+    # asked at interpreter exit.
+    torch.tpu = types.SimpleNamespace(is_initialized=lambda: False,
+                                      pallas=types.SimpleNamespace(jax_op=lambda name, fn, **kw: fn))
+    from difflet.backends.tpu.ops_impl import collective_matmul as cm
+    from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kernel
+
+    devices = jax.devices()
+    n = len(devices)
+    mesh = jax.sharding.Mesh(np.array(cm._snake(devices)), ("tp",))
+    cm._STATE = {"mesh": mesh, "tp": n, "shard_of_pos": tuple(range(n)), "ids": None, "ops": {}}
+
+    def put(a, spec, dtype=jnp.bfloat16):
+        return jax.device_put(jnp.asarray(a, dtype), NamedSharding(mesh, spec))
+
+    def timed(f, *xs):
+        out = f(*xs)
+        jax.block_until_ready(out)
+        for _ in range(3):
+            jax.block_until_ready(f(*xs))
+        t = time.perf_counter()
+        for _ in range(args.iters):
+            out = f(*xs)
+        jax.block_until_ready(out)
+        return out, (time.perf_counter() - t) / args.iters * 1e3
+
+    print(f"jax {jax.__version__} {devices[0].device_kind} x{n}, iters {args.iters}", flush=True)
+    ids = put(np.arange(n), P("tp"), jnp.int32)
+    rng = np.random.default_rng(0)
+    for name in args.shapes.split(","):
+        rows, local_in_features, out_features = SHAPES[name]
+        operands = []
+        for k in local_in_features:
+            operands.append(put(rng.standard_normal((1, rows * n, k * n), dtype=np.float32),
+                                P(None, None, "tp")))
+            operands.append(put(rng.standard_normal((out_features, k * n), dtype=np.float32)
+                                / (k * n) ** 0.5, P(None, "tp")))
+        results = {}
+        for mode in ("jax", "pallas"):
+            os.environ[cm.SP_RS_KERNEL_ENV] = mode
+            try:
+                op = jax.jit(cm._op("rs", len(local_in_features)))
+                results[mode] = timed(op, *operands, ids)
+            except Exception as exc:  # noqa: BLE001 -- report and go on to the next shape
+                print(f"{name}: {mode} failed: {type(exc).__name__}: {str(exc)[:800]}", flush=True)
+        if len(results) < 2:
+            continue
+        ring = np.asarray(results["jax"][0].astype(jnp.float32))
+        kern = np.asarray(results["pallas"][0].astype(jnp.float32))
+        diff = float(np.abs(kern - ring).max() / np.abs(ring).max())
+        if args.ablate:
+            pairs = len(local_in_features)
+
+            def ablated(ablation, pairs=pairs, column_tile_limit=kernel.DEFAULT_COLUMN_TILE_LIMIT):
+                def body(*args):
+                    xs, ws, ids = args[0:2 * pairs:2], args[1:2 * pairs:2], args[-1]
+                    chunk_shards, neighbors = kernel.reduce_scatter_tables(ids[0], tuple(range(n)), n)
+                    return kernel.matmul_reduce_scatter_body(
+                        xs, ws, chunk_shards, neighbors, ring_size=n, profile_ablation=ablation,
+                        column_tile_limit=column_tile_limit)
+                return jax.jit(jax.shard_map(
+                    body, mesh=mesh, in_specs=(P(None, None, "tp"), P(None, "tp")) * pairs + (P("tp"),),
+                    out_specs=P(None, "tp", None), check_vma=False))
+
+            print(f"{name:13s} ablation (ms): " + ", ".join(
+                f"no-{a} {timed(ablated(a), *operands, ids)[1]:.3f}"
+                for a in ("communication", "compute")), flush=True)
+            for limit in [int(v) for v in args.column_tile_limits.split(",") if v]:
+                print(f"{name:13s} column tile <= {limit}: "
+                      f"{timed(ablated('', column_tile_limit=limit), *operands, ids)[1]:.3f} ms",
+                      flush=True)
+        print(f"{name:13s} rows {rows} K/dev {local_in_features} N {out_features}: "
+              f"ring {results['jax'][1]:.3f} ms  kernel {results['pallas'][1]:.3f} ms  "
+              f"(max|d|/max {diff:.1e})", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/difflet/backends/tpu/ops_impl/collective_matmul.py b/difflet/backends/tpu/ops_impl/collective_matmul.py
index 891030c0..42cd07f2 100644
--- a/difflet/backends/tpu/ops_impl/collective_matmul.py
+++ b/difflet/backends/tpu/ops_impl/collective_matmul.py
@@ -103,6 +103,24 @@ def sp_kernel_mode() -> str:
     return mode
 
 
+#: What runs the matmul + reduce-scatter. ``jax`` (default): the ring below.
+#: ``pallas``: ``collective_matmul_pallas.matmul_reduce_scatter_body`` (partial
+#: sums computed while the accumulators travel VMEM to VMEM, no partial product
+#: in HBM); shapes it does not take fall back to the ring with a notice. Off by
+#: default: op-level (v5e-4, 2026-11-11) Wan attention out 0.518 -> 0.504 ms and
+#: LTX-2's 0.214 -> 0.212 (ICI-bound), but Wan FFN down 0.928 -> 1.040 and FLUX's
+#: single-block pair 0.627 -> 0.663 (the kernel's matmul is slower than XLA's);
+#: Wan's step did not move (262.2 -> 262.7 ms).
+SP_RS_KERNEL_ENV = "DIFFLET_TPU_SP_RS_KERNEL"
+
+
+def sp_rs_kernel_mode() -> str:
+    mode = os.environ.get(SP_RS_KERNEL_ENV, "").strip().lower() or "jax"
+    if mode not in ("jax", "pallas"):
+        raise ValueError(f"{SP_RS_KERNEL_ENV} must be 'jax' or 'pallas', got {mode!r}")
+    return mode
+
+
 def _source_tag() -> str:
     """A short hash of this module's and the Pallas kernel's source, put in every
     op name. torch_tpu's compilation cache keys a ``jax_op`` by its name and
@@ -209,7 +227,7 @@ def _op(kind: str, n_weights: int):
     pairs whose products are summed, 'ag' over ``n_weights`` weights sharing one
     gathered input."""
     key = (kind, n_weights,
-           (ag_assemble_mode(), sp_kernel_mode()) if kind == "ag" else None)
+           (ag_assemble_mode(), sp_kernel_mode()) if kind == "ag" else sp_rs_kernel_mode())
     ops = _STATE["ops"]
     if key in ops:
         return ops[key]
@@ -223,7 +241,7 @@ def _op(kind: str, n_weights: int):
     mesh, n = _STATE["mesh"], _STATE["tp"]
     shard_table = list(_STATE["shard_of_pos"])
     assemble = ag_assemble_mode()
-    use_kernel = kind == "ag" and sp_kernel_mode() == "pallas"
+    use_kernel = (sp_kernel_mode() if kind == "ag" else sp_rs_kernel_mode()) == "pallas"
     cw = [(i, (i + 1) % n) for i in range(n)]
     ccw = [(i, (i - 1) % n) for i in range(n)]
     shard_map = jax.shard_map if hasattr(jax, "shard_map") else \
@@ -239,6 +257,18 @@ def body(*args):
             xs, ws = pairs[0::2], pairs[1::2]
             shard_of = jnp.asarray(shard_table, dtype=jnp.int32)
             j = ids[0]
+            if use_kernel:
+                from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kernel
+
+                reason = kernel.reduce_scatter_unsupported_reason(
+                    [x.shape for x in xs], [w.shape for w in ws], xs[0].dtype.itemsize, n)
+                if reason is None:
+                    # The ring position stays a runtime value (rule above).
+                    chunk_shards, neighbors = kernel.reduce_scatter_tables(j, shard_table, n)
+                    return kernel.matmul_reduce_scatter_body(xs, ws, chunk_shards, neighbors,
+                                                             ring_size=n)
+                print(f"[difflet tpu] {SP_RS_KERNEL_ENV}=pallas: rs x{len(xs)} "
+                      f"{tuple(xs[0].shape)} falls back to the JAX ring ({reason})", flush=True)
             rows = xs[0].shape[1] // n
             half = ws[0].shape[0] // 2
 
@@ -435,7 +465,8 @@ def supports(seq_len: int, tp: int) -> bool:
     return seq_len % (2 * tp) == 0
 
 
-__all__ = ["DEFAULT_OPS", "SP_ENV", "SP_KERNEL_ENV", "all_gather_matmul", "matmul2_reduce_scatter",
+__all__ = ["DEFAULT_OPS", "SP_ENV", "SP_KERNEL_ENV", "SP_RS_KERNEL_ENV", "all_gather_matmul",
+           "matmul2_reduce_scatter",
            "set_kernel_default",
            "matmul_reduce_scatter", "prepare",
            "require_compiled", "requested", "supports"]
diff --git a/difflet/backends/tpu/ops_impl/collective_matmul_pallas.py b/difflet/backends/tpu/ops_impl/collective_matmul_pallas.py
index c8fff541..07e503c5 100644
--- a/difflet/backends/tpu/ops_impl/collective_matmul_pallas.py
+++ b/difflet/backends/tpu/ops_impl/collective_matmul_pallas.py
@@ -314,3 +314,313 @@ def ring_tables(ring_position, shard_of_position, ring_size: int):
     neighbors = jnp.stack([(ring_position + 1) % ring_size,
                            (ring_position - 1) % ring_size]).astype(jnp.int32)
     return owner_shards, neighbors
+
+
+# ---------------------------------------------------------------------------
+# Row-parallel matmul + reduce-scatter
+# ---------------------------------------------------------------------------
+
+#: Receive slots per direction for the travelling accumulators. Accumulators move
+#: tile by tile, so a neighbour can run up to two hops ahead in one direction: it
+#: writes a slot again only three hops later, after it has received this device's
+#: other-direction tiles of the hop in between, which this device sends only after
+#: it has read the slot. The CPU test's race detector catches two slots.
+REDUCE_SCATTER_SLOTS = 3
+
+#: Row chunks are DMA'd straight from the (batch, sequence, in_features) input when
+#: their first row is a multiple of this (Mosaic only slices tiled dims in whole
+#: tiles); otherwise the input is reshaped to (batch, ring_size, rows, ...) first.
+_ROW_TILE = 8
+
+
+def matmul_reduce_scatter_body(xs, weights, chunk_shards, neighbors, *, ring_size: int,
+                               axis: str = "tp",
+                               column_tile_limit: int = DEFAULT_COLUMN_TILE_LIMIT,
+                               vmem_limit_bytes: int = VMEM_BUDGET_BYTES,
+                               interpret=False, profile_ablation: str = ""):
+    """Per-device row-parallel matmul over the (x_i, weight_i) pairs, summed, then
+    reduce-scattered along the ring by sequence.
+
+    The same bidirectional ring as ``collective_matmul``'s JAX ``rs`` op: the
+    output's left half of the columns accumulates clockwise, the right half
+    counter-clockwise. At each hop a device multiplies one row chunk of its
+    inputs by its weights (the partial sum for the chunk's owner), adds the
+    accumulator that arrived from upstream and sends the sum on; after
+    ``ring_size`` hops each device holds the full sum for its own rows.
+
+    - Tile by tile: each output column tile is multiplied, the upstream
+      accumulator's same tile is added (it was sent about a hop earlier, so it
+      has arrived) and the sum is sent on at once. The add runs in the matmul's
+      epilogue, the transfers overlap the following tiles' matmuls.
+    - Accumulators travel VMEM to VMEM; the last hop writes its sums straight to
+      the output. No partial product, accumulator or concatenation goes to HBM.
+    - The weights stream through VMEM one column tile at a time, one
+      double-buffered stream over every (hop, direction, column tile); the input
+      row chunks of the next hop load while the current hop computes.
+    - Rounding as the JAX ring: each pair's product is rounded to the input
+      dtype, the pairs are summed in the input dtype, then the accumulator.
+
+    xs:           sequence of (batch, ring_size * rows, in_features_i): the full
+                  sequence, this device's slice of the contracted features.
+    weights:      sequence of (out_features, in_features_i), one per x; the same
+                  out_features, even, its half divisible by one column tile (a
+                  multiple of 128, at most ``column_tile_limit``).
+    chunk_shards: int32 (2, ring_size) -- chunk_shards[0, hop] / [1, hop] is the
+                  shard (row chunk) whose left / right half of the columns this
+                  device computes at ``hop``. From :func:`reduce_scatter_tables`.
+    neighbors:    int32 (2,) -- mesh index along ``axis`` of the clockwise
+                  neighbour (receives the left halves) and the counter-clockwise one.
+    profile_ablation: profiling only, wrong values -- ``"communication"`` skips
+                  the transfers and the adds, ``"compute"`` skips the matmuls.
+
+    Returns (batch, rows, out_features): this device's rows of the sum.
+    """
+    import jax
+    import jax.numpy as jnp
+    from jax.experimental import pallas as pl
+    from jax.experimental.pallas import tpu as pltpu
+
+    if profile_ablation not in ("", "communication", "compute"):
+        raise ValueError(f"profile_ablation must be '', 'communication' or 'compute', "
+                         f"got {profile_ablation!r}")
+    num_pairs = len(xs)
+    if num_pairs != len(weights) or num_pairs == 0:
+        raise ValueError(f"need one weight per input, got {num_pairs} inputs and "
+                         f"{len(weights)} weights")
+    if ring_size < 2:
+        raise ValueError(f"ring_size must be at least 2, got {ring_size}")
+    batch, sequence, _ = xs[0].shape
+    if sequence % ring_size:
+        raise ValueError(f"sequence {sequence} does not split into ring_size={ring_size} chunks")
+    rows = sequence // ring_size
+    out_features = weights[0].shape[0]
+    dtype = xs[0].dtype
+    in_features = []
+    for x, weight in zip(xs, weights):
+        if x.shape[:2] != (batch, sequence) or x.dtype != dtype:
+            raise ValueError(f"inputs differ in batch, sequence or dtype: {x.shape} {x.dtype}")
+        if weight.shape != (out_features, x.shape[2]):
+            raise ValueError(f"weight {weight.shape} does not match input {x.shape} "
+                             f"and out_features={out_features}")
+        in_features.append(x.shape[2])
+    if out_features % 2:
+        raise ValueError(f"out_features must be even (two ring halves), got {out_features}")
+    half_columns = out_features // 2
+    tile_width = _common_column_tile([half_columns], column_tile_limit)
+    if tile_width is None:
+        raise ValueError(f"no column tile (a multiple of 128) divides half the columns "
+                         f"({half_columns})")
+    tiles_per_direction = half_columns // tile_width
+    tiles_per_hop = 2 * tiles_per_direction
+    total_tiles = ring_size * tiles_per_hop
+    slots = REDUCE_SCATTER_SLOTS
+    aligned_rows = rows % _ROW_TILE == 0
+    skip_communication = profile_ablation == "communication"
+    skip_compute = profile_ablation == "compute"
+
+    def kernel(chunk_shards_ref, neighbors_ref, *refs):
+        x_refs = refs[:num_pairs]
+        weight_refs = refs[num_pairs:2 * num_pairs]
+        output_ref = refs[2 * num_pairs]
+        scratch = refs[2 * num_pairs + 1:]
+        x_chunks = scratch[:num_pairs]
+        weight_tiles = scratch[num_pairs:2 * num_pairs]
+        (partial_sums, received, chunk_semaphores, weight_semaphores, send_semaphores,
+         receive_semaphores, output_semaphores) = scratch[2 * num_pairs:]
+
+        # Both neighbours must be inside this kernel (their VMEM slots belong to
+        # it) before anything is copied into them.
+        barrier = pltpu.get_barrier_semaphore()
+        for direction in range(2):
+            pl.semaphore_signal(barrier, 1, device_id={axis: neighbors_ref[direction]})
+        pl.semaphore_wait(barrier, 2)
+
+        def columns_of(column_tile):
+            return pl.ds(pl.multiple_of(column_tile * tile_width, tile_width), tile_width)
+
+        def start_chunk_loads(hop, parity):
+            copies = []
+            for pair in range(num_pairs):
+                for direction in range(2):
+                    shard = chunk_shards_ref[direction, hop]
+                    source = (x_refs[pair].at[:, pl.ds(pl.multiple_of(shard * rows, _ROW_TILE),
+                                                       rows)]
+                              if aligned_rows else x_refs[pair].at[:, shard])
+                    copy = pltpu.make_async_copy(source, x_chunks[pair].at[direction, parity],
+                                                 chunk_semaphores.at[pair, direction, parity])
+                    copy.start()
+                    copies.append(copy)
+            return copies
+
+        def start_weight_tile_load(tile, buffer_slot):
+            # Tile u of a hop covers output columns [u * tile_width, (u + 1) * tile_width):
+            # the left half's tiles first, then the right half's.
+            column_tile = jnp.asarray(tile, jnp.int32) % tiles_per_hop
+            for pair in range(num_pairs):
+                pltpu.make_async_copy(weight_refs[pair].at[columns_of(column_tile), :],
+                                      weight_tiles[pair].at[buffer_slot],
+                                      weight_semaphores.at[pair, buffer_slot]).start()
+
+        def wait_weight_tile_load(buffer_slot):
+            for pair in range(num_pairs):
+                pltpu.make_async_copy(weight_refs[pair].at[columns_of(0), :],
+                                      weight_tiles[pair].at[buffer_slot],
+                                      weight_semaphores.at[pair, buffer_slot]).wait()
+
+        def tile_transfer(direction, column_tile, slot):
+            """The descriptor of one accumulator tile into the neighbour's ``slot``
+            (also used to wait for the same tile arriving here)."""
+            columns = columns_of(column_tile)
+            return pltpu.make_async_remote_copy(
+                partial_sums.at[direction, :, :, columns],
+                received.at[direction, slot, :, :, columns],
+                send_semaphores.at[direction, column_tile],
+                receive_semaphores.at[direction, slot, column_tile],
+                device_id={axis: neighbors_ref[direction]})
+
+        def output_write(direction, column_tile):
+            columns = columns_of(column_tile)
+            output_columns = pl.ds(
+                pl.multiple_of(direction * half_columns + column_tile * tile_width, tile_width),
+                tile_width)
+            return pltpu.make_async_copy(partial_sums.at[direction, :, :, columns],
+                                         output_ref.at[:, :, output_columns],
+                                         output_semaphores.at[direction, column_tile])
+
+        chunk_loads = start_chunk_loads(0, 0)
+        start_weight_tile_load(0, 0)
+        for hop in range(ring_size):
+            parity = hop % 2
+            for copy in chunk_loads:
+                copy.wait()
+            chunk_loads = start_chunk_loads(hop + 1, 1 - parity) if hop + 1 < ring_size else []
+
+            def process_tile(tile_in_hop, carry, hop=hop, parity=parity):
+                tile = hop * tiles_per_hop + tile_in_hop
+                buffer_slot = tile % 2
+
+                @pl.when(tile + 1 < total_tiles)
+                def _():
+                    start_weight_tile_load(tile + 1, 1 - buffer_slot)  # always one in flight
+
+                wait_weight_tile_load(buffer_slot)
+                direction = tile_in_hop // tiles_per_direction
+                column_tile = tile_in_hop % tiles_per_direction
+                columns = columns_of(column_tile)
+                sums = []
+                for batch_index in range(batch):
+                    total = None
+                    if not skip_compute:
+                        for pair in range(num_pairs):
+                            product = jax.lax.dot_general(
+                                x_chunks[pair][direction, parity, batch_index],
+                                weight_tiles[pair][buffer_slot],
+                                (((1,), (1,)), ((), ())),
+                                preferred_element_type=jnp.float32).astype(dtype)
+                            total = product if total is None else (
+                                total.astype(jnp.float32)
+                                + product.astype(jnp.float32)).astype(dtype)
+                    sums.append(total)
+                communicate = not skip_communication
+                if hop > 0 and communicate:
+                    # This tile of the upstream accumulator (sent about a hop ago).
+                    tile_transfer(direction, column_tile, hop % slots).wait_recv()
+                    # The previous hop's send of this tile has left partial_sums.
+                    tile_transfer(direction, column_tile, hop % slots).wait_send()
+                for batch_index in range(batch):
+                    total = sums[batch_index]
+                    if hop > 0 and communicate:
+                        upstream = received[direction, hop % slots, batch_index, :, columns]
+                        total = upstream if total is None else (
+                            total.astype(jnp.float32) + upstream.astype(jnp.float32)).astype(dtype)
+                    if total is not None:
+                        partial_sums[direction, batch_index, :, columns] = total
+                if hop + 1 < ring_size:
+                    if communicate:
+                        tile_transfer(direction, column_tile, (hop + 1) % slots).start()
+                else:  # the last hop holds this device's rows of the sum
+                    output_write(direction, column_tile).start()
+                return carry
+
+            jax.lax.fori_loop(0, tiles_per_hop, process_tile, 0)
+
+        def finish(tile_in_hop, carry):
+            direction = tile_in_hop // tiles_per_direction
+            column_tile = tile_in_hop % tiles_per_direction
+            output_write(direction, column_tile).wait()
+            return carry
+
+        jax.lax.fori_loop(0, tiles_per_hop, finish, 0)
+
+    any_memory = pl.BlockSpec(memory_space=pl.ANY)
+    scratch_shapes = (
+        [pltpu.VMEM((2, 2, batch, rows, k), dtype) for k in in_features]        # row chunks
+        + [pltpu.VMEM((2, tile_width, k), dtype) for k in in_features]          # weight tiles
+        + [pltpu.VMEM((2, batch, rows, half_columns), dtype),                   # partial sums
+           pltpu.VMEM((2, slots, batch, rows, half_columns), dtype),            # received
+           pltpu.SemaphoreType.DMA((num_pairs, 2, 2)),                  # row chunk loads
+           pltpu.SemaphoreType.DMA((num_pairs, 2)),                     # weight tile loads
+           pltpu.SemaphoreType.DMA((2, tiles_per_direction)),           # remote sends
+           pltpu.SemaphoreType.DMA((2, slots, tiles_per_direction)),    # remote receives
+           pltpu.SemaphoreType.DMA((2, tiles_per_direction))])          # output writes
+    inputs = (xs if aligned_rows
+              else [x.reshape(batch, ring_size, rows, x.shape[2]) for x in xs])
+    return pl.pallas_call(
+        kernel,
+        out_shape=jax.ShapeDtypeStruct((batch, rows, out_features), dtype),
+        grid_spec=pltpu.PrefetchScalarGridSpec(
+            num_scalar_prefetch=2,
+            in_specs=[any_memory] * (2 * num_pairs),
+            out_specs=any_memory,
+            scratch_shapes=scratch_shapes,
+        ),
+        compiler_params=pltpu.CompilerParams(collective_id=0,
+                                             vmem_limit_bytes=vmem_limit_bytes),
+        interpret=interpret,
+    )(chunk_shards, neighbors, *inputs, *weights)
+
+
+def reduce_scatter_unsupported_reason(x_shapes, weight_shapes, itemsize: int, ring_size: int,
+                                      column_tile_limit: int = DEFAULT_COLUMN_TILE_LIMIT
+                                      ) -> str | None:
+    """Why :func:`matmul_reduce_scatter_body` cannot take these shapes, or None.
+    Callers fall back to the JAX ring."""
+    batch, sequence, _ = x_shapes[0]
+    if sequence % ring_size:
+        return f"sequence {sequence} does not split into {ring_size} chunks"
+    rows = sequence // ring_size
+    out_features = weight_shapes[0][0]
+    if out_features % 2:
+        return f"odd out_features ({out_features})"
+    tile_width = _common_column_tile([out_features // 2], column_tile_limit)
+    if tile_width is None:
+        return f"no column tile (a multiple of 128) divides half the columns ({out_features // 2})"
+    in_features = sum(shape[2] for shape in x_shapes)
+    half_columns = out_features // 2
+    needed = itemsize * (2 * 2 * batch * rows * in_features                  # row chunks
+                         + 2 * tile_width * in_features                      # weight tiles
+                         + (1 + REDUCE_SCATTER_SLOTS) * 2 * batch * rows * half_columns)
+    needed += 4 * rows * tile_width                                          # one fp32 product
+    if needed > VMEM_BUDGET_BYTES:
+        return (f"needs {needed / 2**20:.0f} MiB of VMEM "
+                f"(budget {VMEM_BUDGET_BYTES / 2**20:.0f})")
+    return None
+
+
+def reduce_scatter_tables(ring_position, shard_of_position, ring_size: int):
+    """Runtime chunk shards and neighbours for ``ring_position`` (a traced int32).
+
+    At hop h the device at position p computes the left half of the columns for
+    position (p - 1 - h)'s rows (that accumulator travels clockwise) and the right
+    half for position (p + 1 + h)'s; at the last hop both are its own position.
+    Returns (chunk_shards int32 (2, ring_size), neighbors int32 (2,)).
+    """
+    import jax.numpy as jnp
+
+    shard_of = jnp.asarray(shard_of_position, dtype=jnp.int32)
+    hops = jnp.arange(ring_size, dtype=jnp.int32)
+    chunk_shards = jnp.stack([shard_of[(ring_position - 1 - hops) % ring_size],
+                              shard_of[(ring_position + 1 + hops) % ring_size]])
+    neighbors = jnp.stack([(ring_position + 1) % ring_size,
+                           (ring_position - 1) % ring_size]).astype(jnp.int32)
+    return chunk_shards, neighbors
diff --git a/tests/unit/backends/test_tpu_collective_matmul_pallas_rs.py b/tests/unit/backends/test_tpu_collective_matmul_pallas_rs.py
new file mode 100644
index 00000000..dc43eb8c
--- /dev/null
+++ b/tests/unit/backends/test_tpu_collective_matmul_pallas_rs.py
@@ -0,0 +1,87 @@
+"""The Pallas matmul + reduce-scatter ring kernel, in TPU interpret mode on 4 CPU devices.
+
+Checks values against numpy with a non-identity ring position -> shard table,
+one and two (x, weight) pairs of different widths, and batch 1 and 2, with
+Pallas' race detector on: it catches two receive slots instead of three. It
+cannot catch a missing ``wait_send``: in interpret mode a remote copy reads its
+source when it starts, so overwriting the source while the copy is in flight is
+not modelled. Needs ``jax``; skipped without it (it is not a test dependency).
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+import textwrap
+
+import pytest
+
+pytest.importorskip("jax")
+
+SCRIPT = textwrap.dedent('''
+    import os, sys
+    os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4"
+    os.environ["JAX_PLATFORMS"] = "cpu"
+    import jax, jax.numpy as jnp, numpy as np
+    import jax._src.tpu_info as ti
+    from jax.sharding import NamedSharding, PartitionSpec as P
+    from jax.experimental.pallas import tpu as pltpu
+    # The pipeline's tiling inference needs a TPU generation: treat the host as a v5e.
+    ti.registry["cpu"] = lambda: ti.get_tpu_info_for_chip(ti.ChipVersion.TPU_V5E, 1)
+    from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kernel
+
+    ring_size, batch, rows, out_features = 4, int(sys.argv[1]), 16, 512
+    local_in_features = (256,) if sys.argv[2] == "one" else (256, 128)
+    mesh = jax.sharding.Mesh(np.array(jax.devices()[:ring_size]), ("tp",))
+    shard_of_position = (2, 0, 3, 1)  # ring position -> shard, not the identity
+    rng = np.random.default_rng(0)
+    sequence = ring_size * rows
+    xs = [rng.standard_normal((batch, sequence, k * ring_size)).astype(np.float32)
+          for k in local_in_features]
+    weights = [rng.standard_normal((out_features, k * ring_size)).astype(np.float32)
+               / (k * ring_size) ** 0.5 for k in local_in_features]
+    put = lambda a, spec: jax.device_put(jnp.asarray(a), NamedSharding(mesh, spec))
+    ids = put(np.arange(ring_size, dtype=np.int32), P("tp"))
+    num_pairs = len(xs)
+
+    def body(*args):
+        xs, weights, ids = args[:num_pairs], args[num_pairs:2 * num_pairs], args[-1]
+        chunk_shards, neighbors = kernel.reduce_scatter_tables(ids[0], shard_of_position, ring_size)
+        return kernel.matmul_reduce_scatter_body(
+            xs, weights, chunk_shards, neighbors, ring_size=ring_size, column_tile_limit=128,
+            interpret=pltpu.InterpretParams(detect_races=True))
+
+    f = jax.jit(jax.shard_map(body, mesh=mesh,
+                              in_specs=(P(None, None, "tp"),) * num_pairs
+                              + (P(None, "tp"),) * num_pairs + (P("tp"),),
+                              out_specs=P(None, "tp", None), check_vma=False))
+    output = np.asarray(f(*[put(x, P(None, None, "tp")) for x in xs],
+                          *[put(w, P(None, "tp")) for w in weights], ids))
+    full = sum(x @ w.T for x, w in zip(xs, weights))  # (batch, sequence, out_features)
+    # The device at ring position p holds the rows of shard shard_of_position[p].
+    expected = np.concatenate([full[:, s * rows:(s + 1) * rows] for s in shard_of_position], axis=1)
+    np.testing.assert_allclose(output, expected, rtol=2e-5, atol=2e-5)
+    print("OK", batch, sys.argv[2])
+''')
+
+
+@pytest.mark.parametrize("batch,pairs", [(1, "one"), (2, "one"), (2, "two")])
+def test_pallas_matmul_reduce_scatter_matches_numpy(batch, pairs):
+    proc = subprocess.run([sys.executable, "-c", SCRIPT, str(batch), pairs], capture_output=True,
+                          text=True, timeout=600)
+    assert proc.returncode == 0, proc.stdout[-2000:] + proc.stderr[-4000:]
+    assert f"OK {batch} {pairs}" in proc.stdout
+
+
+def test_reduce_scatter_unsupported_reason():
+    from difflet.backends.tpu.ops_impl import collective_matmul_pallas as kernel
+
+    # Wan attention out (K 5120 / 4) and FFN down (13824 / 4) at 1170 rows per device fit.
+    assert kernel.reduce_scatter_unsupported_reason([(1, 4680, 1280)], [(5120, 1280)], 2, 4) is None
+    assert kernel.reduce_scatter_unsupported_reason([(1, 4680, 3456)], [(5120, 3456)], 2, 4) is None
+    assert "split" in kernel.reduce_scatter_unsupported_reason([(1, 4681, 1280)], [(5120, 1280)], 2, 4)
+    assert "odd" in kernel.reduce_scatter_unsupported_reason([(1, 4680, 1280)], [(5121, 1280)], 2, 4)
+    assert "no column tile" in kernel.reduce_scatter_unsupported_reason(
+        [(1, 4680, 1280)], [(5000, 1280)], 2, 4)
+    assert "VMEM" in kernel.reduce_scatter_unsupported_reason(
+        [(1, 10496, 768), (1, 10496, 3072)], [(3072, 768), (3072, 3072)], 2, 4)

From 79f0976d819d19c9cbec1a05a064e3688ef11546 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Wed, 11 Nov 2026 12:20:34 -0800
Subject: [PATCH 20/22] perf(tpu, wan): pad the SP sequence to whole tiles and
 rotate RoPE as 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>
---
 difflet/backends/tpu/wan/sp.py             | 176 ++++++++++++++++++++-
 tests/unit/backends/test_tpu_wan_sp_pad.py |  94 +++++++++++
 2 files changed, 265 insertions(+), 5 deletions(-)
 create mode 100644 tests/unit/backends/test_tpu_wan_sp_pad.py

diff --git a/difflet/backends/tpu/wan/sp.py b/difflet/backends/tpu/wan/sp.py
index e4f41a00..858f0fc9 100644
--- a/difflet/backends/tpu/wan/sp.py
+++ b/difflet/backends/tpu/wan/sp.py
@@ -22,6 +22,34 @@
 The per-token work between them (residual, adaLN, the three LayerNorms) is the
 shared block forward, now on S/tp tokens. The across-heads q/k RMSNorm keeps its
 (B,S,1) all-reduce: it needs every head's sum of squares.
+
+Row padding (``DIFFLET_TPU_WAN_SP_PAD``, default on). The rings and the Pallas kernels split
+each rank's rows into halves and address them as whole HBM tiles; when a half is
+not a multiple of 16 rows (480x832x9: S=4680, 585-row halves) every all-gather's
+input split and output assembly, and the reduce-scatter kernel's input, are
+relayout copies. Padded, the sequence runs as ``Sp`` rows, a multiple of
+``2 * tp * 16`` (4680 -> 4736 at tp=4). The model forward is the shared one with
+its own scatter / gather turned off: the first block pads the patch sequence and
+keeps this rank's rows (an ``index_select`` on a per-rank buffer -- a runtime
+value, so every rank compiles the same program; a constant slice of a padded
+sequence broke the rings, see collective_matmul), the last block gathers and
+cuts the pad. In self-attention the pad rows are not keys (exact) and get the
+identity rotation; as queries they compute values that are cut off.
+480x832x9, v5e-4, 2026-11-11: step 269.2 -> 262.2 ms; the relayouts of the rings'
+input splits and output assembly (138 ms of a 20-step denoise) are gone, the
+cut of the pad keys and the wider final gather cost ~50 ms back.
+
+Split-halves RoPE (``DIFFLET_TPU_WAN_SPLIT_ROPE``, default on). Wan rotates each
+head's channel pairs ``(2i, 2i+1)``: a ``(..., D/2, 2)`` interleave of every q and
+k, which the trace above put at ~90 ms of relayout per denoise. As for FLUX
+(backends/tpu/flux/rope.py): each head's channels of self-attention's ``to_q`` /
+``to_k`` rows and of the q/k RMSNorm weights are permuted once at install to
+``[0, 2, ..., D-2, 1, 3, ..., D-1]``, so the pairs become ``(i, D/2 + i)`` and the
+rotation is a slice and a concat. ``q . k`` per head is unchanged by a permutation
+of both, V and everything after are untouched; only the RMSNorm's summation
+order differs from the interleaved layout. 480x832x9 (padded), v5e-4,
+2026-11-11: step 261.7 -> 249.2 ms; the interleave's reshapes / copies and the
+q/k/v output relayout (~320 ms of a 20-step denoise) are gone.
 """
 
 from __future__ import annotations
@@ -33,6 +61,35 @@
 
 from difflet.backends.tpu.ops_impl import collective_matmul as cm
 
+#: ``DIFFLET_TPU_WAN_SP_PAD`` (default on, ``0`` turns it off): pad the sequence so
+#: each ring half is a multiple of ``ROW_ALIGN`` rows (see the module docstring).
+PAD_ENV = "DIFFLET_TPU_WAN_SP_PAD"
+
+#: Rows per ring half are padded to a multiple of this (bf16 HBM tiles hold 16 rows).
+ROW_ALIGN = 16
+
+
+def _pad_requested() -> bool:
+    import os
+
+    return os.environ.get(PAD_ENV, "1").strip().lower() in ("1", "true", "yes", "on")
+
+
+#: ``DIFFLET_TPU_WAN_SPLIT_ROPE`` (default on, ``0`` turns it off): see the module docstring.
+SPLIT_ROPE_ENV = "DIFFLET_TPU_WAN_SPLIT_ROPE"
+
+
+def _split_rope_requested() -> bool:
+    import os
+
+    return os.environ.get(SPLIT_ROPE_ENV, "1").strip().lower() in ("1", "true", "yes", "on")
+
+
+def padded_length(seq_len: int, tp: int) -> int:
+    """``seq_len`` rounded up to a multiple of ``2 * tp * ROW_ALIGN``."""
+    step = 2 * int(tp) * ROW_ALIGN
+    return (int(seq_len) + step - 1) // step * step
+
 
 def _attention_forward(self, hidden_states, encoder_hidden_states=None, rotary_emb=None):
     from difflet.models.wan.modeling_wan import _attn_kernel, apply_rotary_emb
@@ -53,8 +110,17 @@ def _attention_forward(self, hidden_states, encoder_hidden_states=None, rotary_e
 
     if rotary_emb is not None and not self.is_cross_attention:
         cos, sin = rotary_emb
-        q = apply_rotary_emb(q, cos, sin)
-        k = apply_rotary_emb(k, cos, sin)
+        pad = q.shape[1] - cos.shape[1]
+        if pad > 0:  # padded sequence: the pad rows get the identity rotation
+            cos = torch.cat([cos, cos.new_ones(cos.shape[0], pad, *cos.shape[2:])], dim=1)
+            sin = torch.cat([sin, sin.new_zeros(sin.shape[0], pad, *sin.shape[2:])], dim=1)
+        rotate = (_rotate_split_halves if getattr(self, "_difflet_split_rope", False)
+                  else apply_rotary_emb)
+        q = rotate(q, cos, sin)
+        k = rotate(k, cos, sin)
+    valid_keys = getattr(self, "_difflet_sp_valid_keys", None)
+    if valid_keys is not None and not self.is_cross_attention:
+        k, v = k[:, :valid_keys], v[:, :valid_keys]  # the pad rows are not keys
 
     out = _attn_kernel(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
                        head_dim=self.head_dim)
@@ -63,12 +129,104 @@ def _attention_forward(self, hidden_states, encoder_hidden_states=None, rotary_e
     return self.to_out[1](out)
 
 
+def _rotate_split_halves(x, cos, sin):
+    """Wan's rotary on channels permuted to split halves: pair i is ``(i, D/2 + i)``.
+    ``cos`` / ``sin`` hold one value per pair (:class:`_SplitHalvesRope`); the
+    elementwise math is the interleaved ``apply_rotary_emb``'s."""
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    out = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+    return out.to(x.dtype)
+
+
+class _SplitHalvesRope(torch.nn.Module):
+    """``WanRotaryPosEmbed`` returning one (cos, sin) per channel pair, (1, S, 1, D/2).
+    The stock tables repeat each pair's value (``repeat_interleave_real``); the
+    interleaved rotation reads cos at even and sin at odd channels, so these do too."""
+
+    def __init__(self, inner: torch.nn.Module) -> None:
+        super().__init__()
+        self.inner = inner
+
+    def forward(self, hidden_states):
+        cos, sin = self.inner(hidden_states)
+        return cos[..., 0::2], sin[..., 1::2]
+
+
+def _install_split_rope(module) -> None:
+    from difflet.backends.tpu.flux.rope import pair_split_permutation
+
+    head_dim = int(module.config.attention_head_dim)
+    perm = pair_split_permutation(head_dim)
+
+    def permuted(tensor, rows_per_group):
+        # On the host, then back as a fresh tensor: permuting on the device left
+        # the parameter in a layout that every forward copied back (v5e-4,
+        # 2026-11-11: copy bf16[1280,5120] x 80 per step, 50 ms per denoise).
+        host = tensor.detach().cpu()
+        host = host.view(-1, rows_per_group, *host.shape[1:]).index_select(1, perm)
+        return host.reshape(tensor.shape).contiguous().to(tensor.device)
+
+    with torch.no_grad():
+        for block in module.blocks:
+            attn = block.attn1
+            for linear in (attn.to_q, attn.to_k):  # this rank's heads
+                linear.weight.data = permuted(linear.weight, head_dim)
+                if linear.bias is not None:
+                    linear.bias.data = permuted(linear.bias, head_dim)
+            for norm in (attn.norm_q, attn.norm_k):  # every head (sliced per rank later)
+                weight = getattr(norm, "weight", None)
+                if weight is not None:
+                    weight.data = permuted(weight, head_dim)
+            attn._difflet_split_rope = True
+    module.rope = _SplitHalvesRope(module.rope)
+
+
 def _ffn_forward(self, x):
     (h,) = cm.all_gather_matmul(x, (self.net_in,))
     h = F.gelu(h, approximate="tanh")
     return cm.matmul_reduce_scatter(h, self.net_out)
 
 
+def _padded_block_forward(self, hidden_states, encoder_hidden_states, temb, rotary_emb):
+    """A block of the padded layout: the first pads the full patch sequence and
+    keeps this rank's rows, the last gathers them and cuts the pad."""
+    from difflet.ops import gather_from_sequence_parallel_region
+
+    first, last, seq_len, padded = self._difflet_sp_pad
+    if first:
+        if temb.ndim == 4:
+            raise NotImplementedError(f"{PAD_ENV}=1 does not cover per-token (Ti2V) modulation")
+        pad = padded - hidden_states.shape[1]
+        hidden_states = torch.cat([hidden_states, hidden_states.new_zeros(
+            hidden_states.shape[0], pad, hidden_states.shape[2])], dim=1)
+        hidden_states = hidden_states.index_select(1, self._difflet_sp_rows)
+    hidden_states = type(self).forward(self, hidden_states, encoder_hidden_states, temb,
+                                       rotary_emb)
+    if last:
+        hidden_states = gather_from_sequence_parallel_region(hidden_states, dim=1)
+        hidden_states = hidden_states[:, :seq_len]
+    return hidden_states
+
+
+def _install_padding(module, *, seq_len: int, padded: int, tp: int) -> None:
+    from difflet.backends.tpu.ops_impl import parallel_mesh
+
+    rank, rows = parallel_mesh.get_tp_rank(), padded // tp
+    device = next(module.parameters()).device
+    blocks = list(module.blocks)
+    for i, block in enumerate(blocks):
+        block._difflet_sp_pad = (i == 0, i == len(blocks) - 1, int(seq_len), int(padded))
+        block.forward = types.MethodType(_padded_block_forward, block)
+        block.attn1._difflet_sp_valid_keys = int(seq_len)
+    blocks[0].register_buffer(
+        "_difflet_sp_rows", torch.arange(rank * rows, (rank + 1) * rows, device=device),
+        persistent=False)
+    # The shared forward's own sequence scatter / gather would split and join the
+    # unpadded sequence; the blocks above do it on the padded one instead.
+    module.sp_enabled = False
+
+
 def install(module: torch.nn.Module, *, seq_len: int) -> bool:
     """Swap the fused forwards into a ``WanTransformer3DModel`` built with
     ``sp_enabled``. Returns False (and changes nothing) when the shape does not
@@ -78,7 +236,8 @@ def install(module: torch.nn.Module, *, seq_len: int) -> bool:
     from difflet.backends.tpu.ops_impl import parallel_mesh
 
     tp = parallel_mesh.get_tp_size()
-    if not cm.supports(seq_len, tp):
+    padded = padded_length(seq_len, tp) if _pad_requested() else int(seq_len)
+    if not cm.supports(padded, tp):
         print(f"[difflet tpu] {cm.SP_ENV}: sequence {seq_len} does not split into "
               f"2 x tp={tp}; using the unfused Megatron-SP path", flush=True)
         return False
@@ -86,6 +245,11 @@ def install(module: torch.nn.Module, *, seq_len: int) -> bool:
     # DIFFLET_TPU_SP_KERNEL=jax keeps the ring.
     cm.set_kernel_default("pallas")
     cm.prepare()
+    if padded != seq_len:
+        _install_padding(module, seq_len=seq_len, padded=padded, tp=tp)
+    split_rope = _split_rope_requested()
+    if split_rope:
+        _install_split_rope(module)
     for block in module.blocks:
         for attn in (block.attn1, block.attn2):
             if getattr(attn, "context_parallel_enabled", False):
@@ -93,8 +257,10 @@ def install(module: torch.nn.Module, *, seq_len: int) -> bool:
             attn.forward = types.MethodType(_attention_forward, attn)
         block.ffn.forward = types.MethodType(_ffn_forward, block.ffn)
     print(f"[difflet tpu] Wan sequence-parallel with collective matmuls "
-          f"({len(module.blocks)} blocks, S={seq_len}, tp={tp})", flush=True)
+          f"({len(module.blocks)} blocks, S={seq_len}"
+          f"{f' padded to {padded}' if padded != seq_len else ''}, tp={tp}"
+          f"{', split-halves RoPE' if split_rope else ''})", flush=True)
     return True
 
 
-__all__ = ["install"]
+__all__ = ["PAD_ENV", "SPLIT_ROPE_ENV", "install", "padded_length"]
diff --git a/tests/unit/backends/test_tpu_wan_sp_pad.py b/tests/unit/backends/test_tpu_wan_sp_pad.py
new file mode 100644
index 00000000..a118eb5d
--- /dev/null
+++ b/tests/unit/backends/test_tpu_wan_sp_pad.py
@@ -0,0 +1,94 @@
+"""CPU check of Wan's padded sequence-parallel layout (backends/tpu/wan/sp.py,
+``DIFFLET_TPU_WAN_SP_PAD=1``).
+
+The collective matmuls are JAX rings that only run on the chips; here they are
+replaced by their tp=1 meaning (all-gather = identity, reduce-scatter = the
+matmul plus the bias). A 24-token sequence is padded to 32: the first block
+pads and selects the rank's rows, attention ignores the pad keys, the last block
+cuts the pad. With split-halves RoPE the q/k channels are permuted at install and
+rotated as halves. Every combination must match the plain model's output.
+
+The q/k RMSNorm is swapped for a float64 one: the stock one computes in float32,
+where the permuted channels' different summation order shows at ~1e-7.
+"""
+
+from __future__ import annotations
+
+import importlib
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from difflet.pipeline.parallel_mesh import MeshSpec
+
+
+@pytest.fixture(autouse=True)
+def _mesh(monkeypatch):
+    # The shared Wan layers resolve their implementation per backend; the TPU one
+    # needs no Neuron SDK.
+    monkeypatch.setenv("DIFFLET_BACKEND", "tpu")
+    pm = importlib.import_module("difflet.backends.tpu.ops_impl.parallel_mesh")
+    pm.destroy_parallel_mesh()
+    pm.init_parallel_mesh(MeshSpec(tp=1))
+    yield
+    pm.destroy_parallel_mesh()
+
+
+def _fake_rings(monkeypatch, cm):
+    def ag(x, linears):
+        return [F.linear(x, lin.weight, lin.bias) for lin in linears]
+
+    def rs(x, linear):
+        return F.linear(x, linear.weight, linear.bias)
+
+    monkeypatch.setattr(cm, "all_gather_matmul", ag)
+    monkeypatch.setattr(cm, "matmul_reduce_scatter", rs)
+    monkeypatch.setattr(cm, "prepare", lambda ops=None: None)
+
+
+def _config(wan, sp_enabled):
+    cfg = wan.WanTransformerConfig(patch_size=(1, 2, 2), num_attention_heads=4,
+                                   attention_head_dim=16, in_channels=4, out_channels=4,
+                                   text_dim=24, freq_dim=32, ffn_dim=48, num_layers=3)
+    cfg.sp_enabled = sp_enabled
+    return cfg
+
+
+def _rms_norm_float64(x, weight, eps):
+    out = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+    return (out * weight if weight is not None else out).to(x.dtype)
+
+
+@pytest.mark.parametrize("pad", ["1", "0"])
+@pytest.mark.parametrize("split_rope", ["1", "0"])
+def test_padded_sp_matches_plain_model(monkeypatch, pad, split_rope):
+    wan = importlib.import_module("difflet.models.wan.modeling_wan")
+    sp = importlib.import_module("difflet.backends.tpu.wan.sp")
+    _fake_rings(monkeypatch, sp.cm)
+    monkeypatch.setenv(sp.PAD_ENV, pad)
+    monkeypatch.setenv(sp.SPLIT_ROPE_ENV, split_rope)
+    monkeypatch.setattr(wan, "rms_norm_fp32", _rms_norm_float64)
+
+    torch.manual_seed(0)
+    plain = wan.WanTransformer3DModel(_config(wan, False)).double().eval()
+    with torch.no_grad():
+        for p in plain.parameters():  # zero-initialised biases / gates would hide bugs
+            p.copy_(torch.randn_like(p) * 0.2)
+    fused = wan.WanTransformer3DModel(_config(wan, True)).double().eval()
+    fused.load_state_dict(plain.state_dict())
+
+    seq_len = 2 * 3 * 4  # frames x (6 / 2) x (8 / 2)
+    assert sp.install(fused, seq_len=seq_len)
+    assert fused.sp_enabled is (pad == "0")  # padded: the blocks scatter and gather
+    if pad == "1":
+        assert sp.padded_length(seq_len, 1) == 32
+
+    g = torch.Generator().manual_seed(1)
+    hidden = torch.randn(1, 4, 2, 6, 8, generator=g, dtype=torch.float64)
+    timestep = torch.tensor([3.0], dtype=torch.float64)
+    text = torch.randn(1, 5, 24, generator=g, dtype=torch.float64)
+    with torch.no_grad():
+        expected = plain(hidden, timestep, text)
+        got = fused(hidden, timestep, text)
+    torch.testing.assert_close(got, expected, rtol=1e-10, atol=1e-10)

From ca6abd35c85228fdecd4bc0813a179e603c02be5 Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Wed, 11 Nov 2026 12:21:01 -0800
Subject: [PATCH 21/22] docs(tpu): the RS kernel, Wan padding and split RoPE,
 the GELU epilogue; trace_diff

The plan gets the afternoon's results: the AG kernel cleanup checked against
868401eb, 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>
---
 benchmark/tpu_profile/trace_diff.py           | 131 ++++++++++++++++++
 ...2026-11-10-tpu-collective-matmul-pallas.md |  64 +++++++++
 2 files changed, 195 insertions(+)
 create mode 100644 benchmark/tpu_profile/trace_diff.py

diff --git a/benchmark/tpu_profile/trace_diff.py b/benchmark/tpu_profile/trace_diff.py
new file mode 100644
index 00000000..68ae96de
--- /dev/null
+++ b/benchmark/tpu_profile/trace_diff.py
@@ -0,0 +1,131 @@
+"""Where two traced TPU denoises differ: device time per HLO op kind, side by side.
+
+    python -m benchmark.tpu_profile.trace_diff <trace_dir_a> <trace_dir_b> [--top 25]
+
+Reads the rank-0 ``jax.profiler`` traces that main's bench writes with
+``DIFFLET_BENCH_TRACE_DIR`` (as trace_breakdown / exposed_comm) and sums the
+device op time ("XLA Ops" on /device:TPU:0) per op kind -- the op name without
+its ``.N`` suffix, plus the profiler's ``hlo_category`` when it has one -- over
+the traced window. Prints the window, the busy time, and the kinds whose totals
+differ most, so a change that should remove one kind of op (for example a
+relayout ``copy``) can be checked against what else moved.
+
+When the JSON export has no device events (it stops at 5M events), the raw
+xplane is read instead (as exposed_comm does); its op names are full HLO lines,
+so ``--by-shape`` can key each op by opcode and output shape.
+"""
+
+from __future__ import annotations
+
+import argparse
+import collections
+import glob
+import gzip
+import json
+import re
+
+SUFFIX = re.compile(r"(\.\d+)+$")
+#: "%name = <shape> <opcode>(" -- an HLO instruction line (xplane op names).
+HLO_LINE = re.compile(r"^%?[\w.-]+ = (.*?)\s([a-z][a-z0-9-]*)\(")
+
+
+def _load_xplane(trace_dir: str):
+    import jax.profiler
+
+    path = sorted(glob.glob(trace_dir + "/plugins/profile/*/*.xplane.pb"))[-1]
+    data = jax.profiler.ProfileData.from_file(path)
+    plane = next(p for p in data.planes if p.name == "/device:TPU:0")
+    lines = {line.name: line for line in plane.lines}
+
+    def events(name):
+        return [{"name": e.name, "ts": e.start_ns / 1e3, "dur": e.duration_ns / 1e3}
+                for e in lines[name].events] if name in lines else []
+
+    return events("XLA Ops"), events("XLA Modules")
+
+
+def _load(trace_dir: str):
+    path = sorted(glob.glob(trace_dir + "/plugins/profile/*/*.trace.json.gz"))[-1]
+    events = json.load(gzip.open(path))["traceEvents"]
+    procs = {e["pid"]: e["args"]["name"] for e in events
+             if e.get("ph") == "M" and e.get("name") == "process_name"}
+    threads = {(e["pid"], e["tid"]): e["args"]["name"] for e in events
+               if e.get("ph") == "M" and e.get("name") == "thread_name"}
+    device = [e for e in events if e.get("ph") == "X" and procs.get(e["pid"]) == "/device:TPU:0"]
+    ops = [e for e in device if threads.get((e["pid"], e["tid"])) == "XLA Ops"]
+    modules = [e for e in device if threads.get((e["pid"], e["tid"])) == "XLA Modules"]
+    if not modules:
+        return _load_xplane(trace_dir)
+    return ops, modules
+
+
+def _busy(intervals) -> float:
+    total, start, end = 0.0, None, None
+    for s, e in sorted(intervals):
+        if end is None or s > end:
+            if end is not None:
+                total += end - start
+            start, end = s, e
+        else:
+            end = max(end, e)
+    return total + ((end - start) if end is not None else 0.0)
+
+
+def _kind(name: str, by_shape: bool) -> str:
+    hlo = HLO_LINE.match(name)
+    if hlo:  # an xplane op: a full HLO line
+        shape, opcode = hlo.group(1), hlo.group(2)
+        return f"{opcode} {shape}" if by_shape else opcode
+    return SUFFIX.sub("", name)
+
+
+def summarize(trace_dir: str, by_shape: bool = False) -> dict:
+    ops, modules = _load(trace_dir)
+    window = (max(m["ts"] + m["dur"] for m in modules) - min(m["ts"] for m in modules)) / 1e3
+    kinds = collections.Counter()
+    counts = collections.Counter()
+    for e in ops:
+        args = e.get("args", {})
+        kind = _kind(e["name"], by_shape)
+        category = args.get("hlo_category") or args.get("category")
+        if category:
+            kind = f"{kind} [{category}]"
+        kinds[kind] += e["dur"] / 1e3
+        counts[kind] += 1
+    return {"window_ms": window,
+            "busy_ms": _busy([(e["ts"], e["ts"] + e["dur"]) for e in ops]) / 1e3,
+            "kinds": kinds, "counts": counts}
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("a")
+    parser.add_argument("b")
+    parser.add_argument("--top", type=int, default=25)
+    parser.add_argument("--by-shape", action="store_true",
+                        help="key ops by opcode and output shape (xplane traces)")
+    parser.add_argument("--only", default="",
+                        help="comma-separated opcodes to list (e.g. copy,reshape)")
+    args = parser.parse_args()
+    a, b = summarize(args.a, args.by_shape), summarize(args.b, args.by_shape)
+    print(f"window ms: a {a['window_ms']:.1f}  b {b['window_ms']:.1f}  "
+          f"(b - a {b['window_ms'] - a['window_ms']:+.1f})")
+    print(f"busy ms:   a {a['busy_ms']:.1f}  b {b['busy_ms']:.1f}  "
+          f"(b - a {b['busy_ms'] - a['busy_ms']:+.1f})")
+    keys = set(a["kinds"]) | set(b["kinds"])
+    if args.only:
+        wanted = tuple(args.only.split(","))
+        keys = {k for k in keys if k.split(" ")[0].split(" [")[0] in wanted}
+    rows = sorted(keys, key=lambda k: -abs(b["kinds"][k] - a["kinds"][k]))[:args.top]
+    width = 110 if args.by_shape else 60
+    print(f"{'op kind':{width}s} {'a ms':>9s} {'b ms':>9s} {'b - a':>9s} {'a n':>7s} {'b n':>7s}")
+    for k in rows:
+        print(f"{k[:width]:{width}s} {a['kinds'][k]:9.1f} {b['kinds'][k]:9.1f} "
+              f"{b['kinds'][k] - a['kinds'][k]:+9.1f} {a['counts'][k]:7d} {b['counts'][k]:7d}")
+    print(f"{'(all listed kinds)' if args.only else '(all kinds)':{width}s} {sum(a['kinds'].values()):9.1f} {sum(b['kinds'].values()):9.1f} "
+          f"{sum(b['kinds'].values()) - sum(a['kinds'].values()):+9.1f}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md b/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
index ee7dde63..94ef47a2 100644
--- a/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
+++ b/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
@@ -304,6 +304,70 @@ no fallback), tree 077ab253f18d:
 
 Not checked: other seeds, a visual review (TP vs kernel max |d| 199/255).
 
+## 2026-11-11 afternoon: cleanup, the RS kernel, Wan padding and RoPE, a GELU epilogue
+
+All on v5e-4 with main's `bench.py`, each pair back to back on one node (granted
+16:16 UTC, released 20:19 UTC); traces compared with
+`benchmark/tpu_profile/trace_diff.py` (per op kind, or per opcode and output
+shape with `--by-shape`).
+
+**AG kernel cleanup.** Only the streamed variant is kept, with full names
+(`all_gather_matmul_body`, `unsupported_reason`, `ring_tables`). Against 868401eb:
+Wan / FLUX / HunyuanVideo outputs sha256-identical, steps within noise. The
+version table above used (K, N) weights and column tiles up to 512; the models
+run (N, K) and up to 256 (Wan FFN up then has 128-wide tiles): op-level 1.187 /
+1.066 / 0.411 / 0.519 ms (Wan q/k/v, FFN up, LTX-2 q/k/v, FFN up), bit-exact.
+
+**RS kernel** (`DIFFLET_TPU_SP_RS_KERNEL=pallas`, off by default). Tile by tile:
+multiply, add the upstream accumulator's tile in the epilogue, send it on;
+accumulators VMEM to VMEM; three receive slots per direction. A first version
+that added each hop's accumulator in a separate pass was slower everywhere.
+Op-level, ms per call:
+
+| shape | JAX ring | kernel | note |
+|---|---|---|---|
+| Wan attention out | 0.518 | 0.504 | ICI-bound (no compute 0.462) |
+| LTX-2 attention out | 0.214 | 0.212 | ICI-bound |
+| Wan FFN down, 1170 rows | 0.928 | 1.040 | ~75 us input relayout (rows not a multiple of 8) |
+| Wan FFN down, 1168 rows | 0.917 | 0.964 | compute-bound; column tile 512: 0.945 |
+| FLUX single pair | 0.627 | 0.663 | compute-bound; column tile 512: 0.649 |
+
+Wan end to end (padded) 262.2 -> 262.7 ms. Unlike the AG ring, the RS ring has no
+large assembly cost to remove, and where it is compute-bound the kernel's matmul
+is ~5% slower than XLA's einsum. In the model one step differs from the ring by
+50.4 dB, four by 48.9 dB (seed 42), as the AG kernel did.
+
+**Wan row padding** (`DIFFLET_TPU_WAN_SP_PAD`, default on): 4680 -> 4736 rows, 1184
+per rank, 592-row halves. Step 269.2 -> 262.2 ms, output vs unpadded 44.7 dB.
+Trace (20 steps): the rings' input-split and output-assembly relayouts, 137.8 ms,
+are gone; cutting the pad keys (4680 is not a tile multiple: copy + slice ~33 ms),
+1.2% more data per hop (+13 ms ICI wait) and the final gather over the hidden
+size (+8 ms) take ~54 back. The "~250 ms relayout" estimated above counted
+every reshape / copy; most of them are attention and RoPE layout changes that
+do not depend on the row count.
+
+**Wan split-halves RoPE** (`DIFFLET_TPU_WAN_SPLIT_ROPE`, default on), as FLUX's:
+q/k rows and RMSNorm weights permuted per head once, rotation as two halves.
+The interleave's reshapes / fusions / copies and, unexpectedly, the 66 ms
+relayout of the AG kernel's q/k/v outputs are gone. Permuting the weights on the
+device with `copy_` left them in a layout every forward copied back
+(`copy bf16[1280,5120]`, 80 per step, 50 ms, plus ~120 ms more device idle);
+permuted on the host, step 262.2 -> 249.2 ms. Output vs TP 42.2 dB.
+
+**Wan, the day's net:** 269.2 -> 249.2 ms per step (-7.4%); TP is 360.7.
+
+**Bias + GELU in the AG kernel's epilogue (tried, reverted).** For Wan's FFN
+up-projection it removes the separate 60 ms `[1,4736,3456]` GELU fusion, but the
+FFN-down RS ring's matmuls then read `h` 10% slower (113 -> 125 us per op,
++65 ms): XLA's memory-space assignment had placed the GELU fusion's output in
+`S(1)` (alternate memory, presumably VMEM), while a Pallas output (`pl.ANY`)
+lives in plain HBM. Step 250.4 -> 249.4 ms, within noise. A Pallas kernel's
+outputs forfeit that placement for their consumers; whether the AG kernel's
+q/k/v outputs pay this in attention is not checked.
+
+Not checked: other seeds, FLUX's RoPE permutation (it uses `copy_` on the device
+too, so may pay the same per-forward weight copy), padding for the other models.
+
 ## Steps
 
 | # | Work | Validation | Estimate |

From 917801b88c75c385469a684791201f7e881795de Mon Sep 17 00:00:00 2001
From: SHIYANG CHENG <shiyangcheng@SHIYANGs-Laptop.local>
Date: Wed, 11 Nov 2026 12:25:39 -0800
Subject: [PATCH 22/22] docs(tpu): how to reproduce the afternoon's A/B
 numbers, and what is 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>
---
 benchmark/tpu_profile/bench_case.sh           | 40 +++++++++++++++++++
 ...2026-11-10-tpu-collective-matmul-pallas.md | 30 ++++++++++++++
 2 files changed, 70 insertions(+)
 create mode 100755 benchmark/tpu_profile/bench_case.sh

diff --git a/benchmark/tpu_profile/bench_case.sh b/benchmark/tpu_profile/bench_case.sh
new file mode 100755
index 00000000..243ead42
--- /dev/null
+++ b/benchmark/tpu_profile/bench_case.sh
@@ -0,0 +1,40 @@
+#!/bin/bash
+# One case of main's bench on a TPU VM, as the A/B tables in
+# docs/plans/2026-11-10-tpu-collective-matmul-pallas.md were measured.
+#
+#   benchmark/tpu_profile/bench_case.sh <tag> <model> [ENV=VALUE ...]
+#
+# e.g. (Wan, padded SP vs not, back to back on one node):
+#   W="DIFFLET_TPU_COMPILE=1 DIFFLET_TPU_DEVICE_VAE=1 DIFFLET_TPU_SP=1"
+#   benchmark/tpu_profile/bench_case.sh wp-pad0 wan_2_1 $W DIFFLET_TPU_WAN_SP_PAD=0
+#   benchmark/tpu_profile/bench_case.sh wp-pad1 wan_2_1 $W DIFFLET_TPU_WAN_SP_PAD=1
+#   python -m benchmark.tpu_profile.image_psnr $OUT/wp-pad0/<frames>.png $OUT/wp-pad1/<frames>.png
+#
+# Add DIFFLET_BENCH_TRACE_DIR=<dir> for a jax.profiler trace (then
+# benchmark.tpu_profile.trace_diff / exposed_comm). Frees the chips first (a
+# previous case may still hold /dev/vfio), runs `benchmark.bench` once
+# (--iters 1 --natural-iters 1: the step time is the mean over the request's
+# denoising steps), saves the outputs under $OUT/<tag> and the log as
+# $OUT/<tag>.log, and prints one summary line.
+#
+# Paths default to the layout of the data disk used on 2026-11-11 (torch 2.14.1+cpu,
+# torch_tpu 0.2.0, libtpu 0.0.51, jax 0.12.0.1 in $VENV); override as needed.
+set -u
+tag=$1; model=$2; shift 2
+ROOT=${DIFFLET_ROOT:-/mnt/models/Difflet-main}
+VENV=${DIFFLET_TPU_VENV:-/mnt/models/tt-venv}
+OUT=${DIFFLET_BENCH_OUT:-/mnt/models/profile}
+cd "$ROOT"
+holders=$(sudo fuser /dev/vfio/* 2>/dev/null | tr -s " " "\n" | grep -E "^[0-9]+$" | sort -u)
+[ -n "$holders" ] && { sudo kill -9 $holders; sleep 3; }
+export PYTHONPATH=$ROOT HF_HOME=${HF_HOME:-/mnt/models/hf} DIFFLET_BACKEND=tpu
+export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-/mnt/models/uvpy/cpython-3.12.14-linux-x86_64-gnu/lib}
+export DIFFLET_BENCH_DEVICE=v5e-$tag DIFFLET_BENCH_SAVE_DIR=$OUT/$tag
+for kv in "$@"; do export "$kv"; done
+echo "tree md5 $(find difflet -name '*.py' | sort | xargs md5sum | md5sum | cut -c1-12)" > "$OUT/$tag.log"
+timeout 2400 "$VENV/bin/python" -m benchmark.bench --model "$model" --backend tpu \
+    --iters 1 --natural-iters 1 >> "$OUT/$tag.log" 2>&1
+status=$?
+step=$("$VENV/bin/python" -c "import json; d=json.load(open('benchmark/v5e-$tag/$model.json')); print(round(d['step_latency']['mean'], 4))" 2>/dev/null)
+echo "== $tag exit=$status $(grep -oE 'status=[a-z_]+' "$OUT/$tag.log" | tail -1) step=$step" \
+     "files=$(ls "$OUT/$tag" 2>/dev/null | tr '\n' ' ')"
diff --git a/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md b/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
index 94ef47a2..4d821541 100644
--- a/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
+++ b/docs/plans/2026-11-10-tpu-collective-matmul-pallas.md
@@ -368,6 +368,36 @@ q/k/v outputs pay this in attention is not checked.
 Not checked: other seeds, FLUX's RoPE permutation (it uses `copy_` on the device
 too, so may pay the same per-forward weight copy), padding for the other models.
 
+### Reproducing the afternoon's numbers
+
+Environment: v5e-4 (us-west1-c), torch 2.14.1+cpu, torch_tpu 0.2.0, libtpu
+0.0.51, jax 0.12.0.1; Wan 2.1 14B T2V at 480x832x9, 20 steps, seed 42 (main's
+`benchmark/models.py` entry `wan_2_1`). Each table row is one
+`benchmark/tpu_profile/bench_case.sh` case; a pair differs only in its
+environment and runs back to back on the same node and tree:
+
+    W="DIFFLET_TPU_COMPILE=1 DIFFLET_TPU_DEVICE_VAE=1 DIFFLET_TPU_SP=1"
+    bench_case.sh wp-pad0 wan_2_1 $W DIFFLET_TPU_WAN_SP_PAD=0 DIFFLET_TPU_WAN_SPLIT_ROPE=0
+    bench_case.sh wp-pad1 wan_2_1 $W DIFFLET_TPU_WAN_SP_PAD=1 DIFFLET_TPU_WAN_SPLIT_ROPE=0
+    bench_case.sh wr-rope1 wan_2_1 $W                       # both defaults on
+    bench_case.sh wp-rs wan_2_1 $W DIFFLET_TPU_SP_RS_KERNEL=pallas
+    # + DIFFLET_BENCH_TRACE_DIR=<dir> for the trace rows, then
+    python -m benchmark.tpu_profile.trace_diff <dir_a> <dir_b> [--by-shape --only copy,reshape]
+    python -m benchmark.tpu_profile.image_psnr <a>/run1_synced_frames_0_4_8.png <b>/...
+    python -m benchmark.tpu_profile.cm_rs_pallas_bench --ablate --column-tile-limits 128,256,512,1024
+    python -m benchmark.tpu_profile.sp_step_drift --model wan_2_1 --steps 1,4 --out <dir> --tag <t>
+
+The 18:42 UTC measurement of "both on" predates the commit only by comments
+(the source hash in the op names differs, so the committed tree compiles anew;
+not re-measured). Run-to-run noise of one configuration on one node: Wan
+267.7-269.2 ms (four runs), FLUX 141.9-144.7 ms with identical outputs.
+
+Not verified: the "vs TP" PSNRs use a TP run from an earlier tree
+(`rp-wan-base`), not one on this tree; split RoPE's 26 dB against the
+interleaved layout after 20 steps is attributed to the RMSNorm's summation order
+by analogy (the RS kernel's 1- and 4-step check), not by its own 1-step run;
+only seed 42; that `S(1)` is VMEM.
+
 ## Steps
 
 | # | Work | Validation | Estimate |