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

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

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)