[feature] Reduce memory usage with SP (#7728)

Signed-off-by: Maciej Bala <mbala@nvidia.com>

docs/design/feature/sequence_parallel.md

@@ -281,6 +281,10 @@ NOTE: be careful to test adequately when refactoring classes that take this styl
 | `expected_dims` | int \| None | Expected tensor rank for validation (optional) |
 | `split_output` | bool | `False`: shard **input** params; `True`: shard **output** tensors |
 | `auto_pad` | bool | Auto-pad if sequence not divisible by world_size (default: `False`) |
+| `clone_shard` | bool | Give multi-rank shards their own contiguous storage (default: `False`); single-rank inputs are unchanged |
+
+With `clone_shard=True`, memory savings require callers to drop their references to the full tensor.
+Other references to the same storage also keep the full allocation alive.
 
 **SequenceParallelOutput Parameters:**
 

tests/diffusion/distributed/test_sp_storage_lifetime.py

@@ -0,0 +1,239 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import weakref
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelConfig,
+    SequenceParallelInput,
+    SequenceParallelPartialInput,
+)
+from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
+from vllm_omni.diffusion.hooks.base import HookRegistry
+from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook, apply_sequence_parallel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion, pytest.mark.parallel]
+
+
+@pytest.fixture
+def parallel_state(monkeypatch):
+    from vllm_omni.diffusion.attention import selector
+    from vllm_omni.diffusion.distributed import parallel_state, sp_sharding
+
+    class MaskBackend:
+        @staticmethod
+        def supports_attention_mask(spec):
+            return True
+
+    def configure(size=4, rank=1):
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: size)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: rank)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+        monkeypatch.setattr(sp_sharding, "get_sequence_parallel_world_size", lambda: size)
+        monkeypatch.setattr(sp_sharding, "get_sequence_parallel_rank", lambda: rank)
+        monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: MaskBackend())
+
+    configure()
+    return configure
+
+
+@pytest.mark.parametrize("sequence_length", [16, 17])
+@pytest.mark.parametrize("batch", [1, 2])
+@pytest.mark.parametrize("rank", [0, 3])
+@torch.inference_mode()
+def test_cosmos3_real_sp_hook_releases_full_embedding(parallel_state, sequence_length, batch, rank):
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GenSPPrepare, Cosmos3VFMTransformer
+
+    parallel_state(rank=rank)
+    model = nn.Module()
+    model.gen_sp_prepare = Cosmos3GenSPPrepare()
+    metadata = Cosmos3VFMTransformer._sp_plan["gen_sp_prepare"]
+    assert metadata[0].clone_shard
+    assert not metadata[1].clone_shard and not metadata[2].clone_shard
+    apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=4), {"gen_sp_prepare": metadata})
+
+    def prepare():
+        hidden = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(batch, sequence_length, 8)
+        full_ref = weakref.ref(hidden)
+        full_pointer = hidden.untyped_storage().data_ptr()
+        cos = torch.ones(batch, sequence_length, 1, 2)
+        sin = torch.zeros_like(cos)
+        return model.gen_sp_prepare(hidden, cos, sin), full_ref, full_pointer
+
+    for _ in range(2):
+        with set_forward_context():
+            (hidden, cos, sin), full_ref, full_pointer = prepare()
+            # Check the real backing allocation, not only Tensor.data_ptr():
+            # nonzero-offset and inference-mode views can hide the shared base.
+            assert full_ref() is None
+            assert hidden.untyped_storage().data_ptr() != full_pointer
+            assert hidden.untyped_storage().nbytes() == hidden.numel() * hidden.element_size()
+            assert hidden.is_contiguous()
+            expected = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(
+                batch, sequence_length, 8
+            )
+            padding = -sequence_length % 4
+            expected = torch.nn.functional.pad(expected, (0, 0, 0, padding)).chunk(4, dim=1)[rank]
+            torch.testing.assert_close(hidden, expected, rtol=0, atol=0)
+            assert cos.shape[1] == sin.shape[1] == hidden.shape[1]
+
+            # Once the next layer replaces the shard, the hook must not keep
+            # even the smaller initial activation alive through later layers.
+            shard_ref = weakref.ref(hidden)
+            hidden = hidden + 1
+            assert shard_ref() is None
+
+
+@pytest.mark.parametrize("clone_shard", [False, True])
+@pytest.mark.parametrize("world_size", [1, 4])
+def test_shard_storage_policy_is_opt_in_and_skips_single_rank(parallel_state, clone_shard, world_size):
+    parallel_state(size=world_size, rank=0)
+    hook = SequenceParallelSplitHook({}, SequenceParallelConfig(ulysses_degree=4))
+    tensor = torch.arange(128).reshape(1, 16, 8)
+    spec = SequenceParallelInput(split_dim=1, clone_shard=clone_shard)
+    with set_forward_context():
+        shard = hook._prepare_sp_input(tensor, spec)
+    torch.testing.assert_close(shard, tensor.chunk(world_size, dim=1)[0])
+    owns_storage = shard.untyped_storage().data_ptr() != tensor.untyped_storage().data_ptr()
+    assert owns_storage == (clone_shard and world_size > 1)
+    if world_size == 1:
+        assert shard is tensor
+
+
+def test_owned_shard_preserves_gradients(parallel_state):
+    parallel_state(rank=1)
+    hook = SequenceParallelSplitHook({}, SequenceParallelConfig(ulysses_degree=4))
+    tensor = torch.randn(2, 16, 8, requires_grad=True)
+    with set_forward_context():
+        shard = hook._prepare_sp_input(tensor, SequenceParallelInput(split_dim=1, clone_shard=True))
+        assert shard.untyped_storage().data_ptr() != tensor.untyped_storage().data_ptr()
+        assert shard.untyped_storage().nbytes() == shard.numel() * shard.element_size()
+        shard.sum().backward()
+    expected = torch.zeros_like(tensor)
+    expected[:, 4:8] = 1
+    torch.testing.assert_close(tensor.grad, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("sequence_length", [16, 17])
+@pytest.mark.parametrize("batch", [1, 2])
+@pytest.mark.parametrize("keyword", [False, True])
+@pytest.mark.parametrize("shard_group", [None, "video"])
+@torch.inference_mode()
+def test_pre_forward_clone_owns_storage_and_preserves_shard_metadata(
+    parallel_state, sequence_length, batch, keyword, shard_group
+):
+    parallel_state(rank=1)
+    full_refs = []
+    full_pointers = []
+
+    class Module(nn.Module):
+        def forward(self, x):
+            assert x.untyped_storage().data_ptr() != full_pointers[-1]
+            assert x.untyped_storage().nbytes() == x.numel() * x.element_size()
+            assert x.is_contiguous()
+            return x
+
+    module = Module()
+    spec = SequenceParallelInput(
+        split_dim=1,
+        split_output=False,
+        auto_pad=True,
+        clone_shard=True,
+        shard_group=shard_group,
+    )
+    apply_sequence_parallel(module, SequenceParallelConfig(ulysses_degree=4), {"": {"x": spec}})
+
+    def run():
+        tensor = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(batch, sequence_length, 8)
+        full_refs.append(weakref.ref(tensor))
+        full_pointers.append(tensor.untyped_storage().data_ptr())
+        return module(x=tensor) if keyword else module(tensor)
+
+    for _ in range(2):
+        with set_forward_context():
+            ctx = get_forward_context()
+            if shard_group is not None:
+                ctx.sp_shard_metadata[shard_group] = 999  # Stale metadata must be replaced.
+            shard = run()
+            assert all(ref() is None for ref in full_refs)
+            expected = torch.arange(batch * sequence_length * 8, dtype=torch.bfloat16).reshape(
+                batch, sequence_length, 8
+            )
+            padding = -sequence_length % 4
+            expected = torch.nn.functional.pad(expected, (0, 0, 0, padding)).chunk(4, dim=1)[1]
+            torch.testing.assert_close(shard, expected, rtol=0, atol=0)
+            assert ctx.sp_shard_metadata == ({} if shard_group is None else {shard_group: sequence_length})
+            assert ctx.sp_original_seq_len == (sequence_length if padding and shard_group is None else None)
+
+
+@pytest.mark.parametrize("phase", ["success", "non_tensor", "forward_error", "split_error"])
+@pytest.mark.parametrize("keyword", [False, True])
+@torch.inference_mode()
+def test_split_hook_does_not_retain_arguments_on_any_exit(parallel_state, phase, keyword):
+    class Module(nn.Module):
+        def forward(self, x):
+            if phase == "forward_error":
+                raise RuntimeError("forward failed")
+            return None if phase == "non_tensor" else x
+
+    module = Module()
+    index = 1 if phase == "split_error" else 0
+    hook = SequenceParallelSplitHook(
+        {index: SequenceParallelInput(split_dim=1, split_output=True, clone_shard=True)},
+        SequenceParallelConfig(ulysses_degree=4),
+    )
+    HookRegistry.get_or_create(module).register_hook("sp", hook)
+    refs = []
+
+    def run():
+        tensor = torch.randn(1, 16, 8)
+        refs.append(weakref.ref(tensor))
+        return module(x=tensor) if keyword else module(tensor)
+
+    with set_forward_context():
+        if phase.endswith("error"):
+            with pytest.raises((RuntimeError, ValueError)):
+                run()
+        else:
+            result = run()
+            assert (result is None) == (phase == "non_tensor")
+    assert all(ref() is None for ref in refs)
+
+
+@pytest.mark.parametrize("keyword", [False, True])
+@pytest.mark.parametrize("tensor_source", [False, True])
+@torch.inference_mode()
+def test_partial_output_split_keeps_only_lengths_across_calls(parallel_state, keyword, tensor_source):
+    parallel_state(rank=1)
+
+    class Module(nn.Module):
+        def forward(self, x, text):
+            return x
+
+    module = Module()
+    hook = SequenceParallelSplitHook(
+        {0: SequenceParallelPartialInput(split_dim=0, text_len_source="text", split_output=True)},
+        SequenceParallelConfig(ulysses_degree=4),
+    )
+    HookRegistry.get_or_create(module).register_hook("sp", hook)
+    refs = []
+
+    def run(text_length):
+        tensor = torch.arange(text_length + 8).reshape(-1, 1)
+        text = torch.zeros(text_length, 8) if tensor_source else text_length
+        refs.append(weakref.ref(tensor))
+        if tensor_source:
+            refs.append(weakref.ref(text))
+        return module(x=tensor, text=text) if keyword else module(tensor, text)
+
+    for text_length in (2, 4):
+        with set_forward_context():
+            result = run(text_length)
+        expected = torch.cat([torch.arange(text_length), torch.arange(text_length + 2, text_length + 4)]).reshape(-1, 1)
+        torch.testing.assert_close(result, expected)
+        assert all(ref() is None for ref in refs)
+        assert hook._text_len_cache == {"text": text_length}

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import weakref
 from types import SimpleNamespace
 from typing import Any
 
@@ -824,6 +825,134 @@ def test_cache_execution_residual_spans_final_gen_norm(monkeypatch: pytest.Monke
         assert not hasattr(model, name)
 
 
+@pytest.mark.parametrize("sequence_length", [16, 17])
+@pytest.mark.parametrize("batch", [1, 2])
+@pytest.mark.parametrize("rank", [0, 3])
+@pytest.mark.parametrize("seacache", [False, True])
+@pytest.mark.parametrize("control_count", [0, 1, 2])
+@pytest.mark.parallel
+@torch.inference_mode()
+def test_sp_releases_full_gen_embedding_during_stack(
+    monkeypatch, sequence_length, batch, rank, seacache, control_count
+) -> None:
+    from vllm_omni.diffusion.attention import selector
+    from vllm_omni.diffusion.cache.seacache import SeaCacheConfig, apply_sea_cache_hook
+    from vllm_omni.diffusion.distributed import parallel_state, sp_sharding
+    from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+    from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
+    from vllm_omni.diffusion.hooks import sequence_parallel
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    full_refs = []
+    full_pointers = []
+    layer_calls = 0
+    total_length = sequence_length * (control_count + 1)
+    uses_sp = control_count <= 1
+
+    class CheckingLayer(nn.Module):
+        def forward(self, hidden, **kwargs):
+            nonlocal layer_calls
+            layer_calls += 1
+            if full_refs:
+                # Inspect lifetime inside the production stack while the
+                # caller's prepared state is still live.
+                assert full_refs[-1]() is None
+                if layer_calls % 2 == 1:
+                    assert hidden.untyped_storage().data_ptr() != full_pointers[-1]
+                assert hidden.shape == (batch, (total_length + 3) // 4, 8)
+                assert hidden.untyped_storage().nbytes() == hidden.numel() * hidden.element_size()
+                assert kwargs["freqs_cos"].shape[1] == kwargs["freqs_sin"].shape[1] == hidden.shape[1]
+            elif not uses_sp:
+                assert hidden.shape == (batch, total_length, 8)
+            return hidden + 1
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    model.gen_layers = nn.ModuleList([CheckingLayer(), CheckingLayer()])
+    model.cached_kv = [(torch.empty(0), torch.empty(0)) for _ in model.gen_layers]
+    inputs = {
+        "hidden_states": torch.randn(batch, 2, 1, 1, sequence_length),
+        "timestep": torch.ones(batch),
+        "text_ids": torch.ones(batch, 2, dtype=torch.long),
+        "text_mask": torch.ones(batch, 2, dtype=torch.long),
+        "video_shape": (1, 1, sequence_length),
+        "control_latents": [torch.randn(batch, 2, 1, 1, sequence_length) for _ in range(control_count)],
+    }
+    gathered_reference = []
+    normalized_reference = []
+    handle = model.gen_sp_gather.register_forward_pre_hook(
+        lambda module, args: gathered_reference.append(args[0].clone())
+    )
+    norm_handle = model.norm_moe_gen.register_forward_hook(
+        lambda module, args, output: normalized_reference.append(output.clone())
+    )
+    expected = model(**inputs)
+    handle.remove()
+    norm_handle.remove()
+
+    class MaskBackend:
+        @staticmethod
+        def supports_attention_mask(spec):
+            return True
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (4, rank, None))
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 4)
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: rank)
+    monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(sp_sharding, "get_sequence_parallel_world_size", lambda: 4)
+    monkeypatch.setattr(sp_sharding, "get_sequence_parallel_rank", lambda: rank)
+    monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: MaskBackend())
+
+    def gather(hidden, dim, validate):
+        shard_length = hidden.shape[1]
+        start = rank * shard_length
+        valid_length = min(shard_length, total_length - start)
+        reference = normalized_reference[0] if seacache else gathered_reference[0]
+        torch.testing.assert_close(hidden[:, :valid_length], reference[:, start : start + valid_length])
+        # Supply the other ranks' outputs without retaining the original input.
+        return torch.nn.functional.pad(reference, (0, 0, 0, -total_length % 4))
+
+    monkeypatch.setattr(sequence_parallel, "sp_gather", gather)
+    sequence_parallel.apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=4), model._sp_plan)
+
+    def record_full_embedding(module, args):
+        full_refs.append(weakref.ref(args[0]))
+        full_pointers.append(args[0].untyped_storage().data_ptr())
+
+    model.gen_sp_prepare.register_forward_pre_hook(record_full_embedding)
+    metadata = SimpleNamespace(step=0)
+    if seacache:
+        sea_hook = apply_sea_cache_hook(
+            model,
+            SeaCacheConfig(threshold=100.0, residual_order=0),
+            current_step_callback=lambda: metadata.step,
+            current_sigma_callback=lambda: 1.0,
+            num_inference_steps_callback=lambda: 4,
+        )
+        sea_hook.state_manager.set_context("cond")
+    layer_calls = 0
+    for step in range(4 if seacache else 2):
+        metadata.step = step
+        with set_forward_context():
+            output = model(**inputs)
+            torch.testing.assert_close(output, expected)
+            assert get_forward_context()._sp_shard_depth == 0
+            if seacache:
+                if uses_sp:
+                    assert full_refs[-1]() is None
+                history = sea_hook.state_manager.get_state().history
+                assert history
+                residual = history[-1][1]
+                assert residual.shape == (batch, (total_length + 3) // 4 if uses_sp else total_length, 8)
+                assert residual.untyped_storage().nbytes() == residual.numel() * residual.element_size()
+    assert layer_calls == 4
+    assert len(full_refs) == ((4 if seacache else 2) if uses_sp else 0)
+    if seacache:
+        assert sea_hook.full_count == sea_hook.skip_count == 2
+
+
 def test_no_cache_still_runs_final_gen_norm_once(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 
@@ -1394,3 +1523,28 @@ def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> No
     )
     _, offset_gen_pos = rotary.position_ids
     assert offset_gen_pos[0, 0].tolist() == [102, 103, 104, 105, 106, 107]
+
+
+def test_shard_gen_prep_rejects_inputs_already_in_execution_layout(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    prep = model._gen_preprocess(
+        torch.zeros(1, 2, 1, 2, 2),
+        torch.tensor([1.0]),
+        torch.tensor([[1, 2]], dtype=torch.long),
+        torch.ones(1, 2, dtype=torch.long),
+        (1, 2, 2),
+    )
+
+    sharded = model._shard_gen_prep(prep, defer_gather=True)
+    assert sharded.freqs_gen is not None
+    assert sharded.defer_gen_gather
+    with pytest.raises(RuntimeError, match="already in the execution layout"):
+        model._shard_gen_prep(sharded)
+    # A deferred gather without sharding would gather a full-layout tensor.
+    with pytest.raises(RuntimeError, match="_shard_gen_prep"):
+        model._run_gen_stack(prep._replace(defer_gen_gather=True))

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -1620,6 +1620,9 @@ def extract_cosmos3_context(
         control_weights=control_weights,
         transfer_share_vision_temporal_positions=transfer_share_vision_temporal_positions,
     )
+    # Cache residuals stay rank-local: shard before the context and its
+    # closures retain prep; postprocess gathers on both misses and hits.
+    prep = module._shard_gen_prep(prep, defer_gather=True)
 
     def run_transformer_blocks() -> tuple[torch.Tensor, ...]:
         return (module._run_gen_stack(prep),)

vllm_omni/diffusion/distributed/sp_plan.py

@@ -225,6 +225,11 @@ class SequenceParallelInput:
             ForwardContext so models can construct attention masks when needed.
             Note: Ring attention does not support attention mask, so auto_pad
             should only be used with Ulysses SP.
+        clone_shard: If True, give a sharded tensor its own contiguous storage
+            so it does not retain the full input allocation. Applies to both
+            padded and unpadded splits; a single-rank input is returned unchanged.
+            Memory savings require the caller to drop its reference to the full
+            tensor; any other references to its storage also keep it allocated.
         shard_group: Optional key shared by tensors representing the same global
             sequence. Keyed groups track independent padding metadata; omitting
             it preserves the legacy single-sequence padding behavior.
@@ -244,13 +249,15 @@ class SequenceParallelInput:
     expected_dims: int | None = None
     split_output: bool = False
     auto_pad: bool = False
+    clone_shard: bool = False
     shard_group: str | None = None
 
     def __repr__(self) -> str:
         return (
             f"SequenceParallelInput(split_dim={self.split_dim}, "
             f"expected_dims={self.expected_dims}, split_output={self.split_output}, "
-            f"auto_pad={self.auto_pad}, shard_group={self.shard_group!r})"
+            f"auto_pad={self.auto_pad}, clone_shard={self.clone_shard}, "
+            f"shard_group={self.shard_group!r})"
         )
 
 

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -260,9 +260,12 @@ class SequenceParallelSplitHook(ModelHook):
             else:
                 raise ValueError(f"Failed to update parameter '{name}' after sharding.")
 
-        # Store kwargs for post_forward to resolve text lengths
-        self._last_kwargs = kwargs
-        self._last_args = tuple(args_list)
+        # Output splitting only needs integer text lengths. Saving args/kwargs
+        # here would keep full hidden-state inputs alive after output sharding,
+        # including when the module raises before post_forward can run.
+        for spm in self.metadata.values():
+            if isinstance(spm, SequenceParallelPartialInput) and spm.split_output:
+                self._resolve_text_len(spm, args_list, kwargs)
 
         return tuple(args_list), kwargs
 
@@ -291,7 +294,7 @@ class SequenceParallelSplitHook(ModelHook):
                 raise ValueError(f"Index {index} out of bounds for output of length {len(output_list)}.")
 
             original = output_list[index]
-            output_list[index] = self._prepare_sp_input(original, spm, self._last_args, self._last_kwargs)
+            output_list[index] = self._prepare_sp_input(original, spm)
             if output_list[index] is not original:
                 actually_sharded = True
                 equal_rank_seq_lens &= isinstance(spm, SequenceParallelInput) and spm.auto_pad
@@ -392,13 +395,22 @@ class SequenceParallelSplitHook(ModelHook):
         if isinstance(sp_input, SequenceParallelInput):
             # Full split with optional auto-padding
             if sp_input.auto_pad:
-                return self._shard_with_auto_pad(
+                shard = self._shard_with_auto_pad(
                     x,
                     sp_input.split_dim,
                     sp_input.shard_group,
                 )
-            _maybe_validate_strict_divisibility(dim=sp_input.split_dim, seq_len=x.size(sp_input.split_dim))
-            return sp_shard(x, sp_input.split_dim, validate=False)
+            else:
+                _maybe_validate_strict_divisibility(dim=sp_input.split_dim, seq_len=x.size(sp_input.split_dim))
+                shard = sp_shard(x, sp_input.split_dim, validate=False)
+            # contiguous() may be a no-op for a single-batch sequence slice;
+            # clone explicitly to release its full (possibly padded) backing.
+            if sp_input.clone_shard:
+                from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size
+
+                if get_sequence_parallel_world_size() > 1:
+                    return shard.clone(memory_format=torch.contiguous_format)
+            return shard
         elif isinstance(sp_input, SequenceParallelPartialInput):
             # Partial split: keep text portion, split image portion
             text_len = self._resolve_text_len(sp_input, args, kwargs)

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -1182,6 +1182,9 @@ class _GenPrepared(NamedTuple):
     use_multi_control_attention: bool
     multi_control_token_sizes: tuple[int, ...] | None
     multi_control_weights: tuple[float, ...] | None
+    # Execution-layout state; set only by Cosmos3VFMTransformer._shard_gen_prep.
+    freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None
+    defer_gen_gather: bool = False
 
 
 class Cosmos3VFMTransformer(nn.Module):
@@ -1242,7 +1245,13 @@ class Cosmos3VFMTransformer(nn.Module):
 
     _sp_plan = {
         "gen_sp_prepare": {
-            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+            0: SequenceParallelInput(
+                split_dim=1,
+                expected_dims=3,
+                split_output=True,
+                auto_pad=True,
+                clone_shard=True,
+            ),
             1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
             2: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
         },
@@ -1769,6 +1778,8 @@ class Cosmos3VFMTransformer(nn.Module):
         self,
         hidden_gen: torch.Tensor,
         *,
+        freqs_gen: tuple[torch.Tensor, torch.Tensor],
+        gather_output: bool,
         s_video: int,
         s_control: int,
         s_action: int,
@@ -1781,13 +1792,10 @@ class Cosmos3VFMTransformer(nn.Module):
         multi_control_token_sizes: tuple[int, ...] | None,
         multi_control_weights: tuple[float, ...] | None,
     ) -> torch.Tensor:
-        """Run the complete GEN decoder between full-layout boundaries."""
-        if self.cached_kv is None or self.cached_freqs_gen is None:
+        """Run the complete GEN decoder on inputs already in the execution layout."""
+        if self.cached_kv is None:
             raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.")
-        freqs_cos, freqs_sin = self.cached_freqs_gen
-        if not use_multi_control_attention:
-            hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin)
-        freqs_gen = (freqs_cos, freqs_sin)
+        freqs_cos, freqs_sin = freqs_gen
 
         if len(self.gen_layers) == len(self.cached_kv):
             for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True):
@@ -1816,7 +1824,7 @@ class Cosmos3VFMTransformer(nn.Module):
                 if isinstance(hidden_gen, tuple):
                     hidden_gen = hidden_gen[0]
 
-        if not use_multi_control_attention:
+        if gather_output and not use_multi_control_attention:
             hidden_gen = self.gen_sp_gather(hidden_gen)
         return hidden_gen
 
@@ -1863,8 +1871,33 @@ class Cosmos3VFMTransformer(nn.Module):
             control_weights=control_weights,
             transfer_share_vision_temporal_positions=transfer_share_vision_temporal_positions,
         )
