[Bugfix] Fix Helios USP: per-component split for correct sequence parallelism (#6930)

Signed-off-by: yancaocn <yancaochn@163.com>
Co-authored-by: yancaocn <yancaochn@163.com>

tests/diffusion/models/helios/test_helios_usp_2rank.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Two-process USP smoke test for the Helios transformer.
+
+Pins the multi-rank mechanism that ``test_helios_usp_shard.py`` (unit tests
+of ``_sp_split_seq`` alone) cannot see:
+
+  * the manual ``_sp_shard_depth`` bump in ``forward()`` must make attn1 run
+    the Ulysses all-to-all (not ``NoParallelAttention``);
+  * attn2 cross-attention must stay rank-local (``skip_sequence_parallel=
+    True``): if replicated text K/V were fed through Ulysses they would be
+    ws-fold duplicated and outputs would diverge from the single-rank run;
+  * end-to-end correctness: ws=2 output equals the ws=1 baseline for the same
+    weights/inputs (divisible shapes, so the padding path is not involved).
+
+Runs ``HeliosTransformer3DModel.forward`` with tiny random weights via
+``torch.multiprocessing.spawn``: first a baseline (world_size=1), then an
+SP run (world_size=2, ulysses_degree=2); rank-0 outputs are compared.
+
+Requires >= 2 devices; skipped otherwise.
+"""
+
+import os
+import tempfile
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelConfig,
+    get_sp_plan_from_model,
+)
+from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
+from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
+from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
+from vllm_omni.platforms import current_omni_platform
+
+_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, *_L4_TWO_GPU]
+
+RANDOM_SEED = 42
+DTYPE = torch.float32  # keep both runs on the SDPA fallback for tight tolerance
+
+# Tiny but architecturally faithful config:
+#   inner_dim = 4 * 32 = 128, rope_dim sums to head_dim (32),
+#   current chunk 2x4x6 latent -> (1,2,2) patch -> 2*2*3 = 12 tokens,
+#   short history 1x4x6 latent -> (1,2,2) patch -> 6 tokens.
+# All component lengths are even, so ws=2 needs no padding replicas.
+TINY_MODEL_KWARGS = dict(
+    patch_size=(1, 2, 2),
+    num_attention_heads=4,
+    attention_head_dim=32,
+    in_channels=4,
+    out_channels=4,
+    text_dim=64,
+    freq_dim=32,
+    ffn_dim=128,
+    num_layers=2,
+    rope_dim=(12, 10, 10),
+    rope_theta=10000.0,
+    guidance_cross_attn=True,
+    zero_history_timestep=True,
+    has_multi_term_memory_patch=True,
+    is_amplify_history=False,
+)
+
+BATCH = 1
+CUR_FRAMES, LAT_H, LAT_W = 2, 4, 6  # -> 12 current tokens
+HIST_FRAMES = 1  # short history -> 6 tokens
+TEXT_LEN = 5
+
+
+def _update_env(envs_dict: dict[str, str]) -> None:
+    for k, v in envs_dict.items():
+        os.environ[k] = v
+
+
+def _make_inputs(device: torch.device) -> dict[str, torch.Tensor]:
+    torch.manual_seed(RANDOM_SEED + 1)
+    return {
+        "hidden_states": torch.randn(BATCH, 4, CUR_FRAMES, LAT_H, LAT_W, device=device, dtype=DTYPE),
+        "timestep": torch.full((BATCH,), 500.0, device=device),
+        "encoder_hidden_states": torch.randn(
+            BATCH, TEXT_LEN, TINY_MODEL_KWARGS["text_dim"], device=device, dtype=DTYPE
+        ),
+        "indices_latents_history_short": torch.tensor([[0]], device=device),
+        "latents_history_short": torch.randn(BATCH, 4, HIST_FRAMES, LAT_H, LAT_W, device=device, dtype=DTYPE),
+    }
+
+
+def _spy_strategies(model: HeliosTransformer3DModel) -> dict[str, list[str]]:
+    """Record the parallel strategy each Attention resolves to during forward."""
+    seen: dict[str, list[str]] = {"attn1": [], "attn2": []}
+    for name, attn in (("attn1", model.blocks[0].attn1.attn), ("attn2", model.blocks[0].attn2.attn)):
+        original = attn._get_active_parallel_strategy
+
+        def spy(original=original, name=name):
+            strategy = original()
+            seen[name].append(getattr(strategy, "name", type(strategy).__name__))
+            return strategy
+
+        attn._get_active_parallel_strategy = spy
+    return seen
+
+
+def helios_forward_on_model(
+    local_rank: int,
+    world_size: int,
+    ulysses_degree: int,
+    output_file: str,
+    model_state_file: str,
+    is_baseline: bool,
+) -> None:
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+    torch.set_default_device(device)
+    torch.set_default_dtype(DTYPE)
+
+    _update_env(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": "12361",
+        }
+    )
+    init_distributed_environment()
+
+    parallel_config = DiffusionParallelConfig(
+        pipeline_parallel_size=1,
+        data_parallel_size=1,
+        tensor_parallel_size=1,
+        sequence_parallel_size=ulysses_degree,
+        ulysses_degree=ulysses_degree,
+        ring_degree=1,
+        allgather_degree=1,
+        cfg_parallel_size=1,
+    )
+    od_config = OmniDiffusionConfig.from_kwargs(
+        model="test_model",
+        dtype=DTYPE,
+        parallel_config=parallel_config,
+        diffusion_attention_backend="TORCH_SDPA",
+    )
+    initialize_model_parallel(
+        data_parallel_size=1,
+        cfg_parallel_size=1,
+        sequence_parallel_size=ulysses_degree,
+        ulysses_degree=ulysses_degree,
+        ring_degree=1,
+        allgather_degree=1,
+        tensor_parallel_size=1,
+        pipeline_parallel_size=1,
+    )
+
+    with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config):
+        torch.manual_seed(RANDOM_SEED)
+        model = HeliosTransformer3DModel(**TINY_MODEL_KWARGS).to(device).to(DTYPE)
+
+        if is_baseline and local_rank == 0:
+            with open(model_state_file, "wb") as f:
+                torch.save({k: v.cpu() for k, v in model.state_dict().items()}, f)
+
+        if world_size > 1:
+            torch.distributed.barrier()
+
+        with open(model_state_file, "rb") as f:
+            model_state = torch.load(f)
+        model.load_state_dict({k: v.to(device).to(DTYPE) for k, v in model_state.items()})
+
+        if not is_baseline:
+            # Mirror registry._apply_sequence_parallel_if_enabled: install the
+            # proj_out gather hook from the model's own _sp_plan and flip
+            # sp_plan_hooks_applied so sp_active follows _sp_shard_depth
+            # (i.e. the manual bump in Helios forward is what enables Ulysses).
+            plan = get_sp_plan_from_model(model)
+            assert plan is not None, "HeliosTransformer3DModel must define _sp_plan"
+            apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=ulysses_degree), plan)
+            get_forward_context().sp_plan_hooks_applied = True
+
+        seen = _spy_strategies(model)
+        inputs = _make_inputs(device)
+        output = model(**inputs).sample
+
+        assert output.shape == (BATCH, 4, CUR_FRAMES, LAT_H, LAT_W), f"unexpected output shape {output.shape}"
+
+        if is_baseline:
+            assert seen["attn1"] and all(n == "none" for n in seen["attn1"]), (
+                f"baseline attn1 must use NoParallelAttention ('none'), saw {seen['attn1']}"
+            )
+        else:
+            # attn1: SP communication must be active inside the block loop
+            # (regression here = the manual _sp_shard_depth bump was lost).
+            assert seen["attn1"] and all(n == "ulysses" for n in seen["attn1"]), (
+                f"attn1 must run Ulysses under USP (manual _sp_shard_depth bump), saw {seen['attn1']}"
+            )
+            # attn2: replicated text K/V must stay rank-local
+            # (regression here = the ws-fold duplicated-keys bug).
+            assert seen["attn2"] and all(n == "none" for n in seen["attn2"]), (
+                f"attn2 must skip sequence parallel (replicated text K/V), saw {seen['attn2']}"
+            )
+
+        if local_rank == 0:
+            with open(output_file, "wb") as f:
+                torch.save(output.detach().cpu().float(), f)
+
+        destroy_distributed_env()
+
+
+def test_helios_usp2_matches_single_rank():
+    """ws=2 (Ulysses) must reproduce the ws=1 output for identical weights/inputs."""
+    available_devices = current_omni_platform.get_device_count()
+    if available_devices < 2:
+        pytest.skip(f"Test requires 2 devices but only {available_devices} available")
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
+        baseline_output_file = f.name
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
+        sp_output_file = f.name
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f:
+        model_state_file = f.name
+
+    try:
+        torch.multiprocessing.spawn(
+            helios_forward_on_model,
+            args=(1, 1, baseline_output_file, model_state_file, True),
+            nprocs=1,
+        )
+        torch.multiprocessing.spawn(
+            helios_forward_on_model,
+            args=(2, 2, sp_output_file, model_state_file, False),
+            nprocs=2,
+        )
+
+        with open(baseline_output_file, "rb") as f:
+            baseline = torch.load(f)
+        with open(sp_output_file, "rb") as f:
+            sp = torch.load(f)
+
+        assert baseline.shape == sp.shape
+        max_abs_diff = (baseline - sp).abs().max().item()
+        assert torch.allclose(baseline, sp, atol=1e-3, rtol=1e-3), (
+            f"USP=2 output diverges from single-rank baseline: max_abs_diff={max_abs_diff:.6e}"
+        )
+    finally:
+        for path in [baseline_output_file, sp_output_file, model_state_file]:
+            if os.path.exists(path):
+                os.remove(path)

tests/diffusion/models/helios/test_helios_usp_shard.py

@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for Helios per-component sequence splitting.
+
+These tests verify the manual sharding logic that replaces the broken
+rope/blocks.0 hooks in the Helios ``_sp_plan``.  They exercise the *real*
+``HeliosTransformer3DModel._sp_split_seq`` (the parallel-state getters it
+imports lazily inside the method body are monkeypatched), so the model-side
+helper cannot drift silently from what is tested.  For the multi-rank
+mechanism around the helper (manual ``_sp_shard_depth`` bump, attn1/attn2
+parallel strategy, ws=2 vs ws=1 output equivalence) see
+``test_helios_usp_2rank.py``.
+
+Coverage:
+  1. USP disabled (ws == 1) returns the tensor unchanged.
+  2. Divisible split gives each rank the expected slice.
+  3. Non-divisible split pads with replicated tokens (no crash).
+  4. Per-component split: history + current tokens are distributed so
+     that *both* ranks receive some current tokens (no mosaic regression).
+  5. Edge cases: 1-D tensor, empty seq, 4-D tensor.
+"""
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.distributed.parallel_state as sp_parallel_state
+from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@contextmanager
+def _fake_sp_state(ws: int, rank: int) -> Iterator[None]:
+    """Patch the parallel-state getters used by the real helper.
+
+    ``_sp_split_seq`` imports these names from the module *inside* its
+    function body, so replacing the module attributes takes effect.
+    """
+    prev_ws = sp_parallel_state.get_sequence_parallel_world_size
+    prev_rank = sp_parallel_state.get_sequence_parallel_rank
+    sp_parallel_state.get_sequence_parallel_world_size = lambda: ws
+    sp_parallel_state.get_sequence_parallel_rank = lambda: rank
+    try:
+        yield
+    finally:
+        sp_parallel_state.get_sequence_parallel_world_size = prev_ws
+        sp_parallel_state.get_sequence_parallel_rank = prev_rank
+
+
+def _sp_split_seq(x: torch.Tensor, ws: int, rank: int) -> torch.Tensor:
+    """Call the real model-side helper without constructing a model."""
+    with _fake_sp_state(ws, rank):
+        return HeliosTransformer3DModel._sp_split_seq(None, x)
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestSpSplitSeqDisabled:
+    """When USP is off (ws == 1) the helper must be a no-op."""
+
+    def test_ws1_returns_unchanged(self):
+        x = torch.randn(2, 100, 64)
+        out = _sp_split_seq(x, ws=1, rank=0)
+        assert torch.equal(out, x)
+        assert out is x
+
+
+class TestSpSplitSeqDivisible:
+    """Divisible sequences: each rank gets exactly seq_len / ws tokens."""
+
+    def test_each_rank_gets_correct_slice(self):
+        for ws in [2, 4, 8]:
+            seq_len = ws * 10
+            x = torch.arange(seq_len * 4).float().reshape(1, seq_len, 4)
+            for rank in range(ws):
+                out = _sp_split_seq(x, ws=ws, rank=rank)
+                n = seq_len // ws
+                expected = x[:, rank * n : (rank + 1) * n, :]
+                assert out.shape == (1, n, 4)
+                assert torch.equal(out, expected)
+
+    def test_output_is_contiguous(self):
+        x = torch.randn(1, 20, 8)
+        out = _sp_split_seq(x, ws=2, rank=0)
+        assert out.is_contiguous()
+
+
+class TestSpSplitSeqNonDivisible:
+    """Non-divisible sequences are padded with replicated tokens."""
+
+    def test_non_divisible_pads_instead_of_crash(self):
+        """400x272 → 9*25*17 = 3825 tokens, USP=2 → not divisible.
+        Should pad to 3826 and split, not crash."""
+        x = torch.randn(1, 3825, 8)
+        for rank in range(2):
+            out = _sp_split_seq(x, ws=2, rank=rank)
+            # 3825 → pad 1 → 3826 → /2 = 1913
+            assert out.shape[1] == 1913
+
+    def test_non_divisible_pads_with_replicated_last_token(self):
+        """Verify the padding is the replicated last token, not zeros."""
+        x = torch.randn(1, 7, 4)  # 7 % 2 = 1, pad 1
+        last_token = x[:, -1:, :].clone()
+        out = _sp_split_seq(x, ws=2, rank=0)
+        # rank 0 gets tokens [0:4] (first 4 of the padded 8)
+        assert out.shape == (1, 4, 4)
+        # rank 1 gets tokens [4:8] — token 7 (the padded one) = last token
+        out1 = _sp_split_seq(x, ws=2, rank=1)
+        assert out1.shape == (1, 4, 4)
+        assert torch.equal(out1[:, -1:, :], last_token)
+
+    def test_non_divisible_7_4(self):
+        """7 tokens, ws=4 → pad 1 → 8 → /4 = 2 per rank."""
+        x = torch.randn(1, 7, 8)
+        for rank in range(4):
+            out = _sp_split_seq(x, ws=4, rank=rank)
+            assert out.shape[1] == 2
+
+    def test_non_divisible_15_4(self):
+        """15 tokens, ws=4 → pad 1 → 16 → /4 = 4 per rank."""
+        x = torch.randn(1, 15, 8)
+        for rank in range(4):
+            out = _sp_split_seq(x, ws=4, rank=rank)
+            assert out.shape[1] == 4
+
+
+class TestSpSplitSeqPerComponent:
+    """Simulate Helios forward: split current + history components so that
+    *both* ranks receive some current tokens (no mosaic regression)."""
+
+    def test_both_ranks_get_current_tokens(self):
+        """history=2400, current=540, ws=2.
+        After per-component split, each rank must have > 0 current tokens."""
+        current = torch.randn(1, 540, 64)
+        history = torch.randn(1, 2400, 64)
+
+        for rank in range(2):
+            cur_local = _sp_split_seq(current, ws=2, rank=rank)
+            hist_local = _sp_split_seq(history, ws=2, rank=rank)
+
+            assert cur_local.shape[1] > 0, f"Rank {rank} received 0 current tokens"
+            assert cur_local.shape[1] == 270  # 540 / 2
+            assert hist_local.shape[1] == 1200  # 2400 / 2
+
+            original_context_length = cur_local.shape[1]
+            assert original_context_length > 0  # no mosaic regression
+
+    def test_gather_restores_full_current(self):
+        """After splitting current into 2 ranks, simulating proj_out
+        gather must restore the full current sequence length so that
+        the original unpatchify reshape succeeds."""
+        current = torch.randn(1, 540, 64)
+        local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)]
+        gathered = torch.cat(local_chunks, dim=1)
+        assert gathered.shape[1] == 540  # full current restored
+
+    def test_gather_with_padding_truncates_correctly(self):
+        """Non-divisible current (3825) → pad to 3826 → split →
+        gather → 3826. Slice to 3825 for unpatchify."""
+        current = torch.randn(1, 3825, 64)
+        local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)]
+        gathered = torch.cat(local_chunks, dim=1)
+        assert gathered.shape[1] == 3826  # padded size
+        # Simulate the slice before unpatchify
+        expected_seq = 3825
+        if gathered.shape[1] > expected_seq:
+            gathered = gathered[:, :expected_seq, :]
+        assert gathered.shape[1] == 3825  # correct after slice
+
+    def test_rotary_emb_aligned_with_hidden_states(self):
+        """rotary_emb and hidden_states must be split at the same
+        positions so that tokens align inside the attention layer."""
+        hidden = torch.randn(1, 540, 5120)
+        rotary = torch.randn(1, 540, 256)
+        for rank in range(2):
+            h_local = _sp_split_seq(hidden, ws=2, rank=rank)
+            r_local = _sp_split_seq(rotary, ws=2, rank=rank)
+            assert h_local.shape[1] == r_local.shape[1]
+
+    def test_whole_split_causes_mosaic_but_per_component_does_not(self):
+        """Demonstrates why per-component split is needed:
+        with whole split (simulated), rank 0 gets 0 current tokens;
+        with per-component split, both ranks get current tokens."""
+        history = torch.randn(1, 2400, 64)
+        current = torch.randn(1, 540, 64)
+        full = torch.cat([history, current], dim=1)  # 2940
+
+        # Whole split (the OLD broken behavior)
+        whole_n = full.shape[1] // 2  # 1470
+        rank0_whole = full[:, :whole_n, :]
+        rank0_current = rank0_whole.shape[1] - 2400  # negative → 0 current
+        assert rank0_current <= 0, "Whole split should give rank 0 no current"
+
+        # Per-component split (the NEW correct behavior)
+        for rank in range(2):
+            cur_local = _sp_split_seq(current, ws=2, rank=rank)
+            assert cur_local.shape[1] == 270  # both ranks have current
+
+
+class TestSpSplitSeqEdgeCases:
+    """Edge cases that should not crash or silently misbehave."""
+
+    def test_1d_tensor_skipped(self):
+        """1-D tensors must be returned unchanged (dim < 2 guard)."""
+        x = torch.randn(100)
+        out = _sp_split_seq(x, ws=2, rank=0)
+        assert torch.equal(out, x)
+
+    def test_empty_seq_skipped(self):
+        """Empty sequence dimension must be returned unchanged."""
+        x = torch.randn(1, 0, 64)
+        out = _sp_split_seq(x, ws=2, rank=0)
+        assert out.shape == x.shape
+
+    def test_4d_tensor_split_on_dim1(self):
+        """4-D tensor [B, S, H, D] should split along dim=1."""
+        x = torch.randn(1, 20, 8, 32)
+        out = _sp_split_seq(x, ws=2, rank=1)
+        assert out.shape == (1, 10, 8, 32)

