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

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

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)