[diffusion] model: add cosmos3 transfer capability (#34747)

Co-authored-by: Kedi Wu <kediw@nvidia.com>
Co-authored-by: Kedi Wu <31940276+kediwu0331@users.noreply.github.com>

python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py

@@ -113,7 +113,7 @@ class Cosmos3Config(PipelineConfig):
     vae_precision: str = "bf16"
 
     # Pipeline-level (not sampling) knobs.
-    max_sequence_length: int = 512
+    max_sequence_length: int = 4096
     use_duration_template: bool = True
     use_system_prompt: bool = False
 

python/sglang/multimodal_gen/configs/sample/cosmos3.py

@@ -10,7 +10,7 @@ so the file extension and decode path agree.
 """
 
 from dataclasses import dataclass, field
-from typing import Any
+from typing import Any, ClassVar
 
 from sglang.multimodal_gen.configs.sample.sampling_params import (
     DataType,
@@ -74,6 +74,61 @@ class Cosmos3SamplingParams(SamplingParams):
     condition_frame_indexes: list[int] | None = None
     condition_video_keep: str = "first"
 
+    # Transfer (control-video) conditioning. ``control_path`` points to one or
+    # more pre-computed control videos (e.g. edge / blur / depth / seg / wsm
+    # maps). When set, each control clip is VAE-encoded and packed as clean
+    # vision tokens that prefix the target clip in the GEN sequence; multiple
+    # paths drive multi-hint transfer (e.g. edge + depth). Control clips reuse
+    # ``proj_in``, so every Cosmos3 checkpoint supports transfer.
+    control_path: str | list[str] | None = None
+
+    # Optional hint type(s) parallel to ``control_path`` (one of
+    # ``edge`` / ``blur`` / ``depth`` / ``seg`` / ``wsm``). Used only to apply
+    # tuned per-hint defaults (``guidance`` / ``control_guidance`` / ``shift``)
+    # when exactly one control input is given and the user left those unset.
+    control_hint: str | list[str] | None = None
+
+    # Control-CFG scale for transfer. ``1.0`` (default) disables the extra
+    # control-dropped forward; values > 1.0 amplify the control map's influence
+    # by blending the with-control and without-control predictions on the
+    # generated span: ``cond_nc + control_guidance * (cond_full - cond_nc)``.
+    control_guidance: float = 1.0
+
+    # Optional timestep window ``(lo, hi)`` restricting where control-CFG is
+    # applied (analogous to ``guidance_interval`` for text CFG). ``None`` applies
+    # it at every step.
+    control_guidance_interval: tuple[float, float] | None = None
+
+    # Long-video transfer controls. Chunks overlap by
+    # ``num_conditional_frames`` pixel frames; overlap frames from the previous
+    # decoded chunk are kept clean in the next chunk.
+    num_video_frames_per_chunk: int = 93
+    num_conditional_frames: int = 1
+    num_first_chunk_conditional_frames: int = 0
+    max_frames: int = 5000
+    show_control_condition: bool = False
+    show_input: bool = False
+    share_vision_temporal_positions: bool = True
+
+    # Tuned per-hint defaults applied when exactly one control input is given
+    # and the corresponding field was not set explicitly (mirrors the
+    # cosmos-framework ``_TRANSFER_DEFAULTS`` table). ``shift`` maps to
+    # ``flow_shift``. Multi-hint transfer keeps the request's own values.
+    _TRANSFER_DEFAULTS: ClassVar[dict[str, dict[str, float | int]]] = {
+        "edge": {"guidance": 3.0, "control_guidance": 1.5, "shift": 10.0},
+        "blur": {"guidance": 3.0, "control_guidance": 1.5, "shift": 10.0},
+        "depth": {"guidance": 3.0, "control_guidance": 1.5, "shift": 10.0},
+        "seg": {"guidance": 3.0, "control_guidance": 2.0, "shift": 10.0},
+        "wsm": {
+            "guidance": 1.0,
+            "control_guidance": 3.0,
+            "shift": 10.0,
+            "num_frames": 101,
+            "fps": 10,
+            "num_video_frames_per_chunk": 101,
+        },
+    }
+
     supported_resolutions: list[tuple[int, int]] | None = field(
         default_factory=lambda: [
             (1280, 720),
@@ -103,6 +158,108 @@ class Cosmos3SamplingParams(SamplingParams):
     action_stats_path: str | None = None
     action_normalization: str = "quantile"
 
+    @classmethod
+    def video_request_extra_fields(cls) -> frozenset[str]:
+        return frozenset(
+            {
+                "generate_sound",
+                "sound_duration",
+                "condition_frame_indexes",
+                "condition_frame_indexes_vision",
+                "condition_video_keep",
+                "control_path",
+                "control_hint",
+                "control_guidance",
+                "control_guidance_interval",
+                "num_video_frames_per_chunk",
+                "num_conditional_frames",
+                "num_first_chunk_conditional_frames",
+                "max_frames",
+                "show_control_condition",
+                "show_input",
+                "share_vision_temporal_positions",
+                "action_mode",
+                "domain_id",
+                "domain_name",
+                "raw_action_dim",
+                "action_fps",
+                "action",
+                "action_view_point",
+                "action_normalization",
+            }
+        )
+
+    def _resolve_control_paths(self) -> list[str]:
+        cp = self.control_path
+        if cp is None:
+            return []
+        if isinstance(cp, str):
+            return [cp] if cp else []
+        return [p for p in cp if isinstance(p, str) and p]
+
+    def _resolve_control_hints(self) -> list[str]:
+        hint = self.control_hint
+        if hint is None:
+            return []
+        hints = [hint] if isinstance(hint, str) else list(hint)
+        hints = [h for h in hints if h]
+        for h in hints:
+            if h not in self._TRANSFER_DEFAULTS:
+                raise ValueError(
+                    f"Unknown control_hint {h!r}; expected one of "
+                    f"{sorted(self._TRANSFER_DEFAULTS)}"
+                )
+        return hints
+
+    def _apply_transfer_hint_defaults(self) -> None:
+        """Fill tuned per-hint defaults for a single, typed control input.
+
+        Mirrors cosmos-framework: defaults apply only when there is exactly one
+        control input with a known hint type, and only to fields the user did
+        not pass explicitly (tracked via ``_explicit_fields``). Multi-hint
+        transfer keeps the request's own ``guidance`` / ``control_guidance`` /
+        ``flow_shift``.
+        """
+        if len(self._resolve_control_paths()) != 1:
+            return
+        hints = self._resolve_control_hints()
+        if len(hints) != 1:
+            return
+        defaults = self._TRANSFER_DEFAULTS.get(hints[0])
+        if defaults is None:
+            return
+        explicit = getattr(self, "_explicit_fields", None) or set()
+        if "control_guidance" not in explicit:
+            self.control_guidance = defaults["control_guidance"]
+        if "guidance_scale" not in explicit:
+            self.guidance_scale = defaults["guidance"]
+        if "flow_shift" not in explicit and self.flow_shift is None:
+            self.flow_shift = defaults["shift"]
+        for field_name in (
+            "num_frames",
+            "fps",
+            "num_video_frames_per_chunk",
+        ):
+            if field_name in defaults and field_name not in explicit:
+                setattr(self, field_name, defaults[field_name])
+
+    @classmethod
+    def lower_video_request_kwargs(
+        cls, request: Any, kwargs: dict[str, Any]
+    ) -> dict[str, Any]:
+        """Apply defaults that the generic video endpoint pre-resolves."""
+        hint = kwargs.get("control_hint")
+        paths = kwargs.get("control_path")
+        hints = [hint] if isinstance(hint, str) else list(hint or [])
+        control_paths = [paths] if isinstance(paths, str) else list(paths or [])
+        if len(control_paths) == 1 and hints == ["wsm"]:
+            defaults = cls._TRANSFER_DEFAULTS["wsm"]
+            if getattr(request, "num_frames", None) is None:
+                kwargs["num_frames"] = defaults["num_frames"]
+            if getattr(request, "fps", None) is None:
+                kwargs["fps"] = defaults["fps"]
+        return kwargs
+
     def _adjust(self, server_args) -> None:
         # adjust distil and edge args — read from the pre-computed config fields
         # so no checkpoint download happens at request time.
@@ -130,6 +287,37 @@ class Cosmos3SamplingParams(SamplingParams):
                 )
             action_output = self.action_mode != "forward_dynamics"
 
+        # Apply transfer per-hint defaults before the base resolves remaining
+        # fields (e.g. flow_shift per mode), so an unset flow_shift can pick up
+        # the hint's tuned shift.
+        self._apply_transfer_hint_defaults()
+        control_paths = self._resolve_control_paths()
+        if control_paths:
+            if pipeline_config.distilled_sigmas is not None:
+                raise ValueError(
+                    "Cosmos3 distilled checkpoints do not support transfer inference"
+                )
+            if pipeline_config.is_edge:
+                raise ValueError(
+                    "Cosmos3 Edge checkpoints do not support transfer inference"
+                )
+            if self.num_frames == 1:
+                raise ValueError(
+                    "Cosmos3 transfer inference is supported only for video outputs"
+                )
+            if self.image_path is not None:
+                raise ValueError(
+                    "Cosmos3 transfer accepts control videos and an optional source "
+                    "video, not an image input"
+                )
+            if self.action_mode is not None:
+                raise ValueError(
+                    "Cosmos3 transfer cannot be combined with action generation"
+                )
+            if float(self.sound_duration or 0.0) > 0.0:
+                raise ValueError(
+                    "Cosmos3 transfer cannot be combined with sound generation"
+                )
         super()._adjust(server_args)
 
         # Policy and inverse dynamics produce actions. Forward dynamics consumes
@@ -142,6 +330,39 @@ class Cosmos3SamplingParams(SamplingParams):
             self.output_file_name = None
             self.output_compression = 0
 
+    def _validate(self) -> None:
+        super()._validate()
+        paths = self._resolve_control_paths()
+        hints = self._resolve_control_hints()
+        if hints and len(hints) != len(paths):
+            raise ValueError(
+                "control_hint must contain exactly one entry per control_path "
+                f"(got {len(hints)} hint(s) for {len(paths)} path(s))"
+            )
+        if self.control_guidance_interval is not None:
+            if len(self.control_guidance_interval) != 2:
+                raise ValueError(
+                    "control_guidance_interval must contain exactly two values"
+                )
+            lo, hi = self.control_guidance_interval
+            if float(lo) > float(hi):
+                raise ValueError(
+                    "control_guidance_interval must be ordered as (low, high)"
+                )
+        if self.num_video_frames_per_chunk <= 0:
+            raise ValueError("num_video_frames_per_chunk must be positive")
+        if self.num_conditional_frames < 0:
+            raise ValueError("num_conditional_frames must be non-negative")
+        if self.num_conditional_frames >= self.num_video_frames_per_chunk:
+            raise ValueError(
+                "num_conditional_frames must be smaller than "
+                "num_video_frames_per_chunk"
+            )
+        if self.num_first_chunk_conditional_frames < 0:
+            raise ValueError("num_first_chunk_conditional_frames must be non-negative")
+        if self.max_frames <= 0:
+            raise ValueError("max_frames must be positive")
+
     def _guidance_is_explicit(self) -> bool:
         explicit = getattr(self, "_explicit_fields", None)
         return explicit is not None and "guidance_scale" in explicit

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -117,19 +117,6 @@ _MULTIPART_EXTRA_FORM_FIELDS = (
     "guardrails",
     "video_path",
     "video_url",
-    "generate_sound",
-    "sound_duration",
-    "condition_frame_indexes",
-    "action_mode",
-    "domain_id",
-    "domain_name",
-    "raw_action_dim",
-    "action_fps",
-    "action",
-    "action_view_point",
-    "action_normalization",
-    "condition_frame_indexes_vision",
-    "condition_video_keep",
     "quality",
 )
 
@@ -265,6 +252,49 @@ def _coerce_optional_int_list(value: Any) -> list[int] | None:
     return [int(value)]
 
 
+def _coerce_optional_float_list(value: Any) -> list[float] | None:
+    value = _parse_form_extra_value(value)
+    if value is None:
+        return None
+    if isinstance(value, str) and not value.strip():
+        return None
+    if isinstance(value, (list, tuple)):
+        return [float(item) for item in value]
+    return [float(value)]
+
+
+def _coerce_optional_bool(value: Any) -> bool | None:
+    value = _parse_form_extra_value(value)
+    if value is None or (isinstance(value, str) and not value.strip()):
+        return None
+    if isinstance(value, str):
+        normalized = value.strip().lower()
+        if normalized in {"1", "true", "yes", "on"}:
+            return True
+        if normalized in {"0", "false", "no", "off"}:
+            return False
+        raise ValueError(f"Invalid boolean value: {value!r}")
+    return bool(value)
+
+
+def _coerce_optional_str_list(value: Any) -> str | list[str] | None:
+    """Coerce a control_path/control_hint value to str or list[str].
+
+    Accepts a JSON list (``["edge.mp4", "depth.mp4"]``), a plain string, or a
+    native list. Empty values resolve to ``None`` so unset fields don't override
+    sampling-param defaults.
+    """
+    value = _parse_form_extra_value(value)
+    if value is None:
+        return None
+    if isinstance(value, (list, tuple)):
+        items = [str(item) for item in value if str(item).strip()]
+        return items or None
+    if isinstance(value, str):
+        return value if value.strip() else None
+    return str(value)
+
+
 def _resolve_video_path(req: VideoGenerationsRequest) -> str | None:
     video_path = _request_value(req, "video_path") or _request_value(req, "video_url")
     if video_path:
@@ -324,6 +354,41 @@ def _cosmos3_sampling_param_kwargs(
     if condition_frame_indexes is not None:
         kwargs["condition_frame_indexes"] = condition_frame_indexes
 
+    # Transfer (control-video) conditioning.
+    control_path = _coerce_optional_str_list(_request_value(req, "control_path"))
+    if control_path is not None:
+        kwargs["control_path"] = control_path
+    control_hint = _coerce_optional_str_list(_request_value(req, "control_hint"))
+    if control_hint is not None:
+        kwargs["control_hint"] = control_hint
+    control_guidance = _request_value(req, "control_guidance")
+    if control_guidance is not None:
+        kwargs["control_guidance"] = float(control_guidance)
+    control_guidance_interval = _coerce_optional_float_list(
+        _request_value(req, "control_guidance_interval")
+    )
+    if control_guidance_interval is not None:
+        kwargs["control_guidance_interval"] = tuple(control_guidance_interval)
+
+    for name in (
+        "num_video_frames_per_chunk",
+        "num_conditional_frames",
+        "num_first_chunk_conditional_frames",
+        "max_frames",
+    ):
+        value = _parse_form_extra_value(_request_value(req, name))
+        if value is not None and value != "":
+            kwargs[name] = int(value)
+
+    for name in (
+        "show_control_condition",
+        "show_input",
+        "share_vision_temporal_positions",
+    ):
+        value = _coerce_optional_bool(_request_value(req, name))
+        if value is not None:
+            kwargs[name] = value
+
     for name in (
         "condition_video_keep",
         "action_mode",

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -206,6 +206,13 @@ def _maybe_dequantize_fp8(
     return full_tensor
 
 
+def _move_to_device_preserving_meta(model: nn.Module, device: torch.device) -> None:
+    # Buffers absent from the checkpoint (e.g. cosmos3's RoPE inv_freq) are
+    # still on the meta device here and .to() cannot copy out of meta; leave
+    # them for the model's post_load_weights() to rebuild on the real device.
+    model._apply(lambda t: t if t.is_meta else t.to(device))
+
+
 def register_fsdp_entrypoints(model: torch.nn.Module) -> None:
     """Let FSDP2 unshard around forward passes that bypass ``__call__``.
 
@@ -435,7 +442,7 @@ def maybe_load_fsdp_model(
     # 3. postprocessing
     if weight_postprocess_device is not None:
         # move to device to perform postprocessing
-        model.to(weight_postprocess_device)
+        _move_to_device_preserving_meta(model, weight_postprocess_device)
 
     for _, module in model.named_modules():
         quant_method = getattr(module, "quant_method", None)

python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py

@@ -7,7 +7,7 @@ cross-attends from noisy visual tokens to that cache at every denoising step.
 """
 
 import math
-from collections.abc import Iterable, Iterator
+from collections.abc import Iterable, Iterator, Sequence
 from typing import Any
 
 import torch
@@ -1333,12 +1333,30 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
         action_frames: int = 0,
         action_fps: float | None = None,
         action_start_frame_offset: int = 1,
+        control_frames: int | Sequence[int] = 0,
+        share_vision_temporal_positions: bool = True,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Compute mRoPE position IDs for UND text and GEN visual + action + sound tokens."""
+        """Compute mRoPE position IDs for UND text and GEN tokens.
+
+        The GEN sequence is ordered ``[control, video, action, sound]``.
+        Control and target videos either share matching temporal coordinates or
+        occupy consecutive temporal ranges, according to
+        ``share_vision_temporal_positions``.
+        """
         B = text_mask.shape[0]
         S_text = text_mask.shape[1]
         text_lengths = text_mask.sum(dim=1).long()
         effective_fps = fps if fps is not None and T > 1 else None
+        control_frame_counts = (
+            [control_frames]
+            if isinstance(control_frames, int) and control_frames > 0
+            else (
+                [int(count) for count in control_frames if int(count) > 0]
+                if not isinstance(control_frames, int)
+                else []
+            )
+        )
+        vision_frame_counts = [*control_frame_counts, T]
 
         text_pos_list = []
         vis_pos_list = []
@@ -1348,16 +1366,28 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
                 real_len, temporal_offset=0, device=device
             )
             media_offset = t_offset + self.temporal_margin
-            v_pos, _ = compute_mrope_position_ids_vision(
-                T,
-                Hp,
-                Wp,
-                temporal_offset=media_offset,
-                device=device,
-                fps=effective_fps,
-                base_fps=self.base_fps,
-                temporal_compression_factor=self.temporal_compression_factor,
-            )
+            vision_offset = media_offset
+            vision_pos_blocks = []
+            for frame_count in vision_frame_counts:
+                temporal_offset = (
+                    media_offset if share_vision_temporal_positions else vision_offset
+                )
+                vision_pos, vision_offset = compute_mrope_position_ids_vision(
+                    frame_count,
+                    Hp,
+                    Wp,
+                    temporal_offset=temporal_offset,
+                    device=device,
+                    fps=effective_fps,
+                    base_fps=self.base_fps,
+                    temporal_compression_factor=self.temporal_compression_factor,
+                )
+                vision_pos_blocks.append(vision_pos)
+
+            pos_dtype = vision_pos_blocks[0].dtype
+            for pos in vision_pos_blocks[1:]:
+                pos_dtype = torch.promote_types(pos_dtype, pos.dtype)
+            v_pos = torch.cat([pos.to(pos_dtype) for pos in vision_pos_blocks], dim=1)
             if action_frames > 0:
                 a_pos, _ = compute_mrope_position_ids_action(
                     action_frames,
@@ -1394,9 +1424,8 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
             text_pos_list.append(t_pos)
             vis_pos_list.append(v_pos)
 
-        text_pos_ids = torch.stack(text_pos_list, dim=1).to(device)  # [3, B, S_text]
-        vis_pos_ids = torch.stack(vis_pos_list, dim=1).to(device)  # [3, B, S_gen]
-
+        text_pos_ids = torch.stack(text_pos_list, dim=1).to(device)
+        vis_pos_ids = torch.stack(vis_pos_list, dim=1).to(device)
         return text_pos_ids, vis_pos_ids
 
     def reset_cache(self, cache_key: str | None = None):
@@ -1443,6 +1472,8 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
         action_noisy_mask: torch.Tensor | None = None,
         action_fps: float | None = None,
         action_start_frame_offset: int = 1,
+        control_latents: torch.Tensor | list[torch.Tensor] | None = None,
+        transfer_share_vision_temporal_positions: bool = True,
         **kwargs,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Forward pass for denoising.
@@ -1473,6 +1504,11 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
                 Defaults to the video fps when None.
             action_start_frame_offset: Temporal offset applied to action
                 position IDs relative to the video's media_offset (default 1).
+            control_latents: Optional [B, C, T_ctrl, H, W] control-video latents
+                (transfer / control-net conditioning). They are patchified and
+                projected with the shared ``proj_in``, prepended to the GEN
+                sequence as clean (noise-free) tokens that share the video's
+                temporal positions, and excluded from the output projection.
 
         Returns:
             [B, C, T, H, W] velocity prediction, or a tuple
@@ -1505,8 +1541,44 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
                     device=action_latents.device,
                 )
 
+        # Transfer / control-video conditioning: one or more control clips
+        # (e.g. edge + depth) are packed as clean vision tokens that prefix the
+        # target clip in the GEN sequence. Each clip reuses ``proj_in`` and the
+        # shared transformer; blocks are concatenated in input order.
+        control_clips: list[torch.Tensor] = []
+        if control_latents is not None:
+            control_clips = (
+                list(control_latents)
+                if isinstance(control_latents, (list, tuple))
+                else [control_latents]
+            )
+        control_frame_counts: list[int] = []
+        hidden_control_blocks: list[torch.Tensor] = []
+        control_token_len = 0
+        for clip in control_clips:
+            _, _, c_frames, Hc, Wc = clip.shape
+            if (Hc, Wc) != (H, W):
+                raise ValueError(
+                    "control_latents spatial dims "
+                    f"{(Hc, Wc)} must match hidden_states {(H, W)}"
+                )
+            block, _ = self.proj_in(
+                self.patchify(clip.to(hidden_states.dtype), c_frames, Hc, Wc)
+            )
+            hidden_control_blocks.append(block)
+            control_frame_counts.append(c_frames)
+            control_token_len += block.shape[1]
+        has_control = len(hidden_control_blocks) > 0
+        hidden_control = (
+            torch.cat(hidden_control_blocks, dim=1) if has_control else None
+        )
+
         extra_frames = action_frames + sound_frames
         sequence_shard_enabled = self.sp_size > 1
+        # When a control clip is present we always assemble the combined GEN
+        # stream (control prefix + video [+ action] [+ sound]) instead of the
+        # video-only fast path.
+        use_assembly_path = extra_frames > 0 or has_control
 
         # Add timestep embedding (computed in float32 for numerical stability, then cast back)
         time_embed = self.time_embedder(timestep.float())
@@ -1531,7 +1603,7 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
                 .to(hidden_gen.dtype)
             )
 
