[diffusion] feat: support action output for cosmos3 (#27168)

Co-authored-by: Kedi Wu <kediw@nvidia.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

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

@@ -25,23 +25,26 @@ def _build_cosmos3_param_names_mapping() -> dict:
         norm_moe_gen.weight                                    -> norm_moe_gen.weight
         time_embedder.linear_{1,2}.weight                      -> (pass-through)
         proj_in.weight, proj_out.weight                        -> (pass-through)
+        vae2llm.weight                                         -> proj_in.weight  (FP8 ckpt alias)
+        llm2vae.weight                                         -> proj_out.weight (FP8 ckpt alias)
 
     GEN patterns (`*_moe_gen`, `add_*`, `to_add_out`, `norm_added_*`) must
     precede the UND catch-all so the catch-all can't claim GEN keys.
     `norm.weight` and `lm_head.weight` are inherited from Qwen3-VL
-    pretraining and not used at inference; audio/action keys are reserved
+    pretraining and not used at inference; both are skipped.
-    for a future modality extension — all skipped via empty-string replacement.
+    Audio and action keys (``audio_proj_*``, ``action_proj_*``, modality
+    embeds) pass through unchanged.
     """
     return {
         # Inherited from Qwen3-VL pretraining; unused at diffusion inference.
         r"^lm_head\.weight$": "",
         r"^norm\.weight$": "",
-        # Audio / action modalities — not yet wired; skip to avoid load warnings.
-        r"^audio_.*$": "",
-        r"^action_.*$": "",
         # Top-level norms / embeddings.
         r"^norm_moe_gen\.(.*)$": r"norm_moe_gen.\1",
         r"^embed_tokens\.(.*)$": r"language_model.embed_tokens.\1",
+        # FP8 checkpoint aliases for the latent projection layers.
+        r"^vae2llm\.(.*)$": r"proj_in.\1",
+        r"^llm2vae\.(.*)$": r"proj_out.\1",
         # GEN pathway: per-layer (must run before the UND catch-all below).
         # Q/K/V merge into MergedColumnParallelLinear to_qkv (concat order: Q, K, V).
         r"^layers\.(\d+)\.self_attn\.add_q_proj\.(.*)$": (
@@ -151,6 +154,17 @@ class Cosmos3VideoArchConfig(DiTArchConfig):
     temporal_compression_factor: int = 4
     unified_3d_mrope_temporal_modality_margin: int = 15000
 
+    # Audio (sound) modality
+    sound_gen: bool = False
+    sound_dim: int = 64
+    sound_latent_fps: float = 25.0
+    temporal_compression_factor_sound: int = 1
+
+    # Action modality
+    action_gen: bool = False
+    action_dim: int = 64
+    num_embodiment_domains: int = 32
+
     # Timestep embedding
     timestep_scale: float = 0.001
     frequency_embedding_size: int = 256
@@ -167,6 +181,11 @@ class Cosmos3VideoArchConfig(DiTArchConfig):
     )
     reverse_param_names_mapping: dict = field(default_factory=dict)
     lora_param_names_mapping: dict = field(default_factory=dict)
+    # FP8 checkpoint quantization_config.ignore uses checkpoint module names;
+    # translate them to model names so is_layer_excluded matches correctly.
+    quant_ignore_remap: dict = field(
+        default_factory=lambda: {"vae2llm": "proj_in", "llm2vae": "proj_out"}
+    )
 
     def __post_init__(self):
         super().__post_init__()

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

@@ -195,6 +195,7 @@ class PipelineConfig:
     cfg_policy: CFGPolicy = field(default_factory=CFGPolicy)
     generator_device: str | None = None
     flow_shift: float | None = None
+    scheduler_class_override: str | None = None
     disable_autocast: bool = False
 
     # Model configuration
@@ -711,6 +712,13 @@ class PipelineConfig:
             default=PipelineConfig.flow_shift,
             help="Flow shift parameter",
         )
+        parser.add_argument(
+            f"--{prefix_with_dot}scheduler-class-override",
+            type=str,
+            dest=f"{prefix_with_dot.replace('-', '_')}scheduler_class_override",
+            default=PipelineConfig.scheduler_class_override,
+            help="Override the scheduler class from scheduler_config.json.",
+        )
         parser.add_argument(
             f"--{prefix_with_dot}resolution",
             type=int,

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

@@ -40,7 +40,11 @@ class Cosmos3Config(PipelineConfig):
     vae_tiling: bool = False
     vae_sp: bool = False
 
-    # Sourced from scheduler_config.json in the checkpoint.
+    # Cosmos3 reference inference uses FlowUniPC even when the checkpoint
+    # scheduler_config.json advertises a different scheduler class.
+    scheduler_class_override: str | None = "FlowUniPCMultistepScheduler"
+
+    # Per-request mode defaults are applied in Cosmos3TimestepPreparationStage.
     flow_shift: float | None = None
 
     precision: str = "bf16"
@@ -51,6 +55,11 @@ class Cosmos3Config(PipelineConfig):
     use_duration_template: bool = True
     use_system_prompt: bool = False
 
+    # Filesystem path to dataset-derived action stats (JSON) for action
+    # (de)normalization. Set at server launch rather than per request, since it
+    # names a server-side file. ``None`` disables normalization.
+    action_stats_path: str | None = None
+
     def __post_init__(self):
         self.vae_config.arch_config.z_dim = 48
         # Encoder is needed for I2V; T2V/T2I never invoke it.

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

@@ -97,6 +97,12 @@ class WanT2V480PConfig(PipelineConfig):
             auto_dit_layerwise_offload=True,
         )
 
+    def get_pos_prompt_embeds(self, batch):
+        return batch.prompt_embeds[0]
+
+    def get_neg_prompt_embeds(self, batch):
+        return batch.negative_prompt_embeds[0]
+
 
 @dataclass
 class TurboWanT2V480PConfig(WanT2V480PConfig):

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

@@ -1,13 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 """Cosmos3 sampling parameters.
 
-A single ``SamplingParams`` class serves T2V, I2V, and T2I — the per-request
+A single ``SamplingParams`` class serves T2V, I2V, V2V, T2I, and
-mode is dispatched in the pipeline from ``num_frames`` (``== 1`` → T2I) and
+action-conditioned variants.  Per-request mode is dispatched in the pipeline
-``image_path`` (set → I2V). For ``num_frames == 1`` the output ``data_type``
+from ``num_frames`` (``== 1`` → T2I), ``image_path`` (set → I2V),
-flips to ``IMAGE`` so the file extension and decode path agree.
+``video_path`` (set → V2V), and ``action_mode`` (set → action-conditioned).
+For ``num_frames == 1`` the output ``data_type`` flips to ``IMAGE``
+so the file extension and decode path agree.
 """
 
 from dataclasses import dataclass, field
+from typing import Any
 
 from sglang.multimodal_gen.configs.sample.sampling_params import (
     DataType,
@@ -17,7 +20,7 @@ from sglang.multimodal_gen.configs.sample.sampling_params import (
 
 @dataclass
 class Cosmos3SamplingParams(SamplingParams):
-    """Cosmos3 sampling parameters (T2V defaults; also used for I2V / T2I)."""
+    """Cosmos3 sampling parameters (T2V defaults; also used for I2V / V2V / T2I)."""
 
     height: int = 720
     width: int = 1280
@@ -30,9 +33,16 @@ class Cosmos3SamplingParams(SamplingParams):
     negative_prompt: str = ""
 
     # Optional CFG window — T2I requests typically pass e.g. ``(400, 1000)`` to
-    # skip guidance at low noise levels. T2V / I2V leave it unset.
+    # skip guidance at low noise levels. T2V / I2V / V2V leave it unset.
     guidance_interval: tuple[float, float] | None = None
 
+    # V2V conditioning: which latent-frame indices stay locked to the input
+    # video. ``None`` resolves to ``[0]`` for I2V (single frame) and ``[0, 1]``
+    # for V2V. ``condition_video_keep`` controls whether the first or last
+    # source frames are used when the input video is longer than needed.
+    condition_frame_indexes: list[int] | None = None
+    condition_video_keep: str = "first"
+
     supported_resolutions: list[tuple[int, int]] | None = field(
         default_factory=lambda: [
             (1280, 720),
@@ -43,6 +53,24 @@ class Cosmos3SamplingParams(SamplingParams):
         ]
     )
 
+    # Action modality (requires action_gen=True in the model checkpoint)
+    # action_mode: "forward_dynamics" | "policy" | "inverse_dynamics"
+    action_mode: str | None = None
+    domain_id: int | None = None
+    domain_name: str | None = None
+    raw_action_dim: int | None = None
+    action_fps: float | None = None
+    # Action data for forward_dynamics: [T, D] nested list (API) or JSON string
+    # (CLI via --action). Ignored by the other action modes.
+    action: Any = None
+    # Viewpoint phrasing for the structured action caption.
+    action_view_point: str = "ego_view"
+    # Optional dataset-derived action stats (JSON) for (de)normalization. When
+    # set, input actions are normalized and predicted actions de-normalized
+    # into physical units with ``action_normalization``.
+    action_stats_path: str | None = None
+    action_normalization: str = "quantile"
+
     def _set_output_file_name(self) -> None:
         # The pipeline config's ``task_type=TI2V`` drives ``data_type`` to
         # VIDEO, but a single-frame request is a T2I and must pick the IMAGE

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

@@ -103,6 +103,9 @@ class SamplingParams:
     # Image inputs
     image_path: str | list[str] | None = None
 
+    # Video inputs (video-to-video conditioning)
+    video_path: str | list[str] | None = None
+
     # Text inputs
     prompt: str | list[str] | None = field(
         default=None, metadata={"batch_sig_exclude": True}
@@ -158,6 +161,9 @@ class SamplingParams:
         default=None, metadata={"batch_sig_exclude": True}
     )  # None means all resolutions allowed
 
+    # Output audio duration in seconds (models without an audio modality ignore this).
+    sound_duration: float = 0.0
+
     # Denoising parameters
     num_inference_steps: int = None
     guidance_scale: float = 1.0
@@ -841,6 +847,11 @@ class SamplingParams:
             type=int,
             help="Number of frames to generate",
         )
+        add_argument(
+            "--sound-duration",
+            type=float,
+            help="Duration of generated audio in seconds; 0 disables audio output (audio-capable models only)",
+        )
         add_argument(
             "--height",
             type=int,
@@ -952,9 +963,10 @@ class SamplingParams:
             "--action",
             type=str,
             help=(
-                "SANA-WM WASD/IJKL action DSL, e.g. "
+                "Action input. SANA-WM uses a WASD/IJKL DSL, e.g. "
-                "'w-80,jw-40,w-40,lw-60,w-100'. Model-specific fields are "
+                "'w-80,jw-40,w-40,lw-60,w-100'. Cosmos3 uses a JSON array of "
-                "ignored by other pipelines."
+                "shape [T, D], e.g. '[[0.1, 0.2, ...], ...]'. Model-specific "
+                "fields are ignored by other pipelines."
             ),
         )
         add_argument(
@@ -978,6 +990,57 @@ class SamplingParams:
             dest="pitch_limit_deg",
             help="SANA-WM action DSL absolute pitch clamp in degrees.",
         )
+        add_argument(
+            "--video-path",
+            type=str,
+            nargs="+",
+            help=(
+                "Path(s) to input video(s) for video-to-video generation. "
+                "The first/last frames of the video become the conditioning "
+                "frames for the generated output."
+            ),
+        )
+        add_argument(
+            "--action-mode",
+            type=str,
+            dest="action_mode",
+            help=(
+                "Cosmos3 action mode: 'forward_dynamics' (predict next frame "
+                "from action), 'policy' (predict action from frame), or "
+                "'inverse_dynamics' (predict action from two frames)."
+            ),
+        )
+        add_argument(
+            "--domain-id",
+            type=int,
+            dest="domain_id",
+            help="Action embodiment domain ID (integer). Overrides --domain-name.",
+        )
+        add_argument(
+            "--domain-name",
+            type=str,
+            dest="domain_name",
+            help="Action embodiment domain name (e.g. 'av', 'camera_pose', 'umi').",
+        )
+        add_argument(
+            "--raw-action-dim",
+            type=int,
+            dest="raw_action_dim",
+            help=(
+                "Number of active action dimensions to predict; remaining "
+                "dimensions are zero-padded. Required for 'policy' and "
+                "'inverse_dynamics' modes."
+            ),
+        )
+        add_argument(
+            "--action-fps",
+            type=float,
+            dest="action_fps",
+            help=(
+                "Frame rate used for action token temporal mRoPE positions. "
+                "Defaults to the video fps when not set."
+            ),
+        )
         add_argument(
             "--moba-config-path",
             type=str,

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -436,6 +436,7 @@ class DiffGenerator:
             generation_time=generation_time,
             peak_memory_mb=output_batch.peak_memory_mb,
             metrics=metrics.to_dict() if metrics else {},
+            action=output_batch.action_pred,
             trajectory_latents=output_batch.trajectory_latents,
             trajectory_timesteps=output_batch.trajectory_timesteps,
             rollout_trajectory_data=output_batch.rollout_trajectory_data,

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

@@ -86,6 +86,7 @@ class VideoResponse(BaseModel):
     num_outputs: Optional[int] = None
     peak_memory_mb: Optional[float] = None
     inference_time_s: Optional[float] = None
+    action: Optional[Dict[str, Any]] = None
 
 
 class VideoGenerationsRequest(BaseModel):
@@ -94,6 +95,8 @@ class VideoGenerationsRequest(BaseModel):
     prompt: str
     input_reference: Optional[str] = None
     reference_url: Optional[str] = None
+    video_path: Optional[str] = None
+    video_url: Optional[str] = None
     model: Optional[str] = None
     n: Optional[int] = 1
     num_outputs_per_prompt: Optional[int] = None

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

@@ -432,6 +432,17 @@ def add_common_data_to_response(
 
     response["id"] = request_id
 
+    if result.action_pred is not None:
+        t = result.action_pred
+        response["action"] = {
+            "data": t.tolist(),
+            "shape": list(t.shape),
+            "dtype": str(t.dtype).replace("torch.", ""),
+            "raw_action_dim": result.action_raw_action_dim,
+            "action_mode": result.action_mode,
+            "domain_id": result.action_domain_id,
+        }
+
     return response
 
 

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

@@ -50,11 +50,30 @@ from sglang.srt.observability.trace import extract_trace_headers
 logger = init_logger(__name__)
 router = APIRouter(prefix="/v1/videos", tags=["videos"])
 
+_VIDEO_EXTENSIONS = {
+    ".avi",
+    ".gif",
+    ".m4v",
+    ".mkv",
+    ".mov",
+    ".mp4",
+    ".mpeg",
+    ".mpg",
+    ".webm",
+}
+
 
 def _extra_value(request: VideoGenerationsRequest, name: str) -> Any:
     return (request.model_extra or {}).get(name)
 
 
+def _request_value(request: VideoGenerationsRequest, name: str) -> Any:
+    value = getattr(request, name, None)
+    if value is not None:
+        return value
+    return _extra_value(request, name)
+
+
 def _parse_form_extra_value(value: Any) -> Any:
     if not isinstance(value, str):
         return value
@@ -64,15 +83,146 @@ def _parse_form_extra_value(value: Any) -> Any:
         return value
 
 
+def _is_probably_video_source(source: Any) -> bool:
+    content_type = (getattr(source, "content_type", "") or "").lower()
+    if content_type.startswith("video/"):
+        return True
+
+    if isinstance(source, str):
+        if source.lower().startswith("data:video"):
+            return True
+        source_name = source
+    else:
+        source_name = getattr(source, "filename", None)
+
+    if not source_name:
+        return False
+    source_name = str(source_name).split("?", 1)[0].split("#", 1)[0]
+    return os.path.splitext(source_name)[1].lower() in _VIDEO_EXTENSIONS
+
+
+def _is_cosmos3_server(server_args) -> bool:
+    from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config
+
+    return isinstance(server_args.pipeline_config, Cosmos3Config)
+
+
+def _normalize_optional_string(value: Any) -> Any:
+    if isinstance(value, str) and not value.strip():
+        return None
+    return value
+
+
+def _coerce_optional_int_list(value: Any) -> list[int] | 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 [int(item) for item in value]
+    return [int(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:
+        return str(video_path)
+
+    input_reference = _request_value(req, "input_reference")
+    if _is_probably_video_source(input_reference):
+        return str(input_reference)
+
+    reference_url = _request_value(req, "reference_url")
+    if _is_probably_video_source(reference_url):
+        return str(reference_url)
+
+    return None
+
+
+def _resolve_image_path(
+    req: VideoGenerationsRequest, video_path: str | None
+) -> str | None:
+    image_path = _request_value(req, "input_reference")
+    if video_path and image_path == video_path:
+        return None
+    if _is_probably_video_source(image_path):
+        return None
+    return image_path
+
+
+def _resolve_sound_duration(
+    req: VideoGenerationsRequest, *, num_frames: int, fps: int
+) -> float | None:
+    generate_sound = _request_value(req, "generate_sound")
+    sound_duration = _request_value(req, "sound_duration")
+
+    if generate_sound is False:
+        return 0.0
+    if sound_duration is not None:
+        return float(sound_duration)
+    if generate_sound is True:
+        return float(num_frames) / float(fps)
+    return None
+
+
+def _cosmos3_sampling_param_kwargs(
+    req: VideoGenerationsRequest, *, num_frames: int, fps: int
+) -> Dict[str, Any]:
+    """Map HTTP/API aliases to Cosmos3SamplingParams field names."""
+    kwargs: Dict[str, Any] = {}
+
+    sound_duration = _resolve_sound_duration(req, num_frames=num_frames, fps=fps)
+    if sound_duration is not None:
+        kwargs["sound_duration"] = sound_duration
+
+    condition_frame_indexes = _request_value(req, "condition_frame_indexes")
+    if condition_frame_indexes is None:
+        condition_frame_indexes = _request_value(req, "condition_frame_indexes_vision")
+    condition_frame_indexes = _coerce_optional_int_list(condition_frame_indexes)
+    if condition_frame_indexes is not None:
+        kwargs["condition_frame_indexes"] = condition_frame_indexes
+
+    for name in (
+        "condition_video_keep",
+        "action_mode",
+        "domain_id",
+        "domain_name",
+        "raw_action_dim",
+        "action_fps",
+        "action",
+        "action_view_point",
+        "action_normalization",
+    ):
+        value = _parse_form_extra_value(_request_value(req, name))
+        value = _normalize_optional_string(value)
+        if value is not None:
+            kwargs[name] = value
+
+    return kwargs
+
+
 def _build_video_sampling_params(request_id: str, request: VideoGenerationsRequest):
     """Resolve video-specific defaults (fps, seconds → num_frames) then
     delegate to the shared build_sampling_params."""
+    server_args = get_global_server_args()
     seconds = request.seconds if request.seconds is not None else DEFAULT_VIDEO_SECONDS
     fps = request.fps if request.fps is not None else DEFAULT_FPS
     num_frames = request.num_frames if request.num_frames is not None else fps * seconds
     num_outputs = request.num_outputs_per_prompt
     if num_outputs is None:
         num_outputs = request.n or 1
+    video_path = _resolve_video_path(request)
+    image_path = _resolve_image_path(request, video_path)
+    cosmos3_kwargs = {}
+    if _is_cosmos3_server(server_args):
+        cosmos3_kwargs = _cosmos3_sampling_param_kwargs(
+            request, num_frames=num_frames, fps=fps
+        )
+        if server_args.pipeline_config.action_stats_path is not None:
+            cosmos3_kwargs["action_stats_path"] = (
+                server_args.pipeline_config.action_stats_path
+            )
 
     return build_sampling_params(
         request_id,
@@ -83,7 +233,8 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
         height=request.height,
         num_frames=num_frames,
         fps=fps,
-        image_path=request.input_reference,
+        image_path=image_path,
+        video_path=video_path,
         output_file_name=request_id,
         seed=request.seed,
         generator_device=request.generator_device,
@@ -110,30 +261,7 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
         output_quality=request.output_quality,
         perf_dump_path=request.perf_dump_path,
         diffusers_kwargs=request.diffusers_kwargs,
-    )
+        **cosmos3_kwargs,
-
-
-def _reject_unsupported_cosmos3_modes(
-    req: VideoGenerationsRequest, model_path: str | None
-) -> None:
-    if "cosmos3" not in (model_path or "").lower():
-        return
-
-    extra = req.model_extra or {}
-    if extra.get("generate_sound"):
-        raise HTTPException(
-            status_code=400,
-            detail="Cosmos3 video-with-sound is not supported by SGLang yet; omit generate_sound for video-only generation.",
-        )
-    if extra.get("action_mode"):
-        raise HTTPException(
-            status_code=400,
-            detail="Cosmos3 action generation is not supported by SGLang yet.",
-        )
-    if extra.get("condition_frame_indexes_vision") or extra.get("condition_video_keep"):
-        raise HTTPException(
-            status_code=400,
-            detail="Cosmos3 video-to-video conditioning is not supported by SGLang yet.",
     )
 
 
@@ -234,6 +362,9 @@ async def create_video(
     prompt: Optional[str] = Form(None),
     input_reference: Optional[UploadFile] = File(None),
     reference_url: Optional[str] = Form(None),
+    video_reference: Optional[UploadFile] = File(None),
+    video_url: Optional[str] = Form(None),
+    video_path: Optional[str] = Form(None),
     model: Optional[str] = Form(None),
     n: Optional[int] = Form(1),
     num_outputs_per_prompt: Optional[int] = Form(None),
@@ -287,13 +418,38 @@ async def create_video(
     if "multipart/form-data" in content_type:
         if not prompt:
             raise HTTPException(status_code=400, detail="prompt is required")
-        # Validate image input based on model task type
+
+        video_input_path = None
         image_sources = merge_image_input_list(input_reference, reference_url)
+        if video_reference is not None:
+            video_input_path = await _save_first_input_image(
+                video_reference,
+                request_id,
+                uploads_dir,
+                prefer_remote_source=server_args.input_save_path is None,
+            )
+        elif video_path or video_url:
+            video_input_path = video_path or video_url
+        elif input_reference is not None and _is_probably_video_source(input_reference):
+            video_input_path = await _save_first_input_image(
+                input_reference,
+                request_id,
+                uploads_dir,
+                prefer_remote_source=server_args.input_save_path is None,
+            )
+            image_sources = merge_image_input_list(reference_url)
+        elif reference_url and _is_probably_video_source(reference_url):
+            video_input_path = reference_url
+            image_sources = merge_image_input_list(input_reference)
+
+        # Validate image input based on model task type
         if task_type.requires_image_input() and not image_sources:
             raise HTTPException(
                 status_code=400,
                 detail="input_reference or reference_url is required for image-to-video generation",
             )
+        input_path = None
+        if image_sources:
             try:
                 input_path = await _save_first_input_image(
                     image_sources,
@@ -322,7 +478,8 @@ async def create_video(
                 pass
 
         def form_value(name: str, value: Any) -> Any:
-            return value if value is not None else extra_from_form.get(name)
+            selected = value if value is not None else extra_from_form.get(name)
+            return _parse_form_extra_value(selected)
 
         raw_form = await request.form()
         for key in (
@@ -331,9 +488,19 @@ async def create_video(
             "use_system_prompt",
             "use_guardrails",
             "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",
         ):
@@ -353,6 +520,8 @@ async def create_video(
         req = VideoGenerationsRequest(
             prompt=prompt,
             input_reference=input_path,
+            video_path=form_value("video_path", video_input_path),
+            video_url=form_value("video_url", video_url),
             model=form_value("model", model),
             n=form_value("n", n),
             num_outputs_per_prompt=form_value(
@@ -414,8 +583,19 @@ async def create_video(
                 payload.update(flatten_extra_params(extra_json))
             flatten_extra_params(payload)
             # Validate image input based on model task type
-            has_image_input = payload.get("reference_url") or payload.get(
+            if payload.get("video_url") and not payload.get("video_path"):
-                "input_reference"
+                payload["video_path"] = payload["video_url"]
+            if _is_probably_video_source(payload.get("reference_url")):
+                payload.setdefault("video_path", payload.get("reference_url"))
+            if _is_probably_video_source(payload.get("input_reference")):
+                payload.setdefault("video_path", payload.get("input_reference"))
+
+            has_image_input = (
+                payload.get("reference_url")
+                and not _is_probably_video_source(payload.get("reference_url"))
+            ) or (
+                payload.get("input_reference")
+                and not _is_probably_video_source(payload.get("input_reference"))
             )
             if task_type.requires_image_input() and not has_image_input:
                 raise HTTPException(
@@ -423,7 +603,9 @@ async def create_video(
                     detail="input_reference or reference_url is required for image-to-video generation",
                 )
             # for non-multipart/form-data type
-            if payload.get("reference_url"):
+            if payload.get("reference_url") and not _is_probably_video_source(
+                payload.get("reference_url")
+            ):
                 try:
                     input_path = await _save_first_input_image(
                         payload.get("reference_url"),
@@ -454,8 +636,6 @@ async def create_video(
 
     logger.debug(f"Server received from create_video endpoint: req={req}")
 
-    _reject_unsupported_cosmos3_modes(req, server_args.model_path)
-
     try:
         sampling_params = _build_video_sampling_params(request_id, req)
     except (ValueError, TypeError) as e:

python/sglang/multimodal_gen/runtime/entrypoints/utils.py

@@ -117,6 +117,7 @@ class GenerationResult:
     samples: Any = None
     frames: Any = None
     audio: Any = None
+    action: Any = None  # [T, raw_action_dim] predicted action (policy/inverse_dynamics)
     prompt: str | None = None
     size: tuple | None = None  # (height, width, num_frames)
     generation_time: float = 0.0

python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py

@@ -74,7 +74,11 @@ class ModelOptFp8Config(QuantizationConfig):
         return []
 
     @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> ModelOptFp8Config:
+    def from_config(
+        cls,
+        config: Dict[str, Any],
+        ignore_remap: Optional[Dict[str, str]] = None,
+    ) -> ModelOptFp8Config:
         quant_algo = config.get("quant_algo")
         if quant_algo is None:
             raise ValueError(
@@ -85,6 +89,8 @@ class ModelOptFp8Config(QuantizationConfig):
                 f"ModelOptFp8Config only supports FP8, got quant_algo={quant_algo!r}."
             )
         ignore = config.get("ignore", [])
+        if ignore_remap and ignore:
+            ignore = [ignore_remap.get(pattern, pattern) for pattern in ignore]
         return cls(is_checkpoint_fp8_serialized=True, ignore=ignore)
 
     def _is_layer_ignored(self, prefix: str) -> bool:
@@ -122,7 +128,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
     """Linear method for ModelOpt static per-tensor FP8 quantization.
 
     Uses ``torch._scaled_mm`` (or CUTLASS FP8 GEMM when available) for
-    the FP8 matrix multiply — the same kernels used by the LLM runtime.
+    the FP8 matrix multiply - the same kernels used by the LLM runtime.
     """
 
     def __init__(self, quant_config: ModelOptFp8Config):
@@ -177,8 +183,8 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
         max_w_scale = layer.weight_scale.max()
 
         # Transpose weight to [in, out] column-major layout for
-        # apply_fp8_linear / CUTLASS fp8_scaled_mm.  Do NOT call
+        # apply_fp8_linear / CUTLASS fp8_scaled_mm. Do not call .contiguous();
-        # .contiguous() — the kernel requires column-major stride.
+        # the kernel requires column-major stride.
         layer.weight = torch.nn.Parameter(layer.weight.data.t(), requires_grad=False)
 
         if self.cutlass_fp8_supported:

python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py

@@ -187,7 +187,11 @@ class ModelOptFp8Config(ModelOptQuantConfig):
         return 89
 
     @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> ModelOptFp8Config:
+    def from_config(
+        cls,
+        config: Dict[str, Any],
+        ignore_remap: Optional[Dict[str, str]] = None,
+    ) -> ModelOptFp8Config:
         quant_method = config.get("quant_algo")
         exclude_modules = config.get("ignore")
         if quant_method is None:
@@ -205,6 +209,9 @@ class ModelOptFp8Config(ModelOptQuantConfig):
                 "ModelOptFp8Config only supports static FP8 quantization in SGLang diffusion."
             )
 
+        if ignore_remap and exclude_modules:
+            exclude_modules = [ignore_remap.get(p, p) for p in exclude_modules]
+
         return cls(
             is_checkpoint_fp8_serialized=True,
             exclude_modules=exclude_modules,

python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py

@@ -23,11 +23,22 @@ class SchedulerLoader(ComponentLoader):
         """Load the scheduler based on the model path, and inference args."""
         config = get_diffusers_component_config(component_path=component_model_path)
 
-        class_name = config.pop("_class_name")
+        checkpoint_class_name = config.pop("_class_name", None)
+        class_name = (
+            getattr(server_args.pipeline_config, "scheduler_class_override", None)
+            or checkpoint_class_name
+        )
         assert (
             class_name is not None
         ), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
 
+        if checkpoint_class_name is not None and class_name != checkpoint_class_name:
+            logger.info(
+                "Overriding scheduler class from %s to %s",
+                checkpoint_class_name,
+                class_name,
+            )
+
         scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)
 
         scheduler = scheduler_cls(**config)

python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+from safetensors.torch import load_file as safetensors_load_file
+
+from sglang.multimodal_gen.configs.models import ModelConfig
+from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
+    ComponentLoader,
+)
+from sglang.multimodal_gen.runtime.loader.utils import (
+    _list_safetensors_files,
+    set_default_torch_dtype,
+    skip_init_modules,
+)
+from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+    get_diffusers_component_config,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class SoundTokenizerLoader(ComponentLoader):
+    component_names = ["sound_tokenizer"]
+    expected_library = "diffusers"
+
+    def should_offload(
+        self, server_args: ServerArgs, model_config: ModelConfig | None = None
+    ) -> bool:
+        return server_args.vae_cpu_offload
+
+    def load_customized(
+        self, component_model_path: str, server_args: ServerArgs, component_name: str
+    ):
+        config = get_diffusers_component_config(component_path=component_model_path)
+        class_name = config.pop("_class_name", None) or self.component_architecture
+        assert (
+            class_name is not None
+        ), "Sound tokenizer class name must be available from component config."
+
+        server_args.model_paths[component_name] = component_model_path
+
+        try:
+            precision = server_args.pipeline_config.vae_precision
+        except AttributeError:
+            precision = "bf16"
+        dtype = PRECISION_TO_TYPE[precision]
+        target_device = self.target_device(self.should_offload(server_args))
+
+        with set_default_torch_dtype(dtype), skip_init_modules():
+            model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+            model = model_cls(config).to(target_device)
+
+        safetensors_list = _list_safetensors_files(component_model_path)
+        assert (
+            len(safetensors_list) == 1
+        ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}"
+        loaded = safetensors_load_file(safetensors_list[0])
+        incompatible = model.load_state_dict(loaded, strict=False)
+        missing = getattr(incompatible, "missing_keys", [])
+        # The tokenizer is decoder-only; the checkpoint's encoder weights are
+        # expected leftovers, so they're excluded from the load warning.
+        unexpected = [
+            k
+            for k in getattr(incompatible, "unexpected_keys", [])
+            if not k.startswith("encoder.")
+        ]
+        if missing or unexpected:
+            logger.warning(
+                "Loaded sound_tokenizer with missing_keys=%d unexpected_keys=%d",
+                len(missing),
+                len(unexpected),
+            )
+        model.eval()
+        return model

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

@@ -615,10 +615,12 @@ def _resolve_quant_config(
     reverse_param_names_mapping_dict = getattr(
         arch_config, "reverse_param_names_mapping", None
     )
+    quant_ignore_remap_dict = getattr(arch_config, "quant_ignore_remap", None)
     quant_config = get_quant_config(
         hf_config,
         component_model_path,
         reverse_param_names_mapping=reverse_param_names_mapping_dict,
+        quant_ignore_remap=quant_ignore_remap_dict,
     )
     quant_config_name = _get_quant_config_name(quant_config)
     inferred_nvfp4_config = None

python/sglang/multimodal_gen/runtime/managers/scheduler.py

@@ -732,6 +732,9 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag
                     output_batch.audio, start, end, total_items
                 ),
                 audio_sample_rate=output_batch.audio_sample_rate,
+                action_pred=self._slice_batched_value(
+                    output_batch.action_pred, start, end, total_items
+                ),
                 trajectory_timesteps=self._slice_batched_value(
                     output_batch.trajectory_timesteps, start, end, total_items
                 ),

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

@@ -45,6 +45,9 @@ from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.srt.utils import add_prefix
@@ -84,6 +87,8 @@ def compute_mrope_position_ids_vision(
     fps: float | None = None,
     base_fps: float = 24.0,
     temporal_compression_factor: int = 4,
+    base_temporal_compression_factor: int | None = None,
+    start_frame_offset: int = 0,
 ) -> tuple[torch.Tensor, int | float]:
     """Generate 3D mRoPE position IDs for vision tokens.
 
@@ -91,6 +96,12 @@ def compute_mrope_position_ids_vision(
     per vision segment (Qwen3VL-style).
     Flattened in T-major order.
 
+    When the token rate (``temporal_compression_factor``) differs from the
+    base-fps rate (``base_temporal_compression_factor``), the two no longer
+    cancel: action tokens run at frame rate (factor 1) while their temporal
+    positions are scaled by the video factor. ``base_temporal_compression_factor``
+    defaults to ``temporal_compression_factor`` so vision/sound are unchanged.
+
     Returns:
         (position_ids [3, grid_t * grid_h * grid_w], next_temporal_offset)
     """
@@ -98,18 +109,28 @@ def compute_mrope_position_ids_vision(
 
     if fps_modulation:
         tps = fps / temporal_compression_factor
-        base_tps = base_fps / temporal_compression_factor
+        effective_base_tcf = (
+            base_temporal_compression_factor
+            if base_temporal_compression_factor is not None
+            else temporal_compression_factor
+        )
+        base_tps = base_fps / effective_base_tcf
         frame_indices = torch.arange(grid_t, dtype=torch.float32, device=device)
         t_index = (
-            (frame_indices / tps * base_tps + temporal_offset)
+            ((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset)
             .view(-1, 1)
             .expand(-1, grid_h * grid_w)
             .flatten()
         )
     else:
-        t_index = torch.arange(grid_t, dtype=torch.long, device=device).view(
+        t_index = (
-            -1, 1
+            torch.arange(grid_t, dtype=torch.long, device=device)
-        ).expand(-1, grid_h * grid_w).flatten() + int(temporal_offset)
+            .view(-1, 1)
+            .expand(-1, grid_h * grid_w)
+            .flatten()
+            + int(temporal_offset)
+            + start_frame_offset
+        )
 
     h_index = (
         torch.arange(grid_h, dtype=torch.long, device=device)
@@ -135,6 +156,58 @@ def compute_mrope_position_ids_vision(
     return mrope_ids, next_offset
 
 
+def compute_mrope_position_ids_sound(
+    grid_t: int,
+    temporal_offset: int | float,
+    sound_latent_fps: float,
+    device: torch.device,
+    base_fps: float = 24.0,
+    temporal_compression_factor_sound: int = 1,
+) -> tuple[torch.Tensor, int | float]:
+    """mRoPE position IDs for sound tokens: a (T, 1, 1) grid."""
+    return compute_mrope_position_ids_vision(
+        grid_t=grid_t,
+        grid_h=1,
+        grid_w=1,
+        temporal_offset=temporal_offset,
+        device=device,
+        fps=sound_latent_fps,
+        base_fps=base_fps,
+        temporal_compression_factor=temporal_compression_factor_sound,
+    )
+
+
+def compute_mrope_position_ids_action(
+    grid_t: int,
+    temporal_offset: int | float,
+    action_fps: float | None,
+    device: torch.device,
+    base_fps: float = 24.0,
+    base_temporal_compression_factor: int = 4,
+    start_frame_offset: int = 1,
+) -> tuple[torch.Tensor, int | float]:
+    """mRoPE position IDs for action tokens: a (T, 1, 1) grid.
+
+    Action tokens run at frame rate (``temporal_compression_factor=1``) but
+    their positions are scaled by the video's ``base_temporal_compression_factor``
+    so they share the video's temporal coordinate frame. ``start_frame_offset=1``
+    (default) shifts them one frame ahead so they align with the video frame
+    they condition on.
+    """
+    return compute_mrope_position_ids_vision(
+        grid_t=grid_t,
+        grid_h=1,
+        grid_w=1,
+        temporal_offset=temporal_offset,
+        device=device,
+        fps=action_fps,
+        base_fps=base_fps,
+        temporal_compression_factor=1,
+        base_temporal_compression_factor=base_temporal_compression_factor,
+        start_frame_offset=start_frame_offset,
+    )
+
+
 # -----------------------------------------------------------------------------
 # Qwen3-style RoPE functions
 # -----------------------------------------------------------------------------
@@ -195,6 +268,44 @@ def _apply_qwen3_qk_norm_rope_split(
     return _apply_qwen3_rope_from_cache(q, k, cos_sin_cache)
 
 
+# -----------------------------------------------------------------------------
+# Action domain-aware projection
+# -----------------------------------------------------------------------------
+
+
+class DomainAwareLinear(nn.Module):
+    """Per-domain linear projection for action conditioning.
+
+    Maintains one weight matrix and bias per embodiment domain via embedding
+    tables, enabling multi-domain robot action generation from a shared
+    backbone.
+    """
+
+    def __init__(self, input_size: int, output_size: int, num_domains: int) -> None:
+        super().__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+        self.num_domains = num_domains
+        self.fc = nn.Embedding(num_domains, output_size * input_size)
+        self.bias = nn.Embedding(num_domains, output_size)
+        nn.init.xavier_uniform_(
+            self.fc.weight.view(num_domains, output_size, input_size)
+        )
+        nn.init.zeros_(self.bias.weight)
+
+    def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor:
+        if domain_id.ndim == 0:
+            domain_id = domain_id.unsqueeze(0)
+        domain_id = domain_id.to(device=x.device, dtype=torch.long).reshape(-1)
+        weight = self.fc(domain_id).view(
+            domain_id.shape[0], self.input_size, self.output_size
+        )
+        bias = self.bias(domain_id).view(domain_id.shape[0], self.output_size)
+        if x.ndim == 2:
+            return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias
+        return torch.bmm(x, weight) + bias.unsqueeze(1)
+
+
 # -----------------------------------------------------------------------------
 # Cosmos3 Timestep Embedder
 # -----------------------------------------------------------------------------
@@ -797,7 +908,7 @@ class Cosmos3LanguageModel(nn.Module):
 # -----------------------------------------------------------------------------
 
 
-class Cosmos3OmniTransformer(CachableDiT):
+class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
     """Cosmos3 Omni transformer.
 
     Dual-pathway architecture:
@@ -837,6 +948,8 @@ class Cosmos3OmniTransformer(CachableDiT):
         self.base_fps = arch.base_fps
         self.temporal_compression_factor = arch.temporal_compression_factor
         self.temporal_margin = arch.unified_3d_mrope_temporal_modality_margin
+        self.sound_latent_fps = arch.sound_latent_fps
+        self.temporal_compression_factor_sound = arch.temporal_compression_factor_sound
         self.rms_norm_eps = arch.rms_norm_eps
 
         # Ulysses sequence parallelism. When CFG-parallel is also enabled
@@ -881,6 +994,40 @@ class Cosmos3OmniTransformer(CachableDiT):
             prefix="proj_out",
         )
 
+        self.sound_gen = arch.sound_gen
+        self.sound_dim = arch.sound_dim
+        if arch.sound_gen:
+            self.audio_proj_in = ReplicatedLinear(
+                self.sound_dim,
+                self.hidden_size,
+                bias=True,
+                quant_config=quant_config,
+                prefix="audio_proj_in",
+            )
+            self.audio_proj_out = ReplicatedLinear(
+                self.hidden_size,
+                self.sound_dim,
+                bias=True,
+                quant_config=quant_config,
+                prefix="audio_proj_out",
+            )
+            self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
+
+        if arch.action_gen:
+            self.action_dim = arch.action_dim
+            self.num_embodiment_domains = arch.num_embodiment_domains
+            self.action_proj_in = DomainAwareLinear(
+                self.action_dim,
+                self.hidden_size,
+                self.num_embodiment_domains,
+            )
+            self.action_proj_out = DomainAwareLinear(
+                self.hidden_size,
+                self.action_dim,
+                self.num_embodiment_domains,
+            )
+            self.action_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
+
         # Timestep embedder
         self.time_embedder = Cosmos3TimestepEmbedder(
             hidden_size=self.hidden_size,
@@ -920,6 +1067,8 @@ class Cosmos3OmniTransformer(CachableDiT):
 
         self.__post_init__()
 
+        self.layer_names = ["gen_layers", "language_model.layers"]
+
     def _pad_to_patch_size(self, H: int, W: int) -> tuple[int, int, int, int]:
         """Compute padded spatial dims aligned to patch_size."""
         p = self.latent_patch_size
@@ -964,8 +1113,12 @@ class Cosmos3OmniTransformer(CachableDiT):
         Wp: int,
         fps: float | None,
         device: torch.device,
+        sound_frames: int = 0,
+        action_frames: int = 0,
+        action_fps: float | None = None,
+        action_start_frame_offset: int = 1,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Compute mRoPE position IDs for UND text and GEN visual tokens."""
+        """Compute mRoPE position IDs for UND text and GEN visual + action + sound tokens."""
         B = text_mask.shape[0]
         S_text = text_mask.shape[1]
         text_lengths = text_mask.sum(dim=1).long()
@@ -978,16 +1131,40 @@ class Cosmos3OmniTransformer(CachableDiT):
             t_pos, t_offset = compute_mrope_position_ids_text(
                 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=t_offset + self.temporal_margin,
+                temporal_offset=media_offset,
                 device=device,
                 fps=effective_fps,
                 base_fps=self.base_fps,
                 temporal_compression_factor=self.temporal_compression_factor,
             )
+            if action_frames > 0:
+                a_pos, _ = compute_mrope_position_ids_action(
+                    action_frames,
+                    temporal_offset=media_offset,
+                    action_fps=action_fps,
+                    device=device,
+                    base_fps=self.base_fps,
+                    base_temporal_compression_factor=self.temporal_compression_factor,
+                    start_frame_offset=action_start_frame_offset,
+                )
+                pos_dtype = torch.promote_types(v_pos.dtype, a_pos.dtype)
+                v_pos = torch.cat([v_pos.to(pos_dtype), a_pos.to(pos_dtype)], dim=1)
+            if sound_frames > 0:
+                s_pos, _ = compute_mrope_position_ids_sound(
+                    sound_frames,
+                    temporal_offset=media_offset,
+                    sound_latent_fps=self.sound_latent_fps,
+                    device=device,
+                    base_fps=self.base_fps,
+                    temporal_compression_factor_sound=self.temporal_compression_factor_sound,
+                )
+                pos_dtype = torch.promote_types(v_pos.dtype, s_pos.dtype)
+                v_pos = torch.cat([v_pos.to(pos_dtype), s_pos.to(pos_dtype)], dim=1)
             if real_len < S_text:
                 t_pos = torch.cat(
                     [
@@ -1002,7 +1179,7 @@ class Cosmos3OmniTransformer(CachableDiT):
             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_vis]
+        vis_pos_ids = torch.stack(vis_pos_list, dim=1).to(device)  # [3, B, S_gen]
 
         return text_pos_ids, vis_pos_ids
 
@@ -1044,8 +1221,14 @@ class Cosmos3OmniTransformer(CachableDiT):
         cache_key: str = "default",
         noisy_frame_mask: torch.Tensor | None = None,
         max_text_seq_len: int | None = None,
+        sound_latents: torch.Tensor | None = None,
+        action_latents: torch.Tensor | None = None,
+        action_domain_ids: torch.Tensor | None = None,
+        action_noisy_mask: torch.Tensor | None = None,
+        action_fps: float | None = None,
+        action_start_frame_offset: int = 1,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Forward pass for denoising.
 
         Args:
@@ -1064,9 +1247,20 @@ class Cosmos3OmniTransformer(CachableDiT):
                 ``None`` means every frame is noisy (T2V / T2I).
             max_text_seq_len: Real text length already computed during
                 tokenization. When omitted it is derived from ``text_mask``.
+            action_latents: Optional [B, T_action, D_action] noisy action
+                latents for action generation.
+            action_domain_ids: [B] embodiment domain IDs (0=no-action default).
+            action_noisy_mask: [B, T_action, 1] where 1=noisy, 0=conditioned;
+                controls which action tokens receive the timestep embedding.
+                ``None`` means all tokens are noisy.
+            action_fps: Frame rate for action token temporal mRoPE scaling.
+                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).
 
         Returns:
-            [B, C, T, H, W] velocity prediction
+            [B, C, T, H, W] velocity prediction, or a tuple
+            (video_pred, ...) with extra tensors when action/sound are active.
         """
         if text_ids is None or text_mask is None:
             raise ValueError("Cosmos3 requires text_ids and text_mask to be passed")
@@ -1079,9 +1273,31 @@ class Cosmos3OmniTransformer(CachableDiT):
             text_ids = text_ids[:, :max_text_seq_len]
             text_mask = text_mask[:, :max_text_seq_len]
 
-        # Check if sequence parallelism is enabled
+        sound_frames = sound_latents.shape[-1] if sound_latents is not None else 0
+
+        action_frames = 0
+        if action_latents is not None:
+            if self.sp_size > 1:
+                raise NotImplementedError(
+                    "Cosmos3 action generation does not support sequence parallelism yet"
+                )
+            action_frames = action_latents.shape[1]
+            if action_domain_ids is None:
+                action_domain_ids = torch.zeros(
+                    action_latents.shape[0],
+                    dtype=torch.long,
+                    device=action_latents.device,
+                )
+
+        extra_frames = action_frames + sound_frames
         sequence_shard_enabled = self.sp_size > 1
 
+        # Add timestep embedding (computed in float32 for numerical stability, then cast back)
+        time_embed = self.time_embedder(timestep.float())
+        time_embed = time_embed.to(
+            hidden_states.dtype
+        )  # Cast to match hidden_gen dtype
+
         # Patchify and project to hidden dim
         hidden_gen, _ = self.proj_in(self.patchify(hidden_states, T, H, W))
         seq_len_orig = hidden_gen.shape[1]
@@ -1099,7 +1315,9 @@ class Cosmos3OmniTransformer(CachableDiT):
                 .to(hidden_gen.dtype)
             )
 
-        # Shard sequence across GPUs if SP enabled
+        if extra_frames == 0:
+            # Video-only: shard the visual tokens, then add the timestep
+            # embedding on the local shard.
             if sequence_shard_enabled:
                 if seq_len_orig % self.sp_size != 0:
                     seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size)
