[Model] Add sequence parallelism to SANA-Video 2B (#5940)

Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>

docs/user_guide/diffusion_features.md

@@ -164,7 +164,7 @@ The following tables show which models support each feature:
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2.5**                  |     ❌     | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) |         ❓         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |    ❓ (FP8)     |        ❌         |
-|  **SANA-Video-2B T2V I2V**   |     ❌     |     ❌      |          ❌           |       ✅       |  ✅ (TP=2 only)   |         ❌         |   ❌   |            ❌             |          ❌          |       ❌       |        ❌        |
+|  **SANA-Video-2B T2V I2V**   |     ❌     |     ❌      |   ✅ (`--usp`, frame-sharded)   |       ✅       |  ✅ (TP=2 only)   |         ❌         |   ❌   |            ❌             |          ❌          |       ❌       |        ❌        |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ✅*        |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |

recipes/NVIDIA/SANA-Video-2B.md

@@ -121,8 +121,8 @@ bash examples/online_serving/text_to_video/run_curl_sana_video.sh
 
 ##### Parallel native serving
 
-The native pipelines support tensor parallelism (up to 2 GPUs) and CFG
-parallelism. Tensor parallelism on 2 GPUs:
+The native pipelines support tensor parallelism (up to 2 GPUs), sequence
+parallelism (2 or 4 GPUs) and CFG parallelism. Tensor parallelism on 2 GPUs:
 
 ```bash
 vllm serve Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
@@ -133,15 +133,32 @@ vllm serve Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
   --port 8091
 ```
 
-For CFG parallelism on 2 GPUs, replace `--tensor-parallel-size 2` with
-`--cfg-parallel-size 2`; it splits the guided and unguided branches across the
-two GPUs and only helps when `guidance_scale` is above 1. Passing both flags
-combines the two on 4 GPUs.
+Sequence parallelism on 2 GPUs (replace `--tensor-parallel-size 2` with):
 
-Whether tensor parallelism lowers latency depends on the interconnect. Each
-transformer block adds an all-reduce, so it speeds up generation only on fast
-GPU links such as NVLink and can be slower than a single GPU on PCIe-only
-systems. Measure on your hardware before enabling it.
+```bash
+  --usp 2
+```
+
+The `--usp` flag carries the sequence parallel degree. On SANA-Video the
+implementation is not Ulysses all-to-all: video frames are sharded across
+ranks and the linear-attention state is combined with one small all-reduce
+per block, which is far less communication than tensor parallelism needs.
+Ring and AllGather-KV sequence parallel do not apply to SANA's linear
+attention and are rejected.
+
+For CFG parallelism on 2 GPUs, use `--cfg-parallel-size 2` instead; it splits
+the guided and unguided branches across the two GPUs and only helps when
+`guidance_scale` is above 1. CFG parallelism combines with either tensor
+parallelism or sequence parallelism on 4 GPUs. Tensor and sequence parallelism
+also combine; on the measured 4-GPU NVLink box tp2 x sp2 reached 2.55x, below
+the pure sp4 (3.24x) and sp2 x cfg2 (3.56x) configurations, so prefer those
+when four GPUs serve one request.
+
+Whether these lower latency depends on the interconnect. Each transformer
+block adds an all-reduce under tensor parallelism, so it speeds up generation
+only on fast GPU links such as NVLink and can be slower than a single GPU on
+PCIe-only systems. Sequence parallelism communicates far less per block, but
+was only measured on NVLink. Measure on your hardware before enabling either.
 
 To run the black-box compatibility backend for T2V, replace the server script
 with `run_server_sana_video_diffusers.sh`. The same `/v1/videos` request
@@ -287,9 +304,9 @@ for the tuning knobs.
   The denoising loop intentionally retains the checkpoint-compatible
   Diffusers `DPMSolverMultistepScheduler`.
 - Known limitations:
-    - Tensor parallelism (up to 2 GPUs) and CFG parallelism are supported for
-    the native pipeline. Sequence parallelism, TeaCache, and step execution
-    are not validated.
+    - Tensor parallelism (up to 2 GPUs), sequence parallelism (2 or 4 GPUs)
+    and CFG parallelism are supported for the native pipeline. TeaCache and
+    step execution are not validated.
     - Cache-DiT and CPU offload are limited to TP1/CFG1/SP1. Cache-DiT with
     distributed layerwise offload is not supported by the native pipeline.
     - Cache-DiT speedup and offload memory numbers above are single-GPU A800

tests/diffusion/distributed/test_sana_video_sp_numeric.py