-        if extra_frames == 0:
+        if not use_assembly_path:
             # Video-only: shard the visual tokens, then add the timestep
             # embedding on the local shard.
             if sequence_shard_enabled:
@@ -1567,15 +1639,21 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
                 hidden_gen = hidden_gen + time_embed.unsqueeze(1)
         else:
             # Multi-modal: assemble the full GEN sequence
-            # (video[, action][, sound]) with timestep embeddings, then shard
-            # the combined stream so sequence parallelism splits every modality
-            # evenly. The per-modality output heads run after the post-loop
-            # all-gather reassembles the sequence.
+            # ([control,] video[, action][, sound]) with timestep embeddings,
+            # then shard the combined stream so sequence parallelism splits
+            # every modality evenly. The per-modality output heads run after the
+            # post-loop all-gather reassembles the sequence.
             if token_noisy_mask is not None:
                 hidden_gen = hidden_gen + time_embed.unsqueeze(1) * token_noisy_mask
             else:
                 hidden_gen = hidden_gen + time_embed.unsqueeze(1)
 
+            # Control tokens are clean conditioning: prepend them WITHOUT a
+            # timestep embedding so the GEN tokens can attend to the raw control
+            # map. They are stripped before the output projection.
+            if has_control:
+                hidden_gen = torch.cat([hidden_control, hidden_gen], dim=1)
+
             if action_latents is not None:
                 hidden_action = self.action_proj_in(
                     action_latents.to(hidden_gen.dtype), action_domain_ids
@@ -1648,6 +1726,10 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
                 action_frames=action_frames,
                 action_fps=action_fps if action_fps is not None else fps,
                 action_start_frame_offset=action_start_frame_offset,
+                control_frames=control_frame_counts,
+                share_vision_temporal_positions=(
+                    transfer_share_vision_temporal_positions
+                ),
             )
             # UND K/V cache is kept FULL on all ranks (not sharded). Text
             # sequence is short, so memory impact is minimal, and the GEN