@@ -1115,7 +1333,9 @@ class Cosmos3OmniTransformer(CachableDiT):
                             dtype=token_noisy_mask.dtype,
                             device=token_noisy_mask.device,
                         )
-                    token_noisy_mask = torch.cat([token_noisy_mask, mask_pad], dim=1)
+                        token_noisy_mask = torch.cat(
+                            [token_noisy_mask, mask_pad], dim=1
+                        )
                 local_seq_len = hidden_gen.shape[1] // self.sp_size
                 hidden_gen = hidden_gen.view(
                     batch_size, self.sp_size, local_seq_len, hidden_gen.shape[2]
@@ -1125,16 +1345,60 @@ class Cosmos3OmniTransformer(CachableDiT):
                     token_noisy_mask = token_noisy_mask.view(
                         batch_size, self.sp_size, local_seq_len, 1
                     )[:, self.sp_rank, :, :]
-
-        # Add timestep embedding (computed in float32 for numerical stability, then cast back)
-        time_embed = self.time_embedder(timestep.float())
-        time_embed = time_embed.to(
-            hidden_states.dtype
-        )  # Cast to match hidden_gen dtype
             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)
+        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.
+            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)
+
+            if action_latents is not None:
+                hidden_action = self.action_proj_in(
+                    action_latents.to(hidden_gen.dtype), action_domain_ids
+                )
+                hidden_action = hidden_action + self.action_modality_embed.to(
+                    hidden_action.dtype
+                )
+                if action_noisy_mask is None:
+                    hidden_action = hidden_action + time_embed.unsqueeze(1)
+                else:
+                    hidden_action = hidden_action + time_embed.unsqueeze(
+                        1
+                    ) * action_noisy_mask.to(hidden_action.dtype)
+                hidden_gen = torch.cat([hidden_gen, hidden_action], dim=1)
+
+            if sound_latents is not None:
+                packed_sound = sound_latents.permute(0, 2, 1).to(hidden_gen.dtype)
+                hidden_sound, _ = self.audio_proj_in(packed_sound)
+                hidden_sound = hidden_sound + self.audio_modality_embed.to(
+                    hidden_sound.dtype
+                )
+                hidden_sound = hidden_sound + time_embed.unsqueeze(1)
+                hidden_gen = torch.cat([hidden_gen, hidden_sound], dim=1)
+
+            seq_len_orig = hidden_gen.shape[1]
+            if sequence_shard_enabled:
+                if seq_len_orig % self.sp_size != 0:
+                    seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size)
+                    pad = torch.zeros(
+                        (batch_size, seq_shard_pad, hidden_gen.shape[2]),
+                        dtype=hidden_gen.dtype,
+                        device=hidden_gen.device,
+                    )
+                    hidden_gen = torch.cat([hidden_gen, pad], dim=1)
+                local_seq_len = hidden_gen.shape[1] // self.sp_size
+                hidden_gen = hidden_gen.view(
+                    batch_size, self.sp_size, local_seq_len, hidden_gen.shape[2]
+                )
+                hidden_gen = hidden_gen[:, self.sp_rank, :, :]
 
         self._ensure_cache_dicts()
 
