Fix Ulysses sequence-parallel crash in self-forcing causal Wan DiT

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>

fastvideo/models/dits/causal_wanvideo.py

@@ -18,7 +18,10 @@ import torch.distributed as dist
 import fastvideo.envs as envs
 from fastvideo.attention import (DistributedAttention, LocalAttention)
 from fastvideo.configs.models.dits import WanVideoConfig
-from fastvideo.distributed.parallel_state import get_sp_world_size
+from fastvideo.distributed.communication_op import (
+    sequence_model_parallel_all_gather)
+from fastvideo.distributed.parallel_state import (get_sp_parallel_rank,
+                                                  get_sp_world_size)
 from fastvideo.forward_context import get_forward_context
 from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual,
                                         ScaleResidualLayerNormScaleShift)
@@ -98,6 +101,29 @@ class CausalWanSelfAttention(nn.Module):
             roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v)
             roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v)
 
+        # Ulysses sequence-parallel (SP): the causal pipeline replicates the full
+        # (unsharded) token sequence on every rank, so RoPE, the block mask and
+        # the KV cache are all full-length here. We parallelize the attention
+        # across the head dimension (à la Ulysses): each rank computes attention
+        # for its own slice of heads over the full sequence, then all-gathers the
+        # per-head outputs to reconstruct all heads. This is numerically identical
+        # to single-GPU attention (heads are independent) and avoids a sequence
+        # all-to-all that would be incompatible with the causal KV-cache windowing
+        # (expressed in global token coordinates). Applied uniformly to all three
+        # attention paths below (flex, relativistic window, and the standard
+        # KV-cache path).
+        sp_world_size = get_sp_world_size()
+        sp_rank = get_sp_parallel_rank()
+        if sp_world_size > 1:
+            assert self.num_heads % sp_world_size == 0, (
+                f"num_heads ({self.num_heads}) must be divisible by sequence "
+                f"parallel size ({sp_world_size})")
+            heads_per_rank = self.num_heads // sp_world_size
+            head_slice = slice(sp_rank * heads_per_rank,
+                               (sp_rank + 1) * heads_per_rank)
+        else:
+            head_slice = slice(None)
+
         if kv_cache is None:
             # Padding for flex attention
             padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
@@ -117,10 +143,13 @@ class CausalWanSelfAttention(nn.Module):
                 [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], device=v.device, dtype=v.dtype)],
                 dim=1)
 
-            x = flex_attention(query=padded_roped_query.transpose(2, 1),
-                               key=padded_roped_key.transpose(2, 1),
-                               value=padded_v.transpose(2, 1),
+            x = flex_attention(query=padded_roped_query[:, :, head_slice].transpose(2, 1),
+                               key=padded_roped_key[:, :, head_slice].transpose(2, 1),
+                               value=padded_v[:, :, head_slice].transpose(2, 1),
                                block_mask=block_mask)[:, :, :-padded_length].transpose(2, 1)
+            if sp_world_size > 1:
+                # Gather per-rank head outputs back to the full head dimension.
+                x = sequence_model_parallel_all_gather(x.contiguous(), dim=2)
         else:
             current_end = current_start + q.shape[1]
             sink_tokens = self.sink_size * frame_seqlen
@@ -178,7 +207,11 @@ class CausalWanSelfAttention(nn.Module):
                                                 is_neox_style=False).type_as(v)
                 key_window = _apply_rotary_emb(key_window, cos[:window_len], sin[:window_len],
                                                is_neox_style=False).type_as(v)
-            x = self.attn(roped_query, key_window, value_window)
+            x = self.attn(roped_query[:, :, head_slice], key_window[:, :, head_slice],
+                          value_window[:, :, head_slice])
+            if sp_world_size > 1:
+                # Gather per-rank head outputs back to the full head dimension.
+                x = sequence_model_parallel_all_gather(x.contiguous(), dim=2)
             if isinstance(kv_cache["global_end_index"], torch.Tensor):
                 kv_cache["global_end_index"].fill_(current_end)
             else:
@@ -581,10 +614,16 @@ class CausalWanTransformer3DModel(BaseDiT):
             # fixed table over [0, max_attention_frames); attention slices it per step
             max_attention_frames = (GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES
                                     if self.local_attn_size == -1 else self.local_attn_size)
-            rope_num_frames = max_attention_frames * get_sp_world_size()
+            # No multiplication by the sequence-parallel world size: the causal
+            # pipeline replicates the full (unsharded) sequence on every rank, so
+            # RoPE must cover the full local sequence length. Multiplying by the
+            # SP world size produced a 2x-too-long RoPE and crashed
+            # _apply_rotary_emb under SP. Attention is parallelized across heads
+            # instead (see CausalWanSelfAttention.forward).
+            rope_num_frames = max_attention_frames
             rope_start_frame = 0
         else:
-            rope_num_frames = post_patch_num_frames * get_sp_world_size()
+            rope_num_frames = post_patch_num_frames
             rope_start_frame = start_frame  # 0 when kv_cache is None
         freqs_cos, freqs_sin = get_rotary_pos_embed((rope_num_frames, post_patch_height, post_patch_width),
                                                     self.hidden_size,
@@ -682,7 +721,12 @@ class CausalWanTransformer3DModel(BaseDiT):
         d = self.hidden_size // self.num_attention_heads
         rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
         freqs_cos, freqs_sin = get_rotary_pos_embed(
-            (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width),
+            # No multiplication by the sequence-parallel world size: the causal
+            # pipeline replicates the full (unsharded) sequence on every rank, so
+            # RoPE must cover the full local sequence length. Attention is
+            # parallelized across heads instead (see
+            # CausalWanSelfAttention.forward).
+            (post_patch_num_frames, post_patch_height, post_patch_width),
             self.hidden_size,
             self.num_attention_heads,
             rope_dim_list,