+        prep = self._shard_gen_prep(prep)
         return self._gen_postprocess(self._run_gen_stack(prep), prep)
 
+    def _shard_gen_prep(self, prep: _GenPrepared, *, defer_gather: bool = False) -> _GenPrepared:
+        """Move prepared GEN inputs into the execution layout.
+
+        Callers must rebind their ``prep`` to the result before running the
+        stack: keeping the original would pin the full embedding even though
+        the sharded ``hidden_gen`` owns separate storage.
+
+        ``defer_gather`` keeps the stack output rank-local and moves the SP
+        gather into ``_gen_postprocess``, so cache residuals stay sharded.
+        """
+        if prep.freqs_gen is not None or prep.defer_gen_gather:
+            raise RuntimeError("Cosmos3 GEN inputs are already in the execution layout.")
+        if self.cached_freqs_gen is None:
+            raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.")
+        if prep.use_multi_control_attention:
+            # Multi-control attention runs unsharded, so there is nothing to gather.
+            return prep._replace(freqs_gen=self.cached_freqs_gen)
+        hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(prep.hidden_gen, *self.cached_freqs_gen)
+        return prep._replace(
+            hidden_gen=hidden_gen,
+            freqs_gen=(freqs_cos, freqs_sin),
+            defer_gen_gather=defer_gather,
+        )
+
     def _gen_preprocess(
         self,
         hidden_states: torch.Tensor,
@@ -2132,9 +2165,13 @@ class Cosmos3VFMTransformer(nn.Module):
             )
 
     def _run_gen_stack(self, prep: _GenPrepared) -> torch.Tensor:
-        """Execute the cacheable full-layout GEN stack, including final norm."""
+        """Execute the cacheable GEN stack, including final norm."""
+        if prep.freqs_gen is None:
+            raise RuntimeError("Cosmos3 GEN inputs must go through _shard_gen_prep before the stack.")
         hidden_gen = self._run_gen_layers(
             prep.hidden_gen,
+            freqs_gen=prep.freqs_gen,
+            gather_output=not prep.defer_gen_gather,
             s_video=prep.s_video,
             s_control=prep.s_control,
             s_action=prep.s_action,
@@ -2155,6 +2192,8 @@ class Cosmos3VFMTransformer(nn.Module):
         prep: _GenPrepared,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Project an already-normalized packed GEN state to model outputs."""
+        if prep.defer_gen_gather:
+            hidden_gen = self.gen_sp_gather(hidden_gen)
         if not prep.has_action and not prep.has_sound and not prep.has_control:
             return self.unpatchify(self.proj_out(hidden_gen), prep.t, prep.h, prep.w)