[diffusion] model: support cosmos3 edge and distilled (#31590)

Co-authored-by: Kedi Wu <31940276+kediwu0331@users.noreply.github.com> Co-authored-by: Kedi Wu <kediw@nvidia.com>
2026-08-12 10:29:40 +08:00
parent 81c88da1ab
commit f7b6800b22
13 changed files with 1036 additions and 104 deletions
@@ -6,7 +6,7 @@ from dataclasses import dataclass, field
 from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig


-def _build_cosmos3_param_names_mapping() -> dict:
+def _build_cosmos3_param_names_mapping(gated_mlp: bool = True) -> dict:
    """Map diffusers-format Cosmos3 weights to the sglang model namespace.

    Source keys (diffusers transformer ckpt) → target keys (sglang model):
@@ -15,23 +15,26 @@ def _build_cosmos3_param_names_mapping() -> dict:
        layers.X.input_layernorm_moe_gen.weight                -> gen_layers.X.input_layernorm.weight
        layers.X.self_attn.{to_q,to_k,to_v}.weight            -> language_model.layers.X.self_attn.to_qkv.weight (concat dim 0)
        layers.X.self_attn.{add_q,add_k,add_v}_proj.weight    -> gen_layers.X.cross_attention.to_qkv.weight (concat dim 0)
-        layers.X.mlp.{gate,up}_proj.weight                     -> language_model.layers.X.mlp.gate_up_proj.weight (concat dim 0)
-        layers.X.mlp_moe_gen.{gate,up}_proj.weight             -> gen_layers.X.mlp.gate_up_proj.weight (concat dim 0)
        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)

+    ``gated_mlp`` selects the MLP weight layout: SwiGLU checkpoints merge
+    ``mlp.{gate,up}_proj`` into a single ``gate_up_proj``; dense (squared-ReLU)
+    checkpoints ship only ``mlp.{up,down}_proj``, which flow through the
+    catch-all unchanged.
+
    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; both are skipped.
-    Audio and action keys (``audio_proj_*``, ``action_proj_*``, modality
-    embeds) pass through unchanged.
+    `norm.weight` and `lm_head.weight` are inherited from the text-pretrained
+    backbone and not used at inference; both are skipped. Audio and action
+    keys (``audio_proj_*``, ``action_proj_*``, modality embeds) pass through
+    unchanged.
    """
-    return {
-        # Inherited from Qwen3-VL pretraining; unused at diffusion inference.
+    mapping = {
+        # Inherited from text pretraining; unused at diffusion inference.
        r"^lm_head\.weight$": "",
        r"^norm\.weight$": "",
        # Top-level norms / embeddings.
@@ -62,52 +65,61 @@ def _build_cosmos3_param_names_mapping() -> dict:
        r"^layers\.(\d+)\.self_attn\.norm_added_k\.(.*)$": r"gen_layers.\1.cross_attention.norm_k.\2",
        r"^layers\.(\d+)\.input_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.input_layernorm.\2",
        r"^layers\.(\d+)\.post_attention_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.post_attention_layernorm.\2",
+    }
+
+    if gated_mlp:
        # GEN MLP gate/up merge into MergedColumnParallelLinear gate_up_proj.
        # Must precede the mlp_moe_gen catch-all below.
-        r"^layers\.(\d+)\.mlp_moe_gen\.gate_proj\.(.*)$": (
+        mapping[r"^layers\.(\d+)\.mlp_moe_gen\.gate_proj\.(.*)$"] = (
            r"gen_layers.\1.mlp.gate_up_proj.\2",
            0,
            2,
-        ),
-        r"^layers\.(\d+)\.mlp_moe_gen\.up_proj\.(.*)$": (
+        )
+        mapping[r"^layers\.(\d+)\.mlp_moe_gen\.up_proj\.(.*)$"] = (
            r"gen_layers.\1.mlp.gate_up_proj.\2",
            1,
            2,
-        ),
-        r"^layers\.(\d+)\.mlp_moe_gen\.(.*)$": r"gen_layers.\1.mlp.\2",
-        # UND pathway: Q/K/V merge into to_qkv; remaining attention keys
-        # (to_out, norm_q, norm_k) and layernorms pass through the catch-all.
-        r"^layers\.(\d+)\.self_attn\.to_q\.(.*)$": (
-            r"language_model.layers.\1.self_attn.to_qkv.\2",
-            0,
-            3,
-        ),
-        r"^layers\.(\d+)\.self_attn\.to_k\.(.*)$": (
-            r"language_model.layers.\1.self_attn.to_qkv.\2",
-            1,
-            3,
-        ),
-        r"^layers\.(\d+)\.self_attn\.to_v\.(.*)$": (
-            r"language_model.layers.\1.self_attn.to_qkv.\2",
-            2,
-            3,
-        ),
+        )
+
+    # GEN MLP catch-all: dense up/down_proj pass through unchanged.
+    mapping[r"^layers\.(\d+)\.mlp_moe_gen\.(.*)$"] = r"gen_layers.\1.mlp.\2"
+
+    # UND pathway: Q/K/V merge into to_qkv; remaining attention keys
+    # (to_out, norm_q, norm_k) and layernorms pass through the catch-all.
+    mapping[r"^layers\.(\d+)\.self_attn\.to_q\.(.*)$"] = (
+        r"language_model.layers.\1.self_attn.to_qkv.\2",
+        0,
+        3,
+    )
+    mapping[r"^layers\.(\d+)\.self_attn\.to_k\.(.*)$"] = (
+        r"language_model.layers.\1.self_attn.to_qkv.\2",
+        1,
+        3,
+    )
+    mapping[r"^layers\.(\d+)\.self_attn\.to_v\.(.*)$"] = (
+        r"language_model.layers.\1.self_attn.to_qkv.\2",
+        2,
+        3,
+    )
+
+    if gated_mlp:
        # UND MLP gate/up merge into MergedColumnParallelLinear gate_up_proj.
        # Must precede the layers catch-all below.
-        r"^layers\.(\d+)\.mlp\.gate_proj\.(.*)$": (
+        mapping[r"^layers\.(\d+)\.mlp\.gate_proj\.(.*)$"] = (
            r"language_model.layers.\1.mlp.gate_up_proj.\2",
            0,
            2,
-        ),
-        r"^layers\.(\d+)\.mlp\.up_proj\.(.*)$": (
+        )
+        mapping[r"^layers\.(\d+)\.mlp\.up_proj\.(.*)$"] = (
            r"language_model.layers.\1.mlp.gate_up_proj.\2",
            1,
            2,
-        ),
-        # UND pathway: layernorms + remaining attention/mlp keys pass through
-        # under language_model.layers namespace.
-        r"^layers\.(\d+)\.(.*)$": r"language_model.layers.\1.\2",
-    }
+        )
+
+    # UND pathway: layernorms + remaining attention/mlp keys pass through
+    # under language_model.layers namespace.
+    mapping[r"^layers\.(\d+)\.(.*)$"] = r"language_model.layers.\1.\2"
+    return mapping


@dataclass
@@ -130,6 +142,12 @@ class Cosmos3VideoArchConfig(DiTArchConfig):
    num_key_value_heads: int = 8  # GQA
    head_dim: int = 128
    intermediate_size: int = 12288
+    # "relu2" selects a dense squared-ReLU MLP; anything else is SwiGLU.
+    hidden_act: str = "silu"
+    # Per-head QK-norm on the UND self-attention (text) pathway.
+    qk_norm_for_text: bool = True
+    # Normalize the UND keys before they feed the GEN cross-attention.
+    use_und_k_norm_for_gen: bool = False

    # Latent space configuration
    latent_patch_size: int = 2
@@ -180,6 +198,10 @@ class Cosmos3VideoArchConfig(DiTArchConfig):
        default_factory=lambda: {"vae2llm": "proj_in", "llm2vae": "proj_out"}
    )

+    @property
+    def gated_mlp(self) -> bool:
+        return self.hidden_act != "relu2"
+
    def __post_init__(self):
        super().__post_init__()
        # Diffusers configs nest the mrope sizes under `rope_scaling`; lift it.
@@ -189,6 +211,9 @@ class Cosmos3VideoArchConfig(DiTArchConfig):
        self.num_channels_latents = self.out_channels
        # Patch latent dimension: (patch_size^2) * latent_channel
        self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel
