Under Ulysses SP (num_gpus>=2), self-forcing causal Wan models
(causal_wanvideo.py) crash in `_apply_rotary_emb`:
RuntimeError: The size of tensor a (4680) must match tensor b (9360)
at non-singleton dimension 1
Cause: the DiT builds RoPE for `post_patch_num_frames * get_sp_world_size()`
(sp x too long) but the causal pipeline *replicates* the full token sequence and
a full-size KV cache on every rank (it does not shard the sequence), so the
attention applies that over-long RoPE to the full local query.
Fix:
- Remove `* get_sp_world_size()` from the RoPE length (relativistic, standard,
and train branches) so RoPE matches the full local sequence. No-op at sp=1.
- Parallelize the causal attention across HEADS (Ulysses over heads): each rank
computes attention for its head slice `[rank*H/sp:(rank+1)*H/sp]` over the full
sequence, then `sequence_model_parallel_all_gather(dim=2)` reconstructs all
heads. Applied uniformly to all three attention paths (flex_attention,
KV-cache `self.attn`, and the relativistic window). A sequence all-to-all a la
wanvideo.py is incompatible with the causal KV-cache windowing (global token
coords), so head-parallel is the exact-matching approach; heads are
independent so it is numerically identical to single-GPU.
Verified: `wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers`, 480x832, 21 frames (7
autoregressive KV-cache blocks), seed 1234, on 2x GPU vs 1x GPU:
torch.equal=True, max_abs_diff=0.0 (byte-identical). 2-GPU previously crashed.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