vllm_omni/diffusion/models/helios/helios_transformer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Adapted from Helios (https://github.com/BestWishYsh/Helios)
 
 import math
@@ -27,10 +27,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
-from vllm_omni.diffusion.distributed.sp_plan import (
-    SequenceParallelInput,
-    SequenceParallelOutput,
-)
+from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelOutput
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import (
@@ -448,6 +445,12 @@ class HeliosCrossAttention(nn.Module):
             num_kv_heads=self.num_heads,
             softmax_scale=1.0 / (head_dim**0.5),
             causal=False,
+            # Text K/V is replicated on every SP rank, so Ulysses must not be
+            # used here: its all-to-all would treat the per-rank replicas as
+            # sequence shards and ws-fold duplicate the keys. Sharded Q over
+            # full replicated K/V is correct locally and needs no
+            # communication (same as Wan2.2 cross-attn).
+            skip_sequence_parallel=True,
         )
 
     def project_kv(self, encoder_hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
@@ -624,16 +627,50 @@ class HeliosTransformer3DModel(nn.Module):
     _hsdp_shard_conditions = [_is_transformer_block]
 
     _sp_plan = {
-        "rope": {
-            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
-            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
-        },
-        "blocks.0": {
-            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True),
-        },
+        # "rope" and "blocks.0" are intentionally omitted.
+        # The rope hook splits along dim=1 of the 5D output [B, D, T, H, W],
+        # which is the frequency dimension, NOT the sequence dimension.
+        # The blocks.0 hook splits the concatenated hidden_states as a whole,
+        # which can put all history tokens in one rank (original_context_length=0).
+        # Instead, hidden_states and rotary_emb are split per-component in
+        # forward() after flatten+transpose, ensuring each rank gets half
+        # of each component (history and current).
         "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
     }
 
+    def _sp_split_seq(self, x: torch.Tensor) -> torch.Tensor:
+        """Split tensor along sequence dim (dim=1) for Ulysses SP.
+
+        If seq_len is not divisible by world_size, replicate the last
+        token(s) to pad to a divisible length.  This avoids crashing
+        on non-divisible history components (which use different patch
+        sizes and thus produce different token counts).
+
+        Residual caveat: the replicated padding tokens attend unmasked
+        in attn1 on the last rank (at most ws-1 replicas per component).
+        The impact is bounded, but this is the first place to look if
+        odd resolutions ever show quality drift.
+        """
+        from vllm_omni.diffusion.distributed.parallel_state import (
+            get_sequence_parallel_rank,
+            get_sequence_parallel_world_size,
+        )
+
+        ws = get_sequence_parallel_world_size()
+        if ws > 1 and x.dim() >= 2 and x.shape[1] > 0:
+            seq_len = x.shape[1]
+            remainder = seq_len % ws
+            if remainder != 0:
+                pad = ws - remainder
+                # Replicate the last token to fill the padding
+                last = x[:, -1:, ...].expand(-1, pad, *([-1] * (x.dim() - 2)))
+                x = torch.cat([x, last], dim=1)
+                seq_len = x.shape[1]
+            r = get_sequence_parallel_rank()
+            n = seq_len // ws
+            x = x[:, r * n : (r + 1) * n, ...].contiguous()
+        return x
+
     def __init__(
         self,
         patch_size: tuple[int, ...] = (1, 2, 2),
@@ -860,6 +897,9 @@ class HeliosTransformer3DModel(nn.Module):
             device=hidden_states.device,
         )
         rotary_emb = rotary_emb.flatten(2).transpose(1, 2)
+        # USP: per-component split (each rank gets half of current frames)
+        hidden_states = self._sp_split_seq(hidden_states)
+        rotary_emb = self._sp_split_seq(rotary_emb)
         original_context_length = hidden_states.shape[1]
 
         # 2. Process short history latents
@@ -876,6 +916,9 @@ class HeliosTransformer3DModel(nn.Module):
                 device=latents_history_short.device,
             )
             rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose(1, 2)