+        # The MLP weight layout depends on the activation, which the checkpoint
+        # may override; rebuild the mapping after arch values are applied.
+        self.param_names_mapping = _build_cosmos3_param_names_mapping(self.gated_mlp)


@dataclass
@@ -8,6 +8,8 @@ in the stages from ``num_frames`` and ``image_path``; T2I overrides
 ``data_type`` to ``IMAGE`` in :meth:`SamplingParams._adjust`.
 """

+import functools
+import os
 from dataclasses import dataclass, field

 from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig
@@ -18,6 +20,63 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
    PipelineConfig,
 )

+COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense"
+
+
+@functools.lru_cache(maxsize=None)
+def is_edge_checkpoint(model_path: str) -> bool:
+    """Whether the checkpoint is the Edge (dense) variant.
+
+    Read from the transformer config rather than the loaded arch so the answer
+    is available before the weights are on device (e.g. when resolving sampling
+    defaults in the client process).
+    """
+    from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+        get_diffusers_component_config,
+    )
+
+    config = get_diffusers_component_config(
+        component_path=os.path.join(model_path, "transformer")
+    )
+    return (
+        config.get("backbone_type") == COSMOS3_EDGE_BACKBONE_TYPE
+        or config.get("hidden_act") == "relu2"
+    )
+
+
+@functools.lru_cache(maxsize=None)
+def _distilled_sampler_config(model_path: str) -> dict | None:
+    """The fixed-step sampler config for a distilled checkpoint, else ``None``.
+
+    Distillation is a scheduler-only change: the checkpoint ships a
+    ``FlowMatchEulerDiscreteScheduler`` with an explicit fixed-step sigma
+    schedule instead of the multi-step FlowUniPC the other variants use.
+    """
+    from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
+        get_diffusers_component_config,
+    )
+
+    config = get_diffusers_component_config(
+        component_path=os.path.join(model_path, "scheduler")
+    )
+    if config.get("_class_name") != "FlowMatchEulerDiscreteScheduler":
+        return None
+    sampler = config.get("fixed_step_sampler_config")
+    if not sampler or not sampler.get("t_list"):
+        return None
+    return sampler
+
+
+def is_distilled_checkpoint(model_path: str) -> bool:
+    """Whether the checkpoint is a few-step distilled variant."""
+    return _distilled_sampler_config(model_path) is not None
+
+
+def get_distilled_sigmas(model_path: str) -> list[float] | None:
+    """The explicit fixed-step sigma schedule for a distilled checkpoint."""
+    sampler = _distilled_sampler_config(model_path)
+    return list(sampler["t_list"]) if sampler is not None else None
+

@dataclass
 class Cosmos3Config(PipelineConfig):
@@ -60,6 +119,11 @@ class Cosmos3Config(PipelineConfig):
    # names a server-side file. ``None`` disables normalization.
    action_stats_path: str | None = None

+    # Pre-computed once in update_config_from_dict from the resolved model_path.
+    # None until that point (e.g. in unit-test mocks that never call update_config_from_dict).
+    is_edge: bool | None = None
+    distilled_sigmas: list[float] | None = None
+
    def __post_init__(self):
        self.vae_config.arch_config.z_dim = 48
        # Encoder is needed for I2V; T2V/T2I never invoke it.
@@ -70,6 +134,17 @@ class Cosmos3Config(PipelineConfig):
        self.vae_config.use_parallel_encode = False
        self.vae_config.use_parallel_decode = True

+    def update_config_from_dict(self, args, prefix: str = "") -> None:
+        super().update_config_from_dict(args, prefix)
+        # model_path is only populated here, after construction. Compute
+        # checkpoint variant flags once so per-request code reads them from the
+        # config rather than re-downloading the scheduler subfolder each time.
+        if self.model_path:
+            self.distilled_sigmas = get_distilled_sigmas(self.model_path)
+            self.is_edge = is_edge_checkpoint(self.model_path)
+            if self.distilled_sigmas is not None:
+                self.scheduler_class_override = None
+
    def adjust_num_frames(self, num_frames: int) -> int:
        """Round ``num_frames`` so ``(n - 1) % 4 == 0`` for the VAE.

@@ -17,17 +17,48 @@ from sglang.multimodal_gen.configs.sample.sampling_params import (
    SamplingParams,
 )

+COSMOS3_DEFAULT_GUIDANCE_SCALE = 4.0
+COSMOS3_EDGE_T2I_GUIDANCE_SCALE = 7.0
+COSMOS3_EDGE_T2V_GUIDANCE_SCALE = 5.0
+COSMOS3_EDGE_T2V_WIDTH = 832
+COSMOS3_EDGE_T2V_HEIGHT = 480
+COSMOS3_EDGE_T2I_SIZE = 640
+
+# Image generation applies guidance only over a high-noise window; guiding the
+# low-noise steps degrades sample quality (Kynkaeaenniemi et al. 2024). Video
+# modes guide at every step.
+COSMOS3_T2I_GUIDANCE_INTERVAL = (400.0, 1000.0)
+
+# Edge is trained at 256p/480p only; larger frames push the spatial mRoPE grid
+# past its trained range and shatter the output.
+COSMOS3_EDGE_SUPPORTED_RESOLUTIONS = [
+    (832, 480),
+    (480, 832),
+    (640, 480),
+    (480, 640),
+    (480, 480),
+    (640, 640),
+    (448, 256),
+    (256, 448),
+    (256, 256),
+]
+

@dataclass
 class Cosmos3SamplingParams(SamplingParams):
-    """Cosmos3 sampling parameters (T2V defaults; also used for I2V / V2V / T2I)."""
+    """Cosmos3 sampling parameters (T2V defaults; also used for I2V / V2V / T2I).

-    height: int = 720
-    width: int = 1280
+    ``height``/``width`` default to ``None`` so the variant (Edge vs. base) can
+    pick the right resolution at request time in
+    :meth:`_resolve_variant_defaults`.
+    """
+
+    height: int | None = None
+    width: int | None = None
    num_frames: int = 81
    fps: int = 24

-    guidance_scale: float = 4.0
+    guidance_scale: float = COSMOS3_DEFAULT_GUIDANCE_SCALE
    num_inference_steps: int = 35

    negative_prompt: str = ""
@@ -50,6 +81,7 @@ class Cosmos3SamplingParams(SamplingParams):
            (832, 480),
            (480, 832),
            (1024, 1024),
+            (640, 640),
        ]
    )

@@ -72,6 +104,18 @@ class Cosmos3SamplingParams(SamplingParams):
    action_normalization: str = "quantile"

    def _adjust(self, server_args) -> None:
+        # adjust distil and edge args — read from the pre-computed config fields
+        # so no checkpoint download happens at request time.
+        pipeline_config = server_args.pipeline_config
+        distilled_sigmas = pipeline_config.distilled_sigmas
+        if distilled_sigmas is not None:
+            self.num_inference_steps = len(distilled_sigmas)
+        self._resolve_variant_defaults(
+            bool(pipeline_config.is_edge),
+            is_distilled=distilled_sigmas is not None,
+        )
+
+        # adjust action args
        action_output = False
        if self.action_mode is not None:
            self.action_mode = str(self.action_mode).strip().lower()
@@ -98,6 +142,41 @@ class Cosmos3SamplingParams(SamplingParams):
            self.output_file_name = None
            self.output_compression = 0