@@ -1145,7 +1409,16 @@ class Cosmos3OmniTransformer(CachableDiT):
             or cache_key not in self.cached_gen_rope_inputs
         ):
             text_pos_ids, vis_pos_ids = self._compute_rope_position_ids(
-                text_mask, T, Hp, Wp, fps, hidden_states.device
+                text_mask,
+                T,
+                Hp,
+                Wp,
+                fps,
+                hidden_states.device,
+                sound_frames=sound_frames,
+                action_frames=action_frames,
+                action_fps=action_fps if action_fps is not None else fps,
+                action_start_frame_offset=action_start_frame_offset,
             )
             # UND K/V cache is kept FULL on all ranks (not sharded). Text
             # sequence is short, so memory impact is minimal, and the GEN
@@ -1193,15 +1466,43 @@ class Cosmos3OmniTransformer(CachableDiT):
         # this cuts the post-loop SP collective bandwidth ~21x.
         hidden_gen = hidden_gen + residual
         hidden_gen = self.norm_moe_gen(hidden_gen)
-        output, _ = self.proj_out(hidden_gen)
 
+        if extra_frames == 0:
+            # 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
+            # collective bandwidth ~21x.
+            output, _ = self.proj_out(hidden_gen)
             if sequence_shard_enabled:
                 output = sequence_model_parallel_all_gather(output, dim=1)
                 if seq_shard_pad > 0:
                     output = output[:, :seq_len_orig, :]
