[Model] Add TP and CFG parallelism to SANA-Video 2B (#5861)

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

docs/user_guide/diffusion_features.md

@@ -162,7 +162,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**   |     ❓     |     ❓      |           ❓           |       ❓        |         ❓         |         ❓         |   ❓    |             ❓             |          ❓           |       ❓        |        ❌         |
+|  **SANA-Video-2B T2V I2V**   |     ❌     |     ❌      |          ❌           |       ✅       |  ✅ (TP=2 only)   |         ❌         |   ❌   |            ❌             |          ❌          |       ❌       |        ❌        |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ✅*        |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |

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

@@ -119,6 +119,30 @@ MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
 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:
+
+```bash
+vllm serve Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
+  --omni \
+  --model-class-name SanaVideoPipeline \
+  --tensor-parallel-size 2 \
+  --dtype bfloat16 \
+  --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.
+
+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.
+
 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
 works; `num_frames` is adapted to Diffusers' `frames` argument. The script
@@ -194,8 +218,9 @@ validated compatibility/reference backend.
   The denoising loop intentionally retains the checkpoint-compatible
   Diffusers `DPMSolverMultistepScheduler`.
 - Known limitations:
-    - Sequence/tensor/CFG parallelism, Cache-DiT, TeaCache, and step execution
-    are not validated for the native pipeline.
+    - Tensor parallelism (up to 2 GPUs) and CFG parallelism are supported for
+    the native pipeline. Sequence parallelism, Cache-DiT, TeaCache, and step
+    execution are not validated.
     - The Diffusers backend is a compatibility path and does not provide native
     vLLM-Omni parallelism or continuous batching.
     - Native describes pipeline and Transformer ownership, not a zero-Diffusers

tests/diffusion/diffusion_backend/test_diffusers_backend.py

@@ -143,7 +143,16 @@ class TestPipelineArgumentsHandling:
 
     @pytest.mark.parametrize(
         "feature_id",
-        ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "unsupported_quantization"],
+        [
+            "cfg_parallel",
+            "tensor_parallel",
+            "ulysses",
+            "ring",
+            "teacache",
+            "cache_dit",
+            "enforce_eager",
+            "unsupported_quantization",
+        ],
     )
     def test_adapter_guard_unsupported_feature(self, feature_id):
         if feature_id == "cfg_parallel":
@@ -151,6 +160,11 @@ class TestPipelineArgumentsHandling:
                 parallel_config=DiffusionParallelConfig(cfg_parallel_size=2, sequence_parallel_size=1),
                 cache_backend="none",
             )
+        elif feature_id == "tensor_parallel":
+            od_config = _make_od_config(
+                parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, sequence_parallel_size=1),
+                cache_backend="none",
+            )
         elif feature_id == "ulysses":
             od_config = _make_od_config(
                 parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, ulysses_degree=2),

tests/diffusion/models/sana_video/test_pipeline_sana_video_cfg.py

@@ -0,0 +1,582 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""CFG dispatch tests for the native SANA-Video pipelines.
+
+These exercise the CFG world-size-1 batch path, the CFG world-size-2 dispatch,
+and the fp32 combine semantics on CPU / mocked distributed. The real all-gather
+path is covered by the multi-process suite, not here.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import SanaVideoPipeline
+from vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v import SanaImageToVideoPipeline
+from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+    SanaVideoTransformerOutput,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_predict_noise_calls_transformer_with_return_dict_false():
+    """predict_noise must call the transformer with return_dict=False.
+
+    The transformer returns a SanaVideoTransformerOutput dataclass when
+    return_dict is True (its default), which is not subscriptable. The CFG
+    mixin's default predict_noise does ``result[0]`` and would therefore
+    raise TypeError on SANA; the pipeline override must force return_dict=False
+    so the tuple form is returned and indexed.
+    """
+    captured = {}
+    expected = torch.full((1, 2, 3), 7.0)
+
+    class _SpyTransformer:
+        def __call__(self, *args, return_dict=True, **kwargs):
+            captured["return_dict"] = return_dict
+            # Mirror the real transformer: tuple on False, dataclass on True.
+            return (expected,) if not return_dict else SanaVideoTransformerOutput(sample=expected)
+
+    pipe = SanaVideoPipeline.__new__(SanaVideoPipeline)  # bypass heavy __init__
+    pipe.transformer = _SpyTransformer()
+
+    out = pipe.predict_noise(
+        torch.zeros(1, 2, 3),
+        encoder_hidden_states=torch.zeros(1, 1, 4),
+        encoder_attention_mask=None,
+        timestep=torch.zeros(1),
+    )
+
+    assert captured["return_dict"] is False
+    assert torch.equal(out, expected)
+
+
+class _IdentityScheduler:
+    """Scheduler stub whose step is a no-op, so diffuse tests isolate the CFG
+    dispatch from real scheduler math."""
+
+    order = 1
+
+    def step(self, noise_pred, t, latents, **kwargs):
+        return (latents,)
+
+
+def _set_cfg_world_size(monkeypatch, n):
+    """Pipelines always run inside an initialized CFG group; unit tests stub its
+    world size instead of spinning up torch.distributed. Patch both pipeline
+    modules since each resolves the symbol in its own namespace."""
+    import vllm_omni.diffusion.distributed.cfg_parallel as cfg
+    import vllm_omni.diffusion.models.sana_video.pipeline_sana_video as t2v
+    import vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v as i2v
+
+    monkeypatch.setattr(t2v, "get_classifier_free_guidance_world_size", lambda: n, raising=False)
+    monkeypatch.setattr(i2v, "get_classifier_free_guidance_world_size", lambda: n, raising=False)
+    monkeypatch.setattr(cfg, "get_classifier_free_guidance_world_size", lambda: n, raising=False)
+    monkeypatch.setattr(t2v, "model_parallel_is_initialized", lambda: True, raising=False)
+    monkeypatch.setattr(i2v, "model_parallel_is_initialized", lambda: True, raising=False)
+
+
+def _bare_t2v_pipeline(transformer, guidance_scale):
+    pipe = SanaVideoPipeline.__new__(SanaVideoPipeline)  # bypass heavy __init__
+    pipe.transformer = transformer
+    pipe.scheduler = _IdentityScheduler()
+    pipe._interrupt = False
+    pipe._guidance_scale = guidance_scale
+    return pipe
+
+
+def test_cfg_world_size_1_does_single_batch2_forward_per_step(monkeypatch):
+    """CFG world size 1 with guidance must do exactly one batch-2 transformer
+    forward per step (the batch path), not the mixin's sequential two-forward
+    path."""
+    _set_cfg_world_size(monkeypatch, 1)
+    batch_dims = []
+
+    class _SpyTransformer:
+        dtype = torch.float32
+
+        def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
+            batch_dims.append(hidden_states.shape[0])
+            # out_channels == in_channels here (no learned-sigma split).
+            return (torch.zeros_like(hidden_states),)
+
+    b = 1
+    pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0)
+
+    out = pipe.diffuse(
+        latents=torch.randn(b, 4, 2, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=None,
+    )
+
+    assert len(batch_dims) == 1, f"expected 1 forward this step, got {len(batch_dims)}"
+    assert batch_dims[0] == 2 * b, f"expected batch-2 (neg+pos) forward, got batch {batch_dims[0]}"
+    assert out.shape == (b, 4, 2, 4, 4)
+
+
+class _DeterministicTransformer:
+    """Learned-sigma transformer whose output depends on the row's encoder
+    embeds and timestep, so uncond != text and the CFG combine is non-trivial."""
+
+    dtype = torch.float32
+
+    def __init__(self, latent_channels):
+        self.out_channels = 2 * latent_channels  # learned sigma
+
+    def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
+        b, c = hidden_states.shape[0], hidden_states.shape[1]
+        sig = encoder_hidden_states.float().mean(dim=(1, 2)).view(b, 1, 1, 1, 1)
+        base = hidden_states.float() + sig + timestep.float().view(b, 1, 1, 1, 1)
+        return (base.repeat(1, self.out_channels // c, 1, 1, 1),)
+
+
+class _LinearScheduler:
+    """Deterministic stepping scheduler so multi-step divergence compounds."""
+
+    order = 1
+
+    def step(self, noise_pred, t, latents, return_dict=False, **kwargs):
+        return (latents - 0.1 * noise_pred,)
+
+
+def _pr1_reference_loop(
+    transformer,
+    scheduler,
+    latents,
+    timesteps,
+    prompt_embeds,
+    prompt_mask,
+    negative_embeds,
+    negative_mask,
+    guidance_scale,
+    latent_channels,
+    dtype,
+):
+    """The original CFG1 denoise math, verbatim, as the parity oracle."""
+    prompt_embeds = torch.cat([negative_embeds, prompt_embeds], dim=0)
+    prompt_mask = torch.cat([negative_mask, prompt_mask], dim=0)
+    for t in timesteps:
+        latent_model_input = torch.cat([latents] * 2)
+        timestep = t.expand(latent_model_input.shape[0])
+        noise_pred = transformer(
+            latent_model_input.to(dtype=dtype),
+            encoder_hidden_states=prompt_embeds.to(dtype=dtype),
+            encoder_attention_mask=prompt_mask,
+            timestep=timestep,
+            return_dict=False,
+        )[0]
+        noise_pred = noise_pred.float()
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        if transformer.out_channels // 2 == latent_channels:
+            noise_pred = noise_pred.chunk(2, dim=1)[0]
+        latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+    return latents
+
+
+def test_cfg1_refactored_diffuse_is_bit_identical_to_pr1_loop(monkeypatch):
+    """The refactored diffuse must reproduce the original CFG1 latents bit-for-bit
+    across multiple steps, preserving the TP1/CFG1 golden output."""
+    _set_cfg_world_size(monkeypatch, 1)
+    torch.manual_seed(0)
+    b, c = 2, 4
+    latents0 = torch.randn(b, c, 2, 4, 4)
+    timesteps = torch.tensor([900.0, 600.0, 300.0])
+    prompt_embeds = torch.randn(b, 3, 8)
+    negative_embeds = torch.randn(b, 3, 8)
+    prompt_mask = torch.ones(b, 3)
+    negative_mask = torch.ones(b, 3)
+    guidance_scale = 6.0
+
+    transformer = _DeterministicTransformer(latent_channels=c)
+    scheduler = _LinearScheduler()
+
+    ref = _pr1_reference_loop(
+        transformer,
+        scheduler,
+        latents0.clone(),
+        timesteps,
+        prompt_embeds,
+        prompt_mask,
+        negative_embeds,
+        negative_mask,
+        guidance_scale,
+        latent_channels=c,
+        dtype=torch.float32,
+    )
+
+    pipe = _bare_t2v_pipeline(transformer, guidance_scale=guidance_scale)
+    pipe.scheduler = scheduler
+    got = pipe.diffuse(
+        latents=latents0.clone(),
+        timesteps=timesteps,
+        prompt_embeds=prompt_embeds,
+        prompt_attention_mask=prompt_mask,
+        negative_prompt_embeds=negative_embeds,
+        negative_prompt_attention_mask=negative_mask,
+        guidance_scale=guidance_scale,
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=c,
+    )
+
+    assert torch.equal(got, ref)
+
+
+def test_combine_cfg_noise_upcasts_to_fp32_before_combine():
+    """combine_cfg_noise must upcast positive/negative to fp32 before the CFG
+    formula, so the CFG2 (bf16 all-gather) combine is bit-identical to CFG1's
+    fp32 combine."""
+    pipe = SanaVideoPipeline.__new__(SanaVideoPipeline)
+    torch.manual_seed(0)
+    pos = torch.randn(1, 4, 2, 4, 4).bfloat16()
+    neg = torch.randn(1, 4, 2, 4, 4).bfloat16()
+    scale = 6.0
+
+    out = pipe.combine_cfg_noise(pos, neg, scale, cfg_normalize=False)
+
+    expected = neg.float() + scale * (pos.float() - neg.float())
+    assert out.dtype == torch.float32
+    assert torch.equal(out, expected)
+
+
+def test_combine_cfg_noise_accepts_tuple_predictions_from_gather():
+    """The CFG-parallel gather path hands combine_cfg_noise single-element
+    tuples, not bare tensors; the override must upcast and combine them."""
+    pipe = SanaVideoPipeline.__new__(SanaVideoPipeline)
+    torch.manual_seed(0)
+    pos = torch.randn(1, 4, 2, 4, 4).bfloat16()
+    neg = torch.randn(1, 4, 2, 4, 4).bfloat16()
+    scale = 6.0
+
+    out = pipe.combine_cfg_noise((pos,), (neg,), scale, cfg_normalize=False)
+
+    expected = neg.float() + scale * (pos.float() - neg.float())
+    assert out.dtype == torch.float32
+    assert torch.equal(out, expected)
+
+
+def test_cfg2_dispatch_builds_independent_branch_kwargs(monkeypatch):
+    """In CFG world size 2, diffuse must build separate batch-B positive/negative
+    kwargs (no pre-cat), and call the mixin with cfg_normalize=False and the
+    learned-sigma output_slice."""
+    _set_cfg_world_size(monkeypatch, 2)
+
+    b, c = 1, 4
+    captured = {}
+
+    def _spy_cfg(**kwargs):
+        captured.update(kwargs)
+        return torch.zeros(b, c, 2, 4, 4)
+
+    prompt_embeds = torch.randn(b, 3, 8)
+    negative_embeds = torch.randn(b, 3, 8)
+    dtype = torch.float32
+
+    pipe = _bare_t2v_pipeline(transformer=None, guidance_scale=6.0)
+    pipe.predict_noise_maybe_with_cfg = _spy_cfg
+
+    pipe.diffuse(
+        latents=torch.randn(b, c, 2, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=prompt_embeds,
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=negative_embeds,
+        negative_prompt_attention_mask=torch.zeros(b, 3),
+        guidance_scale=6.0,
+        extra_step_kwargs={},
+        dtype=dtype,
+        output_slice=c,
+    )
+
+    assert captured, "predict_noise_maybe_with_cfg was not called in CFG world size 2"
+    assert captured["do_true_cfg"] is True
+    assert captured["true_cfg_scale"] == 6.0
+    assert captured["cfg_normalize"] is False
+    assert captured["output_slice"] == c
+
+    pos = captured["positive_kwargs"]
+    neg = captured["negative_kwargs"]
+    assert pos["hidden_states"].shape[0] == b, "positive branch must be batch-B, not pre-concatenated"
+    assert neg["hidden_states"].shape[0] == b, "negative branch must be batch-B, not pre-concatenated"
+    assert pos["timestep"].shape[0] == b
+    assert torch.equal(pos["encoder_hidden_states"], prompt_embeds.to(dtype))
+    assert torch.equal(neg["encoder_hidden_states"], negative_embeds.to(dtype))
+
+
+# ── I2V CFG dispatch ──
+
+
+class _I2VSpyScheduler:
+    """Records the frame count it steps on and mutates non-first frames so
+    first-frame preservation is observable."""
+
+    order = 1
+
+    def __init__(self):
+        self.stepped_frame_counts = []
+
+    def step(self, noise_pred, t, latents, return_dict=False, **kwargs):
+        self.stepped_frame_counts.append(latents.shape[2])
+        return (latents + 1.0,)
+
+
+def _bare_i2v_pipeline(transformer, scheduler, guidance_scale):
+    pipe = SanaImageToVideoPipeline.__new__(SanaImageToVideoPipeline)
+    pipe.transformer = transformer
+    pipe.scheduler = scheduler
+    pipe._interrupt = False
+    pipe._guidance_scale = guidance_scale
+    return pipe
+
+
+def _conditioning_mask(b, frames):
+    mask = torch.zeros(b, 1, frames, 1, 1)
+    mask[:, :, 0] = 1
+    return mask
+
+
+def test_i2v_cfg1_preserves_first_frame_and_steps_only_tail(monkeypatch):
+    """I2V CFG1 must keep the conditioning first frame bit-exact and only run the
+    scheduler on frames [1:]."""
+    _set_cfg_world_size(monkeypatch, 1)
+    b, c, frames = 1, 4, 3
+
+    class _SpyTransformer:
+        dtype = torch.float32
+        out_channels = 2 * c
+
+        def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
+            return (torch.zeros(hidden_states.shape[0], 2 * c, frames, 4, 4),)
+
+    scheduler = _I2VSpyScheduler()
+    pipe = _bare_i2v_pipeline(_SpyTransformer(), scheduler, guidance_scale=6.0)
+
+    latents0 = torch.randn(b, c, frames, 4, 4)
+    out = pipe.diffuse(
+        latents=latents0.clone(),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        conditioning_mask=_conditioning_mask(b, frames),
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=c,
+    )
+
+    assert torch.equal(out[:, :, :1], latents0[:, :, :1]), "first (conditioning) frame must be unchanged"
+    assert torch.equal(out[:, :, 1:], latents0[:, :, 1:] + 1.0), "tail frames must be the scheduler output"
+    assert scheduler.stepped_frame_counts == [frames - 1], "scheduler must step only on frames [1:]"
+
+
+def test_i2v_cfg2_dispatch_builds_masked_branch_kwargs(monkeypatch):
+    """I2V CFG world size 2 must build batch-B per-branch kwargs with the
+    conditioning-frame timestep zeroed (mask not concatenated) and preserve the
+    first frame."""
+    _set_cfg_world_size(monkeypatch, 2)
+    b, c, frames = 1, 4, 3
+    captured = {}
+
+    def _spy_cfg(**kwargs):
+        captured.update(kwargs)
+        return torch.zeros(b, c, frames, 4, 4)
+
+    scheduler = _I2VSpyScheduler()
+    pipe = _bare_i2v_pipeline(transformer=None, scheduler=scheduler, guidance_scale=6.0)
+    pipe.predict_noise_maybe_with_cfg = _spy_cfg
+
+    prompt_embeds = torch.randn(b, 3, 8)
+    negative_embeds = torch.randn(b, 3, 8)
+    latents0 = torch.randn(b, c, frames, 4, 4)
+    mask = _conditioning_mask(b, frames)
+
+    out = pipe.diffuse(
+        latents=latents0.clone(),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=prompt_embeds,
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=negative_embeds,
+        negative_prompt_attention_mask=torch.zeros(b, 3),
+        guidance_scale=6.0,
+        conditioning_mask=mask,
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=c,
+    )
+
+    assert captured["do_true_cfg"] is True
+    assert captured["cfg_normalize"] is False
+    assert captured["output_slice"] == c
+    pos = captured["positive_kwargs"]
+    neg = captured["negative_kwargs"]
+    assert pos["hidden_states"].shape[0] == b, "positive branch must be batch-B"
+    assert neg["hidden_states"].shape[0] == b, "negative branch must be batch-B"
+    assert tuple(pos["timestep"].shape) == tuple(mask.shape), "mask must not be concatenated in CFG2"
+    assert torch.equal(pos["timestep"][:, :, 0], torch.zeros(b, 1, 1, 1)), "conditioning frame timestep must be 0"
+    assert torch.equal(pos["encoder_hidden_states"], prompt_embeds)
+    assert torch.equal(neg["encoder_hidden_states"], negative_embeds)
+    assert torch.equal(out[:, :, :1], latents0[:, :, :1]), "first frame must be preserved"
+
+
+def test_diffuse_runs_cfg_parallel_validity_check(monkeypatch, mocker):
+    """diffuse must run the CFG-parallel validity check so a scale<=1 or missing
+    negative prompt warns the user instead of silently wasting a rank."""
+    _set_cfg_world_size(monkeypatch, 2)
+
+    b = 1
+    pipe = _bare_t2v_pipeline(transformer=None, guidance_scale=6.0)
+    pipe.predict_noise_maybe_with_cfg = lambda **kwargs: torch.zeros(b, 4, 2, 4, 4)
+    check = mocker.patch.object(SanaVideoPipeline, "check_cfg_parallel_validity", return_value=True)
+
+    pipe.diffuse(
+        latents=torch.randn(b, 4, 2, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=None,
+    )
+
+    check.assert_called_once_with(6.0, True)
+
+
+def test_i2v_diffuse_runs_cfg_parallel_validity_check(monkeypatch, mocker):
+    """I2V has its own diffuse; it must also run the CFG-parallel validity check."""
+    _set_cfg_world_size(monkeypatch, 2)
+    b, c, frames = 1, 4, 3
+
+    pipe = _bare_i2v_pipeline(None, _I2VSpyScheduler(), guidance_scale=6.0)
+    pipe.predict_noise_maybe_with_cfg = lambda **kwargs: torch.zeros(b, c, frames, 4, 4)
+    check = mocker.patch.object(SanaVideoPipeline, "check_cfg_parallel_validity", return_value=True)
+
+    pipe.diffuse(
+        latents=torch.randn(b, c, frames, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        conditioning_mask=_conditioning_mask(b, frames),
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=c,
+    )
+
+    check.assert_called_once_with(6.0, True)
+
+
+def test_diffuse_runs_without_initialized_parallel_groups():
+    """Single-process runs (unit tests, offline scripts) never initialize the
+    distributed groups; diffuse must fall back to the serial path instead of
+    tripping the parallel-state assertion. No world-size patching on purpose."""
+
+    class _SpyTransformer:
+        dtype = torch.float32
+
+        def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
+            return (torch.zeros_like(hidden_states),)
+
+    b = 1
+    pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0)
+
+    out = pipe.diffuse(
+        latents=torch.randn(b, 4, 2, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=None,
+    )
+
+    assert out.shape == (b, 4, 2, 4, 4)
+
+
+def test_diffuse_propagates_group_error_when_parallel_initialized(monkeypatch):
+    """Inside an initialized parallel run a failing CFG-group lookup is a real
+    error and must propagate instead of silently flipping to the serial path."""
+    import vllm_omni.diffusion.models.sana_video.pipeline_sana_video as t2v
+
+    monkeypatch.setattr(t2v, "model_parallel_is_initialized", lambda: True, raising=False)
+
+    def _broken_group_lookup():
+        raise AssertionError("classifier_free_guidance group is not initialized")
+
+    monkeypatch.setattr(t2v, "get_classifier_free_guidance_world_size", _broken_group_lookup)
+
+    b = 1
+    pipe = _bare_t2v_pipeline(None, guidance_scale=6.0)
+
+    with pytest.raises(AssertionError, match="classifier_free_guidance"):
+        pipe.diffuse(
+            latents=torch.randn(b, 4, 2, 4, 4),
+            timesteps=torch.tensor([1000.0]),
+            prompt_embeds=torch.randn(b, 3, 8),
+            prompt_attention_mask=torch.ones(b, 3),
+            negative_prompt_embeds=torch.randn(b, 3, 8),
+            negative_prompt_attention_mask=torch.ones(b, 3),
+            guidance_scale=6.0,
+            extra_step_kwargs={},
+            dtype=torch.float32,
+            output_slice=None,
+        )
+
+
+def test_t2v_forwards_eta_and_generator_into_scheduler(monkeypatch):
+    """T2V must keep passing eta/generator (extra_step_kwargs) into scheduler.step."""
+    _set_cfg_world_size(monkeypatch, 1)
+    captured = {}
+
+    class _KwargScheduler:
+        order = 1
+
+        def step(self, noise_pred, t, latents, **kwargs):
+            captured.update(kwargs)
+            return (latents,)
+
+    class _SpyTransformer:
+        dtype = torch.float32
+
+        def __call__(self, hidden_states, **kwargs):
+            return (torch.zeros_like(hidden_states),)
+
+    b = 1
+    pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0)
+    pipe.scheduler = _KwargScheduler()
+    generator = torch.Generator().manual_seed(0)
+
+    pipe.diffuse(
+        latents=torch.randn(b, 4, 2, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        extra_step_kwargs={"eta": 0.5, "generator": generator},
+        dtype=torch.float32,
+        output_slice=None,
+    )
+
+    assert captured.get("eta") == 0.5
+    assert captured.get("generator") is generator
+    assert captured.get("return_dict") is False

tests/diffusion/models/sana_video/test_sana_video_topology.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""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.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import SanaVideoPipeline
+from vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v import SanaImageToVideoPipeline
+from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+    validate_sana_video_parallel_config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _pc(
+    tp: int = 1,
+    cfg: int = 1,
+    sp: int | None = None,
+    text_tp: int = 1,
+    pp: int = 1,
+    hsdp: bool = False,
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        tensor_parallel_size=tp,
+        cfg_parallel_size=cfg,
+        sequence_parallel_size=sp,
+        text_encoder_tp_size=text_tp,
+        pipeline_parallel_size=pp,
+        use_hsdp=hsdp,
+    )
+
+
+@pytest.mark.parametrize("tp", [1, 2])
+def test_supported_tp_sizes_pass(tp: int) -> None:
+    validate_sana_video_parallel_config(_pc(tp=tp))
+
+
+@pytest.mark.parametrize("tp", [3, 4, 8])
+def test_unsupported_tp_sizes_rejected(tp: int) -> None:
+    with pytest.raises(NotImplementedError, match="tensor_parallel_size"):
+        validate_sana_video_parallel_config(_pc(tp=tp))
+
+
+def test_cfg3_rejected() -> None:
+    with pytest.raises(NotImplementedError, match="cfg_parallel_size"):
+        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))
+
+
+def test_pipeline_parallel_rejected() -> None:
+    with pytest.raises(NotImplementedError, match="pipeline parallel"):
+        validate_sana_video_parallel_config(_pc(pp=2))
+
+
+def test_hsdp_rejected() -> None:
+    with pytest.raises(NotImplementedError, match="HSDP"):
+        validate_sana_video_parallel_config(_pc(hsdp=True))
+
+
+def test_text_encoder_tp_rejected() -> None:
+    with pytest.raises(NotImplementedError, match="text encoder"):
+        validate_sana_video_parallel_config(_pc(text_tp=2))
+
+
+def test_tp2_cfg2_combo_passes() -> None:
+    validate_sana_video_parallel_config(_pc(tp=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."""
+    load = mocker.patch.object(SanaVideoPipeline, "_load_components")
+    od_config = SimpleNamespace(parallel_config=_pc(tp=3))
+
+    with pytest.raises(NotImplementedError, match="tensor_parallel_size"):
+        pipeline_cls(od_config=od_config)
+
+    load.assert_not_called()

tests/diffusion/models/sana_video/test_transformer_sana_video.py

@@ -1,11 +1,37 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import os
+
 import pytest
 import torch
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
+
+@pytest.fixture(autouse=True)
+def _init_distributed(monkeypatch):
+    """The native transformer uses vLLM parallel linear layers, which require a
+    tensor-parallel group; initialize a single-process group for CPU tests."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+    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()
+
+
 _TINY_CONFIG = {
     "in_channels": 4,
     "out_channels": 4,
@@ -323,6 +349,7 @@ def test_tiny_transformer_matches_diffusers_and_frozen_output():
     from vllm_omni.diffusion.attention.layer import Attention
     from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
     from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+        SanaDistributedRMSNorm,
         SanaLinearAttention,
         SanaRMSNorm,
         SanaVideoTransformerConfig,
@@ -342,10 +369,10 @@ def test_tiny_transformer_matches_diffusers_and_frozen_output():
     assert isinstance(block.attn1, SanaLinearAttention)
     assert isinstance(block.attn2.attn, Attention)
     assert block.attn2.attn.role == "cross"
-    assert isinstance(block.attn1.norm_q, SanaRMSNorm)
-    assert isinstance(block.attn1.norm_k, SanaRMSNorm)
-    assert isinstance(block.attn2.norm_q, SanaRMSNorm)
-    assert isinstance(block.attn2.norm_k, SanaRMSNorm)
+    assert isinstance(block.attn1.norm_q, SanaDistributedRMSNorm)
+    assert isinstance(block.attn1.norm_k, SanaDistributedRMSNorm)
+    assert isinstance(block.attn2.norm_q, SanaDistributedRMSNorm)
+    assert isinstance(block.attn2.norm_k, SanaDistributedRMSNorm)
     assert isinstance(model.caption_norm, SanaRMSNorm)
 
     torch.manual_seed(11)

tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py

@@ -0,0 +1,445 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""TP unit tests for SANA-Video: distributed RMSNorm, sharded attention and
+sharded GLUMB temporal conv.
+
+CPU-only. The TP world size and the collective are mocked so a single process
+simulates the shard/reduce logic a real multi-rank run depends on.
+"""
+
+import os
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+    GLUMBTempConv,
+    SanaCrossAttention,
+    SanaDistributedRMSNorm,
+    SanaLinearAttention,
+    SanaRMSNorm,
+)
+
+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,
+    )
+
+
+def test_distributed_rms_norm_tp1_matches_sana_rms_norm(mocker) -> None:
+    """TP1 distributed norm must be bit-identical to the checkpoint SanaRMSNorm."""
+    dim = 24
+    torch.manual_seed(0)
+
+    dense = SanaRMSNorm(dim, eps=1e-5, elementwise_affine=True)
+    dense.weight.data.normal_()
+
+    dist = SanaDistributedRMSNorm(dim, eps=1e-5)
+    dist.weight.data.copy_(dense.weight.data)
+
+    x = torch.randn(2, 5, dim)
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1)
+    out = dist(x)
+    ref = dense(x)
+
+    torch.testing.assert_close(out, ref, rtol=0, atol=0)
+
+
+def test_distributed_rms_norm_tp1_bf16_weight_matches_sana_rms_norm(mocker) -> None:
+    """bf16 weight must follow SanaRMSNorm's cast-to-weight-dtype branch.
+
+    An unconditional cast back to the input dtype would diverge here; this pins
+    the fp16/bf16-only cast that SanaRMSNorm performs.
+    """
+    dim = 24
+    torch.manual_seed(0)
+
+    dense = SanaRMSNorm(dim, eps=1e-5, elementwise_affine=True)
+    dense.weight.data.normal_()
+    dense.weight.data = dense.weight.data.to(torch.bfloat16)
+
+    dist = SanaDistributedRMSNorm(dim, eps=1e-5)
+    dist.weight.data = dense.weight.data.clone()
+
+    x = torch.randn(2, 5, dim, dtype=torch.bfloat16)
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1)
+    out = dist(x)
+    ref = dense(x)
+
+    assert out.dtype == torch.bfloat16
+    torch.testing.assert_close(out, ref, rtol=0, atol=0)
+
+
+def test_distributed_rms_norm_tp2_shards_aggregate_to_dense(mocker) -> None:
+    """Two TP2 shards must reproduce dense SanaRMSNorm.
+
+    The RMS must use the full (unsharded) hidden size: each rank reduces its
+    local sum-of-squares across ranks, not a per-shard mean over its half.
+    """
+    full_dim, tp_size = 24, 2
+    half = full_dim // tp_size
+    torch.manual_seed(0)
+
+    dense = SanaRMSNorm(full_dim, eps=1e-5, elementwise_affine=True)
+    dense.weight.data.normal_()
+
+    dist0 = SanaDistributedRMSNorm(half, eps=1e-5)
+    dist1 = SanaDistributedRMSNorm(half, eps=1e-5)
+    dist0.weight.data.copy_(dense.weight.data[:half])
+    dist1.weight.data.copy_(dense.weight.data[half:])
+
+    x = torch.randn(2, 5, full_dim)
+    x0, x1 = x[..., :half], x[..., half:]
+
+    # A real all-reduce returns the sum of every rank's local sum-of-squares.
+    global_sum_sq = x.to(torch.float32).pow(2).sum(-1, keepdim=True)
+
+    def fake_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
+        return global_sum_sq
+
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size)
+    mocker.patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=fake_all_reduce)
+    out0 = dist0(x0)
+    out1 = dist1(x1)
+
+    out = torch.cat([out0, out1], dim=-1)
+    ref = dense(x)
+
+    torch.testing.assert_close(out, ref, rtol=0, atol=0)
+
+
+def test_distributed_rms_norm_weight_loader_shards_dim0(mocker) -> None:
+    """Each rank loads its own dim-0 slice of the global affine weight."""
+    full_dim, tp_size, rank = 24, 2, 1
+    half = full_dim // tp_size
+
+    norm = SanaDistributedRMSNorm(half, eps=1e-5)
+    full_weight = torch.randn(full_dim)
+
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size)
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank)
+    norm.weight.weight_loader(norm.weight, full_weight)
+
+    torch.testing.assert_close(norm.weight.data, full_weight[rank * half : (rank + 1) * half])
+
+
+def test_distributed_rms_norm_weight_loader_tp1_copies_full() -> None:
+    dim = 24
+    norm = SanaDistributedRMSNorm(dim, eps=1e-5)
+    full_weight = torch.randn(dim)
+
+    norm.weight.weight_loader(norm.weight, full_weight)
+
+    torch.testing.assert_close(norm.weight.data, full_weight)
+
+
+_SELF_ATTN_GOLDEN = torch.tensor(
+    [
+        5.0944157,
+        -16.9981823,
+        55.4457436,
+        42.5174408,
+        -9.3055859,
+        19.3744793,
+        -47.4345360,
+        -5.8110709,
+        -14.0679598,
+        -22.3440742,
+        -8.3834534,
+        8.9109459,
+    ]
+)
+
+
+def _fill_by_name(module: torch.nn.Module) -> None:
+    """Deterministic per-parameter fill that is stable across TP refactors."""
+    torch.manual_seed(0)
+    for _, param in sorted(module.named_parameters()):
+        param.data.normal_()
+
+
+def _self_attn_inputs(seq_len: int, dim: int, head_dim: int):
+    torch.manual_seed(1)
+    x = torch.randn(2, seq_len, dim)
+    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_self_attention_tp1_matches_golden_and_keeps_param_names(tp1_group, force_default_gemm) -> None:
+    """TP1 self-attention stays bit-identical to PR1 and preserves checkpoint keys."""
+    dim, num_heads, head_dim, seq_len = 24, 4, 6, 8
+    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()
+    _fill_by_name(attn)
+    x, rotary_emb = _self_attn_inputs(seq_len, dim, head_dim)
+
+    with torch.no_grad():
+        out = attn(x, rotary_emb=rotary_emb)
+
+    assert {name for name, _ in attn.named_parameters()} == {
+        "to_q.weight",
+        "to_k.weight",
+        "to_v.weight",
+        "norm_q.weight",
+        "norm_k.weight",
+        "to_out.0.weight",
+        "to_out.0.bias",
+    }
+    torch.testing.assert_close(out.flatten()[:12], _SELF_ATTN_GOLDEN, rtol=1e-5, atol=1e-5)
+
+
+def _mock_tp2(mocker):
+    """Mock a TP world size of 2 for module construction (no real group)."""
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2)
+    mocker.patch(
+        "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
+        return_value=2,
+    )
+    tp_group = mocker.MagicMock()
+    tp_group.world_size = 2
+    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=tp_group)
+
+
+def test_self_attention_tp2_shards_projections_and_heads(mocker) -> None:
+    """TP2 splits heads column-wise (q/k/v) and the output projection row-wise."""
+    dim, num_heads, head_dim = 24, 4, 6
+    inner_dim = num_heads * head_dim
+    _mock_tp2(mocker)
+
+    attn = SanaLinearAttention(
+        dim=dim, num_heads=num_heads, head_dim=head_dim, dropout=0.0, bias=False, qk_norm="rms_norm_across_heads"
+    )
+
+    assert attn.heads == num_heads // 2
+    assert attn.to_q.weight.shape == (inner_dim // 2, dim)
+    assert attn.to_k.weight.shape == (inner_dim // 2, dim)
+    assert attn.to_v.weight.shape == (inner_dim // 2, dim)
+    assert attn.to_out[0].weight.shape == (dim, inner_dim // 2)
+    assert attn.norm_q.weight.shape == (inner_dim // 2,)
+    assert attn.norm_k.weight.shape == (inner_dim // 2,)
+
+
+def test_cross_attention_tp2_shards_projections_and_heads(mocker) -> None:
+    """TP2 cross-attention shards q/k/v column-wise, to_out row-wise, heads locally."""
+    dim, cross_dim, num_heads, head_dim = 24, 16, 4, 6
+    inner_dim = num_heads * head_dim
+    _mock_tp2(mocker)
+
+    attn = SanaCrossAttention(
+        dim=dim,
+        cross_attention_dim=cross_dim,
+        num_heads=num_heads,
+        head_dim=head_dim,
+        dropout=0.0,
+        bias=False,
+        out_bias=True,
+        qk_norm="rms_norm_across_heads",
+        prefix="attn2",
+    )
+
+    assert attn.heads == num_heads // 2
+    assert attn.to_q.weight.shape == (inner_dim // 2, dim)
+    assert attn.to_k.weight.shape == (inner_dim // 2, cross_dim)
+    assert attn.to_v.weight.shape == (inner_dim // 2, cross_dim)
+    assert attn.to_out[0].weight.shape == (dim, inner_dim // 2)
+    assert isinstance(attn.norm_q, SanaDistributedRMSNorm)
+    assert attn.norm_q.weight.shape == (inner_dim // 2,)
+
+
+_TINY_TRANSFORMER_CONFIG = {
+    "in_channels": 4,
+    "out_channels": 4,
+    "num_attention_heads": 2,
+    "attention_head_dim": 12,
+    "num_layers": 1,
+    "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": 32,
+}
+
+
+def test_tp2_load_weights_shards_full_checkpoint_without_missing(mocker) -> None:
+    """The full checkpoint loads under TP2 via weight_loaders with no missing/unexpected keys."""
+    from diffusers import SanaVideoTransformer3DModel as DiffusersTransformer
+
+    from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+
+    _mock_tp2(mocker)
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=0)
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+
+    reference = DiffusersTransformer(**_TINY_TRANSFORMER_CONFIG)
+    model = SanaVideoTransformer3DModel(**_TINY_TRANSFORMER_CONFIG)
+
+    loaded = model.load_weights(list(reference.state_dict().items()))
+
+    assert loaded == set(model.state_dict())
+
+
+# ── GLUMB temporal conv TP ──
+
+# hidden_channels = expand_ratio * in_channels = 8, so conv_inverted emits 2h = 16.
+_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO = 4, 4, 2.0
+
+_GLUMB_GOLDEN = torch.tensor(
+    [
+        51.4757385,
+        -10.9912930,
+        -4.1482553,
+        -26.1580830,
+        -19.7727261,
+        17.9901562,
+        -1.6007471,
+        13.3993073,
+        -8.6459389,
+        -10.4008112,
+        -7.4962354,
+        1.5446399,
+    ]
+)
+
+
+def _glumb_inputs() -> torch.Tensor:
+    torch.manual_seed(1)
+    return torch.randn(1, 2, 3, 3, _GLUMB_IN)  # [B, F, H, W, C]
+
+
+def test_glumb_tp1_matches_golden_and_keeps_param_names(tp1_group) -> None:
+    """TP1 GLUMB stays bit-identical to the PR1 conv output and preserves keys."""
+    glumb = GLUMBTempConv(_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO, norm_type=None, residual_connection=False)
+    glumb.eval()
+    _fill_by_name(glumb)
+    x = _glumb_inputs()
+
+    with torch.no_grad():
+        out = glumb(x)
+
+    assert {name for name, _ in glumb.named_parameters()} == {
+        "conv_inverted.weight",
+        "conv_inverted.bias",
+        "conv_depth.weight",
+        "conv_depth.bias",
+        "conv_point.weight",
+        "conv_temp.weight",
+    }
+    torch.testing.assert_close(out.flatten()[:12], _GLUMB_GOLDEN, rtol=1e-5, atol=1e-5)
+
+
+def _random_glumb_checkpoint(hidden_channels: int) -> dict[str, torch.Tensor]:
+    """Checkpoint-format (nn.Conv2d) GLUMB weights with per-segment asymmetry.
+
+    conv_inverted / conv_depth carry two h-channel segments (value then gate).
+    Filling the gate half from a different distribution than the value half is
+    what exposes a naive contiguous dim-0 split: that bug would place the whole
+    value segment on rank 0 and the whole gate segment on rank 1, so each rank's
+    local chunk(2) would slice value-vs-value instead of value-vs-gate.
+    """
+    torch.manual_seed(0)
+    h = hidden_channels
+    value_w = torch.randn(h, _GLUMB_IN, 1, 1)
+    gate_w = torch.randn(h, _GLUMB_IN, 1, 1) * 3.0 + 5.0
+    value_b = torch.randn(h)
+    gate_b = torch.randn(h) * 3.0 + 5.0
+    value_dw = torch.randn(h, 1, 3, 3)
+    gate_dw = torch.randn(h, 1, 3, 3) * 3.0 + 5.0
+    value_db = torch.randn(h)
+    gate_db = torch.randn(h) * 3.0 + 5.0
+    return {
+        "conv_inverted.weight": torch.cat([value_w, gate_w], dim=0),
+        "conv_inverted.bias": torch.cat([value_b, gate_b], dim=0),
+        "conv_depth.weight": torch.cat([value_dw, gate_dw], dim=0),
+        "conv_depth.bias": torch.cat([value_db, gate_db], dim=0),
+        "conv_point.weight": torch.randn(_GLUMB_OUT, h, 1, 1),
+        "conv_temp.weight": torch.randn(_GLUMB_OUT, _GLUMB_OUT, 3, 1),
+    }
+
+
+def _mock_glumb_tp(mocker, world_size: int, rank: int) -> None:
+    """Mock a TP group of ``world_size`` at ``rank`` with an identity all-reduce,
+    so a single process can build and run one rank's shard in isolation."""
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=world_size)
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank)
+    # Identity conv_point all-reduce returns each rank's local partial for the
+    # test to sum.
+    mocker.patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=lambda t: t)
+
+
+def _load_glumb_checkpoint(glumb: GLUMBTempConv, checkpoint: dict[str, torch.Tensor]) -> None:
+    """Load like AutoWeightsLoader: custom weight_loader when present (the sharded
+    convs), plain copy for replicated params (conv_temp)."""
+    params = dict(glumb.named_parameters())
+    for name, tensor in checkpoint.items():
+        param = params[name]
+        loader = getattr(param, "weight_loader", None)
+        if loader is None:
+            param.data.copy_(tensor)
+        else:
+            loader(param, tensor)
+
+
+def test_glumb_tp2_shards_aggregate_to_dense(mocker) -> None:
+    """The merged value/gate two-segment shard must aggregate to the dense conv.
+
+    Everything after conv_point (temporal conv + residual) is linear and
+    bias-free, so summing each rank's full forward (with the conv_point
+    all-reduce mocked to identity) reproduces the replicated dense output.
+    """
+    hidden_channels = int(_GLUMB_RATIO * _GLUMB_IN)
+    checkpoint = _random_glumb_checkpoint(hidden_channels)
+    x = _glumb_inputs()
+
+    _mock_glumb_tp(mocker, world_size=1, rank=0)
+    dense = GLUMBTempConv(_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO, norm_type=None, residual_connection=False)
+    dense.eval()
+    _load_glumb_checkpoint(dense, checkpoint)
+    with torch.no_grad():
+        ref = dense(x)
+
+    partials = []
+    for rank in (0, 1):
+        _mock_glumb_tp(mocker, world_size=2, rank=rank)
+        shard = GLUMBTempConv(_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO, norm_type=None, residual_connection=False)
+        shard.eval()
+        _load_glumb_checkpoint(shard, checkpoint)
+        with torch.no_grad():
+            partials.append(shard(x))
+
+    aggregated = partials[0] + partials[1]
+    torch.testing.assert_close(aggregated, ref, rtol=1e-5, atol=1e-5)

tests/model_tests/diffusion/case_filtering.py

@@ -76,6 +76,8 @@ def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts
             test_info.supported_tasks,
             test_info.check_multi_output and test_group is None,
             test_info.check_determinism and test_group is None,
+            # The divergence check compares generated outputs, which only the offline test does.
+            *([] if online else [test_info.check_i2v_t2v_divergence]),
             id=f"{model_name}[{'+'.join(test_group)}]" if test_group else model_name,
             marks=get_test_group_marks(model_name, test_group, test_info.marks),
         )

tests/model_tests/diffusion/config_types.py

@@ -76,6 +76,11 @@ class DiffusionModelTestOpts(NamedTuple):
     check_multi_output: bool = True  # Runs multiple generations in one request
     check_determinism: bool = True  # Runs 2 generations with the same seed and check determinism
 
+    # For IMAGE_TO_VIDEO, also run a T2V request and assert the outputs differ, to catch a
+    # silently dropped input image. Only valid for unified text/image pipelines (e.g. LTX2);
+    # set False for pipelines that mandate an image and fail closed without one (SANA-Video I2V).
+    check_i2v_t2v_divergence: bool = True
+
 
 ### Mappings & utils for building offline Omni() instances given a list of enabled accelerations
 ACC_OMNI_KWARGS = {

tests/model_tests/diffusion/diff_model_builders.py

@@ -1,3 +1,4 @@
+import json
 from functools import partial
 from pathlib import Path
 
@@ -49,15 +50,17 @@ def tiny_ltx2_builder() -> str:
     return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline")
 
 
-def tiny_sana_video_builder() -> str:
+def _build_tiny_sana_video(pipeline_name: str) -> str:
     """Build a tiny 480p SANA-Video model without downloading model weights.
 
     The tokenizer is the only component loaded from the upstream repository.
     All modules with weights are initialized locally from intentionally small
-    configs, and the scheduler is constructed from its weight-free config.
+    configs, and the scheduler is constructed from its weight-free config. The
+    I2V pipeline shares this diffusers layout, so it reuses the same components
+    and only rewrites model_index.json's ``_class_name`` to its own class.
     """
     model_id = "Efficient-Large-Model/SANA-Video_2B_480p_diffusers"
-    model_dir = _get_tiny_model_path("SanaVideoPipeline")
+    model_dir = _get_tiny_model_path(pipeline_name)
 
     tokenizer = GemmaTokenizerFast.from_pretrained(model_id, subfolder="tokenizer")
     scheduler = DPMSolverMultistepScheduler(
@@ -112,9 +115,22 @@ def tiny_sana_video_builder() -> str:
         scheduler=scheduler,
     )
     pipeline.to(dtype=torch.bfloat16).save_pretrained(model_dir)
+    if pipeline_name != "SanaVideoPipeline":
+        index_path = Path(model_dir) / "model_index.json"
+        index = json.loads(index_path.read_text())
+        index["_class_name"] = pipeline_name
+        index_path.write_text(json.dumps(index, indent=2))
     return model_dir
 
 
+def tiny_sana_video_builder() -> str:
+    return _build_tiny_sana_video("SanaVideoPipeline")
+
+
+def tiny_sana_video_i2v_builder() -> str:
+    return _build_tiny_sana_video("SanaImageToVideoPipeline")
+
+
 def _shrink_flux_clip_text_encoder(config: dict) -> dict:
     config["num_hidden_layers"] = 2
     return config

tests/model_tests/diffusion/model_settings.py

@@ -53,6 +53,22 @@ DIFFUSION_TEST_SETTINGS = {
         model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
         builder=diff_model_builders.tiny_sana_video_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO],
+        extra_test_groups=[
+            [DiffusionAccs.TENSOR_PARALLEL],
+            [DiffusionAccs.CFG_PARALLEL],
+            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL],
+        ],
+        check_multi_output=False,
+    ),
+    "SanaImageToVideoPipeline": DiffusionModelTestOpts(
+        model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
+        builder=diff_model_builders.tiny_sana_video_i2v_builder,
+        supported_tasks=[DiffusionTasks.IMAGE_TO_VIDEO],
+        extra_test_groups=[
+            [DiffusionAccs.CFG_PARALLEL],
+            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL],
+        ],
+        check_i2v_t2v_divergence=False,
         check_multi_output=False,
     ),
     "QwenImagePipeline": DiffusionModelTestOpts(

tests/model_tests/diffusion/task_runners.py

@@ -220,16 +220,20 @@ def run_and_validate_text_to_video_request(omni: Omni):
     _validate_video(_run_offline_t2v(omni))
 
 
-def run_and_validate_image_to_video_request(omni: Omni):
+def run_and_validate_image_to_video_request(omni: Omni, check_t2v_divergence: bool = True):
     """Run and validate an image to video request.
 
     LTX2Pipeline uses a unified text/image entry and can fall back to T2V if
     the image never reaches conditioning, so also run T2V with the same seed
     and prompt and assert the outputs differ; a regression that silently
     drops the input image would otherwise still pass a "did we get a valid
-    video back" check alone.
+    video back" check alone. Pipelines that mandate an image and fail closed
+    without one (SANA-Video I2V) pass check_t2v_divergence=False: a dropped
+    image raises instead of silently degrading, so the T2V comparison is moot.
     """
     i2v_video = _validate_video_frames(_get_offline_videos(_run_offline_i2v(omni)))[0]
+    if not check_t2v_divergence:
+        return
     t2v_video = _validate_video_frames(_get_offline_videos(_run_offline_t2v(omni)))[0]
     assert not np.array_equal(i2v_video, t2v_video), (
         "I2V output is identical to T2V output with the same seed and prompt; "

tests/model_tests/diffusion/test_alignment.py

@@ -61,12 +61,6 @@ EXCLUDED_MODELS = [
     "HunyuanVideo15Pipeline",
     "HunyuanVideo15ImageToVideoPipeline",
     "LingBotVideoPipeline",
-    # The common harness does not yet define an image-to-video task or supply
-    # an input image. SANA I2V is covered by
-    # tests/diffusion/models/sana_video/test_pipeline_sana_video.py and
-    # tests/e2e/online_serving/test_sana_video_expansion.py. SANA T2V is exercised through
-    # the common tiny-model setting in model_settings.py.
-    "SanaImageToVideoPipeline",
     # Requires camera/action assets and persistent AR state; covered by its
     # dedicated contract tests and real-checkpoint E2E smoke instead.
     "LingBotWorldCausalDMDPipeline",

tests/model_tests/diffusion/test_common_offline.py

@@ -24,7 +24,7 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
 
 
 @pytest.mark.parametrize(
-    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
+    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism,check_i2v_t2v_divergence",
     get_parametrized_options(DIFFUSION_TEST_SETTINGS),
 )
 def test_pipeline_on_supported_tasks(
@@ -33,6 +33,7 @@ def test_pipeline_on_supported_tasks(
     supported_tasks: list[DiffusionTasks],
     check_multioutput: bool,
     check_determinism: bool,
+    check_i2v_t2v_divergence: bool,
     tiny_model_paths: dict[str, str],
     subtests,
 ):
@@ -61,7 +62,7 @@ def test_pipeline_on_supported_tasks(
                 elif task_type == DiffusionTasks.TEXT_TO_VIDEO:
                     run_and_validate_text_to_video_request(omni)
                 elif task_type == DiffusionTasks.IMAGE_TO_VIDEO:
-                    run_and_validate_image_to_video_request(omni)
+                    run_and_validate_image_to_video_request(omni, check_t2v_divergence=check_i2v_t2v_divergence)
                 else:
                     raise ValueError(f"Task type {task_type} is not yet supported")
 

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -230,6 +230,11 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
     def _raise_unsupported_features(self) -> None:
         """Raise an error for incompatible feature switches."""
         pc = self.od_config.parallel_config
+        if pc.tensor_parallel_size > 1:
+            raise NotImplementedError(
+                "Tensor parallel is not supported with the diffusers backend. "
+                "It requires model-specific layer sharding; use a native pipeline for TP."
+            )
         if pc.cfg_parallel_size > 1:
             raise NotImplementedError(
                 "CFG parallel is not supported with the diffusers backend. "

vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py

@@ -37,6 +37,11 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
     DistributedAutoencoderKLLTX2Video,
 )
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _wrap
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_classifier_free_guidance_world_size,
+    model_parallel_is_initialized,
+)
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
@@ -49,7 +54,7 @@ from vllm_omni.diffusion.request import resolve_video_num_frames
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
 from .pipeline_output import SanaVideoPipelineOutput
-from .transformer_sana_video import SanaVideoTransformer3DModel
+from .transformer_sana_video import SanaVideoTransformer3DModel, validate_sana_video_parallel_config
 
 
 def _resolve_vae_class_and_dtype(
@@ -269,6 +274,7 @@ def retrieve_timesteps(
 
 class SanaVideoPipeline(
     nn.Module,
+    CFGParallelMixin,
     ProgressBarMixin,
     DiffusionPipelineProfilerMixin,
     SupportsComponentDiscovery,
@@ -308,6 +314,9 @@ class SanaVideoPipeline(
     ):
         super().__init__()
 
+        if od_config is not None:
+            validate_sana_video_parallel_config(od_config.parallel_config)
+
         self.od_config = od_config
         self.device = get_local_device()
         self.weights_sources: list[DiffusersPipelineLoader.ComponentSource] = []
@@ -724,6 +733,105 @@ class SanaVideoPipeline(
     def interrupt(self):
         return self._interrupt
 
+    def predict_noise(self, *args, **kwargs) -> torch.Tensor:
+        # return_dict=False: the dataclass output is not subscriptable.
+        return self.transformer(*args, **kwargs, return_dict=False)[0]
+
+    def combine_cfg_noise(
+        self,
+        positive_noise_pred,
+        negative_noise_pred,
+        true_cfg_scale,
+        cfg_normalize=False,
+        kwargs=None,
+    ):
+        # fp32 combine so the bf16 all-gather path matches CFG1 bit-for-bit.
+        return super().combine_cfg_noise(
+            tuple(p.float() for p in _wrap(positive_noise_pred)),
+            tuple(n.float() for n in _wrap(negative_noise_pred)),
+            true_cfg_scale,
+            cfg_normalize,
+            kwargs,
+        )
+
+    def diffuse(
+        self,
+        latents: torch.Tensor,
+        timesteps: torch.Tensor,
+        prompt_embeds: torch.Tensor,
+        prompt_attention_mask: torch.Tensor,
+        negative_prompt_embeds: torch.Tensor | None,
+        negative_prompt_attention_mask: torch.Tensor | None,
+        guidance_scale: float,
+        extra_step_kwargs: dict,
+        dtype: torch.dtype,
+        output_slice: int | None,
+    ) -> torch.Tensor:
+        do_true_cfg = self.do_classifier_free_guidance
+        # Single-process runs never initialize the parallel groups.
+        cfg_parallel = model_parallel_is_initialized() and get_classifier_free_guidance_world_size() > 1
+        if cfg_parallel:
+            self.check_cfg_parallel_validity(guidance_scale, negative_prompt_embeds is not None)
+
+        if do_true_cfg and not cfg_parallel:
+            # Concatenate neg/pos for a single batch-2 forward.
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
+
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for t in timesteps:
+                if self.interrupt:
+                    continue
+
+                if cfg_parallel:
+                    timestep = t.expand(latents.shape[0])
+                    latent_model_input = latents.to(dtype)
+                    positive_kwargs = {
+                        "hidden_states": latent_model_input,
+                        "encoder_hidden_states": prompt_embeds.to(dtype),
+                        "encoder_attention_mask": prompt_attention_mask,
+                        "timestep": timestep,
+                    }
+                    negative_kwargs = (
+                        {
+                            "hidden_states": latent_model_input,
+                            "encoder_hidden_states": negative_prompt_embeds.to(dtype),
+                            "encoder_attention_mask": negative_prompt_attention_mask,
+                            "timestep": timestep,
+                        }
+                        if do_true_cfg
+                        else None
+                    )
+                    noise_pred = self.predict_noise_maybe_with_cfg(
+                        do_true_cfg=do_true_cfg,
+                        true_cfg_scale=guidance_scale,
+                        positive_kwargs=positive_kwargs,
+                        negative_kwargs=negative_kwargs,
+                        cfg_normalize=False,
+                        output_slice=output_slice,
+                    ).float()
+                else:
+                    latent_model_input = torch.cat([latents] * 2) if do_true_cfg else latents
+                    timestep = t.expand(latent_model_input.shape[0])
+                    noise_pred = self.transformer(
+                        latent_model_input.to(dtype),
+                        encoder_hidden_states=prompt_embeds.to(dtype),
+                        encoder_attention_mask=prompt_attention_mask,
+                        timestep=timestep,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = noise_pred.float()
+                    if do_true_cfg:
+                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+                    if output_slice is not None:
+                        noise_pred = noise_pred[:, :output_slice]
+
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                progress_bar.update()
+
+        return latents
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         transformer_weights = (
             (name.removeprefix("transformer."), tensor) for name, tensor in weights if name.startswith("transformer.")
@@ -947,10 +1055,6 @@ class SanaVideoPipeline(
             max_sequence_length=max_sequence_length,
             complex_human_instruction=complex_human_instruction,
         )
-        if self.do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
-
         # 4. Prepare timesteps
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler, num_inference_steps, device, timesteps, sigmas
@@ -974,44 +1078,21 @@ class SanaVideoPipeline(
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
 
         # 7. Denoising loop
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
         self._num_timesteps = len(timesteps)
-
-        transformer_dtype = self.transformer.dtype
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
-
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latent_model_input.shape[0])
-
-                # predict noise model_output
-                noise_pred = self.transformer(
-                    latent_model_input.to(dtype=transformer_dtype),
-                    encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype),
-                    encoder_attention_mask=prompt_attention_mask,
-                    timestep=timestep,
-                    return_dict=False,
-                )[0]
-                noise_pred = noise_pred.float()
-
-                # perform guidance
-                if self.do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                # learned sigma
-                if self.transformer.config.out_channels // 2 == latent_channels:
-                    noise_pred = noise_pred.chunk(2, dim=1)[0]
-
-                # compute previous image: x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
+        # learned sigma: transformer predicts 2x latent channels; keep the first half.
+        output_slice = latent_channels if self.transformer.config.out_channels // 2 == latent_channels else None
+        latents = self.diffuse(
+            latents=latents,
+            timesteps=timesteps,
+            prompt_embeds=prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            guidance_scale=guidance_scale,
+            extra_step_kwargs=extra_step_kwargs,
+            dtype=self.transformer.dtype,
+            output_slice=output_slice,
+        )
 
         if output_type == "latent":
             video = latents

vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py

@@ -16,6 +16,10 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
     DistributedAutoencoderKLLTX2Video,
 )
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_classifier_free_guidance_world_size,
+    model_parallel_is_initialized,
+)
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import (
     ASPECT_RATIO_480_BIN,
@@ -225,6 +229,91 @@ class SanaImageToVideoPipeline(SanaVideoPipeline, SupportImageInput):
         )
         return DiffusionOutput(output=video)
 
+    def diffuse(
+        self,
+        latents: torch.Tensor,
+        timesteps: torch.Tensor,
+        prompt_embeds: torch.Tensor,
+        prompt_attention_mask: torch.Tensor,
+        negative_prompt_embeds: torch.Tensor | None,
+        negative_prompt_attention_mask: torch.Tensor | None,
+        guidance_scale: float,
+        conditioning_mask: torch.Tensor,
+        extra_step_kwargs: dict,
+        dtype: torch.dtype,
+        output_slice: int | None,
+    ) -> torch.Tensor:
+        do_true_cfg = self.do_classifier_free_guidance
+        # Single-process runs never initialize the parallel groups.
+        cfg_parallel = model_parallel_is_initialized() and get_classifier_free_guidance_world_size() > 1
+        if cfg_parallel:
+            self.check_cfg_parallel_validity(guidance_scale, negative_prompt_embeds is not None)
+
+        if do_true_cfg and not cfg_parallel:
+            # Concatenate neg/pos and the mask for a single batch-2 forward.
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask])
+            conditioning_mask = torch.cat([conditioning_mask, conditioning_mask])
+
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for timestep_value in timesteps:
+                # Conditioning frames keep timestep 0 so they are not denoised.
+                timestep = timestep_value.expand(conditioning_mask.shape) * (1 - conditioning_mask)
+
+                if cfg_parallel:
+                    latent_model_input = latents.to(dtype)
+                    positive_kwargs = {
+                        "hidden_states": latent_model_input,
+                        "encoder_hidden_states": prompt_embeds.to(dtype),
+                        "encoder_attention_mask": prompt_attention_mask,
+                        "timestep": timestep,
+                    }
+                    negative_kwargs = (
+                        {
+                            "hidden_states": latent_model_input,
+                            "encoder_hidden_states": negative_prompt_embeds.to(dtype),
+                            "encoder_attention_mask": negative_prompt_attention_mask,
+                            "timestep": timestep,
+                        }
+                        if do_true_cfg
+                        else None
+                    )
+                    noise_pred = self.predict_noise_maybe_with_cfg(
+                        do_true_cfg=do_true_cfg,
+                        true_cfg_scale=guidance_scale,
+                        positive_kwargs=positive_kwargs,
+                        negative_kwargs=negative_kwargs,
+                        cfg_normalize=False,
+                        output_slice=output_slice,
+                    ).float()
+                else:
+                    latent_input = torch.cat([latents, latents]) if do_true_cfg else latents
+                    noise_pred = self.transformer(
+                        latent_input.to(dtype),
+                        encoder_hidden_states=prompt_embeds.to(dtype),
+                        encoder_attention_mask=prompt_attention_mask,
+                        timestep=timestep,
+                        return_dict=False,
+                    )[0].float()
+                    if do_true_cfg:
+                        uncond, cond = noise_pred.chunk(2)
+                        noise_pred = uncond + guidance_scale * (cond - uncond)
+                    if output_slice is not None:
+                        noise_pred = noise_pred[:, :output_slice]
+
+                # Only the non-conditioning frames are denoised; the first frame is carried through.
+                denoised = self.scheduler.step(
+                    noise_pred[:, :, 1:],
+                    timestep_value,
+                    latents[:, :, 1:],
+                    **extra_step_kwargs,
+                    return_dict=False,
+                )[0]
+                latents = torch.cat([latents[:, :, :1], denoised], dim=2)
+                progress_bar.update()
+
+        return latents
+
     @torch.no_grad()
     def _generate_i2v(
         self,
@@ -282,10 +371,6 @@ class SanaImageToVideoPipeline(SanaVideoPipeline, SupportImageInput):
             max_sequence_length=max_sequence_length,
             complex_human_instruction=complex_human_instruction,
         )
-        if self.do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_embeds, prompt_embeds])
-            prompt_mask = torch.cat([negative_mask, prompt_mask])
-
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler,
             num_inference_steps=num_inference_steps,
@@ -319,39 +404,24 @@ class SanaImageToVideoPipeline(SanaVideoPipeline, SupportImageInput):
             latents.shape[4] // patch_w,
         )
         conditioning_mask[:, :, 0] = 1
-        if self.do_classifier_free_guidance:
-            conditioning_mask = torch.cat([conditioning_mask, conditioning_mask])
 
         self._num_timesteps = len(timesteps)
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-        transformer_dtype = self.transformer.dtype
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for index, timestep_value in enumerate(timesteps):
-                latent_input = torch.cat([latents, latents]) if self.do_classifier_free_guidance else latents
-                timestep = timestep_value.expand(conditioning_mask.shape) * (1 - conditioning_mask)
-                noise_pred = self.transformer(
-                    latent_input.to(transformer_dtype),
-                    encoder_hidden_states=prompt_embeds.to(transformer_dtype),
-                    encoder_attention_mask=prompt_mask,
-                    timestep=timestep,
-                    return_dict=False,
-                )[0].float()
-                if self.do_classifier_free_guidance:
-                    uncond, cond = noise_pred.chunk(2)
-                    noise_pred = uncond + guidance_scale * (cond - uncond)
-                if self.transformer.config.out_channels // 2 == latent_channels:
-                    noise_pred = noise_pred.chunk(2, dim=1)[0]
-
-                denoised = self.scheduler.step(
-                    noise_pred[:, :, 1:],
-                    timestep_value,
-                    latents[:, :, 1:],
-                    **extra_step_kwargs,
-                    return_dict=False,
-                )[0]
-                latents = torch.cat([latents[:, :, :1], denoised], dim=2)
-                if index == len(timesteps) - 1 or (index + 1) % self.scheduler.order == 0:
-                    progress_bar.update()
+        # learned sigma: transformer predicts 2x latent channels; keep the first half.
+        output_slice = latent_channels if self.transformer.config.out_channels // 2 == latent_channels else None
+        latents = self.diffuse(
+            latents=latents,
+            timesteps=timesteps,
+            prompt_embeds=prompt_embeds,
+            prompt_attention_mask=prompt_mask,
+            negative_prompt_embeds=negative_embeds,
+            negative_prompt_attention_mask=negative_mask,
+            guidance_scale=guidance_scale,
+            conditioning_mask=conditioning_mask,
+            extra_step_kwargs=extra_step_kwargs,
+            dtype=self.transformer.dtype,
+            output_slice=output_slice,
+        )
 
         if output_type == "latent":
             return latents

vllm_omni/diffusion/models/sana_video/transformer_sana_video.py

@@ -19,12 +19,48 @@ from dataclasses import dataclass, fields
 
 import torch
 from torch import nn
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
+)
+from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
 from vllm.model_executor.models.utils import AutoWeightsLoader
+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
 
 
+def validate_sana_video_parallel_config(parallel_config) -> None:
+    """Reject unsupported SANA-Video parallel topologies before weights load."""
+    tp_size = parallel_config.tensor_parallel_size
+    if tp_size not in (1, 2):
+        raise NotImplementedError(
+            f"SANA-Video supports tensor_parallel_size 1 or 2, got {tp_size}. Set --tensor-parallel-size to 1 or 2."
+        )
+    cfg_size = parallel_config.cfg_parallel_size
+    if cfg_size not in (1, 2):
+        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:
+        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."
+        )
+    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:
+        raise NotImplementedError("SANA-Video does not support HSDP. Remove --use-hsdp.")
+    if parallel_config.text_encoder_tp_size > 1:
+        raise NotImplementedError(
+            "SANA-Video does not support text encoder tensor parallel. Set text_encoder_tp_size to 1."
+        )
+
+
 @dataclass
 class SanaVideoTransformerConfig:
     in_channels: int = 16
@@ -109,6 +145,54 @@ class SanaRMSNorm(nn.Module):
         return hidden_states
 
 
+class SanaDistributedRMSNorm(nn.Module):
+    """RMSNorm that computes global RMS across tensor parallel ranks."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        set_weight_attrs(self.weight, {"weight_loader": self.weight_loader})
+
+    def weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        if param.shape == loaded_weight.shape:
+            param.data.copy_(loaded_weight)
+            return
+
+        tp_size = get_tensor_model_parallel_world_size()
+        if loaded_weight.shape[0] % tp_size != 0:
+            raise ValueError(
+                f"Cannot shard RMSNorm weight of shape {tuple(loaded_weight.shape)} across tp_size={tp_size}."
+            )
+
+        shard_size = loaded_weight.shape[0] // tp_size
+        start_idx = get_tensor_model_parallel_rank() * shard_size
+        shard = loaded_weight.narrow(0, start_idx, shard_size)
+        if param.shape != shard.shape:
+            raise ValueError(f"RMSNorm shard shape mismatch: param={tuple(param.shape)}, shard={tuple(shard.shape)}.")
+        param.data.copy_(shard)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        tp_size = get_tensor_model_parallel_world_size()
+        x_float = x.float()
+        sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
+        count = x.shape[-1]
+
+        if tp_size > 1:
+            # Use vLLM's collective (custom all-reduce / symmetric-mem fast path
+            # for small tensors) instead of raw torch.distributed.all_reduce, and
+            # take the return value so the custom-AR path (which may return a new
+            # buffer) is handled correctly. No .clone() needed.
+            sum_sq = tensor_model_parallel_all_reduce(sum_sq)
+            count = count * tp_size
+
+        hidden_states = x_float * torch.rsqrt(sum_sq / count + self.eps)
+        # Only cast to the weight dtype for fp16/bf16 weights; keep fp32 otherwise.
+        if self.weight.dtype in (torch.float16, torch.bfloat16):
+            hidden_states = hidden_states.to(self.weight.dtype)
+        return hidden_states * self.weight
+
+
 def _get_timestep_embedding(
     timesteps: torch.Tensor,
     embedding_dim: int,
@@ -378,11 +462,22 @@ class GLUMBTempConv(nn.Module):
         hidden_channels = int(expand_ratio * in_channels)
         self.norm_type = norm_type
         self.residual_connection = residual_connection
+        self.hidden_channels = hidden_channels
+
+        tp_size = get_tensor_model_parallel_world_size()
+        local_hidden_channels = hidden_channels // tp_size
 
         self.nonlinearity = nn.SiLU()
-        self.conv_inverted = nn.Conv2d(in_channels, hidden_channels * 2, 1, 1, 0)
-        self.conv_depth = nn.Conv2d(hidden_channels * 2, hidden_channels * 2, 3, 1, 1, groups=hidden_channels * 2)
-        self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False)
+        self.conv_inverted = nn.Conv2d(in_channels, 2 * local_hidden_channels, 1, 1, 0)
+        self.conv_depth = nn.Conv2d(
+            2 * local_hidden_channels, 2 * local_hidden_channels, 3, 1, 1, groups=2 * local_hidden_channels
+        )
+        self.conv_point = nn.Conv2d(local_hidden_channels, out_channels, 1, 1, 0, bias=False)
+        # conv_inverted/conv_depth carry [value(h), gate(h)]; each rank must hold a
+        # slice of both segments so the local chunk(2) still splits value from gate.
+        for param in (self.conv_inverted.weight, self.conv_inverted.bias, self.conv_depth.weight, self.conv_depth.bias):
+            set_weight_attrs(param, {"weight_loader": self._segment_weight_loader})
+        set_weight_attrs(self.conv_point.weight, {"weight_loader": self._point_weight_loader})
 
         self.norm = None
         if norm_type == "rms_norm":
@@ -390,6 +485,19 @@ class GLUMBTempConv(nn.Module):
 
         self.conv_temp = nn.Conv2d(out_channels, out_channels, kernel_size=(3, 1), stride=1, padding=(1, 0), bias=False)
 
+    def _segment_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        h = self.hidden_channels
+        local = h // get_tensor_model_parallel_world_size()
+        start = get_tensor_model_parallel_rank() * local
+        value = loaded_weight.narrow(0, start, local)
+        gate = loaded_weight.narrow(0, h + start, local)
+        param.data.copy_(torch.cat([value, gate], dim=0))
+
+    def _point_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        local = self.hidden_channels // get_tensor_model_parallel_world_size()
+        start = get_tensor_model_parallel_rank() * local
+        param.data.copy_(loaded_weight.narrow(1, start, local))
+
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         if self.residual_connection:
             residual = hidden_states
@@ -404,6 +512,8 @@ class GLUMBTempConv(nn.Module):
         hidden_states = hidden_states * self.nonlinearity(gate)
 
         hidden_states = self.conv_point(hidden_states)
+        if get_tensor_model_parallel_world_size() > 1:
+            hidden_states = tensor_model_parallel_all_reduce(hidden_states)
 
         # Temporal aggregation
         hidden_states_temporal = hidden_states.view(batch_size, num_frames, num_channels, height * width).permute(
@@ -440,13 +550,20 @@ class SanaLinearAttention(nn.Module):
     ) -> None:
         super().__init__()
         inner_dim = num_heads * head_dim
-        self.heads = num_heads
-        self.to_q = nn.Linear(dim, inner_dim, bias=bias)
-        self.to_k = nn.Linear(dim, inner_dim, bias=bias)
-        self.to_v = nn.Linear(dim, inner_dim, bias=bias)
-        self.norm_q = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None
-        self.norm_k = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None
-        self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim, bias=True), nn.Dropout(dropout)])
+        tp_size = get_tensor_model_parallel_world_size()
+        self.heads = num_heads // tp_size
+        local_inner_dim = inner_dim // tp_size
+        self.to_q = ColumnParallelLinear(dim, inner_dim, bias=bias, gather_output=False, return_bias=False)
+        self.to_k = ColumnParallelLinear(dim, inner_dim, bias=bias, gather_output=False, return_bias=False)
+        self.to_v = ColumnParallelLinear(dim, inner_dim, bias=bias, gather_output=False, return_bias=False)
+        self.norm_q = SanaDistributedRMSNorm(local_inner_dim, eps=1e-5) if qk_norm is not None else None
+        self.norm_k = SanaDistributedRMSNorm(local_inner_dim, eps=1e-5) if qk_norm is not None else None
+        self.to_out = nn.ModuleList(
+            [
+                RowParallelLinear(inner_dim, dim, bias=True, input_is_parallel=True, return_bias=False),
+                nn.Dropout(dropout),
+            ]
+        )
 
     def forward(
         self,
@@ -632,18 +749,29 @@ class SanaCrossAttention(nn.Module):
     ) -> None:
         super().__init__()
         inner_dim = num_heads * head_dim
-        self.heads = num_heads
+        tp_size = get_tensor_model_parallel_world_size()
+        self.heads = num_heads // tp_size
         self.head_dim = head_dim
-        self.to_q = nn.Linear(dim, inner_dim, bias=bias)
-        self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias)
-        self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias)
-        self.norm_q = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None
-        self.norm_k = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None
-        self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim, bias=out_bias), nn.Dropout(dropout)])
+        local_inner_dim = inner_dim // tp_size
+        self.to_q = ColumnParallelLinear(dim, inner_dim, bias=bias, gather_output=False, return_bias=False)
+        self.to_k = ColumnParallelLinear(
+            cross_attention_dim, inner_dim, bias=bias, gather_output=False, return_bias=False
+        )
+        self.to_v = ColumnParallelLinear(
+            cross_attention_dim, inner_dim, bias=bias, gather_output=False, return_bias=False
+        )
+        self.norm_q = SanaDistributedRMSNorm(local_inner_dim, eps=1e-5) if qk_norm is not None else None
+        self.norm_k = SanaDistributedRMSNorm(local_inner_dim, eps=1e-5) if qk_norm is not None else None
+        self.to_out = nn.ModuleList(
+            [
+                RowParallelLinear(inner_dim, dim, bias=out_bias, input_is_parallel=True, return_bias=False),
+                nn.Dropout(dropout),
+            ]
+        )
         self.attn = OmniAttention(
-            num_heads=num_heads,
+            num_heads=self.heads,
             head_size=head_dim,
-            num_kv_heads=num_heads,
+            num_kv_heads=self.heads,
             softmax_scale=1.0 / (head_dim**0.5),
             causal=False,
             role="cross",