+    def _guidance_is_explicit(self) -> bool:
+        explicit = getattr(self, "_explicit_fields", None)
+        return explicit is not None and "guidance_scale" in explicit
+
+    def _resolve_variant_defaults(
+        self, is_edge: bool, is_distilled: bool = False
+    ) -> None:
+        """Fill unset resolution/guidance with the variant's defaults.
+
+        Base resolution defaulting (``supported_resolutions[0]``) covers the
+        non-Edge path; only Edge and guidance need explicit handling here.
+        """
+        is_t2i = self.num_frames == 1
+        if is_distilled:
+            # Guidance is distilled into the model; run a single forward.
+            self.guidance_scale = 1.0
+        elif is_edge and not self._guidance_is_explicit():
+            self.guidance_scale = (
+                COSMOS3_EDGE_T2I_GUIDANCE_SCALE
+                if is_t2i
+                else COSMOS3_EDGE_T2V_GUIDANCE_SCALE
+            )
+        if is_t2i and not is_distilled and self.guidance_interval is None:
+            self.guidance_interval = COSMOS3_T2I_GUIDANCE_INTERVAL
+        if is_edge:
+            self.supported_resolutions = COSMOS3_EDGE_SUPPORTED_RESOLUTIONS
+            if self.height is None and self.width is None:
+                if is_t2i:
+                    self.width = self.height = COSMOS3_EDGE_T2I_SIZE
+                else:
+                    self.width, self.height = (
+                        COSMOS3_EDGE_T2V_WIDTH,
+                        COSMOS3_EDGE_T2V_HEIGHT,
+                    )
+
    def _set_output_file_name(self) -> None:
        # Action outputs never need a visual filename. This also avoids hashing
        # in-memory observation images while base visual adjustment is running.
@@ -1088,8 +1088,8 @@ def _register_configs():

    # Cosmos3 — single checkpoint serves T2V, I2V, and T2I. Mode is dispatched
    # per-request inside the pipeline from ``num_frames`` and ``image_path``.
-    # Both Nano (16B) and Super (64B) share the same pipeline; arch dimensions
-    # come from ``transformer/config.json`` via ``update_model_arch``.
+    # All variants share the same pipeline; arch dimensions (size, activation,
+    # QK-norm) come from ``transformer/config.json`` via ``update_model_arch``.
    register_configs(
        sampling_param_cls=Cosmos3SamplingParams,
        pipeline_config_cls=Cosmos3Config,
@@ -1099,8 +1099,11 @@ def _register_configs():
            "nvidia/Cosmos3-Super",
            "nvidia/Cosmos3-Super-Text2Image",
            "nvidia/Cosmos3-Super-Image2Video",
+            "nvidia/Cosmos3-Edge",
        ],
-        model_detectors=[lambda hf_id: "cosmos3omnidiffuserspipeline" in hf_id.lower()],
+        # Match both the new ``Cosmos3OmniPipeline`` and the legacy
+        # ``Cosmos3OmniDiffusersPipeline`` ``_class_name`` (diffusers rename).
+        model_detectors=[lambda hf_id: "cosmos3omni" in hf_id.lower()],
    )

    # SANA
@@ -2,8 +2,14 @@

 from __future__ import annotations

+import json
+import os
+from contextlib import nullcontext
+from typing import Any
+
 from fastapi import APIRouter, HTTPException, Request, Response, WebSocket

+from sglang.multimodal_gen.configs.sample.sampling_params import generate_request_id
 from sglang.multimodal_gen.runtime.entrypoints.action.protocol import (
    action_generation_response,
    action_metadata,
@@ -15,11 +21,41 @@ from sglang.multimodal_gen.runtime.entrypoints.action.protocol import (
 from sglang.multimodal_gen.runtime.entrypoints.action.ws_utils import (
    run_action_msgpack_ws,
 )
+from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
+    flatten_extra_params,
+    save_image_to_path,
+    temp_dir_if_disabled,
+)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.srt.utils.json_response import orjson_response

 router = APIRouter(prefix="/v1/actions", tags=["actions"])

+_ACTION_VIDEO_EXTENSIONS = {
+    ".avi",
+    ".gif",
+    ".m4v",
+    ".mkv",
+    ".mov",
+    ".mp4",
+    ".mpeg",
+    ".mpg",
+    ".webm",
+}
+
+_ACTION_FORM_NON_PARAMETER_FIELDS = {
+    "extra_body",
+    "extra_params",
+    "id",
+    "image_reference",
+    "input_reference",
+    "prompt",
+    "reference_url",
+    "request_id",
+    "task",
+    "video_reference",
+}
+

 def _wants_msgpack(request: Request) -> bool:
    content_type = request.headers.get("content-type", "").lower()
@@ -40,26 +76,154 @@ def _prefer_numpy_output(payload: dict) -> None:
    runtime.setdefault("output_format", "numpy")


+def _parse_form_value(value: Any) -> Any:
+    if not isinstance(value, str):
+        return value
+    if not value.strip():
+        return None
+    try:
+        return json.loads(value)
+    except Exception:
+        return value
+
+
+def _parse_extra_params(value: Any, field_name: str) -> dict[str, Any]:
+    if value in (None, ""):
+        return {}
+    try:
+        parsed = json.loads(value) if isinstance(value, str) else value
+    except (json.JSONDecodeError, TypeError) as exc:
+        raise ValueError(f"{field_name} is not valid JSON") from exc
+    if not isinstance(parsed, dict):
+        raise ValueError(f"{field_name} must be a JSON object")
+    return flatten_extra_params(dict(parsed))
+
+
+def _is_form_upload(value: Any) -> bool:
+    return callable(getattr(value, "read", None)) and hasattr(value, "filename")
+
+
+def _is_probably_video_upload(value: Any) -> bool:
+    content_type = (getattr(value, "content_type", "") or "").lower()
+    if content_type.startswith("video/"):
+        return True
+    filename = getattr(value, "filename", None)
+    if not filename:
+        return False
+    filename = str(filename).split("?", 1)[0].split("#", 1)[0]
+    return os.path.splitext(filename)[1].lower() in _ACTION_VIDEO_EXTENSIONS
+
+
+async def _save_action_upload(
+    upload: Any,
+    request_id: str,
+    uploads_dir: str,
+    *,
+    fallback_name: str,
+) -> str:
+    filename = getattr(upload, "filename", None) or fallback_name
+    target_path = os.path.join(uploads_dir, f"{request_id}_{filename}")
+    return await save_image_to_path(upload, target_path)
+
+
+async def _multipart_action_payload(request: Request, uploads_dir: str) -> dict:
+    raw_form = await request.form()
+    request_id = (
+        raw_form.get("request_id") or raw_form.get("id") or generate_request_id()
+    )
+
+    parameters: dict[str, Any] = {}
+    for key, value in raw_form.multi_items():
+        if key in _ACTION_FORM_NON_PARAMETER_FIELDS or _is_form_upload(value):
+            continue
+        parameters[key] = _parse_form_value(value)
+
+    parameters.update(_parse_extra_params(raw_form.get("extra_body"), "extra_body"))
+    parameters.update(_parse_extra_params(raw_form.get("extra_params"), "extra_params"))
+    flatten_extra_params(parameters)
+
+    action_mode = str(parameters.get("action_mode", "policy")).strip().lower()
+    observation: dict[str, Any] = {}
+
+    upload_fields = (
+        ("video_reference", "video"),
+        ("image_reference", "image"),
+        ("input_reference", None),
+    )
+    for field_name, forced_observation_key in upload_fields:
+        upload = raw_form.get(field_name)
+        if not _is_form_upload(upload):
+            continue
+        saved_path = await _save_action_upload(
+            upload,
+            str(request_id),
+            uploads_dir,
+            fallback_name=field_name,
+        )
+        observation_key = forced_observation_key
+        if observation_key is None:
+            observation_key = (
+                "video"
+                if action_mode == "inverse_dynamics"
+                or _is_probably_video_upload(upload)
+                else "input_reference"
+            )
+        observation[observation_key] = saved_path
+        break
+
+    for form_key, observation_key in (
+        ("video_path", "video"),
+        ("video_url", "video"),
+        ("reference_url", "video" if action_mode == "inverse_dynamics" else "image"),
+        (
+            "input_reference",
+            "video" if action_mode == "inverse_dynamics" else "input_reference",
+        ),
+    ):
+        value = raw_form.get(form_key)
+        if isinstance(value, str) and value and observation_key not in observation:
+            observation[observation_key] = value
+
+    return {
+        "request_id": request_id,
+        "input": {
+            "task": raw_form.get("prompt") or raw_form.get("task") or "",
+            "observation": observation,
+        },
+        "parameters": parameters,
+    }
+
+
@router.post("/generations")
 async def create_action_generation(request: Request):
    server_args: ServerArgs = request.app.state.server_args
+    content_type = request.headers.get("content-type", "").lower()
+    is_multipart = "multipart/form-data" in content_type
+    input_context = (
+        temp_dir_if_disabled(getattr(server_args, "input_save_path", None))
+        if is_multipart
+        else nullcontext(None)
+    )
    try:
-        if "msgpack" in request.headers.get("content-type", "").lower():
-            payload = unpack_msgpack(await request.body())
-        else:
-            payload = await request.json()
-        wants_msgpack = _wants_msgpack(request)
-        if wants_msgpack:
-            _prefer_numpy_output(payload)
-        output = await infer_action(payload, server_args)
-        if _response_format(payload) == "raw":
-            response = action_raw_response(output, preserve_numpy=wants_msgpack)
-        else:
-            response = action_generation_response(
-                output,
-                server_args,
-                preserve_numpy=wants_msgpack,
-            )
+        with input_context as uploads_dir:
+            if is_multipart:
+                payload = await _multipart_action_payload(request, uploads_dir)
+            elif "msgpack" in content_type:
+                payload = unpack_msgpack(await request.body())
+            else:
+                payload = await request.json()
+            wants_msgpack = _wants_msgpack(request)
+            if wants_msgpack:
+                _prefer_numpy_output(payload)
+            output = await infer_action(payload, server_args)
+            if _response_format(payload) == "raw":
+                response = action_raw_response(output, preserve_numpy=wants_msgpack)
+            else:
+                response = action_generation_response(
+                    output,
+                    server_args,
+                    preserve_numpy=wants_msgpack,
+                )
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    if wants_msgpack:
@@ -683,6 +683,15 @@ async def create_video(
            selected = value if value is not None else extra_from_form.get(name)
            return _parse_form_extra_value(selected)

+        def form_text_value(name: str, value: Any) -> Any:
+            """Resolve a text field without JSON-decoding it.
+
+            Some models take a serialized JSON object as the prompt text, which
+            ``_parse_form_extra_value`` would turn back into a dict and fail
+            request validation.
+            """
+            return value if value is not None else extra_from_form.get(name)
+
        request_field_names = set(VideoGenerationsRequest.model_fields)
        extra_request_fields = {
            key: value
@@ -708,7 +717,7 @@ async def create_video(
            num_frames=num_frames_val,
            seed=form_value("seed", seed),
            generator_device=form_value("generator_device", generator_device),
-            negative_prompt=form_value("negative_prompt", negative_prompt),
+            negative_prompt=form_text_value("negative_prompt", negative_prompt),
            num_inference_steps=form_value("num_inference_steps", num_inference_steps),
            guidance_scale=form_value("guidance_scale", guidance_scale),
            guidance_scale_2=form_value("guidance_scale_2", guidance_scale_2),
@@ -1,3 +1,5 @@
+import inspect
+
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
    ComponentLoader,
 )
@@ -11,6 +13,26 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)


+def _supported_init_kwargs(scheduler_cls, config: dict) -> dict:
+    """Drop scheduler_config keys the constructor can't accept.
+
+    Some checkpoints (e.g. distilled fixed-step schedulers) carry keys in
+    ``scheduler_config.json`` that are consumed downstream, not by the
+    scheduler ``__init__``. Constructors taking ``**kwargs`` accept anything
+    and are left untouched; otherwise unknown non-private keys would raise
+    ``TypeError``, so they are dropped (private ``_``-keys are handled by
+    ``register_to_config``).
+    """
+    params = inspect.signature(scheduler_cls.__init__).parameters
+    if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()):
+        return config
+    dropped = [k for k in config if k not in params and not k.startswith("_")]
+    if dropped:
+        logger.debug("Ignoring unsupported scheduler config keys: %s", dropped)
+        return {k: v for k, v in config.items() if k not in dropped}
+    return config
+
+
 class SchedulerLoader(ComponentLoader):
    """Loader for scheduler."""

@@ -41,7 +63,7 @@ class SchedulerLoader(ComponentLoader):

        scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)