-
             return self.unpatchify(output, T, H, W)
 
+        # Multi-modal: gather the full GEN hidden and drop shard padding, then
+        # split per modality so each output head sees its own contiguous tokens.
+        if sequence_shard_enabled:
+            hidden_gen = sequence_model_parallel_all_gather(hidden_gen, dim=1)
+            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, :])
+        video_pred = self.unpatchify(output, T, H, W)
+
+        extra_outputs: list[torch.Tensor] = []
+        idx = 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))
+            idx += action_frames
+        if sound_frames > 0:
+            sound_hidden = hidden_gen[:, idx:, :]
+            sound_output, _ = self.audio_proj_out(sound_hidden)
+            extra_outputs.append(sound_output.permute(0, 2, 1).contiguous())
+
+        return (video_pred, *extra_outputs)
+
     def preprocess_loaded_state_dict(
         self, iterator: Iterable[tuple[str, torch.Tensor]]
     ) -> Iterator[tuple[str, torch.Tensor]]:

python/sglang/multimodal_gen/runtime/models/vaes/cosmos3_avae.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Decoder-only audio tokenizer for the Cosmos3 sound modality."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm
+
+
+class Snake1d(nn.Module):
+    def __init__(self, hidden_dim: int, logscale: bool = True) -> None:
+        super().__init__()
+        self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+        self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1))
+        self.logscale = logscale
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.shape
+        alpha = torch.exp(self.alpha) if self.logscale else self.alpha
+        beta = torch.exp(self.beta) if self.logscale else self.beta
+        x = x.reshape(shape[0], shape[1], -1)
+        x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+        return x.reshape(shape)
+
+
+class OobleckResidualUnit(nn.Module):
+    def __init__(self, dim: int, dilation: int = 1) -> None:
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.snake1 = Snake1d(dim)
+        self.conv1 = weight_norm(
+            nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad)
+        )
+        self.snake2 = Snake1d(dim)
+        self.conv2 = weight_norm(nn.Conv1d(dim, dim, kernel_size=1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.conv1(self.snake1(x))
+        y = self.conv2(self.snake2(y))
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class OobleckDecoderBlock(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        stride: int,
+        output_padding: int,
+    ) -> None:
+        super().__init__()
+        self.snake1 = Snake1d(input_dim)
+        self.conv_t1 = weight_norm(
+            nn.ConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=output_padding,
+            )
+        )
+        self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1)
+        self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3)
+        self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.snake1(x)
+        x = self.conv_t1(x)
+        x = self.res_unit1(x)
+        x = self.res_unit2(x)
+        return self.res_unit3(x)
+
+
+class OobleckDecoder(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        input_channels: int,
+        audio_channels: int,
+        upsampling_ratios: list[int],
+        channel_multiples: list[int],
+    ) -> None:
+        super().__init__()
+        strides = upsampling_ratios
+        mults = [1] + list(channel_multiples)
+
+        self.conv1 = weight_norm(
+            nn.Conv1d(input_channels, channels * mults[-1], kernel_size=7, padding=3)
+        )
+
+        blocks = []
+        for i, stride in enumerate(strides):
+            blocks.append(
+                OobleckDecoderBlock(
+                    input_dim=channels * mults[len(strides) - i],
+                    output_dim=channels * mults[len(strides) - i - 1],
+                    stride=stride,
+                    output_padding=stride % 2,
+                )
+            )
+        self.block = nn.ModuleList(blocks)
+        self.snake1 = Snake1d(channels)
+        self.conv2 = weight_norm(
+            nn.Conv1d(channels, audio_channels, kernel_size=7, padding=3, bias=False)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv1(x)
+        for layer in self.block:
+            x = layer(x)
+        x = self.snake1(x)
+        return self.conv2(x)
+
+
+def _cfg(config: dict[str, Any], *keys: str, default: Any = None) -> Any:
+    for k in keys:
+        v = config.get(k)
+        if v is not None:
+            return v
+    return default
+
+
+class Cosmos3AVAEAudioTokenizer(nn.Module):
+    """Cosmos3 audio tokenizer: latents → waveform via an Oobleck decoder stack."""
+
+    def __init__(self, config: dict[str, Any]) -> None:
+        super().__init__()
+        self.sample_rate = int(
+            _cfg(config, "sampling_rate", "sample_rate", default=48000)
+        )
+        self.audio_channels = int(
+            _cfg(
+                config,
+                "dec_out_channels",
+                "audio_channels",
+                default=2 if bool(config.get("stereo", True)) else 1,
+            )
+        )
+        self.latent_channels = int(
+            _cfg(config, "vocoder_input_dim", "io_channels", "latent_ch", default=64)
+        )
+        dec_strides = [
+            int(s) for s in _cfg(config, "dec_strides", default=[2, 4, 5, 6, 8])
+        ]
+        self.hop_size = int(
+            _cfg(
+                config,
+                "hop_size",
+                default=math.prod(dec_strides) if dec_strides else 1920,
+            )
+        )
+        stride_product = math.prod(dec_strides)
+        if stride_product != self.hop_size:
+            raise ValueError(
+                "Cosmos3 AVAE dec_strides product must equal hop_size: "
+                f"product={stride_product}, hop_size={self.hop_size}."
+            )
+
+        norm = str(_cfg(config, "normalization_type", default="none"))
+        if bool(_cfg(config, "normalize_latents", default=False)) and norm == "none":
+            norm = "tanh"
+        self.normalization_type = norm
+        self.tanh_input_scale = float(_cfg(config, "tanh_input_scale", default=1.5))
+        self.tanh_output_scale = float(_cfg(config, "tanh_output_scale", default=3.5))
+        self.tanh_clamp = float(_cfg(config, "tanh_clamp", default=0.995))
+
+        self.decoder = OobleckDecoder(
+            channels=int(_cfg(config, "dec_dim", default=320)),
+            input_channels=self.latent_channels,
+            audio_channels=self.audio_channels,
+            upsampling_ratios=list(reversed(dec_strides)),
+            channel_multiples=list(
+                _cfg(config, "dec_c_mults", default=[1, 2, 4, 8, 16])
+            ),
+        )
+
+    @property
+    def temporal_compression_factor(self) -> int:
+        return self.hop_size
+
+    def get_latent_num_samples(self, num_audio_samples: int) -> int:
+        return int(num_audio_samples) // self.hop_size
+
+    def get_audio_num_samples(self, num_latent_samples: int) -> int:
+        return int(num_latent_samples) * self.hop_size
+
+    def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
+        if self.normalization_type == "tanh":
+            in_dtype = latent.dtype
+            x = torch.clamp(
+                latent.float() / self.tanh_output_scale,
+                -self.tanh_clamp,
+                self.tanh_clamp,
+            )
+            return (torch.atanh(x) * self.tanh_input_scale).to(in_dtype)
+        if self.normalization_type != "none":
+            raise ValueError(
+                f"Unsupported AVAE normalization_type={self.normalization_type!r}."
+            )
+        return latent
+
+    @torch.no_grad()
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        squeeze = latent.ndim == 2
+        if squeeze:
+            latent = latent.unsqueeze(0)
+        decoder_dtype = next(self.decoder.parameters()).dtype
+        decoder_device = next(self.decoder.parameters()).device
+        z = self._denormalize_latent(latent.to(decoder_device)).to(decoder_dtype)
+        audio = self.decoder(z).clamp(-1.0, 1.0).to(latent.dtype)
+        return audio.squeeze(0) if squeeze else audio
+
+
+EntryClass = Cosmos3AVAEAudioTokenizer

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

@@ -41,8 +41,18 @@ class Cosmos3Pipeline(ComposedPipelineBase):
         "vae",
         "transformer",
         "scheduler",
+        "sound_tokenizer",
     ]
 
+    def load_modules(self, server_args, loaded_modules=None):
+        # Visual-only Cosmos3 checkpoints ship no sound_tokenizer; require it
+        # only when the checkpoint actually provides one.
+        if "sound_tokenizer" not in self._load_config():
+            self._required_config_modules = [
+                m for m in self._required_config_modules if m != "sound_tokenizer"
+            ]
+        return super().load_modules(server_args, loaded_modules)
+
     def create_pipeline_stages(self, server_args: ServerArgs) -> None:
         """Create Cosmos3 pipeline stages.
 
@@ -58,6 +68,7 @@ class Cosmos3Pipeline(ComposedPipelineBase):
         vae = self.get_module("vae")
         transformer = self.get_module("transformer")
         scheduler = self.get_module("scheduler")
+        sound_tokenizer = self.get_module("sound_tokenizer")
 
         guardrails_disabled = (
             os.environ.get("SGLANG_DISABLE_COSMOS3_GUARDRAILS", "0") == "1"
@@ -86,7 +97,11 @@ class Cosmos3Pipeline(ComposedPipelineBase):
         self.add_stage(Cosmos3LatentPreparationStage(vae, transformer))
         self.add_stage(Cosmos3TimestepPreparationStage(scheduler))
         self.add_stage(Cosmos3DenoisingStage(transformer, scheduler, server_args))
-        self.add_stage(Cosmos3DecodingStage(vae, guardrails=guardrails_on))
+        self.add_stage(
+            Cosmos3DecodingStage(
+                vae, guardrails=guardrails_on, sound_tokenizer=sound_tokenizer
+            )
+        )
 
         logger.info(
             "Cosmos3 pipeline stages created successfully (guardrails=%s)",

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -85,6 +85,7 @@ class Req:
     vae_image: torch.Tensor | PIL.Image.Image | None = None
     pixel_values: torch.Tensor | PIL.Image.Image | None = None
     preprocessed_image: torch.Tensor | None = None
+    preprocessed_video: torch.Tensor | None = None
 
     output_file_ext: str | None = None
     # Primary encoder embeddings
@@ -136,6 +137,9 @@ class Req:
     # Audio Parameters
     generate_audio: bool = True
 
+    # Action Latents (Cosmos3 action-conditioned generation)
+    action_latents: torch.Tensor | None = None
+
     raw_latent_shape: torch.Tensor | None = None
     did_sp_shard_latents: bool = False
     sp_video_start_frame: int = 0
@@ -415,6 +419,10 @@ class OutputBatch:
     raw_frame_metadata: dict[str, Any] | None = None
     audio: torch.Tensor | None = None
     audio_sample_rate: int | None = None
+    action_pred: torch.Tensor | None = None
+    action_mode: str | None = None
+    action_domain_id: int | None = None
+    action_raw_action_dim: int | None = None
     trajectory_timesteps: torch.Tensor | None = None
     trajectory_latents: torch.Tensor | None = None
     rollout_trajectory_data: RolloutTrajectoryData | None = None

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

@@ -4,10 +4,13 @@ prep, denoising, decode.
 
 Cosmos3 has no separate text encoder — text is tokenized with Qwen2's chat
 template and embedded inside the transformer's UND pathway. The same
-``Cosmos3Pipeline`` serves T2V, I2V, and T2I; mode is dispatched per-request
+``Cosmos3Pipeline`` serves T2V, I2V, V2V, and T2I; mode is dispatched
-from ``batch.data_type`` and the presence of ``batch.preprocessed_image``.
+per-request from ``batch.data_type`` and the presence of
+``batch.preprocessed_image`` / ``batch.preprocessed_video``.
 """
 
+import copy
+import json
 from typing import Any
 
 import numpy as np
@@ -27,11 +30,23 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_sp_world_size,
 )
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.models.vision_utils import load_video
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
     PipelineStage,
     StageParallelismType,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_action import (
+    ACTION_MODE_FORWARD_DYNAMICS,
+    ACTION_MODE_INVERSE_DYNAMICS,
+    ACTION_MODES,
+    EMBODIMENT_TO_DOMAIN_ID,
+    build_action_prompt,
+    denormalize_action,
+    get_raw_action_dim,
+    load_action_stats,
+    normalize_action,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     StageValidators as V,
 )
@@ -57,12 +72,40 @@ COSMOS3_IMAGE_SYSTEM_PROMPT = (
     "You are a helpful assistant who will generate images from a given prompt."
 )
 
+# Per-mode flow-shift defaults, applied only when the request and pipeline
+# config leave flow_shift unset.
+COSMOS3_T2I_FLOW_SHIFT = 3.0
+COSMOS3_I2V_FLOW_SHIFT = 10.0
+COSMOS3_T2V_FLOW_SHIFT = 10.0
+COSMOS3_V2V_FLOW_SHIFT = 10.0
+COSMOS3_ACTION_FLOW_SHIFT = 10.0
+
+
+def _resize_crop_pil(
+    image: PIL.Image.Image, target_w: int, target_h: int
+) -> PIL.Image.Image:
+    """Aspect-preserving resize then center-crop to ``target_w x target_h``."""
+    scale = max(target_w / image.width, target_h / image.height)
+    resize_w = int(np.ceil(scale * image.width))
+    resize_h = int(np.ceil(scale * image.height))
+    image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS)
+    left = (resize_w - target_w) // 2
+    top = (resize_h - target_h) // 2
+    return image.crop((left, top, left + target_w, top + target_h))
+
+
+def _pil_to_normalized_tensor(image: PIL.Image.Image) -> torch.Tensor:
+    """PIL RGB → ``[3, H, W]`` float32 tensor in ``[-1, 1]``."""
+    arr = np.asarray(image, dtype=np.float32) / 127.5 - 1.0
+    return torch.from_numpy(arr).permute(2, 0, 1).contiguous()
+
 
 class Cosmos3ImagePreprocessStage(PipelineStage):
-    """Load, aspect-resize, and center-crop the I2V conditioning image.
+    """Load, aspect-resize, and center-crop the conditioning input.
 
-    No-op when the request has no image (T2V / T2I). The output is a
+    For I2V: writes ``[1, 3, H, W]`` to ``batch.preprocessed_image``.
-    ``[1, 3, H, W]`` tensor in ``[-1, 1]`` written to ``batch.preprocessed_image``.
+    For V2V: writes ``[1, 3, T_in, H, W]`` to ``batch.preprocessed_video``.
+    No-op for T2V / T2I.
     """
 
     parallelism_type = StageParallelismType.REPLICATED
@@ -74,26 +117,89 @@ class Cosmos3ImagePreprocessStage(PipelineStage):
         image_path = batch.image_path
         if isinstance(image_path, list):
             image_path = image_path[0] if image_path else None
-        if not isinstance(image_path, str) or not image_path:
+        video_path = batch.video_path
-            return batch
+        if isinstance(video_path, list):
+            video_path = video_path[0] if video_path else None
+
+        if image_path and video_path:
+            raise ValueError(
+                "Cosmos3 accepts either --image-path (I2V) or --video-path "
+                "(V2V), not both"
+            )
 
-        image = PIL.Image.open(image_path).convert("RGB")
         target_h, target_w = batch.height, batch.width
-        scale = max(target_w / image.width, target_h / image.height)
-        resize_w = int(np.ceil(scale * image.width))
-        resize_h = int(np.ceil(scale * image.height))
-        image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS)
-        left = (resize_w - target_w) // 2
-        top = (resize_h - target_h) // 2
-        image = image.crop((left, top, left + target_w, top + target_h))
 
-        arr = np.asarray(image, dtype=np.float32) / 127.5 - 1.0
+        if isinstance(image_path, str) and image_path:
-        tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).contiguous()
+            image = PIL.Image.open(image_path).convert("RGB")
-
+            image = _resize_crop_pil(image, target_w, target_h)
-        batch.preprocessed_image = tensor
+            batch.preprocessed_image = _pil_to_normalized_tensor(image).unsqueeze(0)
             self.log_info(f"Preprocessed conditioning image to {target_w}x{target_h}")
             return batch
 
+        if isinstance(video_path, str) and video_path:
+            frames = load_video(video_path)
+            if not frames:
+                raise ValueError(f"No frames decoded from video: {video_path!r}")
+
+            keep = (
+                getattr(batch.sampling_params, "condition_video_keep", "first")
+                or "first"
+            )
+            if keep not in ("first", "last"):
+                raise ValueError(
+                    f"condition_video_keep must be 'first' or 'last', got {keep!r}"
+                )
+            cond_indexes = self._resolve_condition_indexes(batch)
+            # Encode the full output-length video so that the latent positions
+            # we lock match what the decoder will reconstruct at those frame
+            # indices. Encoding only the first ``max_idx*4+1`` frames produces
+            # an out-of-distribution latent for the locked slots and decodes
+            # to noise.
+            num_source_frames = max(cond_indexes) * 4 + 1
+            num_target_frames = batch.num_frames
+            if keep == "last":
+                frames = frames[-num_source_frames:]
+            else:
+                frames = frames[:num_source_frames]
+            if len(frames) < num_source_frames:
+                frames = frames + [frames[-1]] * (num_source_frames - len(frames))
+            if len(frames) < num_target_frames:
+                frames = frames + [frames[-1]] * (num_target_frames - len(frames))
+
+            processed = [
+                _pil_to_normalized_tensor(
+                    _resize_crop_pil(f.convert("RGB"), target_w, target_h)
+                )
+                for f in frames
+            ]
+            video_tensor = torch.stack(processed, dim=1).unsqueeze(0).contiguous()
+            batch.preprocessed_video = video_tensor
+            self.log_info(
+                f"Preprocessed conditioning video to "
+                f"{video_tensor.shape[2]}x{target_h}x{target_w} "
+                f"(keep={keep}, source frames={num_source_frames}, padded to {num_target_frames})"
+            )
+
+        return batch
+
+    @staticmethod
+    def _resolve_condition_indexes(batch: Req) -> list[int]:
+        """Resolve condition_frame_indexes for V2V (default ``[0, 1]``).
+
+        Inverse-dynamics action mode conditions on the whole input video, so
+        every latent frame is locked.
+        """
+        if (
+            getattr(batch.sampling_params, "action_mode", None)
+            == ACTION_MODE_INVERSE_DYNAMICS
+        ):
+            num_latent_frames = (batch.num_frames - 1) // 4 + 1
+            return list(range(num_latent_frames))
+        cond_indexes = getattr(batch.sampling_params, "condition_frame_indexes", None)
+        if not cond_indexes:
+            return [0, 1]
+        return sorted(set(int(i) for i in cond_indexes))
+
 
 class Cosmos3TokenizationStage(PipelineStage):
     """Tokenization stage for Cosmos3.