+            # USP: per-component split
+            latents_history_short = self._sp_split_seq(latents_history_short)
+            rotary_emb_history_short = self._sp_split_seq(rotary_emb_history_short)
 
             hidden_states = torch.cat([latents_history_short, hidden_states], dim=1)
             rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1)
@@ -896,6 +939,9 @@ class HeliosTransformer3DModel(nn.Module):
             rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2))
             rotary_emb_history_mid = center_down_sample_3d(rotary_emb_history_mid, (2, 2, 2))
             rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2)
+            # USP: per-component split
+            latents_history_mid = self._sp_split_seq(latents_history_mid)
+            rotary_emb_history_mid = self._sp_split_seq(rotary_emb_history_mid)
 
             hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1)
             rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1)
@@ -916,6 +962,9 @@ class HeliosTransformer3DModel(nn.Module):
             rotary_emb_history_long = pad_for_3d_conv(rotary_emb_history_long, (4, 4, 4))
             rotary_emb_history_long = center_down_sample_3d(rotary_emb_history_long, (4, 4, 4))
             rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2)
+            # USP: per-component split
+            latents_history_long = self._sp_split_seq(latents_history_long)
+            rotary_emb_history_long = self._sp_split_seq(rotary_emb_history_long)
 
             hidden_states = torch.cat([latents_history_long, hidden_states], dim=1)
             rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1)