-        scheduler = scheduler_cls(**config)
+        scheduler = scheduler_cls(**_supported_init_kwargs(scheduler_cls, config))
        if server_args.pipeline_config.flow_shift is not None:
            scheduler.set_shift(server_args.pipeline_config.flow_shift)

@@ -30,6 +30,7 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
    apply_qk_norm_rope,
 )
 from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
    MergedColumnParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
@@ -432,6 +433,57 @@ class Cosmos3GatedMLP(nn.Module):
        return out


+class Cosmos3DenseMLP(nn.Module):
+    """Dense MLP with a squared-ReLU activation: ``down(relu(up(x)) ** 2)``."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
+    ):
+        super().__init__()
+        self.up_proj = ColumnParallelLinear(
+            hidden_size,
+            intermediate_size,
+            bias=False,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=add_prefix("up_proj", prefix),
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
+            bias=False,
+            input_is_parallel=True,
+            quant_config=quant_config,
+            prefix=add_prefix("down_proj", prefix),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        up, _ = self.up_proj(x)
+        up = F.relu(up)
+        out, _ = self.down_proj(up * up)
+        return out
+
+
+def _build_mlp(
+    hidden_act: str,
+    hidden_size: int,
+    intermediate_size: int,
+    prefix: str,
+    quant_config: QuantizationConfig | None,
+) -> nn.Module:
+    mlp_cls = Cosmos3DenseMLP if hidden_act == "relu2" else Cosmos3GatedMLP
+    return mlp_cls(
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        prefix=prefix,
+        quant_config=quant_config,
+    )
+
+
 # -----------------------------------------------------------------------------
 # Cosmos3 UND Causal Attention
 # -----------------------------------------------------------------------------
@@ -446,6 +498,9 @@ class Cosmos3CausalAttention(nn.Module):
        num_attention_heads: int,
        num_key_value_heads: int,
        head_dim: int,
+        qk_norm: bool = True,
+        use_k_norm_und_for_gen: bool = False,
+        rms_norm_eps: float = 1e-6,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
    ):
@@ -454,6 +509,7 @@ class Cosmos3CausalAttention(nn.Module):
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
+        self.qk_norm = qk_norm
        self.tp_size = get_tp_world_size()
        if num_attention_heads % self.tp_size != 0:
            raise ValueError(
@@ -487,9 +543,16 @@ class Cosmos3CausalAttention(nn.Module):
            prefix=add_prefix("to_out", prefix),
        )

-        # Per-head QK norm.
-        self.norm_q = RMSNorm(head_dim, eps=1e-6)
-        self.norm_k = RMSNorm(head_dim, eps=1e-6)
+        # Per-head QK norm (optional; some backbones omit it on text).
+        if qk_norm:
+            self.norm_q = RMSNorm(head_dim, eps=1e-6)
+            self.norm_k = RMSNorm(head_dim, eps=1e-6)
+
+        # Dense-backbone variants normalize the keys handed to the GEN
+        # cross-attention separately from the reasoner self-attention keys.
+        self.k_norm_und_for_gen = (
+            RMSNorm(head_dim, eps=rms_norm_eps) if use_k_norm_und_for_gen else None
+        )

    def forward(
        self,
@@ -526,14 +589,27 @@ class Cosmos3CausalAttention(nn.Module):
            :,
        ]

-        q = F.rms_norm(
-            q, (self.head_dim,), self.norm_q.weight, self.norm_q.variance_epsilon
-        )
-        k = F.rms_norm(
-            k, (self.head_dim,), self.norm_k.weight, self.norm_k.variance_epsilon
-        )
+        k_und = k
+        if self.qk_norm:
+            q = F.rms_norm(
+                q, (self.head_dim,), self.norm_q.weight, self.norm_q.variance_epsilon
+            )
+            k = F.rms_norm(
+                k, (self.head_dim,), self.norm_k.weight, self.norm_k.variance_epsilon
+            )
        q, k = _apply_qwen3_rope_from_cache(q, k, cos_sin_cache)

+        if self.k_norm_und_for_gen is not None:
+            k_gen = F.rms_norm(
+                k_und,
+                (self.head_dim,),
+                self.k_norm_und_for_gen.weight,
+                self.k_norm_und_for_gen.variance_epsilon,
+            )
+            _, k_gen = _apply_qwen3_rope_from_cache(q, k_gen, cos_sin_cache)
+        else:
+            k_gen = k
+
        out = F.scaled_dot_product_attention(
            q.transpose(1, 2),
            k.transpose(1, 2),
@@ -544,7 +620,7 @@ class Cosmos3CausalAttention(nn.Module):
        out = out.transpose(1, 2).reshape(batch_size, seq_len, -1)

        out, _ = self.to_out(out)
-        return out, k, v
+        return out, k_gen, v


 # -----------------------------------------------------------------------------
@@ -697,6 +773,9 @@ class Cosmos3UndDecoderLayer(nn.Module):
        head_dim: int,
        intermediate_size: int,
        rms_norm_eps: float,
+        hidden_act: str,
+        qk_norm: bool,
+        use_k_norm_und_for_gen: bool,
        layer_idx: int,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
@@ -709,12 +788,16 @@ class Cosmos3UndDecoderLayer(nn.Module):
            num_attention_heads=num_attention_heads,
            num_key_value_heads=num_key_value_heads,
            head_dim=head_dim,
+            qk_norm=qk_norm,
+            use_k_norm_und_for_gen=use_k_norm_und_for_gen,
+            rms_norm_eps=rms_norm_eps,
            prefix=add_prefix("self_attn", prefix),
            quant_config=quant_config,
        )
        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.mlp = Cosmos3GatedMLP(
+        self.mlp = _build_mlp(
+            hidden_act=hidden_act,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            prefix=add_prefix("mlp", prefix),
@@ -763,6 +846,7 @@ class Cosmos3GenDecoderLayer(nn.Module):
        head_dim: int,
        intermediate_size: int,
        rms_norm_eps: float,
+        hidden_act: str,
        layer_idx: int,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
@@ -782,7 +866,8 @@ class Cosmos3GenDecoderLayer(nn.Module):
        )
        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.mlp = Cosmos3GatedMLP(
+        self.mlp = _build_mlp(
+            hidden_act=hidden_act,
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            prefix=add_prefix("mlp", prefix),
@@ -844,6 +929,9 @@ class Cosmos3LanguageModel(nn.Module):
        rms_norm_eps: float,
        rope_theta: float,
        mrope_section: tuple[int, int, int],
+        hidden_act: str,
+        qk_norm: bool,
+        use_k_norm_und_for_gen: bool,
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
@@ -869,6 +957,9 @@ class Cosmos3LanguageModel(nn.Module):
                    head_dim=head_dim,
                    intermediate_size=intermediate_size,
                    rms_norm_eps=rms_norm_eps,
+                    hidden_act=hidden_act,
+                    qk_norm=qk_norm,
+                    use_k_norm_und_for_gen=use_k_norm_und_for_gen,
                    layer_idx=i,
                    prefix=f"layers.{i}",
                    quant_config=quant_config,
@@ -955,6 +1046,14 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
        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
+        self.hidden_act = arch.hidden_act
+        self.rope_theta = arch.rope_theta
+
+        # The checkpoint may override the activation (and thus the MLP weight
+        # layout), so bind the arch-derived mappings on the instance.
+        self.param_names_mapping = arch.param_names_mapping
+        self.reverse_param_names_mapping = arch.reverse_param_names_mapping
+        self.lora_param_names_mapping = arch.lora_param_names_mapping

        # Ulysses sequence parallelism. When CFG-parallel is also enabled
        # the SP group only spans ranks that share a CFG context (cond or
@@ -979,6 +1078,9 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
            rms_norm_eps=arch.rms_norm_eps,
            rope_theta=arch.rope_theta,
            mrope_section=arch.mrope_section,
+            hidden_act=arch.hidden_act,
+            qk_norm=arch.qk_norm_for_text,
+            use_k_norm_und_for_gen=arch.use_und_k_norm_for_gen,
            quant_config=quant_config,
        )

@@ -1051,6 +1153,7 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
                    head_dim=arch.head_dim,
                    intermediate_size=arch.intermediate_size,
                    rms_norm_eps=arch.rms_norm_eps,
+                    hidden_act=arch.hidden_act,
                    layer_idx=i,
                    prefix=f"gen_layers.{i}",
                    quant_config=quant_config,
@@ -1563,6 +1666,11 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
                yield linear_target + ".input_scale", in_t.max()

        for name, tensor in iterator:
+            # ModelOpt calibration buffers are not inference parameters; the
+            # runtime FP8 path reads weight_scale/input_scale instead. Drop them
+            # so they don't reach the fused-linear concat below.
+            if "_quantizer." in name:
+                continue
            target_name, merge_index, num_to_merge = mapping_fn(name)
            if num_to_merge is None:
                yield target_name, tensor
@@ -1606,9 +1714,8 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
        rotary_emb = self.language_model.rotary_emb
        if rotary_emb.inv_freq.is_meta:
            dim = rotary_emb.head_dim
-            rope_theta = 5000000.0  # From config
            inv_freq = 1.0 / (
-                rope_theta
+                self.rope_theta
                ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
            )
            rotary_emb.register_buffer("inv_freq", inv_freq, persistent=False)
@@ -533,7 +533,12 @@ class FlowMatchEulerDiscreteScheduler(
        else:
            if self.config.stochastic_sampling:
                x0 = sample - current_sigma * model_output
-                noise = torch.randn_like(sample)
+                noise = torch.randn(
+                    sample.shape,
+                    generator=generator,
+                    device=sample.device,
+                    dtype=sample.dtype,
+                )
                prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise
            else:
                prev_sample = sample + dt * model_output
@@ -33,7 +33,10 @@ class Cosmos3Pipeline(ComposedPipelineBase):
    stages from ``batch.data_type`` and ``batch.preprocessed_image``.
    """