@@ -207,6 +313,21 @@ class Cosmos3TokenizationStage(PipelineStage):
             COSMOS3_IMAGE_SYSTEM_PROMPT if is_image_gen else COSMOS3_VIDEO_SYSTEM_PROMPT
         )
 
+        # Action mode uses a structured JSON caption with neither a system
+        # prompt nor the duration suffix.
+        if getattr(batch.sampling_params, "action_mode", None) is not None:
+            prompt = build_action_prompt(
+                prompt,
+                getattr(batch.sampling_params, "action_view_point", "ego_view"),
+                num_frames,
+                fps,
+                batch.height,
+                batch.width,
+            )
+            use_system_prompt = False
+            use_duration_template = False
+            self.log_info(f"Action prompt: {prompt}")
+
         # Apply duration template if enabled (no temporal concept for T2I).
         if use_duration_template and not is_image_gen and num_frames > 1:
             duration = num_frames / fps
@@ -250,10 +371,13 @@ class Cosmos3TokenizationStage(PipelineStage):
 class Cosmos3LatentPreparationStage(PipelineStage):
     """Initialize the noisy latent for Cosmos3.
 
-    T2V / T2I produce pure Gaussian noise. I2V VAE-encodes the conditioning
+    T2V / T2I produce pure Gaussian noise. I2V / V2V VAE-encode the
-    image, replaces frame 0 of the latent with the encoded image, and stashes
+    conditioning input, write the resulting latents at the conditioned
-    a per-frame velocity mask plus the clean frame-0 latent for the denoiser
+    frame indexes, and stash a per-frame velocity mask plus the full
-    to re-inject after each scheduler step.
+    condition latent so the denoiser can re-blend after each scheduler step.
+    I2V is the special case of conditioning at frame ``[0]`` with the image
+    expanded across the temporal axis; V2V conditions at ``[0, 1]`` (or a
+    user-supplied list) with frames from the input video.
     """
 
     parallelism_type = StageParallelismType.REPLICATED
@@ -316,26 +440,55 @@ class Cosmos3LatentPreparationStage(PipelineStage):
 
         noise = torch.randn(shape, generator=generator, device=device, dtype=dtype)
 
-        is_i2v = (
+        is_video_gen = batch.data_type == DataType.VIDEO
-            batch.preprocessed_image is not None and batch.data_type == DataType.VIDEO
+        has_image_cond = batch.preprocessed_image is not None and is_video_gen
-        )
+        has_video_cond = batch.preprocessed_video is not None and is_video_gen
 
-        if is_i2v:
+        if has_image_cond or has_video_cond:
             vae_dtype = next(self.vae.parameters()).dtype
-            pixel_video = batch.preprocessed_image.unsqueeze(2).to(
+
+            if has_video_cond:
+                pixel_input = batch.preprocessed_video.to(
                     device=device, dtype=vae_dtype
                 )
-            with torch.no_grad():
+                cond_indexes = Cosmos3ImagePreprocessStage._resolve_condition_indexes(
-                cond_latent = self._vae_encode(pixel_video).to(dtype)
+                    batch
+                )
+            else:
+                pixel_input = batch.preprocessed_image.unsqueeze(2).to(
+                    device=device, dtype=vae_dtype
+                )
+                cond_indexes = [0]
 
+            with torch.no_grad():
+                cond_latent = self._vae_encode(pixel_input).to(dtype)
+
+            max_idx = max(cond_indexes)
+            if max_idx >= num_latent_frames:
+                raise ValueError(
+                    f"condition_frame_indexes={cond_indexes} exceeds the "
+                    f"latent frame count {num_latent_frames} for "
+                    f"num_frames={batch.num_frames}"
+                )
+
+            condition_latents = torch.zeros_like(noise)
             condition_mask = torch.zeros(
                 1, 1, num_latent_frames, 1, 1, device=device, dtype=dtype
             )
-            condition_mask[:, :, 0, :, :] = 1.0
+            for idx in cond_indexes:
-            latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise
+                src = min(idx, cond_latent.shape[2] - 1)
-            batch.image_latent = cond_latent[:, :, 0:1, :, :].clone()
+                condition_latents[:, :, idx, :, :] = cond_latent[:, :, src, :, :]
+                condition_mask[:, :, idx, :, :] = 1.0
+
+            latents = (
+                condition_mask * condition_latents + (1.0 - condition_mask) * noise
+            )
+            batch.extra["condition_latents"] = condition_latents
             batch.extra["velocity_mask"] = 1.0 - condition_mask
-            self.log_info("Prepared I2V latents with frame-0 conditioning")
+            mode = "V2V" if has_video_cond else "I2V"
+            self.log_info(
+                f"Prepared {mode} latents with conditioning at frames {cond_indexes}"
+            )
         else:
             latents = noise
 
@@ -347,8 +500,172 @@ class Cosmos3LatentPreparationStage(PipelineStage):
         batch.extra["vae_scale_factor_spatial"] = vae_scale_factor_spatial
 
         self.log_info(f"Prepared latents with shape {shape}")
+
+        sound_duration = float(getattr(batch, "sound_duration", 0.0) or 0.0)
+        if sound_duration > 0.0:
+            if not getattr(self.transformer, "sound_gen", False):
+                raise ValueError(
+                    "sound generation was requested (sound_duration > 0) but the "
+                    "loaded Cosmos3 checkpoint has no sound modality (sound_gen is "
+                    "False)."
+                )
+            sound_latent_fps = self.transformer.sound_latent_fps
+            sound_latent_frames = max(1, round(sound_duration * sound_latent_fps))
+            sound_shape = (1, self.transformer.sound_dim, sound_latent_frames)
+            batch.audio_latents = torch.randn(
+                sound_shape, generator=generator, device=device, dtype=dtype
+            )
+            self.log_info(f"Prepared sound latents with shape {sound_shape}")
+
+        action_mode = getattr(batch.sampling_params, "action_mode", None)
+        if action_mode is not None:
+            if getattr(self.transformer, "action_dim", None) is None:
+                raise ValueError(
+                    "action_mode is set but the loaded Cosmos3 checkpoint has no "
+                    "action modality (action_gen is False)."
+                )
+            self._prepare_action_latents(batch, generator, device, dtype)
         return batch
 
+    @staticmethod
+    def _resolve_domain_id(batch: Req) -> int:
+        """Resolve action embodiment domain ID; required for action generation."""
+        domain_id = getattr(batch.sampling_params, "domain_id", None)
+        if domain_id is not None:
+            domain_id = int(domain_id)
+            if domain_id < 0:
+                raise ValueError(f"domain_id must be non-negative, got {domain_id}")
+            return domain_id
+        domain_name = getattr(batch.sampling_params, "domain_name", None)
+        if domain_name:
+            key = str(domain_name).strip().lower()
+            if key not in EMBODIMENT_TO_DOMAIN_ID:
+                raise ValueError(
+                    f"Unknown action domain name {domain_name!r}. "
+                    f"Valid names: {sorted(EMBODIMENT_TO_DOMAIN_ID)}"
+                )
+            return EMBODIMENT_TO_DOMAIN_ID[key]
+        raise ValueError(
+            "Cosmos3 action generation requires --domain-id or --domain-name."
+        )
+
+    def _prepare_action_latents(
+        self,
+        batch: Req,
+        generator,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> None:
+        """Prepare action latents and conditioning, writing them onto ``batch``.
+
+        Action tokens run at frame rate (no temporal compression), so the chunk
+        length is ``num_frames - 1`` with ``start_frame_offset=1`` so each action
+        aligns with the frame it drives.
+
+        Three modes:
+        - ``forward_dynamics``: the user supplies the action; all tokens are
+          clean conditioning (velocity mask 0) and the model predicts video.
+        - ``policy`` / ``inverse_dynamics``: actions are denoised from noise
+          (velocity mask 1); ``raw_action_dim`` is required.
+        """
+        sp = batch.sampling_params
+        mode = str(sp.action_mode).strip().lower()
+        if mode not in ACTION_MODES:
+            raise ValueError(
+                f"Unsupported action_mode={sp.action_mode!r}; "
+                f"expected one of {sorted(ACTION_MODES)}"
+            )
+        action_dim = self.transformer.action_dim
+        num_frames = batch.num_frames
+
+        action_chunk_size = num_frames - 1 if num_frames > 1 else 1
+        action_offset = 1 if action_chunk_size == num_frames - 1 else 0
+
+        domain_id = self._resolve_domain_id(batch)
+        raw_action_dim = getattr(sp, "raw_action_dim", None)
+        if raw_action_dim is None:
+            embodiment = getattr(sp, "domain_name", None)
+            if embodiment:
+                raw_action_dim = get_raw_action_dim(embodiment)
+
+        if mode == ACTION_MODE_FORWARD_DYNAMICS:
+            raw = getattr(sp, "action", None)
+            if raw is None:
+                raise ValueError(
+                    "action_mode='forward_dynamics' requires an 'action' array "
+                    "(list[list[float]] of shape [T, D])."
+                )
+            if isinstance(raw, str):
+                raw = json.loads(raw)
+            action = torch.as_tensor(np.asarray(raw), dtype=torch.float32)
+            if action.ndim == 3 and action.shape[0] == 1:
+                action = action.squeeze(0)
+            if action.ndim != 2:
+                raise ValueError(
+                    f"action must have shape [T, D], got {tuple(action.shape)}"
+                )
+            if action.shape[0] < action_chunk_size:
+                pad = action[-1:].repeat(action_chunk_size - action.shape[0], 1)
+                action = torch.cat([action, pad], dim=0)
+            elif action.shape[0] > action_chunk_size:
+                action = action[:action_chunk_size]
+            if raw_action_dim is None:
+                raw_action_dim = int(action.shape[-1])
+            stats_path = getattr(sp, "action_stats_path", None)
+            if stats_path is not None:
+                method = getattr(sp, "action_normalization", "quantile")
+                action = normalize_action(action, method, load_action_stats(stats_path))
+            if action.shape[-1] < action_dim:
+                pad = torch.zeros(action.shape[0], action_dim - action.shape[-1])
+                action = torch.cat([action, pad], dim=-1)
+            clean_action = action.to(device=device, dtype=dtype).unsqueeze(0)
+        else:
+            if raw_action_dim is None:
+                raise ValueError(f"action_mode={mode!r} requires --raw-action-dim.")
+            clean_action = torch.zeros(
+                1, action_chunk_size, action_dim, device=device, dtype=dtype
+            )
+
+        raw_action_dim = int(raw_action_dim)
+        if not 0 < raw_action_dim <= action_dim:
+            raise ValueError(
+                f"raw_action_dim must be in [1, {action_dim}], got {raw_action_dim}"
+            )
+
+        # condition_mask marks clean (given) action tokens. forward_dynamics
+        # conditions on the whole action sequence; the others denoise it fully.
+        condition_mask = torch.zeros(
+            1, action_chunk_size, 1, device=device, dtype=dtype
+        )
+        if mode == ACTION_MODE_FORWARD_DYNAMICS:
+            condition_mask[:] = 1.0
+
+        noise = torch.randn(
+            1,
+            action_chunk_size,
+            action_dim,
+            generator=generator,
+            device=device,
+            dtype=dtype,
+        )
+        noise[:, :, raw_action_dim:] = 0.0
+        clean_action[:, :, raw_action_dim:] = 0.0
+        action_latents = condition_mask * clean_action + (1.0 - condition_mask) * noise
+
+        batch.action_latents = action_latents
+        batch.extra["action_domain_ids"] = torch.tensor(
+            [domain_id], dtype=torch.long, device=device
+        )
+        batch.extra["action_velocity_mask"] = 1.0 - condition_mask
+        batch.extra["action_condition_latents"] = clean_action
+        batch.extra["raw_action_dim"] = raw_action_dim
+        batch.extra["action_start_frame_offset"] = action_offset
+        self.log_info(
+            f"Prepared action latents with shape {tuple(action_latents.shape)} "
+            f"(mode={mode}, domain_id={domain_id}, raw_action_dim={raw_action_dim}, "
+            f"start_frame_offset={action_offset})"
+        )
+
 
 class Cosmos3TimestepPreparationStage(PipelineStage):
     """
@@ -362,9 +679,21 @@ class Cosmos3TimestepPreparationStage(PipelineStage):
     def __init__(self, scheduler):
         super().__init__()
         self.scheduler = scheduler
-        self.default_flow_shift = getattr(
+
-            getattr(scheduler, "config", None), "flow_shift", None
+    def _default_flow_shift_for_mode(self, batch: Req) -> float | None:
-        )
+        """Resolve the per-mode default flow_shift for the request.
+
+        Matches cosmos-framework's built-in per-mode sample defaults.
+        """
+        if getattr(batch.sampling_params, "action_mode", None) is not None:
+            return COSMOS3_ACTION_FLOW_SHIFT
+        if batch.data_type == DataType.IMAGE:
+            return COSMOS3_T2I_FLOW_SHIFT
+        if batch.preprocessed_image is not None:
+            return COSMOS3_I2V_FLOW_SHIFT
+        if batch.preprocessed_video is not None:
+            return COSMOS3_V2V_FLOW_SHIFT
+        return COSMOS3_T2V_FLOW_SHIFT
 
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         """Prepare scheduler timesteps."""
@@ -374,7 +703,7 @@ class Cosmos3TimestepPreparationStage(PipelineStage):
         if flow_shift is None:
             flow_shift = server_args.pipeline_config.flow_shift
         if flow_shift is None:
-            flow_shift = self.default_flow_shift
+            flow_shift = self._default_flow_shift_for_mode(batch)
         if flow_shift is not None and hasattr(self.scheduler, "set_shift"):
             self.scheduler.set_shift(float(flow_shift))
 
@@ -502,7 +831,13 @@ class Cosmos3DenoisingStage(PipelineStage):
         noisy_frame_mask: torch.Tensor | None = None,
         max_text_seq_len: int | None = None,
         current_timestep: int | None = None,
-    ) -> torch.Tensor:
+        sound_latents: torch.Tensor | None = None,
+        action_latents: torch.Tensor | None = None,
+        action_domain_ids: torch.Tensor | None = None,
+        action_noisy_mask: torch.Tensor | None = None,
+        action_fps: float | None = None,
+        action_start_frame_offset: int = 1,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Run transformer forward pass.
 
         Args:
@@ -514,7 +849,7 @@ class Cosmos3DenoisingStage(PipelineStage):
             fps: Video frame rate
             cache_key: Key for the UND K/V cache. Use "cond" for conditional
                 and "uncond" for unconditional to enable cache reuse across steps.
-            noisy_frame_mask: Optional [B, 1, T, 1, 1] I2V conditioning mask.
+            noisy_frame_mask: Optional [B, 1, T, 1, 1] I2V / V2V conditioning mask.
         """
         if current_timestep is None:
             current_timestep = int(timestep.flatten()[0].item())
@@ -529,6 +864,12 @@ class Cosmos3DenoisingStage(PipelineStage):
                 cache_key=cache_key,
                 noisy_frame_mask=noisy_frame_mask,
                 max_text_seq_len=max_text_seq_len,
+                sound_latents=sound_latents,
+                action_latents=action_latents,
+                action_domain_ids=action_domain_ids,
+                action_noisy_mask=action_noisy_mask,
+                action_fps=action_fps,
+                action_start_frame_offset=action_start_frame_offset,
             )
 
     def _manage_device_placement(self, server_args: ServerArgs):
@@ -565,6 +906,14 @@ class Cosmos3DenoisingStage(PipelineStage):
         self._manage_device_placement(server_args)
 
         latents = batch.latents
+        sound_latents = batch.audio_latents
+        action_latents = getattr(batch, "action_latents", None)
+        action_domain_ids = batch.extra.get("action_domain_ids")
+        action_velocity_mask = batch.extra.get("action_velocity_mask")
+        action_condition_latents = batch.extra.get("action_condition_latents")
+        action_raw_dim = batch.extra.get("raw_action_dim")
+        action_start_frame_offset = batch.extra.get("action_start_frame_offset", 1)
+        action_fps = getattr(batch.sampling_params, "action_fps", None)
         timesteps = batch.timesteps
         guidance_scale = batch.guidance_scale
 
@@ -575,12 +924,28 @@ class Cosmos3DenoisingStage(PipelineStage):
         video_shape = batch.extra["video_shape"]
         fps = batch.extra.get("fps", 24.0)
         velocity_mask = batch.extra.get("velocity_mask")
-        image_latent = batch.image_latent
+        condition_latents = batch.extra.get("condition_latents")
         guidance_interval = getattr(batch.sampling_params, "guidance_interval", None)
 
         do_cfg = guidance_scale > 1.0
 
         enable_cfg_parallel = server_args.enable_cfg_parallel and do_cfg
+        if action_latents is not None and enable_cfg_parallel:
+            raise NotImplementedError(
+                "Cosmos3 action generation does not support CFG parallel yet"
+            )
+
+        # Use separate scheduler instances for action/sound: UniPC keeps a
+        # per-call output history sized to the last sample, so video (5D),
+        # action (3D), and sound (3D) steps must not share state.
+        sound_scheduler = None
+        if sound_latents is not None:
+            sound_scheduler = copy.deepcopy(self.scheduler)
+            sound_scheduler.set_timesteps(len(timesteps), device=timesteps.device)
+        action_scheduler = None
+        if action_latents is not None:
+            action_scheduler = copy.deepcopy(self.scheduler)
+            action_scheduler.set_timesteps(len(timesteps), device=timesteps.device)
         cfg_rank = get_classifier_free_guidance_rank() if enable_cfg_parallel else 0
         cfg_world_size = (
             get_classifier_free_guidance_world_size() if enable_cfg_parallel else 1
@@ -644,6 +1009,12 @@ class Cosmos3DenoisingStage(PipelineStage):
                         cond_text_seq_len=batch.extra["cond_text_seq_len"],
                         uncond_text_seq_len=batch.extra["uncond_text_seq_len"],
                         current_timestep=i,
+                        sound_latents=sound_latents,
+                        action_latents=action_latents,
+                        action_domain_ids=action_domain_ids,
+                        action_noisy_mask=action_velocity_mask,
+                        action_fps=action_fps,
+                        action_start_frame_offset=action_start_frame_offset,
                     )
                 elif effective_scale == 1.0:
                     noise_pred = self._run_transformer(
@@ -657,6 +1028,12 @@ class Cosmos3DenoisingStage(PipelineStage):
                         noisy_frame_mask=velocity_mask,
                         max_text_seq_len=batch.extra["cond_text_seq_len"],
                         current_timestep=i,
+                        sound_latents=sound_latents,
+                        action_latents=action_latents,
+                        action_domain_ids=action_domain_ids,
+                        action_noisy_mask=action_velocity_mask,
+                        action_fps=action_fps,
+                        action_start_frame_offset=action_start_frame_offset,
                     )
                 else:
                     noise_pred = self._predict_noise_cfg_batched(
@@ -675,6 +1052,12 @@ class Cosmos3DenoisingStage(PipelineStage):
                             batch.extra["uncond_text_seq_len"],
                         ),
                         current_timestep=i,
+                        sound_latents=sound_latents,
+                        action_latents=action_latents,
+                        action_domain_ids=action_domain_ids,
+                        action_noisy_mask=action_velocity_mask,
+                        action_fps=action_fps,
+                        action_start_frame_offset=action_start_frame_offset,
                     )
             else:
                 noise_pred = self._run_transformer(
@@ -688,11 +1071,30 @@ class Cosmos3DenoisingStage(PipelineStage):
                     noisy_frame_mask=velocity_mask,
                     max_text_seq_len=batch.extra["cond_text_seq_len"],
                     current_timestep=i,
+                    sound_latents=sound_latents,
+                    action_latents=action_latents,
+                    action_domain_ids=action_domain_ids,
+                    action_noisy_mask=action_velocity_mask,
+                    action_fps=action_fps,
+                    action_start_frame_offset=action_start_frame_offset,
                 )
 