@@ -958,6 +1007,18 @@ class HeliosTransformer3DModel(nn.Module):
             timestep_proj = timestep_proj.permute(0, 2, 1, 3)
 
         # 6. Transformer blocks
+        # Manually increment _sp_shard_depth so that attention layers
+        # know SP is active and perform Ulysses All-to-All communication.
+        # Normally this is done by split_output=True hooks (e.g. rope),
+        # but those are removed because they split the wrong dimension.
+        from vllm_omni.diffusion.forward_context import (
+            get_forward_context,
+            is_forward_context_available,
+        )
+
+        if is_forward_context_available():
+            get_forward_context()._sp_shard_depth += 1
+
         hidden_states = hidden_states.contiguous()
         encoder_hidden_states = encoder_hidden_states.contiguous()
         rotary_emb = rotary_emb.contiguous()
@@ -979,6 +1040,11 @@ class HeliosTransformer3DModel(nn.Module):
         hidden_states = self.proj_out(hidden_states)
 
         # 8. Unpatchify
+        # proj_out gather may include padded tokens from non-divisible
+        # sequences. Slice to the expected size before reshape.
+        expected_seq = post_patch_num_frames * post_patch_height * post_patch_width
+        if hidden_states.shape[1] > expected_seq:
+            hidden_states = hidden_states[:, :expected_seq, :]
         hidden_states = hidden_states.reshape(
             batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1
         )