-    pipeline_name = "Cosmos3OmniDiffusersPipeline"
+    # Canonical ``_class_name`` in newer checkpoints' ``model_index.json``.
+    # Older checkpoints declare ``Cosmos3OmniDiffusersPipeline`` (see the
+    # back-compat alias below).
+    pipeline_name = "Cosmos3OmniPipeline"
    is_video_pipeline = True

    _required_config_modules = [
@@ -109,4 +112,15 @@ class Cosmos3Pipeline(ComposedPipelineBase):
        )


-EntryClass = Cosmos3Pipeline
+class Cosmos3OmniDiffusersPipeline(Cosmos3Pipeline):
+    """Back-compat alias for checkpoints whose ``model_index.json`` still
+    declares the legacy ``_class_name`` ``Cosmos3OmniDiffusersPipeline``.
+
+    Registering both names lets old (Nano/Super) and new (Edge) checkpoints
+    resolve to the same native pipeline.
+    """
+
+    pipeline_name = "Cosmos3OmniDiffusersPipeline"
+
+
+EntryClass = [Cosmos3Pipeline, Cosmos3OmniDiffusersPipeline]
@@ -79,6 +79,35 @@ COSMOS3_I2V_FLOW_SHIFT = 10.0
 COSMOS3_T2V_FLOW_SHIFT = 10.0
 COSMOS3_V2V_FLOW_SHIFT = 10.0
 COSMOS3_ACTION_FLOW_SHIFT = 10.0