-            # I2V: zero-velocity at conditioned frames so the scheduler keeps
+            # Unpack multi-modality outputs; ordering is (video[, action][, sound]).
-            # them clean; UniPC's predictor-corrector still rescales the
+            action_noise_pred = None
-            # sample, so we re-inject the clean image latent below.
+            sound_noise_pred = None
+            if isinstance(noise_pred, tuple):
+                out_idx = 1
+                video_noise_pred = noise_pred[0]
+                if action_latents is not None:
+                    action_noise_pred = noise_pred[out_idx]
+                    out_idx += 1
+                if sound_latents is not None:
+                    sound_noise_pred = noise_pred[out_idx]
+                noise_pred = video_noise_pred
+
+            # I2V / V2V: zero-velocity at conditioned frames so the scheduler
+            # keeps them clean; UniPC's predictor-corrector still rescales the
+            # sample, so we re-blend the clean condition latents below.
             if velocity_mask is not None:
                 noise_pred = noise_pred * velocity_mask
 
@@ -703,13 +1105,53 @@ class Cosmos3DenoisingStage(PipelineStage):
                 return_dict=False,
             )[0]
 
-            if image_latent is not None:
+            if action_noise_pred is not None:
-                latents[:, :, 0:1, :, :] = image_latent
+                # Zero the velocity at conditioned (clean) action tokens and at
+                # padding dims so the scheduler only denoises the active slots,
+                # then re-blend the clean condition after the step.
+                if action_velocity_mask is not None:
+                    action_noise_pred = action_noise_pred * action_velocity_mask
+                if (
+                    action_raw_dim is not None
+                    and action_raw_dim < action_noise_pred.shape[-1]
+                ):
+                    action_noise_pred[..., action_raw_dim:] = 0.0
+                action_latents = action_scheduler.step(
+                    action_noise_pred,
+                    t,
+                    action_latents,
+                    return_dict=False,
+                )[0]
+                if (
+                    action_condition_latents is not None
+                    and action_velocity_mask is not None
+                ):
+                    action_latents = (
+                        action_velocity_mask * action_latents
+                        + (1.0 - action_velocity_mask) * action_condition_latents
+                    )
+
+            if sound_noise_pred is not None:
+                sound_latents = sound_scheduler.step(
+                    sound_noise_pred,
+                    t,
+                    sound_latents,
+                    return_dict=False,
+                )[0]
+
+            if condition_latents is not None and velocity_mask is not None:
+                latents = (
+                    velocity_mask * latents + (1.0 - velocity_mask) * condition_latents
+                )
 
             if batch.profile and not batch.is_warmup:
                 self.step_profile()
 
         batch.latents = latents
+        if action_latents is not None:
+            batch.action_latents = action_latents
+        if sound_latents is not None:
+            batch.audio_latents = sound_latents
         self.log_info("Denoising complete")
         return batch
 
@@ -727,7 +1169,13 @@ class Cosmos3DenoisingStage(PipelineStage):
         noisy_frame_mask: torch.Tensor | None = None,
         max_text_seq_len: int | None = None,
         current_timestep: int | None = None,
-    ) -> torch.Tensor:
+        sound_latents: torch.Tensor | None = None,
+        action_latents: torch.Tensor | None = None,
+        action_domain_ids: torch.Tensor | None = None,
+        action_noisy_mask: torch.Tensor | None = None,
+        action_fps: float | None = None,
+        action_start_frame_offset: int = 1,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Run CFG by stacking both branches into a batch_size=2 forward.
 
         Halves the kernel-launch count vs running cond and uncond serially.
@@ -743,8 +1191,28 @@ class Cosmos3DenoisingStage(PipelineStage):
             if noisy_frame_mask is not None
             else None
         )
+        sound_batched = (
+            torch.cat([sound_latents, sound_latents], dim=0)
+            if sound_latents is not None
+            else None
+        )
+        action_batched = (
+            torch.cat([action_latents, action_latents], dim=0)
+            if action_latents is not None
+            else None
+        )
+        action_domain_ids_batched = (
+            torch.cat([action_domain_ids, action_domain_ids], dim=0)
+            if action_domain_ids is not None
+            else None
+        )
+        action_noisy_mask_batched = (
+            torch.cat([action_noisy_mask, action_noisy_mask], dim=0)
+            if action_noisy_mask is not None
+            else None
+        )
 
-        noise_pred = self._run_transformer(
+        out = self._run_transformer(
             latents=latents_batched,
             timestep=timestep_batched,
             text_ids=text_ids_batched,
@@ -755,13 +1223,21 @@ class Cosmos3DenoisingStage(PipelineStage):
             noisy_frame_mask=mask_batched,
             max_text_seq_len=max_text_seq_len,
             current_timestep=current_timestep,
+            sound_latents=sound_batched,
+            action_latents=action_batched,
+            action_domain_ids=action_domain_ids_batched,
+            action_noisy_mask=action_noisy_mask_batched,
+            action_fps=action_fps,
+            action_start_frame_offset=action_start_frame_offset,
         )
 
-        noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2, dim=0)
+        def _cfg_combine(pred: torch.Tensor) -> torch.Tensor:
-        # CFG: uncond + g·(cond − uncond).
+            uncond, cond = pred.chunk(2, dim=0)
-        return noise_pred_uncond + guidance_scale * (
+            return uncond + guidance_scale * (cond - uncond)
-            noise_pred_cond - noise_pred_uncond
+
-        )
+        if isinstance(out, tuple):
+            return tuple(_cfg_combine(p) for p in out)
+        return _cfg_combine(out)
 
     def _predict_noise_cfg_parallel(
         self,
@@ -779,45 +1255,54 @@ class Cosmos3DenoisingStage(PipelineStage):
         cond_text_seq_len: int | None = None,
         uncond_text_seq_len: int | None = None,
         current_timestep: int | None = None,
-    ) -> torch.Tensor:
+        sound_latents: torch.Tensor | None = None,
+        action_latents: torch.Tensor | None = None,
+        action_domain_ids: torch.Tensor | None = None,
+        action_noisy_mask: torch.Tensor | None = None,
+        action_fps: float | None = None,
+        action_start_frame_offset: int = 1,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Run CFG with one branch per CFG rank, combined by all-reduce.
 
         Rank 0 runs the conditional branch and contributes ``g·cond`` to the
         sum; rank 1 runs the unconditional branch and contributes
         ``(1−g)·uncond``. The all-reduce sum is exactly the standard CFG
         result. Each rank keeps its own UND K/V cache (``"cond"`` /
-        ``"uncond"``).
+        ``"uncond"``). When sound/action modalities are present the forward
+        returns a per-modality tuple; each branch scales every modality by its
+        coefficient and the reduction combines them element-wise.
         """
         if cfg_rank == 0:
-            noise_pred = self._run_transformer(
+            text_ids, text_mask, cache_key = cond_text_ids, cond_text_mask, "cond"
-                latents=latents,
+            text_seq_len = cond_text_seq_len
-                timestep=timestep,
+            coeff = guidance_scale
-                text_ids=cond_text_ids,
-                text_mask=cond_text_mask,
-                video_shape=video_shape,
-                fps=fps,
-                cache_key="cond",
-                noisy_frame_mask=noisy_frame_mask,
-                max_text_seq_len=cond_text_seq_len,
-                current_timestep=current_timestep,
-            )
-            partial = guidance_scale * noise_pred
         else:
-            noise_pred = self._run_transformer(
+            text_ids, text_mask, cache_key = uncond_text_ids, uncond_text_mask, "uncond"
+            text_seq_len = uncond_text_seq_len
+            coeff = 1.0 - guidance_scale
+
+        out = self._run_transformer(
             latents=latents,
             timestep=timestep,
-                text_ids=uncond_text_ids,
+            text_ids=text_ids,
-                text_mask=uncond_text_mask,
+            text_mask=text_mask,
             video_shape=video_shape,
             fps=fps,
-                cache_key="uncond",
+            cache_key=cache_key,
             noisy_frame_mask=noisy_frame_mask,
-                max_text_seq_len=uncond_text_seq_len,
+            max_text_seq_len=text_seq_len,
             current_timestep=current_timestep,
+            sound_latents=sound_latents,
+            action_latents=action_latents,
+            action_domain_ids=action_domain_ids,
+            action_noisy_mask=action_noisy_mask,
+            action_fps=action_fps,
+            action_start_frame_offset=action_start_frame_offset,
         )
-            partial = (1.0 - guidance_scale) * noise_pred
 
-        return cfg_model_parallel_all_reduce(partial)
+        if isinstance(out, tuple):
+            return tuple(cfg_model_parallel_all_reduce(coeff * p) for p in out)
+        return cfg_model_parallel_all_reduce(coeff * out)
 
 
 class Cosmos3DecodingStage(PipelineStage):
@@ -830,12 +1315,13 @@ class Cosmos3DecodingStage(PipelineStage):
 
     parallelism_type = StageParallelismType.REPLICATED
 
-    def __init__(self, vae, guardrails: bool = False):
+    def __init__(self, vae, guardrails: bool = False, sound_tokenizer=None):
         super().__init__()
         self.vae = vae
         self._latents_mean = None
         self._latents_std = None
         self._guardrails = guardrails
+        self.sound_tokenizer = sound_tokenizer
         if guardrails:
             from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_guardrails import (
                 _init_guardrails,
@@ -935,7 +1421,50 @@ class Cosmos3DecodingStage(PipelineStage):
         elif not is_image_gen:
             self.log_info(f"Postprocessed video tensor shape: {output.shape}")
 
+        audio = None
+        audio_sample_rate = None
+        if self.sound_tokenizer is not None and batch.audio_latents is not None:
+            if server_args.vae_cpu_offload:
+                self.sound_tokenizer.to(device)
+            with torch.no_grad():
+                decoded_audio = self.sound_tokenizer.decode(
+                    batch.audio_latents.to(device)
+                )
+            audio = decoded_audio.float().cpu()
+            audio_sample_rate = self.sound_tokenizer.sample_rate
+            if server_args.vae_cpu_offload and not getattr(batch, "is_warmup", False):
+                self.sound_tokenizer.to("cpu", non_blocking=True)
+            self.log_info(
+                f"Decoded audio tensor shape: {tuple(audio.shape)} @ {audio_sample_rate} Hz"
+            )
+
+        action_pred = None
+        if getattr(batch, "action_latents", None) is not None:
+            raw_action_dim = batch.extra.get("raw_action_dim")
+            action_pred = batch.action_latents.float().cpu()
+            if raw_action_dim is not None:
+                action_pred = action_pred[:, :, :raw_action_dim]
+            stats_path = getattr(batch.sampling_params, "action_stats_path", None)
+            if stats_path is not None:
+                method = getattr(
+                    batch.sampling_params, "action_normalization", "quantile"
+                )
+                action_pred = denormalize_action(
+                    action_pred, method, load_action_stats(stats_path)
+                )
+            self.log_info(f"Action predictions shape: {tuple(action_pred.shape)}")
+
         return OutputBatch(
             output=output,
+            audio=audio,
+            audio_sample_rate=audio_sample_rate,
+            action_pred=action_pred,
+            action_mode=getattr(batch.sampling_params, "action_mode", None),
+            action_domain_id=getattr(batch.sampling_params, "domain_id", None),
+            action_raw_action_dim=(
+                batch.extra.get("raw_action_dim")
+                if getattr(batch, "extra", None)
+                else None
+            ),
             metrics=batch.metrics if hasattr(batch, "metrics") else None,
         )

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

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 action modality helpers: domain mapping, mode constants, the
+structured JSON caption, and dataset-derived action (de)normalization.
+"""
+
+import json
+import math
+from pathlib import Path
+
+import numpy as np
+import torch
+
+ACTION_MODE_POLICY = "policy"
+ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics"
+ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics"
+ACTION_MODES = {
+    ACTION_MODE_POLICY,
+    ACTION_MODE_FORWARD_DYNAMICS,
+    ACTION_MODE_INVERSE_DYNAMICS,
+}
+
+EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = {
+    "no_action": 0,
+    "av": 1,
+    "camera_pose": 2,
+    "hand_pose": 3,
+    "pusht": 4,
+    "libero": 5,
+    "umi": 6,
+    "bridge_orig_lerobot": 7,
+    "droid_lerobot": 8,
+    "robomind-franka": 8,
+    "galbot": 9,
+    "robomind-franka-dual": 12,
+    "robomind-ur": 13,
+    "agibotworld": 15,
+    "agibot_gear_gripper": 15,
+    "agibot_gear_gripper_ext": 15,
+    "fractal": 20,
+}
+
+# Embodiment -> real (unpadded) action channel count. Channels beyond this are
+# zero-padding up to the model's action_dim.
+EMBODIMENT_TO_RAW_ACTION_DIM: dict[str, int] = {
+    "av": 9,
+    "camera_pose": 9,
+    "pusht": 2,
+    "umi": 10,
+    "bridge_orig_lerobot": 10,
+    "droid_lerobot": 10,
+    "robomind-franka": 10,
+    "robomind-franka-dual": 20,
+    "robomind-ur": 10,
+    "agibotworld": 29,
+    "fractal": 10,
+}
+
+# Canonical (width, height) targets per resolution tier and aspect ratio, used
+# to render the aspect_ratio field of the action caption.
+VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = {
+    "256": {
+        "1,1": (256, 256),
+        "4,3": (320, 256),
+        "3,4": (256, 320),
+        "16,9": (320, 192),
+        "9,16": (192, 320),
+    },
+    "480": {
+        "1,1": (640, 640),
+        "4,3": (736, 544),
+        "3,4": (544, 736),
+        "16,9": (832, 480),
+        "9,16": (480, 832),
+    },
+    "704": {
+        "1,1": (960, 960),
+        "4,3": (1088, 832),
+        "3,4": (832, 1088),
+        "16,9": (1280, 704),
+        "9,16": (704, 1280),
+    },
+    "720": {
+        "1,1": (960, 960),
+        "4,3": (1104, 832),
+        "3,4": (832, 1104),
+        "16,9": (1280, 720),
+        "9,16": (720, 1280),
+    },
+}
+
+VIEWPOINT_TEMPLATES: dict[str, str] = {
+    "ego_view": "This video is captured from a first-person perspective looking at the scene.",
+    "third_person_view": "This video is captured from a third-person perspective looking towards the agent from the front.",
+    "wrist_view": "This video is captured from a wrist-mounted camera.",
+    "concat_view": "This video contains concatenated views from multiple camera perspectives.",
+}
+
+_STAT_KEYS = {"mean", "std", "min", "max", "q01", "q99"}
+
+
+def get_raw_action_dim(embodiment: str) -> int:
+    key = embodiment.lower().strip()
+    if key not in EMBODIMENT_TO_RAW_ACTION_DIM:
+        raise ValueError(
+            f"No raw action dim for Cosmos3 embodiment {embodiment!r}. Expected one "
+            f"of {sorted(EMBODIMENT_TO_RAW_ACTION_DIM)}."
+        )
+    return EMBODIMENT_TO_RAW_ACTION_DIM[key]
+
+
+def canonical_aspect_ratio(width: int, height: int) -> str:
+    """Canonical ``"W,H"`` aspect string for the action caption."""
+    for sizes in VIDEO_RES_SIZE_INFO.values():
+        for aspect, (cand_w, cand_h) in sizes.items():
+            if width == cand_w and height == cand_h:
+                return aspect
+    divisor = math.gcd(width, height)
+    if divisor == 0:
+        raise ValueError(
+            f"width and height must be non-zero, got width={width}, height={height}."
+        )
+    return f"{width // divisor},{height // divisor}"
+
+
+def build_action_prompt(
+    description: str,
+    view_point: str,
+    num_frames: int,
+    fps: float,
+    height: int,
+    width: int,
+) -> str:
+    """Render the structured JSON action caption the action checkpoints expect."""
+    duration_seconds = num_frames / fps
+    minutes, secs = divmod(round(duration_seconds), 60)
+    if description and description[-1] not in ".!?":
+        description = description + "."
+    prompt = {
+        "cinematography": {
+            "framing": VIEWPOINT_TEMPLATES.get(
+                view_point, VIEWPOINT_TEMPLATES["ego_view"]
+            )
+        },
+        "actions": [{"time": f"0:00-{minutes}:{secs:02d}", "description": description}],
+        "duration": f"{int(duration_seconds)}s",
+        "fps": float(fps),
+        "resolution": {"H": int(height), "W": int(width)},
+        "aspect_ratio": canonical_aspect_ratio(int(width), int(height)),
+    }
+    return json.dumps(prompt)
+
+
+def load_action_stats(
+    stats_path: str, stats_key: str = "global"
+) -> dict[str, torch.Tensor]:
+    """Load per-channel action normalization stats from a JSON file."""
+    path = Path(stats_path)
+    if not path.exists():
+        raise FileNotFoundError(
+            f"Action normalization stats not found at {stats_path}."
+        )
+    raw = json.loads(path.read_text())
+    if stats_key in raw:
+        raw = raw[stats_key]
+    return {
+        k: torch.as_tensor(np.array(v, dtype=np.float32))
+        for k, v in raw.items()
+        if k in _STAT_KEYS
+    }
+
+
+def normalize_action(
+    action: torch.Tensor, method: str, stats: dict[str, torch.Tensor]
+) -> torch.Tensor:
+    if method == "quantile":
+        q01, q99 = stats["q01"].to(action), stats["q99"].to(action)
+        return (2.0 * (action - q01) / (q99 - q01).clamp(min=1e-8) - 1.0).clamp(
+            -1.0, 1.0
+        )
+    if method == "meanstd":
+        return (action - stats["mean"].to(action)) / stats["std"].to(action).clamp(
+            min=1e-8
+        )
+    if method == "minmax":
+        lo, hi = stats["min"].to(action), stats["max"].to(action)
+        return (2.0 * (action - lo) / (hi - lo).clamp(min=1e-8) - 1.0).clamp(-1.0, 1.0)
+    raise ValueError(f"Unknown action normalization method {method!r}.")
+
+
+def denormalize_action(
+    action: torch.Tensor, method: str, stats: dict[str, torch.Tensor]
+) -> torch.Tensor:
+    if method == "quantile":
+        q01, q99 = stats["q01"].to(action), stats["q99"].to(action)
+        return (action + 1.0) / 2.0 * (q99 - q01) + q01
+    if method == "meanstd":
+        return action * stats["std"].to(action) + stats["mean"].to(action)
+    if method == "minmax":
+        lo, hi = stats["min"].to(action), stats["max"].to(action)
+        return (action + 1.0) / 2.0 * (hi - lo) + lo
+    raise ValueError(f"Unknown action normalization method {method!r}.")