@@ -1722,7 +1804,7 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
         hidden_gen = hidden_gen + residual
         hidden_gen = self.norm_moe_gen(hidden_gen)
 
-        if extra_frames == 0:
+        if not use_assembly_path:
             # Video-only: project on the local shard and gather the (much
             # smaller) patch-space output. With patch_latent_dim ~=
             # hidden_size / 21 for cosmos3, this cuts the post-loop SP
@@ -1741,12 +1823,15 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
             if seq_shard_pad > 0:
                 hidden_gen = hidden_gen[:, :seq_len_orig, :]
 
-        s_video = seq_len_orig - extra_frames
-        output, _ = self.proj_out(hidden_gen[:, :s_video, :])
+        # Sequence layout: [control prefix | video | action | sound]. Control
+        # tokens are conditioning only and produce no output.
+        s_video = seq_len_orig - extra_frames - control_token_len
+        video_start = control_token_len
+        output, _ = self.proj_out(hidden_gen[:, video_start : video_start + s_video, :])
         video_pred = self.unpatchify(output, T, H, W)
 
         extra_outputs: list[torch.Tensor] = []
-        idx = s_video
+        idx = video_start + s_video
         if action_frames > 0:
             action_hidden = hidden_gen[:, idx : idx + action_frames, :]
             extra_outputs.append(self.action_proj_out(action_hidden, action_domain_ids))