+# Edge uses a single low flow-shift for every video mode (t2v/i2v/v2v).
+COSMOS3_EDGE_VIDEO_FLOW_SHIFT = 3.0
+
+
+def _inject_caption_metadata(
+    prompt: str, num_frames: int, fps: float, height: int, width: int
+) -> str | None:
+    """Add the generation metadata that Cosmos3's structured captions carry.
+
+    Training captions always ship ``resolution``, plus ``duration``/``fps`` for
+    video, so a JSON caption without them is out of distribution. Returns
+    ``None`` when the prompt is not a JSON object; those prompts carry the same
+    metadata as trailing prose via the duration template instead.
+    """
+    try:
+        caption = json.loads(prompt)
+    except (TypeError, ValueError):
+        return None
+    if not isinstance(caption, dict):
+        return None
+
+    caption["resolution"] = {"H": int(height), "W": int(width)}
+    if num_frames > 1:
+        caption["duration"] = f"{int(num_frames / fps) if fps > 0 else 0}s"
+        caption["fps"] = float(fps)
+    else:
+        caption.pop("duration", None)
+        caption.pop("fps", None)
+    return json.dumps(caption)


 def _resize_crop_pil(
@@ -327,6 +356,15 @@ class Cosmos3TokenizationStage(PipelineStage):
            use_system_prompt = False
            use_duration_template = False
            self.log_info(f"Action prompt: {prompt}")
+        else:
+            structured_caption = _inject_caption_metadata(
+                prompt, num_frames, fps, batch.height, batch.width
+            )
+            if structured_caption is not None:
+                # The metadata is already in the caption; appending the prose
+                # template too would state the duration twice.
+                prompt = structured_caption
+                use_duration_template = False

        # Apply duration template if enabled (no temporal concept for T2I).
        if use_duration_template and not is_image_gen and num_frames > 1:
@@ -680,15 +718,14 @@ class Cosmos3TimestepPreparationStage(PipelineStage):
        super().__init__()
        self.scheduler = scheduler

-    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.
-        """
+    def _default_flow_shift_for_mode(self, batch: Req, is_edge: bool) -> float | None:
+        """Resolve the per-mode default flow_shift for the request."""
        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 is_edge:
+            return COSMOS3_EDGE_VIDEO_FLOW_SHIFT
        if batch.preprocessed_image is not None:
            return COSMOS3_I2V_FLOW_SHIFT
        if batch.preprocessed_video is not None:
@@ -698,12 +735,32 @@ class Cosmos3TimestepPreparationStage(PipelineStage):
    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
        """Prepare scheduler timesteps."""
        device = get_local_torch_device()
+
+        pipeline_config = server_args.pipeline_config
+        distilled_sigmas = pipeline_config.distilled_sigmas
+        if distilled_sigmas is not None:
+            # Distilled checkpoints carry an explicit fixed-step sigma schedule
+            # with the shift already baked in; drive the scheduler from it
+            # directly (step count == len(sigmas), num_inference_steps ignored).
+            # Reset shift so set_timesteps does not re-shift the baked-in sigmas.
+            if hasattr(self.scheduler, "set_shift"):
+                self.scheduler.set_shift(1.0)
+            self.scheduler.set_timesteps(sigmas=distilled_sigmas, device=device)
+            batch.timesteps = self.scheduler.timesteps
+            self.log_info(
+                f"Prepared {len(batch.timesteps)} distilled timesteps "
+                f"(sigmas={distilled_sigmas})"
+            )
+            return batch
+
        num_inference_steps = batch.num_inference_steps
        flow_shift = getattr(batch, "flow_shift", None)
        if flow_shift is None:
-            flow_shift = server_args.pipeline_config.flow_shift
+            flow_shift = pipeline_config.flow_shift
        if flow_shift is None:
-            flow_shift = self._default_flow_shift_for_mode(batch)
+            flow_shift = self._default_flow_shift_for_mode(
+                batch, bool(pipeline_config.is_edge)
+            )
        if flow_shift is not None and hasattr(self.scheduler, "set_shift"):
            self.scheduler.set_shift(float(flow_shift))

@@ -742,6 +799,7 @@ class Cosmos3DenoisingStage(PipelineStage):
        self.scheduler = scheduler
        self.server_args = server_args
        self._logged_parallel_config = False
+        self._logged_cfg_split = False

        # Apply torch.compile if enabled
        if server_args is not None:
@@ -917,6 +975,13 @@ class Cosmos3DenoisingStage(PipelineStage):
        timesteps = batch.timesteps
        guidance_scale = batch.guidance_scale

+        # Seed the scheduler's stochastic (SDE) noise from the request seed so it
+        # is identical on every sequence-parallel rank; otherwise each rank draws
+        # its own noise and the sharded latents diverge at the shard boundary.
+        generator = batch.generator
+        if generator is None and batch.seed is not None:
+            generator = torch.Generator(device=latents.device).manual_seed(batch.seed)
+
        cond_text_ids = batch.extra["cond_text_ids"]
        cond_text_mask = batch.extra["cond_text_mask"]
        uncond_text_ids = batch.extra["uncond_text_ids"]
@@ -1036,7 +1101,7 @@ class Cosmos3DenoisingStage(PipelineStage):
                        action_start_frame_offset=action_start_frame_offset,
                    )
                else:
-                    noise_pred = self._predict_noise_cfg_batched(
+                    noise_pred = self._predict_noise_cfg(
                        latents=latents,
                        timestep=timestep,
                        cond_text_ids=cond_text_ids,
@@ -1047,10 +1112,8 @@ class Cosmos3DenoisingStage(PipelineStage):
                        fps=fps,
                        guidance_scale=effective_scale,
                        noisy_frame_mask=velocity_mask,
-                        max_text_seq_len=max(
-                            batch.extra["cond_text_seq_len"],
-                            batch.extra["uncond_text_seq_len"],
-                        ),
+                        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,
@@ -1102,6 +1165,7 @@ class Cosmos3DenoisingStage(PipelineStage):
                noise_pred,
                t,
                latents,
+                generator=generator,
                return_dict=False,
            )[0]

@@ -1155,6 +1219,87 @@ class Cosmos3DenoisingStage(PipelineStage):
        self.log_info("Denoising complete")
        return batch

+    def _predict_noise_cfg(
+        self,
+        latents: torch.Tensor,
+        timestep: torch.Tensor,
+        cond_text_ids: torch.Tensor,
+        cond_text_mask: torch.Tensor,
+        uncond_text_ids: torch.Tensor,
+        uncond_text_mask: torch.Tensor,
+        video_shape: tuple[int, int, int],
+        fps: float,
+        guidance_scale: float,
+        cond_text_seq_len: int,
+        uncond_text_seq_len: int,
+        **kwargs: Any,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        """Run CFG, batching the two branches only when that is lossless.
+
+        A batched forward needs one shared text length, so the shorter prompt
+        gets right-padded. Those pad positions carry all-zero UND K/V, and the
+        GEN cross-attention runs unmasked over the full text K/V — a zero key
+        scores logit 0 and still takes softmax mass while contributing nothing,
+        which weakens the padded branch's conditioning. Fall back to one
+        forward per branch, each trimmed to its own length, whenever the
+        prompts differ in length.
+        """
+        if cond_text_seq_len == uncond_text_seq_len:
+            return self._predict_noise_cfg_batched(
+                latents=latents,
+                timestep=timestep,
+                cond_text_ids=cond_text_ids,
+                cond_text_mask=cond_text_mask,
+                uncond_text_ids=uncond_text_ids,
+                uncond_text_mask=uncond_text_mask,
+                video_shape=video_shape,
+                fps=fps,
+                guidance_scale=guidance_scale,
+                max_text_seq_len=cond_text_seq_len,
+                **kwargs,
+            )
+
+        if not self._logged_cfg_split and not self._current_batch_is_warmup:
+            self._logged_cfg_split = True
+            self.log_info(
+                "Prompt and negative prompt tokenize to different lengths "
+                f"({cond_text_seq_len} vs {uncond_text_seq_len}); running the "
+                "CFG branches in separate forwards to keep padding out of the "
+                "cross-attention"
+            )
+
+        cond = self._run_transformer(
+            latents=latents,
+            timestep=timestep,
+            text_ids=cond_text_ids,
+            text_mask=cond_text_mask,
+            video_shape=video_shape,
+            fps=fps,
+            cache_key="cond",
+            max_text_seq_len=cond_text_seq_len,
+            **kwargs,
+        )
+        uncond = self._run_transformer(
+            latents=latents,
+            timestep=timestep,
+            text_ids=uncond_text_ids,
+            text_mask=uncond_text_mask,
+            video_shape=video_shape,
+            fps=fps,
+            cache_key="uncond",
+            max_text_seq_len=uncond_text_seq_len,
+            **kwargs,
+        )
+
+        def _cfg_combine(
+            cond_pred: torch.Tensor, uncond_pred: torch.Tensor
+        ) -> torch.Tensor:
+            return uncond_pred + guidance_scale * (cond_pred - uncond_pred)
+
+        if isinstance(cond, tuple):
+            return tuple(_cfg_combine(c, u) for c, u in zip(cond, uncond, strict=True))
+        return _cfg_combine(cond, uncond)
+
    def _predict_noise_cfg_batched(
        self,
        latents: torch.Tensor,
@@ -440,13 +440,31 @@ def load_dict(file_path):
        ) from e


+def _split_hf_subfolder(path: str) -> tuple[str, str | None]:
+    """Split 'namespace/repo/subfolder' into (repo_id, subfolder), or return (path, None)."""
+    if os.path.isabs(path):
+        return path, None
+    parts = path.split("/")
+    if len(parts) > 2:
+        return "/".join(parts[:2]), "/".join(parts[2:])
+    return path, None
+
+
 def prepare_diffusers_component_path_for_loading(component_path: str) -> str:
    """Download component repos if needed and patch legacy flat ModelOpt configs."""
-    local_component_path = (
-        maybe_download_model(component_path)
-        if not os.path.exists(component_path)
-        else component_path
-    )
+    if os.path.exists(component_path):
+        local_component_path = component_path
+    else:
+        repo_id, subfolder = _split_hf_subfolder(component_path)
+        if subfolder is not None:
+            # component_path is 'namespace/repo/subfolder' — download only that subfolder
+            local_repo = maybe_download_model(
+                repo_id,
+                allow_patterns=[f"{subfolder}/**", f"{subfolder}/*"],
+            )
+            local_component_path = os.path.join(local_repo, subfolder)
+        else:
+            local_component_path = maybe_download_model(component_path)
    config_path = os.path.join(local_component_path, "config.json")
    if not os.path.exists(config_path):
        return local_component_path
@@ -2,6 +2,7 @@
 """Unit tests for Cosmos3 config, weight mapping, and sampling params."""

 import importlib.util
+import json
 import types
 import unittest
 from unittest import mock
@@ -12,9 +13,14 @@ from sglang.multimodal_gen.configs.models.dits.cosmos3video import (
    _build_cosmos3_param_names_mapping,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config
-from sglang.multimodal_gen.configs.sample.cosmos3 import Cosmos3SamplingParams
+from sglang.multimodal_gen.configs.sample.cosmos3 import (
+    COSMOS3_EDGE_SUPPORTED_RESOLUTIONS,
+    Cosmos3SamplingParams,
+)
 from sglang.multimodal_gen.configs.sample.sampling_params import DataType
 from sglang.multimodal_gen.registry import (
+    _PIPELINE_REGISTRY,
+    _discover_and_register_pipelines,
    _get_config_info,
    get_non_diffusers_pipeline_name,
 )
@@ -48,6 +54,8 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.c
    Cosmos3ImagePreprocessStage,
    Cosmos3LatentPreparationStage,
    Cosmos3TimestepPreparationStage,
+    Cosmos3TokenizationStage,
+    _inject_caption_metadata,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_action import (
    EMBODIMENT_TO_DOMAIN_ID,
@@ -224,6 +232,42 @@ class TestCosmos3ParamNamesMapping(unittest.TestCase):
        self.assertNotIn("language_model", key)


+class TestCosmos3DenseParamNamesMapping(unittest.TestCase):
+    """Dense (squared-ReLU) checkpoints ship no gate_proj; up/down_proj must
+    pass through unmerged."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.fn = staticmethod(
+            get_param_names_mapping(_build_cosmos3_param_names_mapping(gated_mlp=False))
+        )
+
+    def test_und_mlp_up_proj_unmerged(self):
+        key, idx, _ = _apply(self.fn, "layers.1.mlp.up_proj.weight")
+        self.assertEqual(key, "language_model.layers.1.mlp.up_proj.weight")
+        self.assertIsNone(idx)
+
+    def test_und_mlp_down_proj_unmerged(self):
+        key, idx, _ = _apply(self.fn, "layers.1.mlp.down_proj.weight")
+        self.assertEqual(key, "language_model.layers.1.mlp.down_proj.weight")
+        self.assertIsNone(idx)
+
+    def test_gen_mlp_up_proj_unmerged(self):
+        key, idx, _ = _apply(self.fn, "layers.2.mlp_moe_gen.up_proj.weight")
+        self.assertEqual(key, "gen_layers.2.mlp.up_proj.weight")
+        self.assertIsNone(idx)
+
+    def test_gen_mlp_down_proj_unmerged(self):
+        key, idx, _ = _apply(self.fn, "layers.2.mlp_moe_gen.down_proj.weight")
+        self.assertEqual(key, "gen_layers.2.mlp.down_proj.weight")
+        self.assertIsNone(idx)
+
+    def test_qkv_merge_still_applies(self):
+        key, idx, total = _apply(self.fn, "layers.0.self_attn.to_q.weight")
+        self.assertEqual(key, "language_model.layers.0.self_attn.to_qkv.weight")
+        self.assertEqual((idx, total), (0, 3))
+
+
 class TestCosmos3AdjustNumFrames(unittest.TestCase):
    """Verify VAE-aligned frame rounding in Cosmos3Config."""