python/sglang/multimodal_gen/runtime/utils/quantization_utils.py

@@ -162,6 +162,7 @@ def get_quant_config(
     packed_modules_mapping: Dict[str, List[str]] = {},
     reverse_param_names_mapping: Dict[str, List[str]] = {},
     remap_prefix: Dict[str, str] | None = None,
+    quant_ignore_remap: Optional[Dict[str, str]] = None,
 ) -> QuantizationConfig:
     quant_cfg = find_quant_modelslim_config(model_config, component_model_path)
     if quant_cfg is not None:
@@ -191,7 +192,16 @@ def get_quant_config(
         hf_quant_config = getattr(model_config, "compression_config", None)
     if hf_quant_config is not None:
         hf_quant_config["packed_modules_mapping"] = packed_modules_mapping
-        return quant_cls.from_config(hf_quant_config)
+        is_modelopt_fp8 = (
+            hf_quant_config.get("quant_method") == "modelopt"
+            and "FP8" in str(hf_quant_config.get("quant_algo", "")).upper()
+        )
+        extra_kwargs = (
+            {"ignore_remap": quant_ignore_remap}
+            if quant_ignore_remap and is_modelopt_fp8
+            else {}
+        )
+        return quant_cls.from_config(hf_quant_config, **extra_kwargs)
 
     model_name_or_path = model_config["model_path"]
     hf_folder = model_name_or_path

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -483,7 +483,6 @@ COSMOS3_NANO_CI_sampling_params = DiffusionSamplingParams(
         "num_inference_steps": 35,
         "seed": 0,
         "max_sequence_length": 128,
-        "flow_shift": 10.0,
         "extra_args": {
             "guardrails": False,
             "use_resolution_template": False,

python/sglang/multimodal_gen/test/test_utils.py

@@ -34,7 +34,7 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-SGL_TEST_FILES_CI_DATA_REVISION = "77bd016251220fee8917a30ec92e89da03794a8a"
+SGL_TEST_FILES_CI_DATA_REVISION = "9a64abec5a7517a9f2b04ac1b4eab4173adb2d38"
 
 if current_platform.is_npu():
     SGL_TEST_FILES_CI_DATA_REVISION = "6b62f4b6825c76a25fd2ba28248df68f2b400e65"

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

@@ -2,10 +2,11 @@
 """Unit tests for Cosmos3 config, weight mapping, and sampling params."""
 
 import importlib.util
+import types
 import unittest
 from unittest import mock
 
-from fastapi import HTTPException
+import torch
 
 from sglang.multimodal_gen.configs.models.dits.cosmos3video import (
     _build_cosmos3_param_names_mapping,
@@ -22,9 +23,29 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
     VideoGenerationsRequest,
 )
 from sglang.multimodal_gen.runtime.entrypoints.openai.video_api import (
-    _reject_unsupported_cosmos3_modes,
+    _cosmos3_sampling_param_kwargs,
+    _resolve_sound_duration,
+    _resolve_video_path,
+)
+from sglang.multimodal_gen.runtime.loader.component_loaders import scheduler_loader
+from sglang.multimodal_gen.runtime.loader.component_loaders.scheduler_loader import (
+    SchedulerLoader,
 )
 from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+from sglang.multimodal_gen.runtime.models.dits.cosmos3video import (
+    DomainAwareLinear,
+    compute_mrope_position_ids_action,
+    compute_mrope_position_ids_sound,
+    compute_mrope_position_ids_vision,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import (
+    Cosmos3ImagePreprocessStage,
+    Cosmos3LatentPreparationStage,
+    Cosmos3TimestepPreparationStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_action import (
+    EMBODIMENT_TO_DOMAIN_ID,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_guardrails import (
     is_cosmos_guardrail_available,
 )
@@ -54,16 +75,16 @@ class TestCosmos3ParamNamesMapping(unittest.TestCase):
         key, idx, total = _apply(self.fn, "norm.weight")
         self.assertEqual(key, "")
 
-    def test_audio_proj_in_dropped(self):
-        key, *_ = _apply(self.fn, "audio_proj_in.weight")
-        self.assertEqual(key, "")
-
-    def test_action_proj_in_dropped(self):
-        key, *_ = _apply(self.fn, "action_proj_in.weight")
-        self.assertEqual(key, "")
-
     # --- top-level pass-through ---
 
+    def test_audio_proj_in_passthrough(self):
+        key, *_ = _apply(self.fn, "audio_proj_in.weight")
+        self.assertEqual(key, "audio_proj_in.weight")
+
+    def test_action_proj_in_passthrough(self):
+        key, *_ = _apply(self.fn, "action_proj_in.fc.weight")
+        self.assertEqual(key, "action_proj_in.fc.weight")
+
     def test_embed_tokens(self):
         key, *_ = _apply(self.fn, "embed_tokens.weight")
         self.assertEqual(key, "language_model.embed_tokens.weight")
@@ -204,6 +225,87 @@ class TestCosmos3AdjustNumFrames(unittest.TestCase):
         self.assertEqual(self.cfg.adjust_num_frames(2), 1)
 
 
+class TestCosmos3SchedulerConfig(unittest.TestCase):
+    """Verify Cosmos3 scheduler class and flow-shift defaults."""
+
+    def test_config_overrides_checkpoint_scheduler_class(self):
+        cfg = Cosmos3Config()
+        self.assertEqual(cfg.scheduler_class_override, "FlowUniPCMultistepScheduler")
+        self.assertIsNone(cfg.flow_shift)
+
+    def test_scheduler_loader_uses_configured_class_override(self):
+        class FakeScheduler:
+            def __init__(self, **config):
+                self.config = config
+
+        server_args = types.SimpleNamespace(
+            pipeline_config=types.SimpleNamespace(
+                scheduler_class_override="FlowUniPCMultistepScheduler",
+                flow_shift=None,
+            )
+        )
+        with (
+            mock.patch.object(
+                scheduler_loader,
+                "get_diffusers_component_config",
+                return_value={"_class_name": "CheckpointScheduler", "foo": "bar"},
+            ),
+            mock.patch.object(
+                scheduler_loader.ModelRegistry,
+                "resolve_model_cls",
+                return_value=(FakeScheduler, None),
+            ) as resolve,
+        ):
+            scheduler = SchedulerLoader().load_customized("unused", server_args)
+
+        resolve.assert_called_once_with("FlowUniPCMultistepScheduler")
+        self.assertEqual(scheduler.config["foo"], "bar")
+
+    @staticmethod
+    def _stage():
+        stage = Cosmos3TimestepPreparationStage.__new__(Cosmos3TimestepPreparationStage)
+        stage.scheduler = types.SimpleNamespace(
+            config=types.SimpleNamespace(flow_shift=1.0)
+        )
+        return stage
+
+    @staticmethod
+    def _batch(**kwargs):
+        sp_kwargs = kwargs.pop("sp_kwargs", {})
+        return types.SimpleNamespace(
+            sampling_params=Cosmos3SamplingParams(prompt="t", **sp_kwargs),
+            data_type=kwargs.pop("data_type", DataType.VIDEO),
+            preprocessed_image=kwargs.pop("preprocessed_image", None),
+            preprocessed_video=kwargs.pop("preprocessed_video", None),
+        )
+
+    def test_per_mode_flow_shift_defaults(self):
+        stage = self._stage()
+        self.assertEqual(
+            stage._default_flow_shift_for_mode(self._batch(data_type=DataType.IMAGE)),
+            3.0,
+        )
+        self.assertEqual(
+            stage._default_flow_shift_for_mode(
+                self._batch(preprocessed_image=torch.empty(1))
+            ),
+            10.0,
+        )
+        self.assertEqual(
+            stage._default_flow_shift_for_mode(
+                self._batch(preprocessed_video=torch.empty(1))
+            ),
+            10.0,
+        )
+        self.assertEqual(stage._default_flow_shift_for_mode(self._batch()), 10.0)
+        self.assertEqual(
+            stage._default_flow_shift_for_mode(
+                self._batch(sp_kwargs={"action_mode": "policy"})
+            ),
+            10.0,
+        )
+
+
 class TestCosmos3SamplingParamsDataType(unittest.TestCase):
     """Verify num_frames==1 flips data_type to IMAGE before file name derivation."""
 
@@ -246,9 +348,9 @@ class TestCosmos3ModelResolution(unittest.TestCase):
 
 
 class TestCosmos3OpenAIProtocol(unittest.TestCase):
-    """Verify Cosmos3-only knobs stay out of the stable request schema."""
+    """Verify Cosmos3 modality knobs are exposed by the video HTTP schema."""
 
-    def test_cosmos3_private_fields_are_extra_fields(self):
+    def test_cosmos3_template_fields_remain_extra_fields(self):
         for request_cls in (ImageGenerationsRequest, VideoGenerationsRequest):
             with self.subTest(request_cls=request_cls.__name__):
                 self.assertIn("max_sequence_length", request_cls.model_fields)
@@ -258,22 +360,66 @@ class TestCosmos3OpenAIProtocol(unittest.TestCase):
                 self.assertNotIn("use_system_prompt", request_cls.model_fields)
                 self.assertNotIn("use_guardrails", request_cls.model_fields)
 
-        self.assertNotIn("generate_sound", VideoGenerationsRequest.model_fields)
+    def test_cosmos3_modal_fields_pass_through_as_extras(self):
-        self.assertNotIn("sound_duration", VideoGenerationsRequest.model_fields)
+        for field_name in ("video_path", "video_url"):
+            with self.subTest(field_name=field_name):
+                self.assertIn(field_name, VideoGenerationsRequest.model_fields)
 
-    def test_unsupported_cosmos3_modes_allow_falsy_extra_fields(self):
+        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():
+            with self.subTest(field_name=field_name):
+                self.assertNotIn(field_name, VideoGenerationsRequest.model_fields)
+                self.assertEqual(getattr(req, field_name), value)
+
+    def test_cosmos3_http_aliases_map_to_sampling_params(self):
         req = VideoGenerationsRequest(
             prompt="test",
-            generate_sound=False,
+            video_url="https://example.com/input.mp4",
-            action_mode="",
+            generate_sound=True,
-            condition_frame_indexes_vision=[],
+            condition_frame_indexes_vision=[0, 2],
-            condition_video_keep={},
+            condition_video_keep="last",
+            action_mode="policy",
+            domain_name="umi",
+            raw_action_dim=9,
+            action_fps=30.0,
+            action_view_point="ego_view",
         )
-        _reject_unsupported_cosmos3_modes(req, "nvidia/Cosmos3-Nano")
 
-        req = VideoGenerationsRequest(prompt="test", generate_sound=True)
+        self.assertEqual(_resolve_video_path(req), "https://example.com/input.mp4")
-        with self.assertRaises(HTTPException):
+
-            _reject_unsupported_cosmos3_modes(req, "nvidia/Cosmos3-Nano")
+        kwargs = _cosmos3_sampling_param_kwargs(req, num_frames=48, fps=24)
+        self.assertEqual(kwargs["sound_duration"], 2.0)
+        self.assertEqual(kwargs["condition_frame_indexes"], [0, 2])
+        self.assertEqual(kwargs["condition_video_keep"], "last")
+        self.assertEqual(kwargs["action_mode"], "policy")
+        self.assertEqual(kwargs["domain_name"], "umi")
+        self.assertEqual(kwargs["raw_action_dim"], 9)
+        self.assertEqual(kwargs["action_fps"], 30.0)
+        self.assertEqual(kwargs["action_view_point"], "ego_view")
+
+    def test_generate_sound_false_disables_sound_duration(self):
+        req = VideoGenerationsRequest(
+            prompt="test", generate_sound=False, sound_duration=3.0
+        )
+        self.assertEqual(
+            _resolve_sound_duration(req, num_frames=48, fps=24),
+            0.0,
+        )
 
 
 class TestCosmos3Guardrails(unittest.TestCase):
@@ -296,5 +442,306 @@ class TestCosmos3Guardrails(unittest.TestCase):
         self.assertFalse(is_cosmos_guardrail_available())
 
 
+class TestCosmos3MRoPE(unittest.TestCase):
+    """mRoPE position-ID computation for vision / sound / action token grids."""
+
+    DEVICE = torch.device("cpu")
+
+    def test_vision_default_args_unchanged(self):
+        # With base_temporal_compression_factor=None and start_frame_offset=0
+        # the token and base rates cancel, so t-index == frame index.
+        pos, _ = compute_mrope_position_ids_vision(
+            grid_t=21,
+            grid_h=1,
+            grid_w=1,
+            temporal_offset=0,
+            device=self.DEVICE,
+            fps=24.0,
+            base_fps=24.0,
+            temporal_compression_factor=4,
+        )
+        self.assertEqual(tuple(pos.shape), (3, 21))
+        self.assertAlmostEqual(float(pos[0, 0]), 0.0, places=5)
+        self.assertAlmostEqual(float(pos[0, 20]), 20.0, places=5)
+
+    def test_sound_grid_shape_and_scaling(self):
+        pos, _ = compute_mrope_position_ids_sound(
+            grid_t=10,
+            temporal_offset=0,
+            sound_latent_fps=25.0,
+            device=self.DEVICE,
+            base_fps=24.0,
+            temporal_compression_factor_sound=1,
+        )
+        # (T, 1, 1) grid -> spatial axes are all zero.
+        self.assertEqual(tuple(pos.shape), (3, 10))
+        self.assertTrue(torch.all(pos[1] == 0))
+        self.assertTrue(torch.all(pos[2] == 0))
+        # t-index = i / sound_fps * base_fps = i / 25 * 24.
+        self.assertAlmostEqual(float(pos[0, 5]), 5 / 25 * 24, places=4)
+
+    def test_action_uses_video_base_compression_and_offset(self):
+        # Action runs at frame rate (tcf=1) but is scaled by the video's
+        # base_temporal_compression_factor=4, and shifted by start_frame_offset.
+        pos, _ = compute_mrope_position_ids_action(
+            grid_t=16,
+            temporal_offset=0,
+            action_fps=10.0,
+            device=self.DEVICE,
+            base_fps=24.0,
+            base_temporal_compression_factor=4,
+            start_frame_offset=1,
+        )
+        self.assertEqual(tuple(pos.shape), (3, 16))
+        self.assertTrue(torch.all(pos[1] == 0))
+        self.assertTrue(torch.all(pos[2] == 0))
+        # t-index[i] = (i + start_frame_offset) / action_fps * (base_fps / base_tcf)
+        #            = (i + 1) / 10 * (24 / 4) = (i + 1) * 0.6
+        self.assertAlmostEqual(float(pos[0, 0]), 0.6, places=4)
+        self.assertAlmostEqual(float(pos[0, 15]), 16 * 0.6, places=4)
+
+    def test_action_offset_zero(self):
+        pos, _ = compute_mrope_position_ids_action(
+            grid_t=8,
+            temporal_offset=0,
+            action_fps=10.0,
+            device=self.DEVICE,
+            base_fps=24.0,
+            base_temporal_compression_factor=4,
+            start_frame_offset=0,
+        )
+        self.assertAlmostEqual(float(pos[0, 0]), 0.0, places=5)
+
+    def test_action_aligns_with_video_positions(self):
+        # Action frames at frame rate should share the video's temporal frame:
+        # every 4th action token lands on the next video latent-frame position.
+        media_offset = 100
+        vid, _ = compute_mrope_position_ids_vision(
+            grid_t=5,
+            grid_h=1,
+            grid_w=1,
+            temporal_offset=media_offset,
+            device=self.DEVICE,
+            fps=24.0,
+            base_fps=24.0,
+            temporal_compression_factor=4,
+        )
+        act, _ = compute_mrope_position_ids_action(
+            grid_t=16,
+            temporal_offset=media_offset,
+            action_fps=24.0,
+            device=self.DEVICE,
+            base_fps=24.0,
+            base_temporal_compression_factor=4,
+            start_frame_offset=0,
+        )
+        # video latent frame 1 sits at media_offset+1; action frame 4 (4 frames
+        # per latent at tcf=4) lands at the same temporal position.
+        self.assertAlmostEqual(float(vid[0, 1]), float(act[0, 4]), places=4)
+
+
+class TestCosmos3DomainAwareLinear(unittest.TestCase):
+    """Per-domain action projection."""
+
+    def test_rank3_and_rank2_shapes(self):
+        layer = DomainAwareLinear(input_size=7, output_size=64, num_domains=32)
+        x3 = torch.randn(2, 16, 7)
+        out3 = layer(x3, torch.tensor([1, 5]))
+        self.assertEqual(tuple(out3.shape), (2, 16, 64))
+        x2 = torch.randn(3, 7)
+        out2 = layer(x2, torch.tensor([0, 1, 2]))
+        self.assertEqual(tuple(out2.shape), (3, 64))
+
+    def test_distinct_domains_give_distinct_outputs(self):
+        torch.manual_seed(0)
+        layer = DomainAwareLinear(input_size=4, output_size=8, num_domains=4)
+        x = torch.randn(1, 3, 4)
+        out_a = layer(x, torch.tensor([0]))
+        out_b = layer(x, torch.tensor([2]))
+        self.assertFalse(torch.allclose(out_a, out_b))
+
+    def test_scalar_domain_id_promoted(self):
+        layer = DomainAwareLinear(input_size=4, output_size=8, num_domains=4)
+        out = layer(torch.randn(1, 2, 4), torch.tensor(3))
+        self.assertEqual(tuple(out.shape), (1, 2, 8))
+
+
+class TestCosmos3ConditionIndexes(unittest.TestCase):
+    """Vision condition-frame resolution across V2V and action modes."""
+
+    @staticmethod
+    def _batch(num_frames=61, **sp_kwargs):
+        sp = Cosmos3SamplingParams(prompt="t", num_frames=num_frames, **sp_kwargs)
+        return types.SimpleNamespace(sampling_params=sp, num_frames=num_frames)
+
+    def test_v2v_default(self):
+        idx = Cosmos3ImagePreprocessStage._resolve_condition_indexes(self._batch())
+        self.assertEqual(idx, [0, 1])
+
+    def test_v2v_explicit_sorted_unique(self):
+        idx = Cosmos3ImagePreprocessStage._resolve_condition_indexes(
+            self._batch(condition_frame_indexes=[2, 0, 2])
+        )
+        self.assertEqual(idx, [0, 2])
+
+    def test_inverse_dynamics_conditions_all_latent_frames(self):
+        # 61 frames -> (61-1)//4 + 1 = 16 latent frames, all locked.
+        idx = Cosmos3ImagePreprocessStage._resolve_condition_indexes(
+            self._batch(num_frames=61, action_mode="inverse_dynamics")
+        )
+        self.assertEqual(idx, list(range(16)))
+
+
+class TestCosmos3DomainResolution(unittest.TestCase):
+    """Embodiment domain-id resolution for action generation."""
+
+    @staticmethod
+    def _batch(**sp_kwargs):
+        sp = Cosmos3SamplingParams(prompt="t", **sp_kwargs)
+        return types.SimpleNamespace(sampling_params=sp)
+
+    def test_explicit_domain_id(self):
+        self.assertEqual(
+            Cosmos3LatentPreparationStage._resolve_domain_id(self._batch(domain_id=7)),
+            7,
+        )
+
+    def test_domain_name_lookup(self):
+        self.assertEqual(
+            Cosmos3LatentPreparationStage._resolve_domain_id(
+                self._batch(domain_name="av")
+            ),
+            EMBODIMENT_TO_DOMAIN_ID["av"],
+        )
+        self.assertEqual(
+            Cosmos3LatentPreparationStage._resolve_domain_id(
+                self._batch(domain_name="umi")
+            ),
+            EMBODIMENT_TO_DOMAIN_ID["umi"],
+        )
+
+    def test_missing_domain_raises(self):
+        with self.assertRaises(ValueError):
+            Cosmos3LatentPreparationStage._resolve_domain_id(self._batch())
+
+    def test_unknown_domain_name_raises(self):
+        with self.assertRaises(ValueError):
+            Cosmos3LatentPreparationStage._resolve_domain_id(
+                self._batch(domain_name="not_a_robot")
+            )
+
+
+class TestCosmos3ActionLatentPrep(unittest.TestCase):
+    """Action latent / mask preparation per action mode."""
+
+    @classmethod
+    def setUpClass(cls):
+        # Bypass PipelineStage.__init__ (needs global server args); only the
+        # transformer's action_dim and log_info are used by _prepare_action_latents.
+        cls.stage = Cosmos3LatentPreparationStage.__new__(Cosmos3LatentPreparationStage)
+        cls.stage.transformer = types.SimpleNamespace(action_dim=64)
+        cls.stage.log_info = lambda *a, **k: None
+        cls.device = torch.device("cpu")
+        cls.dtype = torch.float32
+
+    def _run(self, num_frames=17, **sp_kwargs):
+        sp = Cosmos3SamplingParams(prompt="t", num_frames=num_frames, **sp_kwargs)
+        batch = types.SimpleNamespace(
+            sampling_params=sp, num_frames=num_frames, extra={}
+        )
+        gen = torch.Generator(device=self.device).manual_seed(0)
+        self.stage._prepare_action_latents(batch, gen, self.device, self.dtype)
+        return batch
+
+    def test_forward_dynamics_clean_conditioning(self):
+        batch = self._run(
+            action_mode="forward_dynamics",
+            domain_name="agibotworld",
+            action=[[0.1] * 29 for _ in range(16)],
+        )
+        # action_chunk_size = num_frames - 1 = 16; padded to action_dim 64.
+        self.assertEqual(tuple(batch.action_latents.shape), (1, 16, 64))
+        # raw_action_dim is derived from the embodiment (agibotworld -> 29).
+        self.assertEqual(batch.extra["raw_action_dim"], 29)
+        self.assertEqual(batch.extra["action_start_frame_offset"], 1)
+        self.assertEqual(
+            int(batch.extra["action_domain_ids"][0]),
+            EMBODIMENT_TO_DOMAIN_ID["agibotworld"],
+        )
+        # forward_dynamics: action is clean conditioning -> velocity mask all zero.
+        self.assertTrue(torch.all(batch.extra["action_velocity_mask"] == 0))
+
+    def test_policy_denoises_from_noise(self):
+        batch = self._run(
+            action_mode="policy", domain_name="droid_lerobot", raw_action_dim=10
+        )
+        self.assertEqual(tuple(batch.action_latents.shape), (1, 16, 64))
+        self.assertEqual(batch.extra["raw_action_dim"], 10)
+        # policy: action fully denoised -> velocity mask all one.
+        self.assertTrue(torch.all(batch.extra["action_velocity_mask"] == 1))
+        # padding dims beyond raw_action_dim start at zero.
+        self.assertTrue(torch.all(batch.action_latents[:, :, 10:] == 0))
+
+    def test_inverse_dynamics_denoises_from_noise(self):
+        batch = self._run(
+            num_frames=61,
+            action_mode="inverse_dynamics",
+            domain_name="av",
+            raw_action_dim=9,
+        )
+        self.assertEqual(tuple(batch.action_latents.shape), (1, 60, 64))
+        self.assertTrue(torch.all(batch.extra["action_velocity_mask"] == 1))
+
+    def test_forward_dynamics_requires_action(self):
+        with self.assertRaises(ValueError):
+            self._run(action_mode="forward_dynamics", domain_name="agibotworld")
+
+    def test_policy_requires_raw_action_dim(self):
+        # policy has no input action to infer from, so it needs raw_action_dim
+        # from either the embodiment or an explicit value. With only a numeric
+        # domain_id (no embodiment name) and no raw_action_dim, it must raise.
+        with self.assertRaises(ValueError):
+            self._run(action_mode="policy", domain_id=0)
+
+    def test_policy_raw_action_dim_from_embodiment(self):
+        # droid_lerobot -> 10, so policy no longer needs an explicit raw dim.
+        batch = self._run(action_mode="policy", domain_name="droid_lerobot")
+        self.assertEqual(batch.extra["raw_action_dim"], 10)
+
+    def test_unknown_action_mode_raises(self):
+        with self.assertRaises(ValueError):
+            self._run(action_mode="teleport", domain_id=0)
+
+
+class TestCosmos3ModalitySamplingParams(unittest.TestCase):
+    """Sound / V2V / action sampling-param fields and defaults."""
+
+    def test_sound_duration_default_and_set(self):
+        self.assertEqual(Cosmos3SamplingParams(prompt="t").sound_duration, 0.0)
+        self.assertEqual(
+            Cosmos3SamplingParams(prompt="t", sound_duration=3.0).sound_duration, 3.0
+        )
+
+    def test_v2v_fields(self):
+        sp = Cosmos3SamplingParams(
+            prompt="t", video_path="in.mp4", condition_frame_indexes=[0, 1]
+        )
+        self.assertEqual(sp.video_path, "in.mp4")
+        self.assertEqual(sp.condition_frame_indexes, [0, 1])
+        self.assertEqual(sp.condition_video_keep, "first")
+
+    def test_action_fields_default_none(self):
+        sp = Cosmos3SamplingParams(prompt="t")
+        for field in (
+            "action_mode",
+            "domain_id",
+            "domain_name",
+            "raw_action_dim",
+            "action_fps",
+            "action",
+        ):
+            self.assertIsNone(getattr(sp, field))
+
+
 if __name__ == "__main__":
     unittest.main()

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

@@ -52,6 +52,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i
 from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config
 from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
     ModelOptFp4Config,
+    ModelOptFp8Config,
     _prepare_nvfp4_weight_bytes,
 )
 from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
@@ -456,6 +457,37 @@ class TestTransformerQuantHelpers(unittest.TestCase):
             ["single_transformer_blocks.*.proj_mlp*"],
         )
 
+    def test_modelopt_fp8_hf_config_uses_general_modelopt_fp8(self):
+        config = get_quant_config(
+            {
+                "quantization_config": {
+                    "quant_method": "modelopt",
+                    "quant_algo": "FP8",
+                    "ignore": ["vae2llm", "llm2vae"],
+                }
+            },
+            "/unused/component/path",
+            quant_ignore_remap={"vae2llm": "proj_in", "llm2vae": "proj_out"},
+        )
+
+        self.assertIsInstance(config, ModelOptFp8Config)
+        self.assertEqual(config.exclude_modules, ["proj_in", "proj_out"])
+
+    def test_modelopt_fp8_explicit_config_uses_general_modelopt_fp8(self):
+        config = get_quant_config(
+            {
+                "quantization_config": {
+                    "quant_method": "modelopt_fp8",
+                    "quant_algo": "FP8",
+                    "ignore": ["proj_out"],
+                }
+            },
+            "/unused/component/path",
+        )
+
+        self.assertIsInstance(config, ModelOptFp8Config)
+        self.assertEqual(config.exclude_modules, ["proj_out"])
+
     @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_rank", return_value=0)
     @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_size", return_value=1)
     @patch(

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

@@ -0,0 +1,19 @@
+from types import SimpleNamespace
+
+import torch
+
+from sglang.multimodal_gen.configs.pipeline_configs.wan import WanT2V480PConfig
+
+
+def test_wan_prompt_embed_accessors_return_transformer_tensor():
+    prompt_embeds = torch.empty(1, 2, 3)
+    negative_prompt_embeds = torch.empty(1, 2, 3)
+    batch = SimpleNamespace(
+        prompt_embeds=[prompt_embeds],
+        negative_prompt_embeds=[negative_prompt_embeds],
+    )
+
+    config = WanT2V480PConfig()
+
+    assert config.get_pos_prompt_embeds(batch) is prompt_embeds
+    assert config.get_neg_prompt_embeds(batch) is negative_prompt_embeds