LTX-2 `_cp_plan` appears copied from LTX-1 and fails under `enable_parallelism()` (crash + one silent-correctness hazard)

### Describe the bug

`LTX2VideoTransformer3DModel._cp_plan` ([`transformer_ltx2.py:1094`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_ltx2.py#L1094)) is textually identical to `LTXVideoTransformer3DModel._cp_plan` ([`transformer_ltx.py:417`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_ltx.py#L417)), but LTX-2 is a dual-stream audio+video model with a different forward. The plan doesn't cover:

- **`timestep`** — per-token on LTX-2, `(B, num_video_tokens)`, so it must be sharded
- **the audio stream** — `audio_hidden_states`, `audio_timestep`, `audio_encoder_hidden_states`, `audio_encoder_attention_mask`, `audio_proj_out`
- **3 of 4 RoPE modules** — `audio_rope`, `cross_attn_rope`, `cross_attn_audio_rope`

`enable_parallelism()` accepts it, then the forward fails at [`transformer_ltx2.py:634`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_ltx2.py#L634) with `size of tensor a (16) must match tensor b (32)` — sharded video tokens against the unsharded per-token timestep.

**The part I'd flag as more important than the crash:** the obvious minimal repair — mirror Wan and shard video only — is silently wrong on LTX-2. Its blocks run bidirectional AV cross-attention; `v2a` has Q=audio, KV=**video**, so with video sharded and no collective on that call each rank attends over `1/P` of the video keys. Shapes line up, nothing raises, output is wrong. Both streams have to be sharded.

Verified against `main` @ `60092255`.

### Reproduction

CPU-only, gloo, 2 ranks, no GPU — spawns its own ranks, so `python repro.py` is all it takes.

repro.py — --mask and --fixed select the arms

```python
"""
python repro.py # A: stock _cp_plan -> crash
python repro.py --mask # B: corrected plan -> attn_mask head-axis bug
python repro.py --fixed # C: corrected plan + fix -> matches single-rank
"""

import argparse, os, traceback
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

import diffusers.models.attention_dispatch as AD
from diffusers.models._modeling_parallel import (
ContextParallelConfig, ContextParallelInput, ContextParallelOutput,
)
from diffusers.models.transformers.transformer_ltx2 import LTX2VideoTransformer3DModel

WORLD = 2

CONFIG = dict(
in_channels=8, out_channels=8, patch_size=1, patch_size_t=1,
num_attention_heads=4, attention_head_dim=16, cross_attention_dim=64,
vae_scale_factors=(8, 32, 32), pos_embed_max_pos=20,
base_height=2048, base_width=2048, gated_attn=True, cross_attn_mod=True,
audio_in_channels=8, audio_out_channels=8, audio_patch_size=1,
audio_patch_size_t=1, audio_num_attention_heads=4, audio_attention_head_dim=8,
audio_cross_attention_dim=32, audio_scale_factor=4, audio_pos_embed_max_pos=20,
audio_gated_attn=True, audio_cross_attn_mod=True,
num_layers=2, caption_channels=32, rope_type="interleaved",
use_prompt_embeddings=True,
)
B, N_VIDEO, N_AUDIO, N_TEXT = 1, 32, 8, 16 # all divisible by WORLD

def make_inputs():
g = torch.Generator().manual_seed(0)
return dict(
hidden_states=torch.randn(B, N_VIDEO, 8, generator=g),
audio_hidden_states=torch.randn(B, N_AUDIO, 8, generator=g),
encoder_hidden_states=torch.randn(B, N_TEXT, 32, generator=g),
audio_encoder_hidden_states=torch.randn(B, N_TEXT, 32, generator=g),
timestep=torch.full((B, N_VIDEO), 500.0), # PER-TOKEN, as the pipelines pass it
audio_timestep=torch.full((B, N_AUDIO), 500.0),
sigma=torch.full((B,), 0.5), # per-batch, must NOT be split
audio_sigma=torch.full((B,), 0.5),
encoder_attention_mask=torch.ones(B, N_TEXT),
audio_encoder_attention_mask=torch.ones(B, N_TEXT),
num_frames=2, height=4, width=4, # 2 * 4 * 4 = 32 video tokens
audio_num_frames=N_AUDIO,
fps=24.0, use_cross_timestep=True, return_dict=False,
)

def _in(dim=1, nd=3, out=False):
return ContextParallelInput(split_dim=dim, expected_dims=nd, split_output=out)

_ROPE = {0: _in(out=True), 1: _in(out=True)}

LTX2_CP_PLAN = {
"": {
"hidden_states": _in(),
"timestep": _in(nd=2),
"audio_hidden_states": _in(),
"audio_timestep": _in(nd=2),
"encoder_hidden_states": _in(),
"encoder_attention_mask": _in(nd=2),
"audio_encoder_hidden_states": _in(),
"audio_encoder_attention_mask": _in(nd=2),
},
"rope": _ROPE,
"audio_rope": _ROPE,
"cross_attn_rope": _ROPE,
"cross_attn_audio_rope": _ROPE,
"proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
"audio_proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
}

_ORIG = AD.TemplatedUlyssesAttention.forward

def _patched(ctx, query, key, value, attn_mask, *args, **kwargs):
parallel_config = kwargs.get("_parallel_config") or args[-1]
if attn_mask is not None and attn_mask.ndim == 4:
cp = parallel_config.context_parallel_config
world_size, local_rank = cp.ulysses_degree, cp._ulysses_local_rank
num_heads = query.shape[2]
if attn_mask.shape[1] == num_heads and num_heads % world_size == 0:
attn_mask = attn_mask.chunk(world_size, dim=1)[local_rank].contiguous()
return _ORIG(ctx, query, key, value, attn_mask, *args, **kwargs)

def worker(rank, mode):
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "29591")
dist.init_process_group("gloo", rank=rank, world_size=WORLD)
torch.manual_seed(0)

model = LTX2VideoTransformer3DModel(**CONFIG).eval()
with torch.no_grad():
ref_video, ref_audio = model(**make_inputs())

if mode == "fixed":
AD.TemplatedUlyssesAttention.forward = staticmethod(_patched)
kwargs = {} if mode == "stock" else {"cp_plan": LTX2_CP_PLAN}
model.enable_parallelism(config=ContextParallelConfig(ulysses_degree=WORLD), **kwargs)

try:
with torch.no_grad():
video, audio = model(**make_inputs())
if rank == 0:
print(f"[{mode}] ran. max|delta| vs single-rank: "
f"video={(video - ref_video).abs().max():.3e} "
f"audio={(audio - ref_audio).abs().max():.3e}")
except Exception as exc:
if rank == 0:
frames = [ln.strip() for ln in traceback.format_exc().splitlines()
if "transformer_ltx2.py" in ln or "attention_dispatch.py" in ln]
print(f"[{mode}] FAILED: {type(exc).__name__}: {exc}")
print(f"[{mode}] at {frames[-1] if frames else '?'}")
dist.destroy_process_group()

if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument("--fixed", action="store_true")
ap.add_argument("--mask", action="store_true")
a = ap.parse_args()
mp.spawn(worker, args=("fixed" if a.fixed else "mask" if a.mask else "stock",),
nprocs=WORLD, join=True)
```

```
$ python repro.py
[stock] FAILED: RuntimeError: The size of tensor a (16) must match the size of tensor b (32) at non-singleton dimension 1
[stock] at File ".../transformer_ltx2.py", line 634, in forward

$ python repro.py --mask
[mask] FAILED: RuntimeError: The size of tensor a (2) must match the size of tensor b (4) at non-singleton dimension 1
[mask] at File ".../attention_dispatch.py", line 851, in _native_attention_forward_op

$ python repro.py --fixed
[fixed] ran. max|delta| vs single-rank: video=4.768e-07 audio=2.384e-07
```

The `--fixed` arm is the 12-entry plan in the repro, which matches single-rank to fp32 accumulation error through the a2v/v2a path.

### Secondary: `TemplatedUlyssesAttention` doesn't shard `attn_mask`'s head axis

Independent of LTX-2, but LTX-2 is the first model to reach it (it applies CP to masked cross-attention; Wan guards that at [`transformer_wan.py:152`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py#L152)). [`attention_dispatch.py:2393`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_dispatch.py#L2393) all-gathers a sharded mask on the KV axis, but q/k/v come back from the all-to-all with `H // world_size` heads while the mask keeps all `H` → broadcast failure in SDPA. Distinct from #13696/#13756, which fixed the *sequence* axis. Fix is a `chunk(world_size, dim=1)[ulysses_local_rank]` alongside the KV gather; same gap in the `ulysses_anything` branch at [`:2609`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_dispatch.py#L2609).

Two smaller notes: `_prepare_cp_input` [warns and skips](https://github.com/huggingface/diffusers/blob/main/src/diffusers/hooks/context_parallel.py#L206) on a dims mismatch rather than raising, so `expected_dims=3` for `"rope"` silently degrades to unsharded at `rope_type="split"` (4-D output); and `video_self_attention_mask` is `(B, T_v, T_v)`, a two-axis shard `ContextParallelInput` can't express, so the IC-LoRA path can't be fixed by a plan alone.

### System Info

`diffusers` `main` @ `60092255` (also on `0.39.0`) · torch 2.13.0 · Python 3.12 · Linux · CPU-only, gloo, 2 ranks

### Who can help?

@sayakpaul @DN6 @yiyixuxu · cc @zhtmike re the Ulysses mask handling

Run the supplied CPU-only two-rank repro.py first, comparing the stock, --mask, and --fixed modes. Read transformer_ltx2.py, attention_dispatch.py, and context_parallel.py around the cited entry points, including _cp_plan, masked attention, and input preparation. Done means LTX-2 parallelism covers both streams and per-token inputs, masked attention runs, and outputs match the single-rank reference within the repro's fp32 tolerance.

Written by the indexing model from the issue text.