@@ -306,29 +350,100 @@ class TestCosmos3SchedulerConfig(unittest.TestCase):
    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)),
+            stage._default_flow_shift_for_mode(
+                self._batch(data_type=DataType.IMAGE), is_edge=False
+            ),
            3.0,
        )
        self.assertEqual(
            stage._default_flow_shift_for_mode(
-                self._batch(preprocessed_image=torch.empty(1))
+                self._batch(preprocessed_image=torch.empty(1)), is_edge=False
            ),
            10.0,
        )
        self.assertEqual(
            stage._default_flow_shift_for_mode(
-                self._batch(preprocessed_video=torch.empty(1))
+                self._batch(preprocessed_video=torch.empty(1)), is_edge=False
            ),
            10.0,
        )
-        self.assertEqual(stage._default_flow_shift_for_mode(self._batch()), 10.0)
+        self.assertEqual(
+            stage._default_flow_shift_for_mode(self._batch(), is_edge=False), 10.0
+        )
        self.assertEqual(
            stage._default_flow_shift_for_mode(
-                self._batch(sp_kwargs={"action_mode": "policy"})
+                self._batch(sp_kwargs={"action_mode": "policy"}), is_edge=False
            ),
            10.0,
        )

+    def test_edge_flow_shift_default(self):
+        stage = self._stage()
+        # Edge uses 3.0 for T2I and every video mode (T2V/I2V/V2V); action stays high.
+        self.assertEqual(
+            stage._default_flow_shift_for_mode(self._batch(), is_edge=True), 3.0
+        )
+        self.assertEqual(
+            stage._default_flow_shift_for_mode(
+                self._batch(data_type=DataType.IMAGE), is_edge=True
+            ),
+            3.0,
+        )
+        self.assertEqual(
+            stage._default_flow_shift_for_mode(
+                self._batch(preprocessed_image=torch.empty(1)), is_edge=True
+            ),
+            3.0,
+        )
+        self.assertEqual(
+            stage._default_flow_shift_for_mode(
+                self._batch(preprocessed_video=torch.empty(1)), is_edge=True
+            ),
+            3.0,
+        )
+        self.assertEqual(
+            stage._default_flow_shift_for_mode(
+                self._batch(sp_kwargs={"action_mode": "policy"}), is_edge=True
+            ),
+            10.0,
+        )
+
+
+class TestCosmos3EdgeSamplingDefaults(unittest.TestCase):
+    """Edge variant fills its own resolution/guidance defaults; base is untouched."""
+
+    def test_edge_t2v_defaults(self):
+        sp = Cosmos3SamplingParams(prompt="t", num_frames=81)
+        sp._resolve_variant_defaults(is_edge=True)
+        self.assertEqual((sp.width, sp.height), (832, 480))
+        self.assertEqual(sp.guidance_scale, 5.0)
+
+    def test_edge_t2i_defaults(self):
+        sp = Cosmos3SamplingParams(prompt="t", num_frames=1)
+        sp._resolve_variant_defaults(is_edge=True)
+        self.assertEqual((sp.width, sp.height), (640, 640))
+        self.assertEqual(sp.guidance_scale, 7.0)
+
+    def test_edge_restricts_supported_resolutions(self):
+        sp = Cosmos3SamplingParams(prompt="t", num_frames=1)
+        sp._resolve_variant_defaults(is_edge=True)
+        self.assertEqual(sp.supported_resolutions, COSMOS3_EDGE_SUPPORTED_RESOLUTIONS)
+        # Base high-res sizes are excluded so they trip the "unsupported" warning.
+        self.assertNotIn((1280, 720), sp.supported_resolutions)
+        self.assertNotIn((1024, 1024), sp.supported_resolutions)
+
+    def test_non_edge_defers_resolution_to_base(self):
+        sp = Cosmos3SamplingParams(prompt="t", num_frames=81)
+        sp._resolve_variant_defaults(is_edge=False)
+        self.assertIsNone(sp.width)
+        self.assertIsNone(sp.height)
+        self.assertEqual(sp.guidance_scale, 4.0)
+
+    def test_explicit_resolution_preserved_for_edge(self):
+        sp = Cosmos3SamplingParams(prompt="t", num_frames=81, width=1024, height=576)
+        sp._resolve_variant_defaults(is_edge=True)
+        self.assertEqual((sp.width, sp.height), (1024, 576))
+

 class TestCosmos3SamplingParamsDataType(unittest.TestCase):
    """Verify num_frames==1 flips data_type to IMAGE before file name derivation."""