@@ -0,0 +1,357 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Real multi-process numeric equivalence gate for SANA-Video sequence parallel.
+
+Runs the native transformer directly (no pipeline, no topology validation) with
+real NCCL groups and compares SP>1 output against a same-weight SP=1 baseline.
+In fp32 the only difference is reduction reordering, so any halo, gather,
+timestep-slice or state-reduction misalignment fails by orders of magnitude.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    get_classifier_free_guidance_rank,
+    get_sequence_parallel_rank,
+    get_sequence_parallel_world_size,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.platforms import current_omni_platform
+
+_TINY_CONFIG = {
+    "in_channels": 4,
+    "out_channels": 4,
+    "num_attention_heads": 2,
+    "attention_head_dim": 12,
+    "num_layers": 2,
+    "num_cross_attention_heads": 2,
+    "cross_attention_head_dim": 12,
+    "cross_attention_dim": 24,
+    "caption_channels": 8,
+    "mlp_ratio": 2.0,
+    "patch_size": (1, 2, 2),
+    "sample_size": 8,
+    "rope_max_seq_len": 512,
+}
+
+_MODEL_SEED = 42
+_INPUT_SEED = 123
+_FP32_REL_L2 = 1e-5
+# bf16 collectives only reorder low-precision sums; informational sanity bound.
+_BF16_REL_L2 = 5e-2
+
+
+def _make_omni_config(sp_size: int, cfg_size: int, dtype: torch.dtype, tp_size: int = 1) -> OmniDiffusionConfig:
+    parallel_config = DiffusionParallelConfig(
+        pipeline_parallel_size=1,
+        data_parallel_size=1,
+        tensor_parallel_size=tp_size,
+        sequence_parallel_size=sp_size,
+        ulysses_degree=sp_size,
+        ring_degree=1,
+        cfg_parallel_size=cfg_size,
+    )
+    return OmniDiffusionConfig(
+        model="test",
+        dtype=dtype,
+        parallel_config=parallel_config,
+        # Backend-agnostic gate; SDPA also avoids platform quirks of the
+        # default backend on older GPUs.
+        diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")),
+    )
+
+
+def _build_model(device: torch.device, dtype: torch.dtype):
+    from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+
+    torch.manual_seed(_MODEL_SEED)
+    model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval()
+    for _, param in sorted(model.named_parameters()):
+        torch.nn.init.normal_(param, mean=0.0, std=0.02)
+    return model.to(device=device, dtype=dtype)
+
+
+def _make_inputs(task: str, num_frames: int, input_seed: int):
+    generator = torch.Generator().manual_seed(input_seed)
+    latent = torch.randn(2, 4, num_frames, 8, 8, generator=generator)
+    encoder_hidden_states = torch.randn(2, 6, 8, generator=generator)
+    encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 1, 0], [1, 1, 1, 1, 1, 1]])
+    if task == "t2v":
+        timestep = torch.tensor([500.0, 700.0])
+    else:
+        timestep = torch.rand(2, 1, num_frames, 4, 4, generator=generator) * 1000.0
+    return latent, encoder_hidden_states, encoder_attention_mask, timestep
+
+
+def _worker(
+    local_rank: int,
+    world_size: int,
+    sp_size: int,
+    cfg_size: int,
+    tp_size: int,
+    dtype: torch.dtype,
+    task: str,
+    num_frames: int,
+    weights_path: str,
+    port: int,
+    result_queue,
+):
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+    os.environ.update(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": str(port),
+        }
+    )
+
+    init_distributed_environment()
+    initialize_model_parallel(
+        cfg_parallel_size=cfg_size,
+        sequence_parallel_size=sp_size,
+        ulysses_degree=sp_size,
+        tensor_parallel_size=tp_size,
+    )
+    assert get_sequence_parallel_world_size() == sp_size
+
+    cfg_rank = get_classifier_free_guidance_rank() if cfg_size > 1 else 0
+    od_config = _make_omni_config(sp_size, cfg_size, dtype, tp_size)
+    with set_current_diffusion_config(od_config):
+        model = _build_model(device, dtype)
+    # Load the baseline's full weights through the real weight loaders so TP
+    # ranks hold proper shards; per-rank random init is not TP-consistent.
+    model.load_weights(iter(torch.load(weights_path, map_location="cpu").items()))
+    # Each CFG branch gets distinct inputs so a CFG/SP group mix-up cannot
+    # cancel out.
+    latent, encoder_hidden_states, encoder_attention_mask, timestep = _make_inputs(
+        task, num_frames, _INPUT_SEED + cfg_rank
+    )
+
+    with torch.no_grad():
+        output = model(
+            latent.to(device=device, dtype=dtype),
+            encoder_hidden_states.to(device=device, dtype=dtype),
+            timestep.to(device=device),
+            encoder_attention_mask=encoder_attention_mask.to(device=device),
+            return_dict=False,
+        )[0]
+
+    result_queue.put((cfg_rank, get_sequence_parallel_rank(), output.float().cpu()))
+    destroy_distributed_env()
+
+
+def _rel_l2(actual: torch.Tensor, expected: torch.Tensor) -> float:
+    return ((actual - expected).double().norm() / expected.double().norm()).item()
+
+
+def _spawn(worker_args: tuple, nprocs: int, result_queue) -> list:
+    torch.multiprocessing.spawn(_worker, args=(*worker_args, result_queue), nprocs=nprocs)
+    return [result_queue.get() for _ in range(nprocs)]
+
+
+def _run_sp_parity(
+    *,
+    sp_size: int,
+    cfg_size: int,
+    dtype: torch.dtype,
+    task: str,
+    num_frames: int,
+    rel_l2_bound: float,
+    tp_size: int = 1,
+) -> float:
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+
+    baselines = {}
+    with tempfile.TemporaryDirectory() as tmpdir:
+        weights_path = os.path.join(tmpdir, "baseline_weights.pt")
+        for cfg_rank in range(cfg_size):
+            queue = manager.Queue()
+            # A world-1 run sees input seed _INPUT_SEED + 0; shift the seed via
+            # the task inputs by spawning one baseline per CFG branch input.
+            torch.multiprocessing.spawn(
+                _baseline_worker,
+                args=(
+                    dtype,
+                    task,
+                    num_frames,
+                    _INPUT_SEED + cfg_rank,
+                    weights_path if cfg_rank == 0 else "",
+                    29531,
+                    queue,
+                ),
+                nprocs=1,
+            )
+            baselines[cfg_rank] = queue.get()
+
+        queue = manager.Queue()
+        world_size = sp_size * cfg_size * tp_size
+        results = _spawn(
+            (world_size, sp_size, cfg_size, tp_size, dtype, task, num_frames, weights_path, 29532), world_size, queue
+        )
+
+    worst = 0.0
+    by_cfg: dict[int, list] = {}
+    for cfg_rank, sp_rank, output in results:
+        by_cfg.setdefault(cfg_rank, []).append((sp_rank, output))
+    assert sorted(by_cfg) == list(range(cfg_size))
+    for cfg_rank, rank_outputs in by_cfg.items():
+        rank_outputs.sort(key=lambda item: item[0])
+        first = rank_outputs[0][1]
+        for sp_rank, output in rank_outputs[1:]:
+            torch.testing.assert_close(
+                output, first, rtol=0, atol=0, msg=f"SP ranks disagree (cfg_rank={cfg_rank}, sp_rank={sp_rank})"
+            )
+        drift = _rel_l2(first, baselines[cfg_rank])
+        worst = max(worst, drift)
+        assert drift <= rel_l2_bound, (
+            f"SP{sp_size} output drifted rel_l2={drift:.3e} from the SP1 baseline "
+            f"(bound {rel_l2_bound:.1e}, cfg_rank={cfg_rank}, task={task}, frames={num_frames})"
+        )
+    return worst
+
+
+def _baseline_worker(
+    local_rank: int,
+    dtype: torch.dtype,
+    task: str,
+    num_frames: int,
+    input_seed: int,
+    weights_path: str,
+    port: int,
+    result_queue,
+):
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+    os.environ.update(
+        {
+            "RANK": "0",
+            "LOCAL_RANK": "0",
+            "WORLD_SIZE": "1",
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": str(port),
+        }
+    )
+
+    init_distributed_environment()
+    initialize_model_parallel()
+
+    od_config = _make_omni_config(1, 1, dtype)
+    with set_current_diffusion_config(od_config):
+        model = _build_model(device, dtype)
+    if weights_path:
+        torch.save({name: param.detach().cpu() for name, param in model.named_parameters()}, weights_path)
+    latent, encoder_hidden_states, encoder_attention_mask, timestep = _make_inputs(task, num_frames, input_seed)
+
+    with torch.no_grad():
+        output = model(
+            latent.to(device=device, dtype=dtype),
+            encoder_hidden_states.to(device=device, dtype=dtype),
+            timestep.to(device=device),
+            encoder_attention_mask=encoder_attention_mask.to(device=device),
+            return_dict=False,
+        )[0]
+
+    result_queue.put(output.float().cpu())
+    destroy_distributed_env()
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@pytest.mark.parametrize("task", ["t2v", "i2v"])
+@pytest.mark.parametrize("num_frames", [21, 5])
+def test_sp2_matches_sp1_fp32(task: str, num_frames: int) -> None:
+    _run_sp_parity(
+        sp_size=2,
+        cfg_size=1,
+        dtype=torch.float32,
+        task=task,
+        num_frames=num_frames,
+        rel_l2_bound=_FP32_REL_L2,
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4"}, num_cards=4)
+@pytest.mark.parametrize("task", ["t2v", "i2v"])
+@pytest.mark.parametrize("num_frames", [21, 5])
+def test_sp4_matches_sp1_fp32(task: str, num_frames: int) -> None:
+    _run_sp_parity(
+        sp_size=4,
+        cfg_size=1,
+        dtype=torch.float32,
+        task=task,
+        num_frames=num_frames,
+        rel_l2_bound=_FP32_REL_L2,
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4"}, num_cards=4)
+def test_sp2_cfg2_matches_sp1_fp32() -> None:
+    """Four processes: CFG branches with distinct inputs each run their own SP
+    group; catches CFG/SP group crossover and deadlocks."""
+    _run_sp_parity(
+        sp_size=2,
+        cfg_size=2,
+        dtype=torch.float32,
+        task="i2v",
+        num_frames=21,
+        rel_l2_bound=_FP32_REL_L2,
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4"}, num_cards=4)
+def test_tp2_sp2_matches_sp1_fp32() -> None:
+    """Four processes: TP shards heads and channels while SP shards frames; the
+    combined mesh must still reproduce the serial output."""
+    _run_sp_parity(
+        sp_size=2,
+        cfg_size=1,
+        tp_size=2,
+        dtype=torch.float32,
+        task="t2v",
+        num_frames=21,
+        rel_l2_bound=_FP32_REL_L2,
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_sp2_bf16_drift_informational() -> None:
+    """Records the bf16 reduction-reorder drift magnitude; the load-bearing
+    numeric gate is the fp32 suite above."""
+    drift = _run_sp_parity(
+        sp_size=2,
+        cfg_size=1,
+        dtype=torch.bfloat16,
+        task="t2v",
+        num_frames=21,
+        rel_l2_bound=_BF16_REL_L2,
+    )
+    print(f"\nSANA-Video SP2 bf16 drift vs SP1: rel_l2={drift:.3e}")

tests/diffusion/models/sana_video/conftest.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.diffusion.models.sana_video import transformer_sana_video
+
+
+@pytest.fixture(autouse=True)
+def sp_world_1(monkeypatch):
+    """L1 runs single-process without a sequence parallel group; keep the
+    transformer on its SP1 path unless a test mocks a larger SP world."""
+    monkeypatch.setattr(transformer_sana_video, "get_sequence_parallel_world_size", lambda: 1)

tests/diffusion/models/sana_video/test_sana_video_topology.py

@@ -3,8 +3,10 @@
 
 """Topology validation for SANA-Video parallelism.
 
-The native pipeline only supports tensor/CFG parallel of size 1 or 2 and rejects
-sequence parallel and text-encoder TP before any large weight is allocated.
+The native pipeline only supports tensor/CFG parallel of size 1 or 2 and
+sequence parallel of degree 1, 2 or 4 (Ulysses flag, frame-sharded manual
+implementation); it rejects ring/AllGather-KV SP, the untested TP x SP combo
+and text-encoder TP before any large weight is allocated.
 """
 
 from types import SimpleNamespace
@@ -24,6 +26,8 @@ def _pc(
     tp: int = 1,
     cfg: int = 1,
     sp: int | None = None,
+    ring: int = 1,
+    allgather: int = 1,
     text_tp: int = 1,
     pp: int = 1,
     hsdp: bool = False,
@@ -32,6 +36,8 @@ def _pc(
         tensor_parallel_size=tp,
         cfg_parallel_size=cfg,
         sequence_parallel_size=sp,
+        ring_degree=ring,
+        allgather_degree=allgather,
         text_encoder_tp_size=text_tp,
         pipeline_parallel_size=pp,
         use_hsdp=hsdp,
@@ -54,9 +60,30 @@ def test_cfg3_rejected() -> None:
         validate_sana_video_parallel_config(_pc(cfg=3))
 
 
-def test_sequence_parallel_rejected() -> None:
-    with pytest.raises(NotImplementedError, match="[Ss]equence parallel"):
-        validate_sana_video_parallel_config(_pc(sp=2))
+@pytest.mark.parametrize("sp", [None, 1, 2, 4])
+def test_supported_sp_sizes_pass(sp: int | None) -> None:
+    validate_sana_video_parallel_config(_pc(sp=sp))
+
+
+@pytest.mark.parametrize("sp", [3, 8])
+def test_unsupported_sp_sizes_rejected(sp: int) -> None:
+    with pytest.raises(NotImplementedError, match="sequence_parallel_size"):
+        validate_sana_video_parallel_config(_pc(sp=sp))
+
+
+def test_ring_sequence_parallel_rejected() -> None:
+    with pytest.raises(NotImplementedError, match="ring"):
+        validate_sana_video_parallel_config(_pc(sp=2, ring=2))
+
+
+def test_allgather_sequence_parallel_rejected() -> None:
+    with pytest.raises(NotImplementedError, match="AllGather"):
+        validate_sana_video_parallel_config(_pc(sp=2, allgather=2))
+
+
+@pytest.mark.parametrize("sp", [2, 4])
+def test_tp2_sp_combos_pass(sp: int) -> None:
+    validate_sana_video_parallel_config(_pc(tp=2, sp=sp))
 
 
 def test_pipeline_parallel_rejected() -> None:
@@ -78,6 +105,10 @@ def test_tp2_cfg2_combo_passes() -> None:
     validate_sana_video_parallel_config(_pc(tp=2, cfg=2))
 
 
+def test_sp2_cfg2_combo_passes() -> None:
+    validate_sana_video_parallel_config(_pc(sp=2, cfg=2))
+
+
 @pytest.mark.parametrize("pipeline_cls", [SanaVideoPipeline, SanaImageToVideoPipeline])
 def test_pipeline_rejects_unsupported_topology_before_loading(pipeline_cls, mocker) -> None:
     """The pipeline must reject an unsupported topology before it loads weights."""

tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py

@@ -0,0 +1,567 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""SP unit tests for SANA-Video: balanced frame split sizes, linear-attention
+state reduction and the uneven frame gather.
+
+CPU-only. The SP world size, rank and the collectives are mocked so a single
+process simulates the shard/reduce logic a real multi-rank run depends on.
+"""
+
+import os
+import threading
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+    GLUMBTempConv,
+    SanaAdaLayerNormSingle,
+    SanaLinearAttention,
+    _sp_frame_split_sizes,
+    _sp_gather_frames,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+_MODULE = "vllm_omni.diffusion.models.sana_video.transformer_sana_video"
+
+
+@pytest.fixture
+def tp1_group():
+    """Real single-process TP group so the parallel linear layers construct and
+    run on CPU at tensor_parallel_size=1."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29501")
+    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture
+def force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for the parallel linear layers."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+# ── frame split sizes ──
+
+
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("num_frames", [5, 6, 9, 11, 21])
+def test_frame_split_sizes_are_balanced_and_cover_all_frames(num_frames, world_size) -> None:
+    """Exactly world_size non-empty near-equal chunks that sum to num_frames.
+
+    torch.chunk's ceil semantics produce fewer than world_size chunks for
+    5/6/9 frames at world size 4, which is why this helper exists.
+    """
+    sizes = _sp_frame_split_sizes(num_frames, world_size)
+
+    assert len(sizes) == world_size
+    assert all(size >= 1 for size in sizes)
+    assert sum(sizes) == num_frames
+    assert max(sizes) - min(sizes) <= 1
+
+
+def test_frame_split_sizes_known_layouts() -> None:
+    assert _sp_frame_split_sizes(21, 2) == [11, 10]
+    assert _sp_frame_split_sizes(11, 2) == [6, 5]
+    assert _sp_frame_split_sizes(9, 4) == [3, 2, 2, 2]
+    assert _sp_frame_split_sizes(5, 4) == [2, 1, 1, 1]
+
+
+def test_frame_split_sizes_rejects_fewer_frames_than_ranks() -> None:
+    with pytest.raises(ValueError, match="latent frame per rank"):
+        _sp_frame_split_sizes(3, 4)
+
+
+# ── linear attention state reduction ──
+
+
+class _FakeSpAllReduce:
+    """Two-pass stand-in for the SP all-reduce: pass 1 records each rank's
+    partial, pass 2 replays their sum to every rank."""
+
+    def __init__(self) -> None:
+        self.partials: list[torch.Tensor] = []
+        self.total: torch.Tensor | None = None
+
+    def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
+        self.partials.append(tensor.clone())
+        if self.total is not None:
+            return self.total.clone()
+        return tensor
+
+
+def _make_attn(dim: int, num_heads: int, head_dim: int) -> SanaLinearAttention:
+    attn = SanaLinearAttention(
+        dim=dim, num_heads=num_heads, head_dim=head_dim, dropout=0.0, bias=False, qk_norm="rms_norm_across_heads"
+    )
+    attn.eval()
+    torch.manual_seed(0)
+    for _, param in sorted(attn.named_parameters()):
+        param.data.normal_()
+    return attn
+
+
+def _attn_inputs(seq_len: int, dim: int, head_dim: int):
+    """Second half of the sequence drawn from a shifted, scaled distribution so
+    per-rank partial sums differ; a missing state reduction cannot cancel out."""
+    torch.manual_seed(1)
+    x = torch.randn(2, seq_len, dim)
+    x[:, seq_len // 2 :] = x[:, seq_len // 2 :] * 3.0 + 1.0
+    freqs_cos = torch.randn(1, seq_len, 1, head_dim)
+    freqs_sin = torch.randn(1, seq_len, 1, head_dim)
+    return x, (freqs_cos, freqs_sin)
+
+
+def test_attn1_sp2_state_reduction_matches_dense(tp1_group, force_default_gemm, mocker) -> None:
+    """Two token shards with a summed packed state must reproduce the dense
+    full-sequence linear attention output."""
+    dim, num_heads, head_dim, seq_len = 24, 4, 6, 12
+    attn = _make_attn(dim, num_heads, head_dim)
+    x, (freqs_cos, freqs_sin) = _attn_inputs(seq_len, dim, head_dim)
+    bounds = [(0, 8), (8, 12)]
+
+    with torch.no_grad():
+        ref = attn(x, rotary_emb=(freqs_cos, freqs_sin))
+
+    fake_group = _FakeSpAllReduce()
+    mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=2)
+    mocker.patch(f"{_MODULE}.get_sp_group", return_value=fake_group)
+
+    def run_ranks() -> list[torch.Tensor]:
+        outs = []
+        for start, stop in bounds:
+            local_rotary = (freqs_cos[:, start:stop], freqs_sin[:, start:stop])
+            with torch.no_grad():
+                outs.append(attn(x[:, start:stop], rotary_emb=local_rotary))
+        return outs
+
+    run_ranks()
+    fake_group.total = fake_group.partials[0] + fake_group.partials[1]
+    out = torch.cat(run_ranks(), dim=1)
+
+    torch.testing.assert_close(out, ref, rtol=1e-5, atol=1e-5)
+
+
+def test_attn1_issues_single_packed_all_reduce(tp1_group, force_default_gemm, mocker) -> None:
+    """One collective per forward, carrying scores and k_sum packed together."""
+    dim, num_heads, head_dim, seq_len = 24, 4, 6, 8
+    attn = _make_attn(dim, num_heads, head_dim)
+    x, rotary_emb = _attn_inputs(seq_len, dim, head_dim)
+
+    fake_group = _FakeSpAllReduce()
+    mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=2)
+    mocker.patch(f"{_MODULE}.get_sp_group", return_value=fake_group)
+
+    with torch.no_grad():
+        attn(x, rotary_emb=rotary_emb)
+
+    assert len(fake_group.partials) == 1
+    packed = fake_group.partials[0]
+    assert packed.shape == (2, num_heads, head_dim, head_dim + 1)
+    assert packed.dtype == torch.float32
+
+
+def test_attn1_sp1_path_is_untouched(tp1_group, force_default_gemm, mocker) -> None:
+    """At SP1 no collective may run: the dense operator sequence must not
+    change shape or route through get_sp_group at all."""
+    dim, num_heads, head_dim, seq_len = 24, 4, 6, 8
+    attn = _make_attn(dim, num_heads, head_dim)
+    x, rotary_emb = _attn_inputs(seq_len, dim, head_dim)
+
+    mock_group = mocker.patch(f"{_MODULE}.get_sp_group")
+
+    with torch.no_grad():
+        attn(x, rotary_emb=rotary_emb)
+
+    mock_group.assert_not_called()
+
+
+# ── uneven frame gather ──
+
+
+class _FakeSpAllGather:
+    """Stand-in for the SP all-gather: returns precomputed equal-shape per-rank
+    contributions, enforcing the equal-shard contract of the real collective."""
+
+    def __init__(self, parts: list[torch.Tensor]) -> None:
+        self.parts = parts
+
+    def all_gather(self, tensor: torch.Tensor, dim: int = 0, separate_tensors: bool = False):
+        assert separate_tensors, "frame gather must request per-rank tensors"
+        assert all(part.shape == tensor.shape for part in self.parts), "all_gather requires equal shards"
+        return [part.clone() for part in self.parts]
+
+
+@pytest.mark.parametrize("num_frames,world_size", [(21, 2), (6, 2), (5, 4)])
+def test_gather_frames_roundtrips_uneven_shards(num_frames, world_size, mocker) -> None:
+    """shard -> pad -> all_gather -> per-rank narrow -> cat must reproduce the
+    full sequence bitwise, with communication pads never reaching the output."""
+    batch, hw, channels = 2, 3, 4
+    torch.manual_seed(0)
+    full = torch.randn(batch, num_frames * hw, channels)
+
+    sizes = _sp_frame_split_sizes(num_frames, world_size)
+    frame_view = full.unflatten(1, (num_frames, hw))
+    locals_ = [t.flatten(1, 2) for t in frame_view.split(sizes, dim=1)]
+
+    # Per-rank padded contributions as the real collective would deliver them;
+    # NaN pads prove a leaked pad frame cannot go unnoticed.
+    parts = []
+    for local, size in zip(locals_, sizes):
+        part = local.unflatten(1, (size, hw))
+        pad = max(sizes) - size
+        if pad:
+            part = torch.cat([part, torch.full((batch, pad, hw, channels), torch.nan)], dim=1)
+        parts.append(part)
+
+    mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=world_size)
+    mocker.patch(f"{_MODULE}.get_sp_group", return_value=_FakeSpAllGather(parts))
+    for rank in range(world_size):
+        mocker.patch(f"{_MODULE}.get_sequence_parallel_rank", return_value=rank)
+
+        gathered = _sp_gather_frames(locals_[rank], sizes)
+
+        torch.testing.assert_close(gathered, full, rtol=0, atol=0)
+
+
+# ── GLUMB conv_temp halo exchange ──
+
+
+class _NoopReq:
+    def wait(self) -> None:
+        pass
+
+
+class _ReplayP2PDist:
+    """Two-pass stand-in for torch.distributed P2P: pass 1 records each rank's
+    boundary sends, pass 2 replays them into the matching neighbour recvs."""
+
+    irecv = "irecv"
+    isend = "isend"
+
+    def __init__(self) -> None:
+        self.rank = 0
+        self.sent: dict[tuple[int, int], torch.Tensor] = {}
+        self.replay = False
+        self.batches = 0
+        self.peers: set[tuple[int, int]] = set()
+
+    def P2POp(self, op, tensor, peer, group):
+        return (op, tensor, peer)
+
+    def batch_isend_irecv(self, ops):
+        self.batches += 1
+        for op, tensor, peer in ops:
+            self.peers.add((self.rank, peer))
+            if op == "isend":
+                self.sent[(self.rank, peer)] = tensor.clone()
+            elif self.replay:
+                tensor.copy_(self.sent[(peer, self.rank)])
+        return [_NoopReq() for _ in ops]
+
+
+class _FakeSpGroup:
+    def __init__(self, world_size: int) -> None:
+        self.ranks = list(range(world_size))
+        self.device_group = None
+
+
+def _make_glumb(mocker) -> GLUMBTempConv:
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1)
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=0)
+    torch.manual_seed(0)
+    glumb = GLUMBTempConv(4, 4, 2.0, norm_type=None, residual_connection=False)
+    glumb.eval()
+    for _, param in sorted(glumb.named_parameters()):
+        param.data.normal_()
+    return glumb
+
+
+@pytest.mark.parametrize("world_size", [2, 4])
+def test_glumb_conv_temp_sp_halo_matches_dense(world_size, mocker) -> None:
+    """Frame shards with neighbour boundary frames exchanged must reproduce the
+    dense temporal conv, including zero padding at the global first/last frame.
+    The exchange must talk to the two neighbouring ranks only."""
+    glumb = _make_glumb(mocker)
+    num_frames = 5
+    torch.manual_seed(1)
+    x = torch.randn(2, num_frames, 3, 3, 4)
+
+    with torch.no_grad():
+        ref = glumb(x)
+
+    sizes = _sp_frame_split_sizes(num_frames, world_size)
+    shards = [shard.contiguous() for shard in x.split(sizes, dim=1)]
+    fake_dist = _ReplayP2PDist()
+    mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=world_size)
+    mocker.patch(f"{_MODULE}.get_sp_group", return_value=_FakeSpGroup(world_size))
+    mocker.patch(f"{_MODULE}.dist", fake_dist)
+
+    def run_ranks() -> list[torch.Tensor]:
+        outs = []
+        for rank in range(world_size):
+            fake_dist.rank = rank
+            mocker.patch(f"{_MODULE}.get_sequence_parallel_rank", return_value=rank)
+            with torch.no_grad():
+                outs.append(glumb(shards[rank]))
+        return outs
+
+    run_ranks()
+    assert fake_dist.batches == world_size
+    assert fake_dist.peers == {(a, b) for a in range(world_size) for b in (a - 1, a + 1) if 0 <= b < world_size}
+    fake_dist.replay = True
+    out = torch.cat(run_ranks(), dim=1)
+
+    torch.testing.assert_close(out, ref, rtol=1e-5, atol=1e-5)
+
+
+def test_glumb_sp_halo_rejects_non_three_frame_kernel(mocker) -> None:
+    """The single-frame halo is only correct for a three-frame temporal kernel;
+    any other width must fail loudly instead of corrupting parity."""
+    glumb = _make_glumb(mocker)
+    glumb.conv_temp = nn.Conv2d(4, 4, kernel_size=(5, 1), stride=1, padding=(2, 0), bias=False)
+    mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=2)
+    mocker.patch(f"{_MODULE}.get_sequence_parallel_rank", return_value=0)
+    mocker.patch(f"{_MODULE}.get_sp_group", return_value=_FakeSpGroup(2))
+    mocker.patch(f"{_MODULE}.dist", _ReplayP2PDist())
+    with pytest.raises(AssertionError, match="temporal kernel"):
+        with torch.no_grad():
+            glumb(torch.randn(1, 3, 3, 3, 4))
+
+
+# ── I2V per-token timestep frame slicing ──
+
+
+def test_i2v_timestep_frame_slice_matches_dense_token_region() -> None:
+    """Slicing the 5D timestep along its frame dim before the per-token
+    time-embed MLP must match the dense modulation for the same token region.
+
+    The sinusoidal projection is bitwise under slicing; the MLP is only equal
+    up to CPU GEMM blocking, which depends on the row count. A wrongly sliced
+    frame produces O(1) differences, far outside this tolerance."""
+    dim = 24
+    torch.manual_seed(0)
+    layer = SanaAdaLayerNormSingle(dim).eval()
+
+    batch, num_frames, height, width = 2, 5, 2, 3
+    timestep = torch.rand(batch, 1, num_frames, height, width) * 1000.0
+
+    with torch.no_grad():
+        dense_mod, dense_emb = layer(timestep.flatten(), batch_size=batch, hidden_dtype=torch.float32)
+    dense_mod = dense_mod.view(batch, -1, dense_mod.size(-1))
+    dense_emb = dense_emb.view(batch, -1, dense_emb.size(-1))
+
+    sizes = _sp_frame_split_sizes(num_frames, 2)
+    start = 0
+    for rank, size in enumerate(sizes):
+        local = timestep.split(sizes, dim=2)[rank]
+        with torch.no_grad():
+            local_mod, local_emb = layer(local.flatten(), batch_size=batch, hidden_dtype=torch.float32)
+        local_mod = local_mod.view(batch, -1, local_mod.size(-1))
+        local_emb = local_emb.view(batch, -1, local_emb.size(-1))
+
+        lo, hi = start * height * width, (start + size) * height * width
+        torch.testing.assert_close(local_mod, dense_mod[:, lo:hi], rtol=1e-5, atol=1e-6)
+        torch.testing.assert_close(local_emb, dense_emb[:, lo:hi], rtol=1e-5, atol=1e-6)
+        start += size
+
+
+# ── full-model lockstep parity ──
+
+
+_TINY_SP_CONFIG = {
+    "in_channels": 4,
+    "out_channels": 4,
+    "num_attention_heads": 2,
+    "attention_head_dim": 12,
+    "num_layers": 2,
+    "num_cross_attention_heads": 2,
+    "cross_attention_head_dim": 12,
+    "cross_attention_dim": 24,
+    "caption_channels": 8,
+    "mlp_ratio": 2.0,
+    "patch_size": (1, 2, 2),
+    "sample_size": 4,
+    "rope_max_seq_len": 64,
+}
+
+
+class _LockstepSpGroup:
+    """Barrier-synchronized collectives so rank threads run true lockstep SP in
+    one process; records per-call inputs so tests can prove ranks computed on
+    different shards."""
+
+    def __init__(self, world_size: int, tls: threading.local) -> None:
+        self.world_size = world_size
+        self.ranks = list(range(world_size))
+        self.device_group = None
+        self.tls = tls
+        self.barrier = threading.Barrier(world_size)
+        self.slots: list[torch.Tensor | None] = [None] * world_size
+        self.reduce_inputs: list[list[torch.Tensor]] = []
+        self.gather_calls = 0
+        self.lock = threading.Lock()
+
+    def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
+        self.slots[self.tls.rank] = tensor
+        self.barrier.wait(timeout=60)
+        total = self.slots[0].clone()
+        for part in self.slots[1:]:
+            total = total + part
+        if self.tls.rank == 0:
+            self.reduce_inputs.append([part.clone() for part in self.slots])
+        self.barrier.wait(timeout=60)
+        return total
+
+    def all_gather(self, tensor: torch.Tensor, dim: int = 0, separate_tensors: bool = False):
+        assert separate_tensors
+        self.slots[self.tls.rank] = tensor
+        self.barrier.wait(timeout=60)
+        assert all(part.shape == tensor.shape for part in self.slots), "all_gather requires equal shards"
+        parts = [part.clone() for part in self.slots]
+        if self.tls.rank == 0:
+            with self.lock:
+                self.gather_calls += 1
+        self.barrier.wait(timeout=60)
+        return parts
+
+
+class _LockstepP2PDist:
+    """Barrier-synchronized torch.distributed P2P stand-in: sends land in a
+    mailbox, all rank threads rendezvous, then recvs read their neighbour's
+    send of the same batch."""
+
+    irecv = "irecv"
+    isend = "isend"
+
+    def __init__(self, world_size: int, tls: threading.local) -> None:
+        self.tls = tls
+        self.barrier = threading.Barrier(world_size)
+        self.mailbox: dict[tuple[int, int], torch.Tensor] = {}
+        self.peers: set[tuple[int, int]] = set()
+        self.batches = 0
+        self.lock = threading.Lock()
+
+    def P2POp(self, op, tensor, peer, group):
+        return (op, tensor, peer)
+
+    def batch_isend_irecv(self, ops):
+        rank = self.tls.rank
+        with self.lock:
+            self.batches += 1
+            for op, tensor, peer in ops:
+                self.peers.add((rank, peer))
+                if op == "isend":
+                    self.mailbox[(rank, peer)] = tensor.clone()
+        self.barrier.wait(timeout=60)
+        for op, tensor, peer in ops:
+            if op == "irecv":
+                tensor.copy_(self.mailbox[(peer, rank)])
+        self.barrier.wait(timeout=60)
+        return [_NoopReq() for _ in ops]
+
+
+@pytest.mark.parametrize("task", ["t2v", "i2v"])
+def test_tiny_transformer_sp2_lockstep_matches_dense(tp1_group, force_default_gemm, mocker, task) -> None:
+    """Two lockstep rank threads over uneven frame shards must reproduce the
+    dense tiny-transformer output, and must actually have sharded the work."""
+    from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+
+    torch.manual_seed(3)
+    model = SanaVideoTransformer3DModel(**_TINY_SP_CONFIG).eval()
+    for _, param in sorted(model.named_parameters()):
+        param.data.normal_()
+    batch, frames, height, width = 2, 5, 4, 4
+    torch.manual_seed(11)
+    latent = torch.randn(batch, 4, frames, height, width)
+    encoder_hidden_states = torch.randn(batch, 6, 8)
+    encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 1, 0], [1, 1, 1, 1, 1, 1]])
+    if task == "t2v":
+        timestep = torch.tensor([500.0, 700.0])
+    else:
+        timestep = torch.rand(batch, 1, frames, height // 2, width // 2) * 1000.0
+
+    with torch.no_grad():
+        ref = model(
+            latent,
+            encoder_hidden_states,
+            timestep,
+            encoder_attention_mask=encoder_attention_mask,
+            return_dict=False,
+        )[0]
+
+    world_size = 2
+    tls = threading.local()
+    group = _LockstepSpGroup(world_size, tls)
+    fake_dist = _LockstepP2PDist(world_size, tls)
+    mocker.patch(f"{_MODULE}.get_sequence_parallel_world_size", return_value=world_size)
+    mocker.patch(f"{_MODULE}.get_sequence_parallel_rank", side_effect=lambda: tls.rank)
+    mocker.patch(f"{_MODULE}.get_sp_group", return_value=group)
+    mocker.patch(f"{_MODULE}.dist", fake_dist)
+
+    def run_rank(rank: int) -> torch.Tensor:
+        tls.rank = rank
+        with torch.no_grad():
+            return model(
+                latent,
+                encoder_hidden_states,
+                timestep,
+                encoder_attention_mask=encoder_attention_mask,
+                return_dict=False,
+            )[0]
+
+    with ThreadPoolExecutor(max_workers=world_size) as pool:
+        outs = list(pool.map(run_rank, range(world_size)))
+
+    num_layers = _TINY_SP_CONFIG["num_layers"]
+    # One attn1 state reduce per block, its inputs distinct per rank (the ranks
+    # really processed different frames), one neighbour halo exchange per block
+    # per rank, and only the final frame gather is a full collective.
+    assert len(group.reduce_inputs) == num_layers
+    for parts in group.reduce_inputs:
+        assert not torch.equal(parts[0], parts[1])
+    assert group.gather_calls == 1
+    assert fake_dist.batches == world_size * num_layers
+    assert fake_dist.peers == {(0, 1), (1, 0)}
+
+    torch.testing.assert_close(outs[1], outs[0], rtol=0, atol=0)
+    # Reduction-reorder noise amplified by the unit-normal test weights; any
+    # sharding misalignment is orders of magnitude above this.
+    torch.testing.assert_close(outs[0], ref, rtol=1e-3, atol=1e-3)
+
+
+def test_transformer_declares_empty_sp_plan() -> None:
+    """Empty plan: the registry takes its normal SP enablement path with zero
+    hooks and no misleading 'no _sp_plan found' warning; sharding is manual."""
+    from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+
+    assert SanaVideoTransformer3DModel._sp_plan == {}
+
+
+def test_dummy_run_num_frames_covers_max_sp_degree() -> None:
+    """Engine warmup must produce enough latent frames for every allowed SP
+    degree on both VAE variants (temporal /4 and /8); one pixel frame would
+    fail the per-rank frame check at startup."""
+    from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline, SanaVideoPipeline
+
+    for temporal_scale in (4, 8):
+        latent_frames = (SanaVideoPipeline.dummy_run_num_frames - 1) // temporal_scale + 1
+        assert latent_frames >= 4
+    assert SanaImageToVideoPipeline.dummy_run_num_frames == SanaVideoPipeline.dummy_run_num_frames

tests/model_tests/diffusion/model_settings.py

@@ -57,6 +57,9 @@ DIFFUSION_TEST_SETTINGS = {
             [DiffusionAccs.TENSOR_PARALLEL],
             [DiffusionAccs.CFG_PARALLEL],
             [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL],
+            [DiffusionAccs.SEQUENCE_PARALLEL],
+            [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CFG_PARALLEL],
+            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.SEQUENCE_PARALLEL],
             [DiffusionAccs.CACHE_DIT],
             [DiffusionAccs.CPU_OFFLOAD],
             [DiffusionAccs.LAYERWISE_OFFLOAD],
@@ -72,6 +75,8 @@ DIFFUSION_TEST_SETTINGS = {
         extra_test_groups=[
             [DiffusionAccs.CFG_PARALLEL],
             [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL],
+            [DiffusionAccs.SEQUENCE_PARALLEL],
+            [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CFG_PARALLEL],
         ],
         check_i2v_t2v_divergence=False,
         check_multi_output=False,

vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py

@@ -333,6 +333,9 @@ class SanaVideoPipeline(
     _encoder_modules = ["text_encoder"]
     _vae_modules = ["vae"]
     supports_step_execution = False
+    # Warmup must yield >= 4 latent frames on both VAE variants (temporal /4
+    # and /8) so every allowed SP degree has one frame per rank.
+    dummy_run_num_frames = 25
     default_num_inference_steps = 50
 
     def __init__(

vllm_omni/diffusion/models/sana_video/transformer_sana_video.py

@@ -18,6 +18,7 @@ from collections.abc import Iterable
 from dataclasses import dataclass, fields
 
 import torch
+import torch.distributed as dist
 from cache_dit import ForwardPattern
 from torch import nn
 from vllm.distributed import (
@@ -32,6 +33,11 @@ from vllm.model_executor.utils import set_weight_attrs
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention as OmniAttention
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_sequence_parallel_rank,
+    get_sequence_parallel_world_size,
+    get_sp_group,
+)
 
 
 def validate_sana_video_parallel_config(parallel_config) -> None:
@@ -46,13 +52,18 @@ def validate_sana_video_parallel_config(parallel_config) -> None:
         raise NotImplementedError(
             f"SANA-Video supports cfg_parallel_size 1 or 2, got {cfg_size}. Set --cfg-parallel-size to 1 or 2."
         )
-    sp_size = parallel_config.sequence_parallel_size
-    if sp_size is not None and sp_size > 1:
+    sp_size = parallel_config.sequence_parallel_size or 1
+    if sp_size not in (1, 2, 4):
         raise NotImplementedError(
-            "Sequence parallel is not supported for SANA-Video: its linear attention and "
-            "GLUMB temporal conv have no SANA-specific SP implementation. Use tensor and/or "
-            "CFG parallel instead."
+            f"SANA-Video supports sequence_parallel_size 1, 2 or 4, got {sp_size}. Set --usp to 1, 2 or 4."
         )
+    if parallel_config.ring_degree > 1:
+        raise NotImplementedError(
+            "SANA-Video does not support ring sequence parallel: its linear attention has no "
+            "log-sum-exp state to merge across ring stages. Use --usp for sequence parallel."
+        )
+    if parallel_config.allgather_degree > 1:
+        raise NotImplementedError("SANA-Video does not support AllGather-KV sequence parallel. Use --usp instead.")
     if parallel_config.pipeline_parallel_size > 1:
         raise NotImplementedError("SANA-Video does not support pipeline parallel. Set --pipeline-parallel-size to 1.")
     if parallel_config.use_hsdp:
@@ -63,6 +74,32 @@ def validate_sana_video_parallel_config(parallel_config) -> None:
         )
 
 
+def _sp_frame_split_sizes(num_frames: int, world_size: int) -> list[int]:
+    """Balanced per-rank frame counts: the first ``num_frames % world_size``
+    ranks take one extra frame, so exactly ``world_size`` non-empty chunks."""
+    if num_frames < world_size:
+        raise ValueError(
+            f"SANA-Video sequence parallel needs at least one latent frame per rank, "
+            f"got {num_frames} frames for sequence parallel size {world_size}. "
+            "Lower the sequence parallel degree or increase num_frames."
+        )
+    base, rem = divmod(num_frames, world_size)
+    return [base + 1 if rank < rem else base for rank in range(world_size)]
+
+
+def _sp_gather_frames(hidden_states: torch.Tensor, sizes: list[int]) -> torch.Tensor:
+    """Gather per-rank frame shards of shape (B, sizes[rank]*hw, C) into the
+    full (B, sum(sizes)*hw, C) sequence. Shards are padded to the widest rank
+    only for the collective; pads are dropped before concatenation."""
+    rank = get_sequence_parallel_rank()
+    frames = hidden_states.unflatten(1, (sizes[rank], -1))
+    pad = max(sizes) - sizes[rank]
+    if pad:
+        frames = torch.cat([frames, frames.new_zeros((frames.shape[0], pad, *frames.shape[2:]))], dim=1)
+    parts = get_sp_group().all_gather(frames, dim=0, separate_tensors=True)
+    return torch.cat([part.narrow(1, 0, size) for part, size in zip(parts, sizes)], dim=1).flatten(1, 2)
+
+
 @dataclass
 class SanaVideoTransformerConfig:
     in_channels: int = 16
@@ -521,7 +558,37 @@ class GLUMBTempConv(nn.Module):
         hidden_states_temporal = hidden_states.view(batch_size, num_frames, num_channels, height * width).permute(
             0, 2, 1, 3
         )
-        hidden_states = hidden_states_temporal + self.conv_temp(hidden_states_temporal)
+        if get_sequence_parallel_world_size() > 1:
+            # The temporal conv reads one frame across each shard boundary, so
+            # each rank exchanges single frames with its neighbours only; the
+            # global first/last rank keeps the dense zero padding on its outer
+            # side.
+            assert self.conv_temp.kernel_size == (3, 1), (
+                "SP halo exchange assumes a temporal kernel of 3 (one frame per side), "
+                f"got {self.conv_temp.kernel_size}"
+            )
+            x = hidden_states_temporal
+            rank = get_sequence_parallel_rank()
+            world_size = get_sequence_parallel_world_size()
+            sp_group = get_sp_group()
+            left = x.new_zeros((x.shape[0], x.shape[1], 1, x.shape[3]))
+            right = torch.zeros_like(left)
+            p2p_ops = []
+            if rank > 0:
+                prev_rank = sp_group.ranks[rank - 1]
+                p2p_ops.append(dist.P2POp(dist.irecv, left, prev_rank, sp_group.device_group))
+                p2p_ops.append(dist.P2POp(dist.isend, x[:, :, :1].contiguous(), prev_rank, sp_group.device_group))
+            if rank < world_size - 1:
+                next_rank = sp_group.ranks[rank + 1]
+                p2p_ops.append(dist.P2POp(dist.isend, x[:, :, -1:].contiguous(), next_rank, sp_group.device_group))
+                p2p_ops.append(dist.P2POp(dist.irecv, right, next_rank, sp_group.device_group))
+            for req in dist.batch_isend_irecv(p2p_ops):
+                req.wait()
+            padded = torch.cat([left, x, right], dim=2)
+            temporal_out = nn.functional.conv2d(padded, self.conv_temp.weight, padding=0)
+        else:
+            temporal_out = self.conv_temp(hidden_states_temporal)
+        hidden_states = hidden_states_temporal + temporal_out
         hidden_states = hidden_states.permute(0, 2, 3, 1).view(batch_size, num_frames, height, width, num_channels)
 
         if self.norm_type == "rms_norm":
@@ -621,9 +688,15 @@ class SanaLinearAttention(nn.Module):
         # Keep the unrotated denominator contraction in the input dtype. This
         # intentionally matches Diffusers 0.38.0's SanaLinearAttnProcessor3_0;
         # only the rotated numerator path is accumulated in FP32.
-        z = 1 / (key.sum(dim=-1, keepdim=True).transpose(-2, -1) @ query + 1e-15)
-
+        k_sum = key.sum(dim=-1, keepdim=True)
         scores = torch.matmul(value, key_rotate.transpose(-1, -2))
+        if get_sequence_parallel_world_size() > 1:
+            # Token sums are the only sequence coupling; one packed all-reduce
+            # turns the per-rank partial state into the global one.
+            packed = torch.cat([scores, k_sum.float()], dim=-1)
+            packed = get_sp_group().all_reduce(packed)
+            scores, k_sum = packed[..., :-1], packed[..., -1:].to(key.dtype)
+        z = 1 / (k_sum.transpose(-2, -1) @ query + 1e-15)
         hidden_states = torch.matmul(scores, query_rotate)
 
         hidden_states = hidden_states * z
@@ -971,6 +1044,10 @@ class SanaVideoTransformer3DModel(nn.Module):
     )
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
 
+    # SP is implemented manually in forward (frame-aligned uneven sharding);
+    # the empty plan enables the registry's SP path with zero generic hooks.
+    _sp_plan = {}
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return AutoWeightsLoader(self).load_weights(weights)
 
@@ -1151,6 +1228,26 @@ class SanaVideoTransformer3DModel(nn.Module):
         hidden_states = self.patch_embedding(hidden_states)
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
+        sp_world_size = get_sequence_parallel_world_size()
+        local_frames = post_patch_num_frames
+        if sp_world_size > 1:
+            # Frame-aligned sharding: every rank keeps full spatial extent, so
+            # the per-frame convs, the GLUMB unflatten and the per-token
+            # modulation stay local. The I2V per-token timestep is sliced
+            # before the time embed so its MLP also runs on local tokens only.
+            sp_sizes = _sp_frame_split_sizes(post_patch_num_frames, sp_world_size)
+            sp_rank = get_sequence_parallel_rank()
+            local_frames = sp_sizes[sp_rank]
+            hidden_states = (
+                hidden_states.unflatten(1, (post_patch_num_frames, -1)).split(sp_sizes, dim=1)[sp_rank].flatten(1, 2)
+            )
+            rotary_emb = tuple(
+                freqs.unflatten(1, (post_patch_num_frames, -1)).split(sp_sizes, dim=1)[sp_rank].flatten(1, 2)
+                for freqs in rotary_emb
+            )
+            if timestep.ndim == 5:
+                timestep = timestep.split(sp_sizes, dim=2)[sp_rank]
+
         if guidance is not None:
             timestep, embedded_timestep = self.time_embed(
                 timestep.flatten(), guidance=guidance, hidden_dtype=hidden_states.dtype
@@ -1176,7 +1273,7 @@ class SanaVideoTransformer3DModel(nn.Module):
                 encoder_hidden_states,
                 encoder_attention_mask,
                 timestep,
-                post_patch_num_frames,
+                local_frames,
                 post_patch_height,
                 post_patch_width,
                 rotary_emb,
@@ -1189,6 +1286,10 @@ class SanaVideoTransformer3DModel(nn.Module):
 
         hidden_states = self.proj_out(hidden_states)
 
+        if sp_world_size > 1:
+            # Gather after proj_out where the channel width is smallest.
+            hidden_states = _sp_gather_frames(hidden_states, sp_sizes)
+
         # 5. Unpatchify
         hidden_states = hidden_states.reshape(
             batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1