@@ -1756,6 +1841,10 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
             sound_output, _ = self.audio_proj_out(sound_hidden)
             extra_outputs.append(sound_output.permute(0, 2, 1).contiguous())
 
+        # Control-only conditioning (no action/sound): keep the bare-tensor
+        # return type identical to the video-only path.
+        if not extra_outputs:
+            return video_pred
         return (video_pred, *extra_outputs)
 
     def preprocess_loaded_state_dict(

python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py

@@ -100,7 +100,11 @@ class Cosmos3Pipeline(LoRAPipeline, ComposedPipelineBase):
             self.add_stage(Cosmos3TextGuardrailStage())
         self.add_stage(Cosmos3LatentPreparationStage(vae, transformer))
         self.add_stage(Cosmos3TimestepPreparationStage(scheduler))
-        self.add_stage(Cosmos3DenoisingStage(transformer, scheduler, server_args))
+        self.add_stage(
+            Cosmos3DenoisingStage(
+                transformer, scheduler, server_args=server_args, vae=vae
+            )
+        )
         self.add_stage(
             Cosmos3DecodingStage(
                 vae, guardrails=guardrails_on, sound_tokenizer=sound_tokenizer

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py

@@ -10,12 +10,17 @@ Enabled by default when available; opt out with
 
 from __future__ import annotations
 
+import hashlib
 import importlib.util
+import os
+import shutil
 from functools import lru_cache
+from pathlib import Path
 
 import numpy as np
 import torch
 
+from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
     PipelineStage,
@@ -35,6 +40,68 @@ def is_cosmos_guardrail_available() -> bool:
     return importlib.util.find_spec("cosmos_guardrail") is not None
 
 
+def _mirror_symlinked_nltk_data() -> None:
+    """Make the guardrail's nltk_data readable under NLTK's hardened opener.
+
+    ``CosmosSafetyChecker`` registers its HF-hub snapshot's
+    ``blocklist/nltk_data`` directory on ``nltk.data.path``. Hub snapshot files
+    are symlinks into the blob store, which NLTK builds that ship the
+    ``pathsec`` hardened opener refuse to follow (O_NOFOLLOW, CWE-59 TOCTOU
+    guard) — every text-safety check then fails with "refusing to follow a
+    symlink at open time". Mirror each symlink-containing search entry to a
+    plain-file copy and register the mirror ahead of the original, so NLTK's
+    data lookup resolves to real files first. No-op for NLTK builds without
+    the hardened opener and for search entries that are already plain files.
+    """
+    try:
+        import nltk.data
+    except ImportError:
+        return
+
+    mirror_root = Path(envs.SGLANG_DIFFUSION_CACHE_ROOT) / "nltk_data_deref"
+    for entry in list(nltk.data.path):
+        try:
+            root = Path(entry)
+            if not root.is_dir():
+                continue
+            if not any(p.is_symlink() for p in root.rglob("*")):
+                continue
+            mirror = mirror_root / hashlib.sha256(str(root).encode()).hexdigest()[:16]
+            if str(mirror) in nltk.data.path:
+                continue
+            if not mirror.is_dir():
+                # Every GPU worker process runs this at pipeline construction on
+                # a shared filesystem, so stage under a per-process name and
+                # publish with an atomic rename; whichever process publishes
+                # first wins and the others adopt its mirror.
+                staging = mirror.with_name(f"{mirror.name}.{os.getpid()}.tmp")
+                shutil.rmtree(staging, ignore_errors=True)
+                # symlinks=False dereferences: the copy holds real file contents.
+                shutil.copytree(root, staging, symlinks=False)
+                try:
+                    staging.rename(mirror)
+                except OSError:
+                    shutil.rmtree(staging, ignore_errors=True)
+                    if not mirror.is_dir():
+                        raise
+            nltk.data.path.insert(nltk.data.path.index(entry), str(mirror))
+            logger.info(
+                "Mirrored symlinked nltk_data %s -> %s (hardened-NLTK compatibility)",
+                root,
+                mirror,
+            )
+        except OSError as exc:
+            # Best-effort: an unwritable cache root must not break guardrail
+            # init. Hardened-NLTK builds may still fail at check time; plain
+            # NLTK builds work fine without the mirror.
+            logger.warning(
+                "Could not mirror symlinked nltk_data %s under %s: %s",
+                entry,
+                mirror_root,
+                exc,
+            )
+
+
 def _init_guardrails(offload_to_cpu: bool = False) -> None:
     global _checker
     if _checker is not None:
@@ -50,6 +117,7 @@ def _init_guardrails(offload_to_cpu: bool = False) -> None:
         "Initializing Cosmos3 guardrails (offload_to_cpu=%s) ...", offload_to_cpu
     )
     _checker = CosmosSafetyChecker()
+    _mirror_symlinked_nltk_data()
     idle_device = "cpu" if offload_to_cpu else current_platform.device_type
     for runner in (_checker.text_guardrail, _checker.video_guardrail):
         if runner is None or not hasattr(runner, "models"):

python/sglang/multimodal_gen/test/unit/test_cosmos3.py

@@ -18,7 +18,10 @@ from sglang.multimodal_gen.configs.sample.cosmos3 import (
     COSMOS3_EDGE_SUPPORTED_RESOLUTIONS,
     Cosmos3SamplingParams,
 )
-from sglang.multimodal_gen.configs.sample.sampling_params import DataType
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
 from sglang.multimodal_gen.registry import (
     _PIPELINE_REGISTRY,
     _discover_and_register_pipelines,
@@ -36,6 +39,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
 )
 from sglang.multimodal_gen.runtime.entrypoints.openai.video_api import (
     _cosmos3_sampling_param_kwargs,
+    _multipart_video_extras,
     _resolve_sound_duration,
     _resolve_video_path,
 )
@@ -45,6 +49,7 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.scheduler_loader imp
 )
 from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
 from sglang.multimodal_gen.runtime.models.dits.cosmos3video import (
+    Cosmos3OmniTransformer,
     DomainAwareLinear,
     _can_enable_t1_fused_qk_norm_rope,
     compute_mrope_position_ids_action,
@@ -59,6 +64,8 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.c
     Cosmos3TimestepPreparationStage,
     Cosmos3TokenizationStage,
     _inject_caption_metadata,
+    _pad_transfer_frames,
+    _resize_center_crop_uint8_cthw,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_action import (
     EMBODIMENT_TO_DOMAIN_ID,
@@ -877,7 +884,7 @@ class TestCosmos3ModelResolution(unittest.TestCase):
 
 
 class TestCosmos3OpenAIProtocol(unittest.TestCase):
-    """Verify Cosmos3 modality knobs are exposed by the video HTTP schema."""
+    """Verify Cosmos3 modality knobs stay model-specific video extras."""
 
     def test_cosmos3_template_fields_remain_extra_fields(self):
         for request_cls in (ImageGenerationsRequest, VideoGenerationsRequest):
@@ -889,31 +896,45 @@ class TestCosmos3OpenAIProtocol(unittest.TestCase):
                 self.assertNotIn("use_system_prompt", request_cls.model_fields)
                 self.assertNotIn("use_guardrails", request_cls.model_fields)
 
-    def test_cosmos3_modal_fields_pass_through_as_extras(self):
-        for field_name in ("video_path", "video_url"):
-            with self.subTest(field_name=field_name):
-                self.assertIn(field_name, VideoGenerationsRequest.model_fields)
-
-        modal_values = {
-            "generate_sound": True,
-            "sound_duration": 3.0,
-            "condition_frame_indexes": [0, 2],
-            "condition_frame_indexes_vision": [0, 2],
-            "condition_video_keep": "last",
-            "action_mode": "policy",
-            "domain_id": 1,
-            "domain_name": "umi",
-            "raw_action_dim": 9,
-            "action_fps": 30.0,
-            "action": [0.0, 1.0],
-            "action_view_point": "ego_view",
-            "action_normalization": "mean_std",
-        }
-        req = VideoGenerationsRequest(prompt="test", **modal_values)
-        for field_name, value in modal_values.items():
+    def test_cosmos3_modal_fields_are_model_specific_video_extras(self):
+        for field_name in (
+            "generate_sound",
+            "sound_duration",
+            "condition_frame_indexes",
+            "condition_frame_indexes_vision",
+            "condition_video_keep",
+            "control_path",
+            "control_hint",
+            "control_guidance",
+            "control_guidance_interval",
+            "num_video_frames_per_chunk",
+            "num_conditional_frames",
+            "num_first_chunk_conditional_frames",
+            "max_frames",
+            "show_control_condition",
+            "show_input",
+            "share_vision_temporal_positions",
+            "action_mode",
+            "domain_id",
+            "domain_name",
+            "raw_action_dim",
+            "action_fps",
+            "action",
+            "action_view_point",
+            "action_normalization",
+        ):
             with self.subTest(field_name=field_name):
                 self.assertNotIn(field_name, VideoGenerationsRequest.model_fields)
-                self.assertEqual(getattr(req, field_name), value)
+                self.assertIn(
+                    field_name, Cosmos3SamplingParams.video_request_extra_fields()
+                )
+
+        self.assertIn("video_path", VideoGenerationsRequest.model_fields)
+        self.assertIn("video_url", VideoGenerationsRequest.model_fields)
+        self.assertNotIn("action_stats_path", VideoGenerationsRequest.model_fields)
+        self.assertNotIn(
+            "action_stats_path", Cosmos3SamplingParams.video_request_extra_fields()
+        )
 
     def test_cosmos3_http_aliases_map_to_sampling_params(self):
         req = VideoGenerationsRequest(
@@ -927,6 +948,17 @@ class TestCosmos3OpenAIProtocol(unittest.TestCase):
             raw_action_dim=9,
             action_fps=30.0,
             action_view_point="ego_view",
+            control_path=["edge.mp4", "depth.mp4"],
+            control_hint=["edge", "depth"],
+            control_guidance=1.5,
+            control_guidance_interval=[0.0, 500.0],
+            num_video_frames_per_chunk=97,
+            num_conditional_frames=5,
+            num_first_chunk_conditional_frames=2,
+            max_frames=1200,
+            show_control_condition="true",
+            show_input="false",
+            share_vision_temporal_positions="false",
         )
 
         self.assertEqual(_resolve_video_path(req), "https://example.com/input.mp4")
@@ -940,6 +972,48 @@ class TestCosmos3OpenAIProtocol(unittest.TestCase):
         self.assertEqual(kwargs["raw_action_dim"], 9)
         self.assertEqual(kwargs["action_fps"], 30.0)
         self.assertEqual(kwargs["action_view_point"], "ego_view")
+        self.assertEqual(kwargs["control_path"], ["edge.mp4", "depth.mp4"])
+        self.assertEqual(kwargs["control_hint"], ["edge", "depth"])
+        self.assertEqual(kwargs["control_guidance"], 1.5)
+        self.assertEqual(kwargs["control_guidance_interval"], (0.0, 500.0))
+        self.assertEqual(kwargs["num_video_frames_per_chunk"], 97)
+        self.assertEqual(kwargs["num_conditional_frames"], 5)
+        self.assertEqual(kwargs["num_first_chunk_conditional_frames"], 2)
+        self.assertEqual(kwargs["max_frames"], 1200)
+        self.assertTrue(kwargs["show_control_condition"])
+        self.assertFalse(kwargs["show_input"])
+        self.assertFalse(kwargs["share_vision_temporal_positions"])
+
+    def test_cosmos3_multipart_extras_are_model_specific(self):
+        raw_form = {
+            "generate_sound": "true",
+            "control_path": '["edge.mp4", "depth.mp4"]',
+            "control_hint": '["edge", "depth"]',
+            "action_mode": "policy",
+            "action_stats_path": "/tmp/action_stats.json",
+        }
+
+        generic = _multipart_video_extras(
+            raw_form,
+            extra_body=None,
+            extra_params=None,
+            sampling_params_cls=SamplingParams,
+        )
+        self.assertNotIn("generate_sound", generic)
+        self.assertNotIn("control_path", generic)
+        self.assertNotIn("action_mode", generic)
+
+        cosmos = _multipart_video_extras(
+            raw_form,
+            extra_body=None,
+            extra_params=None,
+            sampling_params_cls=Cosmos3SamplingParams,
+        )
+        self.assertIs(cosmos["generate_sound"], True)
+        self.assertEqual(cosmos["control_path"], ["edge.mp4", "depth.mp4"])
+        self.assertEqual(cosmos["control_hint"], ["edge", "depth"])
+        self.assertEqual(cosmos["action_mode"], "policy")
+        self.assertNotIn("action_stats_path", cosmos)
 
     def test_generate_sound_false_disables_sound_duration(self):
         req = VideoGenerationsRequest(
@@ -1069,6 +1143,403 @@ class TestCosmos3MRoPE(unittest.TestCase):
         self.assertAlmostEqual(float(vid[0, 1]), float(act[0, 4]), places=4)
 
 
+class TestCosmos3Transfer(unittest.TestCase):
+    """Transfer (control-video) conditioning: rope packing, control-CFG, defaults."""
+
+    DEVICE = torch.device("cpu")
+
+    @staticmethod
+    def _transformer_self() -> Cosmos3OmniTransformer:
+        model = Cosmos3OmniTransformer.__new__(Cosmos3OmniTransformer)
+        model.temporal_margin = 15000
+        model.base_fps = 24.0
+        model.temporal_compression_factor = 4
+        model.sound_latent_fps = 25.0
+        model.temporal_compression_factor_sound = 1
+        return model
+
+    def test_single_control_shares_video_positions(self):
+        model = self._transformer_self()
+        text_mask = torch.ones(1, 4)
+        T, Hp, Wp = 3, 2, 2
+        tpc = T * Hp * Wp
+
+        _, base = model._compute_rope_position_ids(
+            text_mask, T, Hp, Wp, fps=None, device=self.DEVICE, control_frames=0
+        )
+        _, with_ctrl = model._compute_rope_position_ids(
+            text_mask, T, Hp, Wp, fps=None, device=self.DEVICE, control_frames=T
+        )
+        self.assertEqual(tuple(with_ctrl.shape), (3, 1, 2 * tpc))
+        # control prefix == video block, and video block is unchanged.
+        self.assertTrue(torch.equal(with_ctrl[:, :, :tpc], with_ctrl[:, :, tpc:]))
+        self.assertTrue(torch.equal(with_ctrl[:, :, tpc:], base))
+
+    def test_multi_control_prepended_in_order(self):
+        model = self._transformer_self()
+        text_mask = torch.ones(1, 4)
+        T, Hp, Wp = 3, 2, 2
+        tpc = T * Hp * Wp
+
+        _, base = model._compute_rope_position_ids(
+            text_mask, T, Hp, Wp, fps=None, device=self.DEVICE, control_frames=0
+        )
+        _, multi = model._compute_rope_position_ids(
+            text_mask, T, Hp, Wp, fps=None, device=self.DEVICE, control_frames=[T, T]
+        )
+        self.assertEqual(tuple(multi.shape), (3, 1, 3 * tpc))
+        c0 = multi[:, :, :tpc]
+        c1 = multi[:, :, tpc : 2 * tpc]
+        vid = multi[:, :, 2 * tpc :]
+        self.assertTrue(torch.equal(c0, vid))
+        self.assertTrue(torch.equal(c1, vid))
+        self.assertTrue(torch.equal(vid, base))
+
+    def test_control_positions_can_be_sequential(self):
+        model = self._transformer_self()
+        text_mask = torch.ones(1, 4)
+        T, Hp, Wp = 3, 1, 1
+
+        _, positions = model._compute_rope_position_ids(
+            text_mask,
+            T,
+            Hp,
+            Wp,
+            fps=None,
+            device=self.DEVICE,
+            control_frames=[T, T],
+            share_vision_temporal_positions=False,
+        )
+        c0 = positions[0, 0, :T]
+        c1 = positions[0, 0, T : 2 * T]
+        video = positions[0, 0, 2 * T :]
+        self.assertLess(c0.max().item(), c1.min().item())
+        self.assertLess(c1.max().item(), video.min().item())
+
+    def test_transfer_chunk_count_and_reflection_padding(self):
+        self.assertEqual(
+            Cosmos3ImagePreprocessStage._get_transfer_num_chunks(93, 93, 1),
+            (1, 93),
+        )
+        self.assertEqual(
+            Cosmos3ImagePreprocessStage._get_transfer_num_chunks(186, 93, 1),
+            (3, 92),
+        )
+        frames = torch.tensor([0, 1, 2], dtype=torch.uint8).view(1, 1, 3, 1, 1)
+        frames = frames.expand(1, 3, 3, 1, 1)
+        padded = _pad_transfer_frames(frames, 5)
+        self.assertEqual(padded[0, 0, :, 0, 0].tolist(), [0, 1, 2, 2, 1])
+
+    def test_transfer_resize_matches_vllm_omni(self):
+        frames = torch.arange(3 * 2 * 2 * 3, dtype=torch.uint8).reshape(3, 2, 2, 3)
+
+        actual = _resize_center_crop_uint8_cthw(frames, height=3, width=3)
+
+        resized = torch.nn.functional.interpolate(
+            frames.permute(1, 0, 2, 3).float(),
+            size=(3, 5),
+            mode="bilinear",
+            align_corners=False,
+        )
+        expected = (
+            resized[:, :, :, 1:4]
+            .round()
+            .clamp(0, 255)
+            .to(torch.uint8)
+            .permute(1, 0, 2, 3)
+            .contiguous()
+        )
+        self.assertTrue(torch.equal(actual, expected))
+        self.assertEqual(tuple(actual.shape), (3, 2, 3, 3))
+        self.assertEqual(actual.dtype, torch.uint8)
+
+    def test_transfer_chunks_stitch_without_duplicate_overlap(self):
+        class Scheduler:
+            def __init__(self):
+                self.timesteps = torch.tensor([])
+                self.calls = 0
+
+            def set_timesteps(self, steps, device):
+                self.calls += 1
+                self.timesteps = torch.arange(steps, device=device)
+
+        stage = Cosmos3DenoisingStage.__new__(Cosmos3DenoisingStage)
+        stage.vae = torch.nn.Linear(1, 1, bias=False)
+        stage.transformer = torch.nn.Linear(1, 1, bias=False)
+        stage.scheduler = Scheduler()
+        stage._prepare_transfer_chunk = mock.Mock(side_effect=[0, 1])
+        stage._denoise_once = mock.Mock(side_effect=lambda batch, *_a, **_k: batch)
+        stage._decode_transfer_latents = mock.Mock(
+            side_effect=[
+                torch.arange(5).view(1, 1, 5, 1, 1).float() / 10,
+                torch.arange(5, 10).view(1, 1, 5, 1, 1).float() / 10,
+            ]
+        )
+        generator = torch.Generator(device="cpu").manual_seed(7)
+        batch = types.SimpleNamespace(
+            latents=torch.zeros(1),
+            generator=generator,
+            seed=7,
+            num_inference_steps=2,
+            is_warmup=False,
+            extra={
+                "transfer_plan": {
+                    "num_chunks": 2,
+                    "total_frames": 9,
+                }
+            },
+        )
+        server_args = types.SimpleNamespace(vae_cpu_offload=False)
+        stage.server_args = server_args
+
+        stage._forward_transfer(batch, server_args)
+
+        stitched = batch.extra["transfer_decoded_output"].flatten()
+        self.assertTrue(
+            torch.allclose(
+                stitched,
+                torch.tensor([0.0, 0.1, 0.2, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9]),
+            )
+        )
+        self.assertEqual(stage.scheduler.calls, 2)
+        for call in stage._prepare_transfer_chunk.call_args_list:
+            self.assertIs(call.args[3], generator)
+        for call in stage._denoise_once.call_args_list:
+            self.assertIs(call.kwargs["generator"], generator)
+
+    def test_transfer_display_composes_input_controls_and_output(self):
+        output = torch.full((1, 3, 2, 1, 2), 0.5)
+        control = torch.full((1, 3, 2, 1, 2), 255, dtype=torch.uint8)
+        source = torch.zeros((1, 3, 2, 1, 2), dtype=torch.uint8)
+        batch = types.SimpleNamespace(
+            sampling_params=types.SimpleNamespace(
+                show_control_condition=True,
+                show_input=True,
+            ),
+            extra={
+                "transfer_plan": {"total_frames": 2},
+                "preprocessed_control": [control],
+                "preprocessed_transfer_video": source,
+            },
+        )
+
+        composed = Cosmos3DecodingStage._compose_transfer_display(output, batch)
+
+        self.assertEqual(tuple(composed.shape), (1, 3, 2, 1, 6))
+        self.assertTrue(torch.equal(composed[..., :2], torch.zeros_like(output)))
+        self.assertTrue(torch.equal(composed[..., 2:4], torch.ones_like(output)))
+        self.assertTrue(torch.equal(composed[..., 4:], output))
+
+    def test_control_cfg_blend_math(self):
+        stage = Cosmos3DenoisingStage.__new__(Cosmos3DenoisingStage)
+
+        # Per-branch forward values (each run un-batched, bs=1): cond_full
+        # (control in, cond text) -> 20; cond_nc (control dropped) -> 10; uncond
+        # (control in, uncond text) -> 2. The unified executor reduces the
+        # coefficient-weighted branch sum built by _control_cfg_branches.
+        def fake_run(**kw):
+            bs = kw["latents"].shape[0]
+            if kw["control_latents"] is None:
+                return torch.full((bs,), 10.0)  # cond_nc, control dropped
+            if kw["cache_key"] == "uncond":
+                return torch.full((bs,), 2.0)  # uncond, control in
+            return torch.full((bs,), 20.0)  # cond_full, control in
+
+        stage._run_transformer = fake_run
+        cond_text_ids = torch.zeros(1)
+        cond_text_mask = torch.ones(1)
+        uncond_text_ids = torch.zeros(1)
+        uncond_text_mask = torch.ones(1)
+        control_latents = [torch.zeros(1)]
+
+        def _run(text_g, control_g):
+            branches = stage._control_cfg_branches(
+                cond_text_ids,
+                cond_text_mask,
+                uncond_text_ids,
+                uncond_text_mask,
+                cond_text_seq_len=None,
+                uncond_text_seq_len=None,
+                control_latents=control_latents,
+                text_guidance_scale=text_g,
+                control_guidance_scale=control_g,
+            )
+            return stage._predict_noise_cfg(
+                branches,
+                latents=torch.zeros(1),
+                timestep=torch.zeros(1),
+                video_shape=(1, 1, 1),
+                fps=24.0,
+                cfg_rank=0,
+                cfg_world_size=1,
+            )
+
+        # control-CFG only (g=1): cg*cond_full + (1-cg)*cond_nc
+        #   = 2*20 + (-1)*10 = 30
+        out = _run(text_g=1.0, control_g=2.0)
+        self.assertTrue(torch.allclose(out, torch.full((1,), 30.0)))
+
+        # control-CFG + text CFG (g=3): g*cg*cond_full + g*(1-cg)*cond_nc
+        #   + (1-g)*uncond = 6*20 + (-3)*10 + (-2)*2 = 86
+        out2 = _run(text_g=3.0, control_g=2.0)
+        self.assertTrue(torch.allclose(out2, torch.full((1,), 86.0)))
+
+    def test_text_cfg_batched_single_gpu(self):
+        """Single-GPU text CFG batches both branches into one bs=2 forward."""
+        stage = Cosmos3DenoisingStage.__new__(Cosmos3DenoisingStage)
+
+        def fake_run(**kw):
+            self.assertIsNone(kw["control_latents"])
+            self.assertEqual(kw["latents"].shape[0], 2)  # batched [uncond, cond]
+            return torch.tensor([2.0, 20.0])
+
+        stage._run_transformer = fake_run
+        out = stage._predict_noise_cfg_batched(
+            latents=torch.zeros(1),
+            timestep=torch.zeros(1),
+            cond_text_ids=torch.zeros(1),
+            cond_text_mask=torch.ones(1),
+            uncond_text_ids=torch.zeros(1),
+            uncond_text_mask=torch.ones(1),
+            video_shape=(1, 1, 1),
+            fps=24.0,
+            guidance_scale=3.0,
+        )
+        # uncond + g*(cond - uncond) = 2 + 3*(20-2) = 56
+        self.assertTrue(torch.allclose(out, torch.full((1,), 56.0)))
+
+    def test_control_cfg_parallel_distribution(self):
+        """Round-robin branch distribution + all-reduce equals the serial blend."""
+        stage = Cosmos3DenoisingStage.__new__(Cosmos3DenoisingStage)
+        seen = {}
+
+        def make_run(rank):
+            seen[rank] = []
+
+            def fake_run(**kw):
+                seen[rank].append(kw["cache_key"])
+                bs = kw["latents"].shape[0]
+                if kw["control_latents"] is None:
+                    return torch.full((bs,), 10.0)  # cond_nc
+                if kw["cache_key"] == "uncond":
+                    return torch.full((bs,), 2.0)  # uncond
+                return torch.full((bs,), 20.0)  # cond_full
+
+            return fake_run
+
+        def run(rank, world):
+            stage._run_transformer = make_run(rank)
+            branches = stage._control_cfg_branches(
+                torch.zeros(1),
+                torch.ones(1),
+                torch.zeros(1),
+                torch.ones(1),
+                cond_text_seq_len=None,
+                uncond_text_seq_len=None,
+                control_latents=[torch.zeros(1)],
+                text_guidance_scale=3.0,
+                control_guidance_scale=2.0,
+            )
+            return stage._predict_noise_cfg(
+                branches,
+                latents=torch.zeros(1),
+                timestep=torch.zeros(1),
+                video_shape=(1, 1, 1),
+                fps=24.0,
+                cfg_rank=rank,
+                cfg_world_size=world,
+            )
+
+        target = (
+            "sglang.multimodal_gen.runtime.pipelines_core.stages."
+            "model_specific_stages.cosmos3.cfg_model_parallel_all_reduce"
+        )
+        # Stand in for the all-reduce with identity so each rank returns its own
+        # partial; summing them mimics the cross-rank reduction.
+        with mock.patch(target, side_effect=lambda x: x):
+            # 2 ranks: rank 0 runs the two control-in forwards (cond_full +
+            # uncond), rank 1 runs only the control-dropped cond_nc.
+            r0 = run(0, 2)
+            r1 = run(1, 2)
+            self.assertEqual(seen[0], ["cond", "uncond"])
+            self.assertEqual(seen[1], ["cond_nc"])
+            self.assertTrue(torch.allclose(r0 + r1, torch.full((1,), 86.0)))
+
+            seen.clear()
+            # 4 ranks for 3 branches: rank 3 is idle and contributes zeros.
+            parts = [run(r, 4) for r in range(4)]
+            self.assertEqual(seen[0], ["cond"])
+            self.assertEqual(seen[1], ["cond_nc"])
+            self.assertEqual(seen[2], ["uncond"])
+            self.assertEqual(seen[3], [])
+            self.assertTrue(torch.allclose(parts[3], torch.zeros(1)))
+            self.assertTrue(torch.allclose(sum(parts), torch.full((1,), 86.0)))
+
+    def test_single_hint_defaults_applied(self):
+        sp = Cosmos3SamplingParams(
+            prompt="t", control_path="edge.mp4", control_hint="edge"
+        )
+        sp._explicit_fields = {"prompt", "control_path", "control_hint"}
+        sp._apply_transfer_hint_defaults()
+        self.assertEqual(sp.control_guidance, 1.5)
+        self.assertEqual(sp.guidance_scale, 3.0)
+        self.assertEqual(sp.flow_shift, 10.0)
+
+    def test_wsm_uses_transfer_spec_defaults(self):
+        sp = Cosmos3SamplingParams(
+            prompt="t", control_path="wsm.mp4", control_hint="wsm"
+        )
+        sp._explicit_fields = {"prompt", "control_path", "control_hint"}
+        sp._apply_transfer_hint_defaults()
+        self.assertEqual(sp.guidance_scale, 1.0)
+        self.assertEqual(sp.control_guidance, 3.0)
+        self.assertEqual(sp.flow_shift, 10.0)
+        self.assertEqual(sp.num_frames, 101)
+        self.assertEqual(sp.fps, 10)
+        self.assertEqual(sp.num_video_frames_per_chunk, 101)
+
+        lowered = Cosmos3SamplingParams.lower_video_request_kwargs(
+            types.SimpleNamespace(num_frames=None, fps=None),
+            {
+                "control_path": "wsm.mp4",
+                "control_hint": "wsm",
+                "num_frames": 121,
+                "fps": 24,
+            },
+        )
+        self.assertEqual(lowered["num_frames"], 101)
+        self.assertEqual(lowered["fps"], 10)
+
+    def test_explicit_value_overrides_hint_default(self):
+        sp = Cosmos3SamplingParams(
+            prompt="t",
+            control_path="seg.mp4",
+            control_hint="seg",
+            control_guidance=4.2,
+        )
+        sp._explicit_fields = {"control_path", "control_hint", "control_guidance"}
+        sp._apply_transfer_hint_defaults()
+        self.assertEqual(sp.control_guidance, 4.2)
+
+    def test_multi_hint_defaults_not_applied(self):
+        sp = Cosmos3SamplingParams(
+            prompt="t",
+            control_path=["edge.mp4", "depth.mp4"],
+            control_hint=["edge", "depth"],
+        )
+        base_guidance = sp.guidance_scale
+        sp._explicit_fields = {"control_path", "control_hint"}
+        sp._apply_transfer_hint_defaults()
+        self.assertEqual(sp.control_guidance, 1.0)
+        self.assertEqual(sp.guidance_scale, base_guidance)
+
+    def test_unknown_hint_rejected(self):
+        with self.assertRaises(ValueError):
+            Cosmos3SamplingParams(
+                prompt="t", control_path="x.mp4", control_hint="bogus"
+            )
+
+
 class TestCosmos3DomainAwareLinear(unittest.TestCase):
     """Per-domain action projection."""
 

python/sglang/multimodal_gen/test/unit/test_fsdp_load.py

@@ -48,6 +48,23 @@ class _CustomEntrypointModel(_UniformDtypeModel):
         pass
 
 
+class _MetaBufferModel(_UniformDtypeModel):
+    """Mirrors cosmos3: a non-checkpoint buffer stays on meta until
+    post_load_weights() rebuilds it."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.proj = ReplicatedLinear(4, 4, bias=False)
+        self.register_buffer("inv_freq", torch.empty(2), persistent=False)
+
+    def post_load_weights(self) -> None:
+        if self.inv_freq.is_meta:
+            device = next(self.parameters()).device
+            self.register_buffer(
+                "inv_freq", torch.ones(2, device=device), persistent=False
+            )
+
+
 class TestFSDPMixedPrecisionPolicy(unittest.TestCase):
     def test_quant_config_detection_uses_the_runtime_instance(self):
         self.assertTrue(fsdp_load._is_bitsandbytes_quant_config(BitsAndBytesConfig()))
@@ -218,6 +235,35 @@ class TestOrdinaryWeightLoading(unittest.TestCase):
         torch.testing.assert_close(model.proj.weight, checkpoint_weight)
 
 
+class TestDevicePostprocessMove(unittest.TestCase):
+    def test_postprocess_move_preserves_meta_buffers_for_post_load_weights(self):
+        # The pre-postprocess device move must not copy buffers that are
+        # still on meta awaiting post_load_weights() (cosmos3's RoPE inv_freq).
+        load_plan = WeightLoadPlan(
+            checkpoint_load_device=torch.device("cpu"),
+            weight_postprocess_device=torch.device("cpu"),
+        )
+        checkpoint_weight = torch.arange(16, dtype=torch.float32).reshape(4, 4)
+
+        with patch.object(fsdp_load.current_platform, "is_mps", return_value=False):
+            model = fsdp_load.maybe_load_fsdp_model(
+                model_cls=_MetaBufferModel,
+                init_params={},
+                weight_dir_list=[],
+                device=torch.device("cpu"),
+                hsdp_replicate_dim=1,
+                hsdp_shard_dim=1,
+                param_dtype=torch.float32,
+                reduce_dtype=torch.float32,
+                weight_load_plan=load_plan,
+                weights_iterator=iter((("proj.weight", checkpoint_weight),)),
+            )
+
+        self.assertFalse(model.inv_freq.is_meta)
+        torch.testing.assert_close(model.inv_freq, torch.ones(2))
+        torch.testing.assert_close(model.proj.weight, checkpoint_weight)
+
+
 class TestRankLocalSafetensorsRead(unittest.TestCase):
     def _source(
         self,

python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Serialized ModelOpt FP8 checkpoints must postprocess on device even under
+layerwise offload: requantize_with_max_scale() runs scaled_fp8_quant(), a
+CUDA-only kernel, so a CPU-resident postprocess must never come back."""
+
+import unittest
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    maybe_init_distributed_environment_and_model_parallel,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.layers.linear import MergedColumnParallelLinear
+from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
+    ModelOptFp8Config,
+)
+from sglang.multimodal_gen.runtime.loader import fsdp_load
+from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
+    ensure_distributed_env_defaults,
+)
+
+_IN_FEATURES = 32
+_SHARD_OUT = 16
+
+
+class _FusedFp8Model(nn.Module):
+    """Minimal stand-in for a serialized ModelOpt FP8 DiT: one fused linear
+    with per-shard scales, plus a cosmos3-style non-checkpoint meta buffer."""
+
+    param_names_mapping = {}
+    _fsdp_forward_methods: tuple[str, ...] = ()
+
+    def __init__(self, quant_config: ModelOptFp8Config) -> None:
+        super().__init__()
+        self.qkv = MergedColumnParallelLinear(
+            input_size=_IN_FEATURES,
+            output_sizes=[_SHARD_OUT, _SHARD_OUT],
+            bias=False,
+            quant_config=quant_config,
+            prefix="qkv",
+        )
+        self.register_buffer("inv_freq", torch.empty(4), persistent=False)
+
+    def post_load_weights(self) -> None:
+        if self.inv_freq.is_meta:
+            device = next(self.parameters()).device
+            self.register_buffer(
+                "inv_freq",
+                torch.ones(4, dtype=torch.float32, device=device),
+                persistent=False,
+            )
+
+
+def _make_serialized_fp8_checkpoint() -> tuple[dict[str, torch.Tensor], torch.Tensor]:
+    """Quantize a reference weight shard-by-shard, as ModelOpt exports do."""
+    torch.manual_seed(0)
+    weight_ref = torch.randn(2 * _SHARD_OUT, _IN_FEATURES, dtype=torch.float32) * 0.05
+
+    fp8_max = torch.finfo(torch.float8_e4m3fn).max
+    shard_scales = []
+    fp8_shards = []
+    for shard in weight_ref.split(_SHARD_OUT, dim=0):
+        scale = shard.abs().max() / fp8_max
+        shard_scales.append(scale)
+        fp8_shards.append((shard / scale).to(torch.float8_e4m3fn))
+
+    state_dict = {
+        "qkv.weight": torch.cat(fp8_shards, dim=0),
+        "qkv.weight_scale": torch.stack(shard_scales),
+        "qkv.input_scale": torch.tensor([0.5, 0.5], dtype=torch.float32),
+    }
+    return state_dict, weight_ref
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA for scaled_fp8_quant")
+class TestModelOptFp8LayerwiseOffloadLoad(unittest.TestCase):
+    def test_serialized_checkpoint_loads_with_component_starting_on_cpu(self):
+        if not model_parallel_is_initialized():
+            ensure_distributed_env_defaults()
+            maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
+
+        state_dict, weight_ref = _make_serialized_fp8_checkpoint()
+        expected_max_scale = state_dict["qkv.weight_scale"].max()
+
+        # The plan a layerwise-offload component gets when
+        # _needs_device_weight_postprocess() returns True: load and postprocess
+        # on GPU, then defer the CPU placement.
+        load_plan = WeightLoadPlan.for_component(
+            checkpoint_load_device=torch.device("cuda"),
+            needs_device_weight_postprocess=True,
+            component_starts_on_cpu=True,
+        )
+        self.assertTrue(load_plan.defer_cpu_placement)
+
+        model = fsdp_load.maybe_load_fsdp_model(
+            model_cls=_FusedFp8Model,
+            init_params={
+                "quant_config": ModelOptFp8Config(is_checkpoint_fp8_serialized=True)
+            },
+            weight_dir_list=[],
+            device=torch.device("cuda"),
+            hsdp_replicate_dim=1,
+            hsdp_shard_dim=1,
+            param_dtype=torch.bfloat16,
+            reduce_dtype=torch.float32,
+            component_starts_on_cpu=True,
+            weight_load_plan=load_plan,
+            weights_iterator=iter(state_dict.items()),
+        )
+
+        # Postprocess ran: the weight was requantized to the shared max scale
+        # and rebound transposed.
+        weight = model.qkv.weight
+        self.assertEqual(weight.dtype, torch.float8_e4m3fn)
+        self.assertEqual(tuple(weight.shape), (_IN_FEATURES, 2 * _SHARD_OUT))
+        weight_scale = model.qkv.weight_scale
+        torch.testing.assert_close(
+            weight_scale.flatten(),
+            expected_max_scale.expand(weight_scale.numel()),
+            check_device=False,
+        )
+        torch.testing.assert_close(
+            model.qkv.input_scale.flatten().max(), torch.tensor(0.5), check_device=False
+        )
+
+        # The round trip through both quantizations stays close to the source.
+        # Loose on purpose: this guards against garbage (wrong scale, wrong
+        # shard order), not fp8 precision.
+        dequant = weight.t().float().cpu() * expected_max_scale
+        torch.testing.assert_close(
+            dequant, weight_ref, rtol=0.5, atol=float(expected_max_scale) * 8
+        )
+
+        # Layerwise offload contract: the component lands on CPU afterwards,
+        # with the non-checkpoint buffer rebuilt.
+        self.assertEqual(weight.device.type, "cpu")
+        self.assertFalse(model.inv_freq.is_meta)
+        self.assertEqual(model.inv_freq.device.type, "cpu")
+
+
+if __name__ == "__main__":
+    unittest.main()

python/sglang/multimodal_gen/test/unit/test_transformer_quant.py

@@ -895,12 +895,19 @@ class TestTransformerQuantHelpers(unittest.TestCase):
 
         warning.assert_called_once()
 
-    def test_modelopt_fp8_serialized_checkpoint_needs_device_postprocess(self):
+    def test_modelopt_fp8_always_needs_device_weight_postprocess(self):
+        # Even a serialized checkpoint requantizes fused shards through
+        # scaled_fp8_quant(), which cannot process CPU tensors.
         self.assertTrue(
             _needs_device_weight_postprocess(
                 ModelOptFp8Config(is_checkpoint_fp8_serialized=True)
             )
         )
+        self.assertTrue(
+            _needs_device_weight_postprocess(
+                ModelOptFp8Config(is_checkpoint_fp8_serialized=False)
+            )
+        )
 
     def test_online_fp8_needs_device_weight_postprocess(self):
         self.assertTrue(_needs_device_weight_postprocess(Fp8Config()))