@@ -525,6 +640,7 @@ class TestCosmos3ModelResolution(unittest.TestCase):
            "nvidia/Cosmos3-Super",
            "nvidia/Cosmos3-Super-Text2Image",
            "nvidia/Cosmos3-Super-Image2Video",
+            "nvidia/Cosmos3-Edge",
        ):
            with self.subTest(model_path=model_path):
                self.assertIsNone(get_non_diffusers_pipeline_name(model_path))
@@ -533,6 +649,31 @@ class TestCosmos3ModelResolution(unittest.TestCase):
                self.assertIs(config_info.sampling_param_cls, Cosmos3SamplingParams)
                self.assertIs(config_info.pipeline_config_cls, Cosmos3Config)

+    def test_class_name_detection_matches_legacy_and_new(self):
+        """Unregistered checkpoints resolve via ``_class_name``: both the legacy
+        ``Cosmos3OmniDiffusersPipeline`` and the current ``Cosmos3OmniPipeline``
+        map to the native Cosmos3 config."""
+        for idx, class_name in enumerate(
+            ("Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline")
+        ):
+            model_path = f"acme/mystery-ckpt-{idx}"
+            with self.subTest(class_name=class_name):
+                with mock.patch(
+                    "sglang.multimodal_gen.registry.maybe_download_model_index",
+                    return_value={"_class_name": class_name},
+                ):
+                    config_info = _get_config_info(model_path)
+                self.assertIsNotNone(config_info)
+                self.assertIs(config_info.pipeline_config_cls, Cosmos3Config)
+
+    def test_legacy_and_new_pipeline_names_both_registered(self):
+        """Both ``_class_name`` spellings resolve to a native pipeline class so
+        old (Nano/Super) and new (Edge) checkpoints load."""
+        _discover_and_register_pipelines()
+        for pipeline_name in ("Cosmos3OmniPipeline", "Cosmos3OmniDiffusersPipeline"):
+            with self.subTest(pipeline_name=pipeline_name):
+                self.assertIn(pipeline_name, _PIPELINE_REGISTRY)
+

 class TestCosmos3OpenAIProtocol(unittest.TestCase):
    """Verify Cosmos3 modality knobs are exposed by the video HTTP schema."""
@@ -930,5 +1071,130 @@ class TestCosmos3ModalitySamplingParams(unittest.TestCase):
            self.assertIsNone(getattr(sp, field))


+class TestCosmos3CaptionMetadata(unittest.TestCase):
+    """Structured captions get generation metadata; prose prompts opt out."""
+
+    def test_video_caption_gets_resolution_duration_and_fps(self):
+        prompt = json.dumps({"temporal_caption": "a cone melts"})
+        caption = json.loads(
+            _inject_caption_metadata(
+                prompt, num_frames=189, fps=24.0, height=480, width=832
+            )
+        )
+        self.assertEqual(caption["resolution"], {"H": 480, "W": 832})
+        # Whole seconds, matching the caption format the model was trained on.
+        self.assertEqual(caption["duration"], "7s")
+        self.assertEqual(caption["fps"], 24.0)
+        self.assertEqual(caption["temporal_caption"], "a cone melts")
+
+    def test_image_caption_drops_temporal_fields(self):
+        prompt = json.dumps({"subjects": ["hands"], "duration": "5s", "fps": 24.0})
+        caption = json.loads(
+            _inject_caption_metadata(
+                prompt, num_frames=1, fps=24.0, height=768, width=768
+            )
+        )
+        self.assertEqual(caption["resolution"], {"H": 768, "W": 768})
+        self.assertNotIn("duration", caption)
+        self.assertNotIn("fps", caption)
+        self.assertEqual(caption["subjects"], ["hands"])
+
+    def test_request_resolution_overrides_caption(self):
+        prompt = json.dumps({"resolution": {"H": 111, "W": 222}})
+        caption = json.loads(
+            _inject_caption_metadata(
+                prompt, num_frames=1, fps=24.0, height=640, width=640
+            )
+        )
+        self.assertEqual(caption["resolution"], {"H": 640, "W": 640})
+
+    def test_unrelated_caption_fields_are_preserved(self):
+        prompt = json.dumps({"aspect_ratio": "1,1", "subjects": ["a vase"]})
+        caption = json.loads(
+            _inject_caption_metadata(
+                prompt, num_frames=1, fps=24.0, height=640, width=640
+            )
+        )
+        self.assertEqual(caption["aspect_ratio"], "1,1")
+        self.assertEqual(caption["subjects"], ["a vase"])
+
+    def test_zero_fps_does_not_divide_by_zero(self):
+        prompt = json.dumps({"temporal_caption": "a cone melts"})
+        caption = json.loads(
+            _inject_caption_metadata(
+                prompt, num_frames=81, fps=0.0, height=480, width=832
+            )
+        )
+        self.assertEqual(caption["duration"], "0s")
+
+    def test_non_structured_prompts_are_declined(self):
+        for prompt in (
+            "A curious raccoon in a field of sunflowers.",
+            json.dumps(["not", "an", "object"]),
+            json.dumps("a bare string"),
+            "",
+        ):
+            with self.subTest(prompt=prompt):
+                self.assertIsNone(
+                    _inject_caption_metadata(
+                        prompt, num_frames=81, fps=24.0, height=480, width=832
+                    )
+                )
+
+
+class TestCosmos3DurationTemplateSuppression(unittest.TestCase):
+    """A structured caption must not also get the prose duration suffix."""
+
+    def _run_prompt_stage(self, prompt: str, use_duration_template: bool) -> str:
+        """Drive Cosmos3TokenizationStage.forward far enough to see the prompt."""
+        seen = {}
+
+        def fake_tokenize(prompt_text, *args, **kwargs):
+            seen.setdefault("prompt", prompt_text)
+            return (
+                torch.zeros(1, 4, dtype=torch.long),
+                torch.ones(1, 4, dtype=torch.long),
+                4,
+            )
+
+        stage = Cosmos3TokenizationStage.__new__(Cosmos3TokenizationStage)
+        batch = types.SimpleNamespace(
+            prompt=prompt,
+            negative_prompt="bad",
+            max_sequence_length=512,
+            use_duration_template=use_duration_template,
+            use_system_prompt=False,
+            fps=24.0,
+            num_frames=189,
+            height=480,
+            width=832,
+            data_type=DataType.VIDEO,
+            sampling_params=types.SimpleNamespace(action_mode=None),
+            extra={},
+        )
+        server_args = types.SimpleNamespace(
+            pipeline_config=types.SimpleNamespace(
+                use_duration_template=True, use_system_prompt=False
+            )
+        )
+        with (
+            mock.patch.object(stage, "_tokenize_prompt", fake_tokenize),
+            mock.patch.object(stage, "log_info"),
+        ):
+            stage.forward(batch, server_args)
+        return seen["prompt"]
+
+    def test_structured_caption_stays_valid_json(self):
+        prompt = json.dumps({"temporal_caption": "a cone melts"})
+        final = self._run_prompt_stage(prompt, use_duration_template=True)
+        caption = json.loads(final)  # would raise if the suffix were appended
+        self.assertEqual(caption["duration"], "7s")
+        self.assertNotIn("seconds long", final)
+
+    def test_prose_prompt_still_gets_the_duration_suffix(self):
+        final = self._run_prompt_stage("A curious raccoon.", use_duration_template=True)
+        self.assertIn("seconds long", final)
+
+
 if __name__ == "__main__":
    unittest.main()