[diffusion] model: support a new model (#24994)

docs/diffusion/compatibility_matrix.md

@@ -39,6 +39,8 @@ default parameters when initializing and generating videos.
 | Helios Distilled             | `BestWishYsh/Helios-Distilled`                     | 720p                 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |               ❌               |                   ❌                    |             ❌             |
 | LTX-2 (one/two-stage/TI2V)   | `Lightricks/LTX-2`                                | 768×512<br>1536×1024 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |               ❌               |                   ❌                    |             ❌             |
 | LTX-2.3 (one/two-stage/TI2V/HQ) | `Lightricks/LTX-2.3`                           | 768×512<br>1536×1024<br>1920×1088 (HQ default) |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |               ❌               |                   ❌                    |             ❌             |             ❌             |
+| Cosmos3-Nano (T2V / I2V / T2I)  | `nvidia/Cosmos3-Nano`                          | 720p · 480p<br>1024×1024 (T2I)                 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |               ❌               |                   ❌                    |             ❌             |
+| Cosmos3-Super (T2V / I2V / T2I) | `nvidia/Cosmos3-Super`                         | 720p · 480p<br>1024×1024 (T2I)                 |    ❌     |         ❌         |     ❌     |              ❌               |               ❌               |                   ❌                    |             ❌             |               ❌               |                   ❌                    |             ❌             |
 
 **Note**:
 
@@ -56,6 +58,11 @@ default parameters when initializing and generating videos.
    - `resident` usually provides the best latency/throughput but uses much more VRAM.
    - `original` keeps official two-stage semantics without the premerged stage-2 transformer path.
    - Example (one prior run): `original` `154.67s`, `snapshot` `114.05s`, `resident` `75.71s`; peak VRAM trend is `original < snapshot < resident`.
+5. Cosmos3 ships in two sizes — `nvidia/Cosmos3-Nano` (8B) and
+   `nvidia/Cosmos3-Super` (32B). Both share the same pipeline; the only
+   difference is transformer depth and width, picked up from
+   `transformer/config.json` at load time. A single checkpoint serves T2V,
+   I2V (`--image-path`), and T2I (`--num-frames 1`).
 
 ### Image Generation Models
 

docs/diffusion/index.md

@@ -4,7 +4,7 @@ SGLang Diffusion is a high-performance inference framework for image and video g
 
 ## Key Features
 
-- Broad model support across Wan, Hunyuan, Qwen-Image, FLUX, Z-Image, GLM-Image, and more
+- Broad model support across Wan, Hunyuan, Cosmos3, Qwen-Image, FLUX, Z-Image, GLM-Image, and more
 - Fast inference with `sgl-kernel`, JIT kernels, scheduler improvements, and caching acceleration
 - Multiple interfaces: `sglang generate`, `sglang serve`, and an OpenAI-compatible API
 - Multi-platform support for NVIDIA, AMD, Intel XPU, Ascend, Apple Silicon, and Moore Threads

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

@@ -1,5 +1,6 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 
+from sglang.multimodal_gen.configs.models.dits.cosmos3video import Cosmos3VideoConfig
 from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig
 from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig
 from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
@@ -11,6 +12,7 @@ from sglang.multimodal_gen.configs.models.dits.stablediffusion3 import (
 from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig
 
 __all__ = [
+    "Cosmos3VideoConfig",
     "HeliosConfig",
     "HunyuanVideoConfig",
     "WanVideoConfig",

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

@@ -0,0 +1,188 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 video DiT architecture configuration."""
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.configs.models.fsdp import is_module_list_entry_in
+
+
+def is_layers(n: str, m) -> bool:
+    return is_module_list_entry_in(n, ("layers", "gen_layers"))
+
+
+def _build_cosmos3_param_names_mapping() -> dict:
+    """Map diffusers-format Cosmos3 weights to the sglang model namespace.
+
+    Source keys (diffusers transformer ckpt) → target keys (sglang model):
+        model.embed_tokens.weight                              -> language_model.embed_tokens.weight
+        model.layers.X.input_layernorm.weight                  -> language_model.layers.X.input_layernorm.weight
+        model.layers.X.input_layernorm_moe_gen.weight          -> gen_layers.X.input_layernorm.weight
+        model.layers.X.self_attn.{q,k,v}_proj.weight           -> language_model.layers.X.self_attn.to_qkv.weight (concat dim 0)
+        model.layers.X.self_attn.{q,k,v}_proj_moe_gen.weight   -> gen_layers.X.cross_attention.to_qkv.weight (concat dim 0)
+        model.layers.X.mlp.{gate,up}_proj.weight               -> language_model.layers.X.mlp.gate_up_proj.weight (concat dim 0)
+        model.layers.X.mlp_moe_gen.{gate,up}_proj.weight       -> gen_layers.X.mlp.gate_up_proj.weight (concat dim 0)
+        model.norm_moe_gen.weight                              -> norm_moe_gen.weight
+        time_embedder.mlp.{0,2}.weight                         -> time_embedder.linear_{1,2}.weight
+        vae2llm.weight, llm2vae.weight                         -> (pass-through)
+
+    GEN patterns (`*_moe_gen`) must precede the UND catch-all so the
+    catch-all can't claim GEN keys. `model.norm.weight` and `lm_head.weight`
+    are inherited from Qwen3-VL pretraining and not used at inference, so
+    they are skipped via empty-string replacement.
+    """
+    return {
+        # Inherited from Qwen3-VL pretraining; unused at diffusion inference.
+        r"^lm_head\.weight$": "",
+        r"^model\.norm\.weight$": "",
+        # Top-level norms / heads.
+        r"^model\.norm_moe_gen\.(.*)$": r"norm_moe_gen.\1",
+        r"^model\.embed_tokens\.(.*)$": r"language_model.embed_tokens.\1",
+        # Time embedder: mlp.0 -> linear_1, mlp.2 -> linear_2 (SiLU at index 1).
+        r"^time_embedder\.mlp\.0\.(.*)$": r"time_embedder.linear_1.\1",
+        r"^time_embedder\.mlp\.2\.(.*)$": r"time_embedder.linear_2.\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"^model\.layers\.(\d+)\.self_attn\.q_proj_moe_gen\.(.*)$": (
+            r"gen_layers.\1.cross_attention.to_qkv.\2",
+            0,
+            3,
+        ),
+        r"^model\.layers\.(\d+)\.self_attn\.k_proj_moe_gen\.(.*)$": (
+            r"gen_layers.\1.cross_attention.to_qkv.\2",
+            1,
+            3,
+        ),
+        r"^model\.layers\.(\d+)\.self_attn\.v_proj_moe_gen\.(.*)$": (
+            r"gen_layers.\1.cross_attention.to_qkv.\2",
+            2,
+            3,
+        ),
+        r"^model\.layers\.(\d+)\.self_attn\.o_proj_moe_gen\.(.*)$": r"gen_layers.\1.cross_attention.to_out.\2",
+        r"^model\.layers\.(\d+)\.self_attn\.q_norm_moe_gen\.(.*)$": r"gen_layers.\1.cross_attention.norm_q.\2",
+        r"^model\.layers\.(\d+)\.self_attn\.k_norm_moe_gen\.(.*)$": r"gen_layers.\1.cross_attention.norm_k.\2",
+        r"^model\.layers\.(\d+)\.input_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.input_layernorm.\2",
+        r"^model\.layers\.(\d+)\.post_attention_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.post_attention_layernorm.\2",
+        # GEN MLP gate/up merge into MergedColumnParallelLinear gate_up_proj.
+        # Must precede the mlp_moe_gen catch-all below.
+        r"^model\.layers\.(\d+)\.mlp_moe_gen\.gate_proj\.(.*)$": (
+            r"gen_layers.\1.mlp.gate_up_proj.\2",
+            0,
+            2,
+        ),
+        r"^model\.layers\.(\d+)\.mlp_moe_gen\.up_proj\.(.*)$": (
+            r"gen_layers.\1.mlp.gate_up_proj.\2",
+            1,
+            2,
+        ),
+        r"^model\.layers\.(\d+)\.mlp_moe_gen\.(.*)$": r"gen_layers.\1.mlp.\2",
+        # UND pathway: per-layer attention rename (q/k/v_proj -> to_qkv merged,
+        # q_norm/k_norm -> norm_q/k, o_proj -> to_out).
+        r"^model\.layers\.(\d+)\.self_attn\.q_proj\.(.*)$": (
+            r"language_model.layers.\1.self_attn.to_qkv.\2",
+            0,
+            3,
+        ),
+        r"^model\.layers\.(\d+)\.self_attn\.k_proj\.(.*)$": (
+            r"language_model.layers.\1.self_attn.to_qkv.\2",
+            1,
+            3,
+        ),
+        r"^model\.layers\.(\d+)\.self_attn\.v_proj\.(.*)$": (
+            r"language_model.layers.\1.self_attn.to_qkv.\2",
+            2,
+            3,
+        ),
+        r"^model\.layers\.(\d+)\.self_attn\.o_proj\.(.*)$": r"language_model.layers.\1.self_attn.to_out.\2",
+        r"^model\.layers\.(\d+)\.self_attn\.q_norm\.(.*)$": r"language_model.layers.\1.self_attn.norm_q.\2",
+        r"^model\.layers\.(\d+)\.self_attn\.k_norm\.(.*)$": r"language_model.layers.\1.self_attn.norm_k.\2",
+        # UND MLP gate/up merge into MergedColumnParallelLinear gate_up_proj.
+        # Must precede the layers catch-all below.
+        r"^model\.layers\.(\d+)\.mlp\.gate_proj\.(.*)$": (
+            r"language_model.layers.\1.mlp.gate_up_proj.\2",
+            0,
+            2,
+        ),
+        r"^model\.layers\.(\d+)\.mlp\.up_proj\.(.*)$": (
+            r"language_model.layers.\1.mlp.gate_up_proj.\2",
+            1,
+            2,
+        ),
+        # UND pathway: layernorms + remaining mlp keys pass through unchanged.
+        r"^model\.layers\.(\d+)\.(.*)$": r"language_model.layers.\1.\2",
+    }
+
+
+@dataclass
+class Cosmos3VideoArchConfig(DiTArchConfig):
+    """Architecture config for Cosmos3 Omni Transformer.
+
+    Cosmos3 uses a dual-pathway design:
+    - Understanding (UND): Causal self-attention for text tokens
+    - Generation (GEN): Cross-attention from visual to cached UND K/V
+
+    Field names mirror the diffusers ``transformer/config.json`` so values
+    flow through ``update_model_arch`` without translation. Defaults are
+    Qwen3-8B-Instruct-derived and overridden by the checkpoint at load time.
+    """
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_layers])
+
+    # Transformer architecture
+    hidden_size: int = 4096
+    num_hidden_layers: int = 36
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8  # GQA
+    head_dim: int = 128
+    intermediate_size: int = 12288
+
+    # Latent space configuration
+    latent_patch_size: int = 2
+    latent_channel: int = 48
+    out_channels: int = 48
+
+    # RoPE configuration (Qwen3-VL 3D mRoPE: temporal, height, width)
+    mrope_section: tuple[int, int, int] = (24, 20, 20)
+    rope_theta: float = 5000000.0
+    # Populated from rope_scaling in the diffusers config when present.
+    rope_scaling: dict | None = None
+
+    # Temporal configuration
+    base_fps: float = 24.0
+    temporal_compression_factor: int = 4
+    unified_3d_mrope_temporal_modality_margin: int = 15000
+
+    # Timestep embedding
+    timestep_scale: float = 0.001
+    frequency_embedding_size: int = 256
+
+    # Vocab size (Qwen3-VL tokenizer)
+    vocab_size: int = 151936
+
+    # RMSNorm epsilon
+    rms_norm_eps: float = 1e-6
+
+    # Weight mapping from checkpoint to model
+    param_names_mapping: dict = field(
+        default_factory=_build_cosmos3_param_names_mapping
+    )
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+        super().__post_init__()
+        # Diffusers configs nest the mrope sizes under `rope_scaling`; lift it.
+        if isinstance(self.rope_scaling, dict) and "mrope_section" in self.rope_scaling:
+            self.mrope_section = tuple(self.rope_scaling["mrope_section"])
+        self.in_channels = self.latent_channel
+        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
+
+
+@dataclass
+class Cosmos3VideoConfig(DiTConfig):
+    """DiT config wrapper for Cosmos3 Video model."""
+
+    arch_config: DiTArchConfig = field(default_factory=Cosmos3VideoArchConfig)
+    prefix: str = "Cosmos3"

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

@@ -4,6 +4,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     PipelineConfig,
     SlidingTileAttnConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config
 from sglang.multimodal_gen.configs.pipeline_configs.diffusers_generic import (
     DiffusersGenericPipelineConfig,
 )
@@ -43,6 +44,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import (
 from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig
 
 __all__ = [
+    "Cosmos3Config",
     "DiffusersGenericPipelineConfig",
     "HeliosDistilledConfig",
     "HeliosMidConfig",

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

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 pipeline configuration.
+
+A single config serves T2V, I2V, and T2I — same checkpoint, same DiT, same
+VAE. ``task_type`` is ``TI2V`` so the request validator accepts an optional
+``image_path`` (for I2V) without requiring it. Per-modality dispatch happens
+in the stages from ``num_frames`` and ``image_path``; T2I overrides
+``data_type`` to ``IMAGE`` in :meth:`SamplingParams._adjust`.
+"""
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig
+from sglang.multimodal_gen.configs.models.dits.cosmos3video import Cosmos3VideoConfig
+from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig
+from sglang.multimodal_gen.configs.pipeline_configs.base import (
+    ModelTaskType,
+    PipelineConfig,
+)
+
+
+@dataclass
+class Cosmos3Config(PipelineConfig):
+    """Cosmos3 unified pipeline config.
+
+    Cosmos3 reuses the Wan VAE (48 latent channels, 4× temporal, 16× spatial),
+    a UniPC rectified-flow sampler, and a ~15B-parameter dual-pathway DiT.
+    There is no separate text encoder — text is tokenized and embedded inside
+    the transformer.
+    """
+
+    # TI2V (text + image → video) so the request validator accepts ``image_path``
+    # without requiring it. T2V ignores it; I2V uses it; T2I disregards it.
+    task_type: ModelTaskType = ModelTaskType.TI2V
+
+    dit_config: DiTConfig = field(default_factory=Cosmos3VideoConfig)
+
+    # Wan VAE with 48 latent channels (overridden in __post_init__ below).
+    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    # Sourced from scheduler_config.json in the checkpoint.
+    flow_shift: float | None = None
+
+    precision: str = "bf16"
+    vae_precision: str = "bf16"
+
+    # Pipeline-level (not sampling) knobs.
+    max_sequence_length: int = 512
+    use_duration_template: bool = True
+    use_system_prompt: bool = False
+
+    def __post_init__(self):
+        self.vae_config.arch_config.z_dim = 48
+        # Encoder is needed for I2V; T2V/T2I never invoke it.
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        # WanVAE defaults use_parallel_encode/decode to True, which silently
+        # activates an SP-sharded VAE path when sp_world_size > 1 and produces
+        # garbled pixels for cosmos3's latent shape.
+        self.vae_config.use_parallel_encode = False
+        self.vae_config.use_parallel_decode = False
+
+    def adjust_num_frames(self, num_frames: int) -> int:
+        """Round ``num_frames`` so ``(n - 1) % 4 == 0`` for the VAE.
+
+        Skips rounding when ``num_frames == 1`` (T2I path) so the single
+        frame survives untouched.
+        """
+        if num_frames == 1:
+            return 1
+        vae_scale_factor_temporal = 4
+        if (num_frames - 1) % vae_scale_factor_temporal != 0:
+            num_frames = (
+                (num_frames - 1) // vae_scale_factor_temporal
+            ) * vae_scale_factor_temporal + 1
+        return num_frames

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

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 sampling parameters.
+
+A single ``SamplingParams`` class serves T2V, I2V, and T2I — the per-request
+mode is dispatched in the pipeline from ``num_frames`` (``== 1`` → T2I) and
+``image_path`` (set → I2V). 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 sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
+
+
+@dataclass
+class Cosmos3SamplingParams(SamplingParams):
+    """Cosmos3 sampling parameters (T2V defaults; also used for I2V / T2I)."""
+
+    height: int = 720
+    width: int = 1280
+    num_frames: int = 81
+    fps: int = 24
+
+    guidance_scale: float = 4.0
+    num_inference_steps: int = 35
+
+    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.
+    guidance_interval: tuple[float, float] | None = None
+
+    supported_resolutions: list[tuple[int, int]] | None = field(
+        default_factory=lambda: [
+            (1280, 720),
+            (720, 1280),
+            (832, 480),
+            (480, 832),
+            (1024, 1024),
+        ]
+    )
+
+    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
+        # extension. Flip before the base derives the file name.
+        if self.num_frames == 1:
+            self.data_type = DataType.IMAGE
+        super()._set_output_file_name()

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

@@ -511,6 +511,7 @@ class SamplingParams:
             "wan" in pipeline_name_lower
             or "helios" in pipeline_name_lower
             or "joy" in pipeline_name_lower
+            or "cosmos3" in pipeline_name_lower
         ) and (self.enable_sequence_shard is None or self.enable_sequence_shard):
             self.enable_sequence_shard = True
             logger.debug("Automatically enabled enable_sequence_shard")

python/sglang/multimodal_gen/registry.py

@@ -29,6 +29,7 @@ if TYPE_CHECKING:
     from sglang.multimodal_gen.runtime.server_args import Backend
 
 from sglang.multimodal_gen.configs.pipeline_configs import (
+    Cosmos3Config,
     FastHunyuanConfig,
     FluxPipelineConfig,
     HeliosDistilledConfig,
@@ -84,6 +85,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import (
     Wan2_2_T2V_A14B_Config,
     Wan2_2_TI2V_5B_Config,
 )
+from sglang.multimodal_gen.configs.sample.cosmos3 import Cosmos3SamplingParams
 from sglang.multimodal_gen.configs.sample.ernie_image import ErnieImageSamplingParams
 from sglang.multimodal_gen.configs.sample.flux import (
     Flux2KleinBaseSamplingParams,
@@ -102,9 +104,7 @@ from sglang.multimodal_gen.configs.sample.hunyuan import (
     HunyuanSamplingParams,
 )
 from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams
-from sglang.multimodal_gen.configs.sample.joy_image import (
-    JoyImageEditSamplingParams,
-)
+from sglang.multimodal_gen.configs.sample.joy_image import JoyImageEditSamplingParams
 from sglang.multimodal_gen.configs.sample.ltx_2 import (
     LTX2SamplingParams,
     LTX23HQSamplingParams,
@@ -951,6 +951,20 @@ 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 (8B) and Super (32B) share the same pipeline; arch dimensions
+    # come from ``transformer/config.json`` via ``update_model_arch``.
+    register_configs(
+        sampling_param_cls=Cosmos3SamplingParams,
+        pipeline_config_cls=Cosmos3Config,
+        hf_model_paths=[
+            "nvidia/Cosmos3-Nano",
+            "nvidia/Cosmos3-Super",
+        ],
+        model_detectors=[lambda hf_id: "cosmos3omnidiffuserspipeline" in hf_id.lower()],
+    )
+
     # SANA
     register_configs(
         sampling_param_cls=SanaSamplingParams,

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -418,6 +418,7 @@ class USPAttention(nn.Module):
         attn_mask: torch.Tensor | None = None,
         num_replicated_prefix: int = 0,
         num_replicated_suffix: int = 0,
+        num_replicated_kv_prefix: int = 0,
         skip_sequence_parallel_override: bool = False,
     ) -> torch.Tensor:
         """
@@ -432,6 +433,12 @@ class USPAttention(nn.Module):
             num_replicated_suffix: number of trailing tokens in q/k/v that are
                 replicated across all SP ranks, e.g. caption tokens appended
                 after image tokens in Z-Image joint attention.
+            num_replicated_kv_prefix: number of leading tokens in k/v only
+                (not q) that are replicated across all SP ranks. Used for
+                cross-attention where the keys/values include a fully-replicated
+                conditioning prefix (e.g. cached text K/V) followed by a
+                sequence-sharded suffix (image tokens). Q has no replicated
+                portion and is fully sequence-sharded.
 
         Note: Replicated tensors are not supported in this implementation.
         When skip_sequence_parallel=True (set at construction time), all SP
@@ -537,9 +544,19 @@ class USPAttention(nn.Module):
             return out
 
         sp_size = get_ulysses_parallel_world_size()
-        if num_replicated_prefix > 0 and num_replicated_suffix > 0:
+        if (
+            sum(
+                bool(n)
+                for n in (
+                    num_replicated_prefix,
+                    num_replicated_suffix,
+                    num_replicated_kv_prefix,
+                )
+            )
+            > 1
+        ):
             raise ValueError(
-                "USPAttention does not support replicated prefix and suffix at the same time."
+                "USPAttention supports at most one replicated-token mode per call."
             )
         if sp_size > 1 and num_replicated_prefix > 0:
             return self._forward_with_replicated_prefix(
@@ -549,6 +566,10 @@ class USPAttention(nn.Module):
             return self._forward_with_replicated_suffix(
                 q, k, v, ctx_attn_metadata, num_replicated_suffix
             )
+        if sp_size > 1 and num_replicated_kv_prefix > 0:
+            return self._forward_with_replicated_kv_prefix(
+                q, k, v, ctx_attn_metadata, num_replicated_kv_prefix
+            )
 
         # Ulysses-style All-to-All for sequence/head sharding
         if sp_size > 1:
@@ -636,6 +657,48 @@ class USPAttention(nn.Module):
 
         return torch.cat([out_rep, out_shard], dim=1)
 
+    def _forward_with_replicated_kv_prefix(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        ctx_attn_metadata,
+        num_rep: int,
+    ) -> torch.Tensor:
+        """Ulysses cross-attention where only K/V have a replicated prefix.
+
+        Q is sequence-sharded across SP ranks with no replicated portion. K/V
+        carry a fully-replicated prefix (``[:num_rep]``, same on every rank,
+        e.g. cached text K/V) followed by a sequence-sharded suffix (e.g.
+        image tokens) that aligns with Q's sharding.
+
+        Strategy:
+        1. All-to-all Q and the sharded K/V suffix (seq → head shard).
+        2. Locally slice the replicated K/V prefix to the same head shard.
+        3. Concatenate prefix + suffix on the sequence dim and attend.
+        4. All-to-all the output back (head shard → seq shard).
+        """
+        sp_rank = get_sp_parallel_rank()
+
+        k_rep, k_shard = k[:, :num_rep], k[:, num_rep:]
+        v_rep, v_shard = v[:, :num_rep], v[:, num_rep:]
+
+        q = _usp_input_all_to_all(q, head_dim=2)
+        k_shard = _usp_input_all_to_all(k_shard, head_dim=2)
+        v_shard = _usp_input_all_to_all(v_shard, head_dim=2)
+
+        h_kv_local = k_shard.shape[2]
+        h_start = sp_rank * h_kv_local
+        h_end = h_start + h_kv_local
+        k_rep = k_rep[:, :, h_start:h_end, :].contiguous()
+        v_rep = v_rep[:, :, h_start:h_end, :].contiguous()
+
+        k = torch.cat([k_rep, k_shard], dim=1)
+        v = torch.cat([v_rep, v_shard], dim=1)
+
+        out = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
+        return _usp_output_all_to_all(out, head_dim=2)
+
     def _forward_with_replicated_suffix(
         self,
         q: torch.Tensor,

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

@@ -432,7 +432,7 @@ class AutoProcessorLoader(ComponentLoader):
 class TokenizerLoader(ComponentLoader):
     """Loader for tokenizers."""
 
-    component_names = ["tokenizer"]
+    component_names = ["tokenizer", "text_tokenizer"]
     expected_library = "transformers"
 
     def load_customized(

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

@@ -245,6 +245,9 @@ def maybe_load_fsdp_model(
         )
 
     weight_iterator = safetensors_weights_iterator(weight_dir_list)
+    preprocess_loaded_state_dict = getattr(model, "preprocess_loaded_state_dict", None)
+    if preprocess_loaded_state_dict is not None:
+        weight_iterator = preprocess_loaded_state_dict(weight_iterator)
     param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
     load_model_from_full_model_state_dict(
         model,

python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py

@@ -192,6 +192,15 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
             dark samples instead of limiting it to samples with medium brightness. Loosely related to
             [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
+        sigma_min (`float`, *optional*):
+            Override the lower bound of the sigma range when `use_karras_sigmas=True` or
+            `use_exponential_sigmas=True`. If `None`, the bound is derived from the trained beta schedule.
+        sigma_max (`float`, *optional*):
+            Override the upper bound of the sigma range when `use_karras_sigmas=True` or
+            `use_exponential_sigmas=True`. If `None`, the bound is derived from the trained beta schedule.
+        shift_terminal (`float`, *optional*):
+            Forward-compat field accepted from diffusers >=0.38 scheduler configs; consulted only by the
+            dynamic-shift code path, which this scheduler does not yet implement.
     """
 
     _compatibles = [e.name for e in KarrasDiffusionSchedulers]
@@ -226,6 +235,9 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         rescale_betas_zero_snr: bool = False,
         use_dynamic_shifting: bool = False,
         time_shift_type: str = "exponential",
+        sigma_min: float | None = None,
+        sigma_max: float | None = None,
+        shift_terminal: float | None = None,
     ):
         if self.config.use_beta_sigmas and not is_scipy_available():
             raise ImportError(
@@ -410,9 +422,16 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             sigmas = self._convert_to_karras(
                 in_sigmas=sigmas, num_inference_steps=num_inference_steps
             )
-            timesteps = np.array(
-                [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]
-            ).round()
+            if self.config.use_flow_sigmas:
+                # Karras builds sigmas in EDM space; flow-matching models expect
+                # sigmas in [0, 1]. Map EDM -> flow with sigma / (sigma + 1) and
+                # derive timesteps from the flow sigmas (matches diffusers >=0.38).
+                sigmas = sigmas / (sigmas + 1)
+                timesteps = (sigmas * self.config.num_train_timesteps).copy()
+            else:
+                timesteps = np.array(
+                    [self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]
+                ).round()
             if self.config.final_sigmas_type == "sigma_min":
                 sigma_last = sigmas[-1]
             elif self.config.final_sigmas_type == "zero":

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

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 video diffusion pipeline.
+
+Cosmos3 has no separate text encoder — the transformer embeds text directly
+via its Understanding (UND) pathway, and the Generation (GEN) pathway
+cross-attends to the cached UND K/V at each denoising step.
+"""
+
+import os
+
+from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
+    ComposedPipelineBase,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import (
+    Cosmos3DecodingStage,
+    Cosmos3DenoisingStage,
+    Cosmos3ImagePreprocessStage,
+    Cosmos3LatentPreparationStage,
+    Cosmos3TimestepPreparationStage,
+    Cosmos3TokenizationStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class Cosmos3Pipeline(ComposedPipelineBase):
+    """Cosmos3 diffusion pipeline shared by T2V, I2V, and T2I.
+
+    Text is tokenized and embedded directly inside the transformer; there is
+    no separate text encoder. Modality is dispatched per-request inside the
+    stages from ``batch.data_type`` and ``batch.preprocessed_image``.
+    """
+
+    pipeline_name = "Cosmos3OmniDiffusersPipeline"
+    is_video_pipeline = True
+
+    _required_config_modules = [
+        "text_tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
+    ]
+
+    def create_pipeline_stages(self, server_args: ServerArgs) -> None:
+        """Create Cosmos3 pipeline stages.
+
+        Stage order:
+        1. Cosmos3ImagePreprocessStage - Load + aspect-resize the I2V image (no-op otherwise)
+        2. Cosmos3TokenizationStage - Tokenize with Qwen2 chat template
+        3. Cosmos3LatentPreparationStage - Noise latent (or image-conditioned for I2V)
+        4. Cosmos3TimestepPreparationStage - Set up scheduler timesteps
+        5. Cosmos3DenoisingStage - Dual-pathway denoising (UND once, GEN per step)
+        6. Cosmos3DecodingStage - VAE decode to video, or to a single image for T2I
+        """
+        text_tokenizer = self.get_module("text_tokenizer")
+        vae = self.get_module("vae")
+        transformer = self.get_module("transformer")
+        scheduler = self.get_module("scheduler")
+
+        # Guardrails on by default; opt out with SGLANG_DISABLE_COSMOS3_GUARDRAILS=1.
+        guardrails_on = os.environ.get("SGLANG_DISABLE_COSMOS3_GUARDRAILS", "0") != "1"
+
+        self.add_stage(Cosmos3ImagePreprocessStage())
+        self.add_stage(Cosmos3TokenizationStage(tokenizer=text_tokenizer))
+        if guardrails_on:
+            from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_guardrails import (
+                Cosmos3TextGuardrailStage,
+            )
+
+            self.add_stage(Cosmos3TextGuardrailStage())
+        self.add_stage(Cosmos3LatentPreparationStage(vae, transformer))
+        self.add_stage(Cosmos3TimestepPreparationStage(scheduler))
+        self.add_stage(Cosmos3DenoisingStage(transformer, scheduler, server_args))
+        self.add_stage(Cosmos3DecodingStage(vae, guardrails=guardrails_on))
+
+        logger.info(
+            "Cosmos3 pipeline stages created successfully (guardrails=%s)",
+            guardrails_on,
+        )
+
+
+EntryClass = Cosmos3Pipeline

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

@@ -0,0 +1,882 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 pipeline stages: image preprocess, tokenization, latent / timestep
+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
+from ``batch.data_type`` and the presence of ``batch.preprocessed_image``.
+"""
+
+from typing import Any
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn as nn
+from tqdm.auto import tqdm
+
+from sglang.multimodal_gen.configs.sample.sampling_params import DataType
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    cfg_model_parallel_all_reduce,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
+    get_sp_parallel_rank,
+    get_sp_world_size,
+)
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+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.validators import (
+    StageValidators as V,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+    VerificationResult,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.utils.common import get_compiler_backend
+
+logger = init_logger(__name__)
+
+COSMOS3_DEFAULT_NEGATIVE_PROMPT = ""
+COSMOS3_DURATION_TEMPLATE = (
+    "The video is {duration:.1f} seconds long and is of {fps} FPS."
+)
+COSMOS3_VIDEO_SYSTEM_PROMPT = (
+    "You are a helpful assistant who will generate videos from a given prompt."
+)
+COSMOS3_IMAGE_SYSTEM_PROMPT = (
+    "You are a helpful assistant who will generate images from a given prompt."
+)
+
+
+class Cosmos3ImagePreprocessStage(PipelineStage):
+    """Load, aspect-resize, and center-crop the I2V conditioning image.
+
+    No-op when the request has no image (T2V / T2I). The output is a
+    ``[1, 3, H, W]`` tensor in ``[-1, 1]`` written to ``batch.preprocessed_image``.
+    """
+
+    parallelism_type = StageParallelismType.REPLICATED
+
+    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+        return VerificationResult()
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        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:
+            return batch
+
+        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
+        tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).contiguous()
+
+        batch.preprocessed_image = tensor
+        self.log_info(f"Preprocessed conditioning image to {target_w}x{target_h}")
+        return batch
+
+
+class Cosmos3TokenizationStage(PipelineStage):
+    """Tokenization stage for Cosmos3.
+
+    Applies the Qwen2 chat template, appends a duration suffix, and writes
+    ``text_ids`` / ``text_mask`` into ``batch.extra`` for the denoising stage.
+    """
+
+    parallelism_type = StageParallelismType.REPLICATED
+
+    def __init__(self, tokenizer):
+        super().__init__()
+        if tokenizer is None:
+            raise ValueError(
+                "Cosmos3TokenizationStage requires a tokenizer; expected the "
+                "Qwen2 tokenizer loaded from the checkpoint's text_tokenizer/ "
+                "subfolder."
+            )
+        self.tokenizer = tokenizer
+
+    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("prompt", batch.prompt, V.string_or_list_strings)
+        return result
+
+    def _tokenize_prompt(
+        self,
+        text: str,
+        max_sequence_length: int,
+        device: torch.device,
+        use_system_prompt: bool = False,
+        system_prompt: str | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Tokenize a prompt using Qwen2 chat template.
+
+        Returns (input_ids, attention_mask) as [1, S] tensors.
+        """
+        conversations = []
+        if use_system_prompt:
+            conversations.append(
+                {
+                    "role": "system",
+                    "content": system_prompt or COSMOS3_VIDEO_SYSTEM_PROMPT,
+                }
+            )
+        conversations.append({"role": "user", "content": text})
+
+        result = self.tokenizer.apply_chat_template(
+            conversations,
+            tokenize=True,
+            add_generation_prompt=True,
+        )
+
+        # Handle different return types from apply_chat_template
+        # Fast tokenizer returns BatchEncoding, slow tokenizer returns list[int]
+        if hasattr(result, "input_ids"):
+            # BatchEncoding from fast tokenizer
+            token_ids = list(result.input_ids)
+        elif isinstance(result, list):
+            # Already a list from slow tokenizer
+            token_ids = list(result)
+        else:
+            raise TypeError(
+                f"Unexpected return type from apply_chat_template: {type(result)}"
+            )
+
+        # Reserve room for the two special tokens (EOS + vision_start) so the
+        # final length cannot exceed ``max_sequence_length``.
+        token_ids = token_ids[: max_sequence_length - 2]
+
+        # Add EOS and vision_start tokens
+        token_ids.append(self.tokenizer.eos_token_id)
+        vision_start_id = self.tokenizer.convert_tokens_to_ids("<|vision_start|>")
+        if vision_start_id is not None:
+            token_ids.append(vision_start_id)
+
+        seq_len = len(token_ids)
+
+        # Pad to max_sequence_length
+        pad_len = max_sequence_length - seq_len
+        attention_mask = [1] * seq_len + [0] * pad_len
+        pad_token_id = self.tokenizer.pad_token_id or 0
+        token_ids = token_ids + [pad_token_id] * pad_len
+
+        input_ids = torch.tensor([token_ids], dtype=torch.long, device=device)
+        attention_mask = torch.tensor([attention_mask], dtype=torch.long, device=device)
+        return input_ids, attention_mask
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        """Tokenize prompt and negative prompt."""
+        device = get_local_torch_device()
+        prompt = batch.prompt
+        negative_prompt = batch.negative_prompt or COSMOS3_DEFAULT_NEGATIVE_PROMPT
+
+        # Get parameters
+        max_sequence_length = getattr(batch, "max_sequence_length", None) or 512
+        use_duration_template = getattr(batch, "use_duration_template", True)
+        use_system_prompt = getattr(batch, "use_system_prompt", False)
+        fps = batch.fps or 24.0
+        num_frames = batch.num_frames
+        is_image_gen = batch.data_type == DataType.IMAGE
+        system_prompt = (
+            COSMOS3_IMAGE_SYSTEM_PROMPT if is_image_gen else COSMOS3_VIDEO_SYSTEM_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
+            suffix = COSMOS3_DURATION_TEMPLATE.format(duration=duration, fps=fps)
+            prompt = f"{prompt} {suffix}"
+            self.log_info(f"Prompt with duration: '{prompt}'")
+
+        # Tokenize prompts
+        cond_ids, cond_mask = self._tokenize_prompt(
+            prompt, max_sequence_length, device, use_system_prompt, system_prompt
+        )
+        uncond_ids, uncond_mask = self._tokenize_prompt(
+            negative_prompt,
+            max_sequence_length,
+            device,
+            use_system_prompt,
+            system_prompt,
+        )
+
+        # Store in batch.extra for denoising stage
+        batch.extra["cond_text_ids"] = cond_ids
+        batch.extra["cond_text_mask"] = cond_mask
+        batch.extra["uncond_text_ids"] = uncond_ids
+        batch.extra["uncond_text_mask"] = uncond_mask
+        batch.extra["fps"] = fps
+
+        # Mark as processed (even though we don't use standard embeddings)
+        batch.is_prompt_processed = True
+
+        return batch
+
+
+class Cosmos3LatentPreparationStage(PipelineStage):
+    """Initialize the noisy latent for Cosmos3.
+
+    T2V / T2I produce pure Gaussian noise. I2V VAE-encodes the conditioning
+    image, replaces frame 0 of the latent with the encoded image, and stashes
+    a per-frame velocity mask plus the clean frame-0 latent for the denoiser
+    to re-inject after each scheduler step.
+    """
+
+    parallelism_type = StageParallelismType.REPLICATED
+
+    def __init__(self, vae, transformer):
+        super().__init__()
+        self.vae = vae
+        self.transformer = transformer
+
+    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("height", batch.height, V.positive_int_divisible(16))
+        result.add_check("width", batch.width, V.positive_int_divisible(16))
+        result.add_check("num_frames", batch.num_frames, V.positive_int)
+        return result
+
+    def _vae_encode(self, video: torch.Tensor) -> torch.Tensor:
+        """VAE-encode a [B, 3, T, H, W] pixel tensor and normalize the latent.
+
+        WanVAE returns a ``DiagonalGaussianDistribution``; ``mode()`` keeps
+        the encoding deterministic for I2V conditioning.
+        """
+        latent = self.vae.encode(video).mode()
+        mean = (
+            torch.as_tensor(self.vae.config.latents_mean)
+            .view(1, -1, 1, 1, 1)
+            .to(latent.device, latent.dtype)
+        )
+        std = (
+            torch.as_tensor(self.vae.config.latents_std)
+            .view(1, -1, 1, 1, 1)
+            .to(latent.device, latent.dtype)
+        )
+        return (latent - mean) / std
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        """Prepare initial latents (pure noise for T2V/T2I, image-conditioned for I2V)."""
+        device = get_local_torch_device()
+        dtype = torch.bfloat16
+
+        vae_scale_factor_temporal = getattr(self.vae.config, "scale_factor_temporal", 4)
+        vae_scale_factor_spatial = getattr(self.vae.config, "scale_factor_spatial", 16)
+
+        num_channels_latents = self.transformer.latent_channel
+        num_latent_frames = (batch.num_frames - 1) // vae_scale_factor_temporal + 1
+        height_latent = batch.height // vae_scale_factor_spatial
+        width_latent = batch.width // vae_scale_factor_spatial
+
+        shape = (
+            1,
+            num_channels_latents,
+            num_latent_frames,
+            height_latent,
+            width_latent,
+        )
+
+        generator = batch.generator
+        if generator is None and batch.seed is not None:
+            generator = torch.Generator(device=device).manual_seed(batch.seed)
+
+        noise = torch.randn(shape, generator=generator, device=device, dtype=dtype)
+
+        is_i2v = (
+            batch.preprocessed_image is not None and batch.data_type == DataType.VIDEO
+        )
+
+        if is_i2v:
+            vae_dtype = next(self.vae.parameters()).dtype
+            pixel_video = (
+                batch.preprocessed_image.unsqueeze(2)
+                .expand(-1, -1, batch.num_frames, -1, -1)
+                .contiguous()
+                .to(device=device, dtype=vae_dtype)
+            )
+            with torch.no_grad():
+                cond_latent = self._vae_encode(pixel_video).to(dtype)
+
+            condition_mask = torch.zeros(
+                1, 1, num_latent_frames, 1, 1, device=device, dtype=dtype
+            )
+            condition_mask[:, :, 0, :, :] = 1.0
+            latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise
+            batch.image_latent = cond_latent[:, :, 0:1, :, :].clone()
+            batch.extra["velocity_mask"] = 1.0 - condition_mask
+            self.log_info("Prepared I2V latents with frame-0 conditioning")
+        else:
+            latents = noise
+
+        batch.latents = latents
+        batch.raw_latent_shape = shape
+
+        batch.extra["video_shape"] = (num_latent_frames, height_latent, width_latent)
+        batch.extra["vae_scale_factor_temporal"] = vae_scale_factor_temporal
+        batch.extra["vae_scale_factor_spatial"] = vae_scale_factor_spatial
+
+        self.log_info(f"Prepared latents with shape {shape}")
+        return batch
+
+
+class Cosmos3TimestepPreparationStage(PipelineStage):
+    """
+    Timestep preparation stage for Cosmos3.
+
+    Sets up the diffusion scheduler timesteps.
+    """
+
+    parallelism_type = StageParallelismType.REPLICATED
+
+    def __init__(self, scheduler):
+        super().__init__()
+        self.scheduler = scheduler
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        """Prepare scheduler timesteps."""
+        device = get_local_torch_device()
+        num_inference_steps = batch.num_inference_steps
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        batch.timesteps = self.scheduler.timesteps
+
+        self.log_info(f"Prepared {len(batch.timesteps)} timesteps")
+        return batch
+
+
+class Cosmos3DenoisingStage(PipelineStage):
+    """Cosmos3 denoise loop, including CFG and the parallelism modes.
+
+    The UND pathway runs once and its K/V is cached per cache_key (``cond`` /
+    ``uncond`` / ``cfg_batched``); the GEN pathway runs every step.
+
+    Parallelism modes (combine freely):
+    - **CFG-parallel** — splits the conditional and unconditional branches
+      across CFG ranks. Each rank runs one branch, then a single all-reduce
+      combines them via ``g·cond + (1−g)·uncond``. Default 2-GPU recipe.
+    - **Ulysses (sequence parallel)** — shards the visual sequence across an
+      SP group. The cross-attention all-gathers visual K/V inside the
+      kernel; after the last GEN layer we all-gather hidden_gen back to
+      full length.
+    - **CFG + Ulysses** — when both are on, the SP group only contains ranks
+      that share a CFG context, so each context shards independently.
+    """
+
+    parallelism_type = StageParallelismType.REPLICATED
+
+    def __init__(self, transformer, scheduler, server_args: ServerArgs | None = None):
+        super().__init__()
+        self.transformer = transformer
+        self.scheduler = scheduler
+        self.server_args = server_args
+        self._logged_parallel_config = False
+
+        # Apply torch.compile if enabled
+        if server_args is not None:
+            self._maybe_enable_torch_compile(transformer, server_args)
+
+    def _maybe_enable_torch_compile(
+        self, transformer: nn.Module, server_args: ServerArgs
+    ) -> None:
+        """Regional ``torch.compile`` over the GEN decoder blocks.
+
+        Only ``gen_layers`` are compiled — they are the per-step hot path and
+        all share the same module class, so a single compilation amortizes
+        across them. The UND ``language_model`` runs once per prompt and is
+        cached, so compiling it would only pay warmup cost.
+
+        Caveat for Ulysses (``sp_size > 1``): the cross-attention's all-to-all
+        calls into ``torch.distributed.all_to_all_single`` through a Python
+        wrapper that fetches the process group at call time, which graph-breaks
+        Dynamo. Compile still works but loses some speedup on that path. The
+        headline 2-GPU CFG-parallel recipe (``sp_size == 1``) skips the SP
+        branch entirely and compiles cleanly.
+        """
+        if not server_args.enable_torch_compile or not isinstance(
+            transformer, nn.Module
+        ):
+            return
+
+        if current_platform.is_npu():
+            compile_kwargs: dict[str, Any] = {
+                "backend": get_compiler_backend(),
+                "fullgraph": False,
+                "dynamic": False,
+            }
+        else:
+            try:
+                import torch._inductor.config as _inductor_cfg
+
+                _inductor_cfg.reorder_for_compute_comm_overlap = True
+            except ImportError:
+                pass
+            # Lift Dynamo's per-callable cache cap above the default (64).
+            # Each gen_layer is its own compiled object, and several shape
+            # specializations (cond/uncond, with/without residual carry,
+            # SP on/off) can accumulate. 128 leaves headroom without
+            # encouraging unbounded specialization.
+            torch._dynamo.config.cache_size_limit = max(
+                getattr(torch._dynamo.config, "cache_size_limit", 64), 128
+            )
+            compile_kwargs = {
+                "mode": "default",
+                "fullgraph": False,
+                "dynamic": True,
+            }
+
+        gen_layers = getattr(transformer, "gen_layers", None)
+        if gen_layers is not None and isinstance(gen_layers, nn.ModuleList):
+            logger.info(
+                "Compiling %d Cosmos3 gen_layers with %s",
+                len(gen_layers),
+                compile_kwargs,
+            )
+            for i, layer in enumerate(gen_layers):
+                gen_layers[i] = torch.compile(layer, **compile_kwargs)
+        else:
+            logger.warning("Cosmos3 gen_layers not found, skipping torch.compile")
+
+    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("latents", batch.latents, V.is_tensor)
+        result.add_check("timesteps", batch.timesteps, V.is_tensor)
+        return result
+
+    def _run_transformer(
+        self,
+        latents: torch.Tensor,
+        timestep: torch.Tensor,
+        text_ids: torch.Tensor,
+        text_mask: torch.Tensor,
+        video_shape: tuple[int, int, int],
+        fps: float,
+        cache_key: str = "default",
+        noisy_frame_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Run transformer forward pass.
+
+        Args:
+            latents: Noisy latent tensor
+            timestep: Current diffusion timestep
+            text_ids: Tokenized text input
+            text_mask: Attention mask for text
+            video_shape: (T, H, W) in latent space
+            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.
+        """
+        with set_forward_context(
+            current_timestep=int(timestep.flatten()[0].item()),
+            attn_metadata=None,
+        ):
+            return self.transformer(
+                hidden_states=latents,
+                encoder_hidden_states=None,  # Not used by Cosmos3
+                timestep=timestep,
+                text_ids=text_ids,
+                text_mask=text_mask,
+                fps=fps,
+                cache_key=cache_key,
+                noisy_frame_mask=noisy_frame_mask,
+            )
+
+    def _manage_device_placement(self, server_args: ServerArgs):
+        """Move transformer to GPU if CPU offload is enabled."""
+        if not server_args.dit_cpu_offload:
+            return
+
+        # FSDP manages offloading internally
+        if server_args.use_fsdp_inference:
+            return
+
+        device = get_local_torch_device()
+        # Load the model to GPU if it's on CPU
+        if next(self.transformer.parameters()).device.type == "cpu":
+            self.log_info("Moving transformer to GPU for inference")
+            self.transformer.to(device)
+
+    @staticmethod
+    def _cfg_active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool:
+        """Return True iff CFG should be applied at timestep ``t``.
+
+        T2I uses a CFG window (e.g. ``[400, 1000]``) to skip guidance at low
+        noise levels, where it is empirically harmful. T2V/I2V leave this
+        unset and CFG is always on.
+        """
+        if interval is None:
+            return True
+        t_scalar = float(t.item()) if torch.is_tensor(t) else float(t)
+        lo, hi = interval
+        return lo <= t_scalar <= hi
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        """Run the denoising loop with CFG and optional I2V conditioning."""
+        self._manage_device_placement(server_args)
+
+        latents = batch.latents
+        timesteps = batch.timesteps
+        guidance_scale = batch.guidance_scale
+
+        cond_text_ids = batch.extra["cond_text_ids"]
+        cond_text_mask = batch.extra["cond_text_mask"]
+        uncond_text_ids = batch.extra["uncond_text_ids"]
+        uncond_text_mask = batch.extra["uncond_text_mask"]
+        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
+        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
+        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
+        )
+
+        sp_size = get_sp_world_size()
+        sp_rank = get_sp_parallel_rank() if sp_size > 1 else 0
+        ulysses_enabled = sp_size > 1
+
+        if not self._logged_parallel_config:
+            self._logged_parallel_config = True
+            if enable_cfg_parallel and ulysses_enabled:
+                self.log_info(
+                    f"CFG + Ulysses enabled: cfg_size={cfg_world_size}, cfg_rank={cfg_rank}, "
+                    f"sp_size={sp_size}, sp_rank={sp_rank}"
+                )
+            elif enable_cfg_parallel:
+                self.log_info(
+                    f"CFG parallel enabled: cfg_size={cfg_world_size}, cfg_rank={cfg_rank}"
+                )
+            elif ulysses_enabled:
+                self.log_info(f"Ulysses enabled: sp_size={sp_size}, sp_rank={sp_rank}")
+
+        # Drop any cached UND K/V from a previous request — its text differs.
+        self.transformer.reset_cache()
+
+        self.log_info(
+            f"Starting denoising with {len(timesteps)} steps, CFG={do_cfg}, "
+            f"CFG_parallel={enable_cfg_parallel}, cfg_rank={cfg_rank}"
+        )
+
+        progress_bar = tqdm(
+            enumerate(timesteps),
+            total=len(timesteps),
+            desc="Denoising",
+            disable=batch.is_warmup,
+        )
+
+        for i, t in progress_bar:
+            timestep = t.unsqueeze(0) if t.dim() == 0 else t
+            # Outside the CFG window the effective scale collapses to 1.0,
+            # which reduces CFG to the cond branch (cfg-parallel safe).
+            effective_scale = (
+                guidance_scale if self._cfg_active_at(t, guidance_interval) else 1.0
+            )
+
+            if do_cfg:
+                if enable_cfg_parallel:
+                    noise_pred = self._predict_noise_cfg_parallel(
+                        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=effective_scale,
+                        cfg_rank=cfg_rank,
+                        noisy_frame_mask=velocity_mask,
+                    )
+                elif effective_scale == 1.0:
+                    noise_pred = 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",
+                        noisy_frame_mask=velocity_mask,
+                    )
+                else:
+                    noise_pred = 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=effective_scale,
+                        noisy_frame_mask=velocity_mask,
+                    )
+            else:
+                noise_pred = 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",
+                    noisy_frame_mask=velocity_mask,
+                )
+
+            # I2V: zero-velocity at conditioned frames so the scheduler keeps
+            # them clean; UniPC's predictor-corrector still rescales the
+            # sample, so we re-inject the clean image latent below.
+            if velocity_mask is not None:
+                noise_pred = noise_pred * velocity_mask
+
+            latents = self.scheduler.step(
+                noise_pred,
+                t,
+                latents,
+                return_dict=False,
+            )[0]
+
+            if image_latent is not None:
+                latents[:, :, 0:1, :, :] = image_latent
+
+        batch.latents = latents
+        self.log_info("Denoising complete")
+        return batch
+
+    def _predict_noise_cfg_batched(
+        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,
+        noisy_frame_mask: torch.Tensor | None = None,
+    ) -> 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.
+        Order is ``[uncond, cond]`` so the chunk-and-combine math below
+        matches the standard CFG formula.
+        """
+        latents_batched = torch.cat([latents, latents], dim=0)
+        text_ids_batched = torch.cat([uncond_text_ids, cond_text_ids], dim=0)
+        text_mask_batched = torch.cat([uncond_text_mask, cond_text_mask], dim=0)
+        timestep_batched = timestep.expand(2)
+        mask_batched = (
+            torch.cat([noisy_frame_mask, noisy_frame_mask], dim=0)
+            if noisy_frame_mask is not None
+            else None
+        )
+
+        noise_pred = self._run_transformer(
+            latents=latents_batched,
+            timestep=timestep_batched,
+            text_ids=text_ids_batched,
+            text_mask=text_mask_batched,
+            video_shape=video_shape,
+            fps=fps,
+            cache_key="cfg_batched",
+            noisy_frame_mask=mask_batched,
+        )
+
+        noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2, dim=0)
+        # CFG: uncond + g·(cond − uncond).
+        return noise_pred_uncond + guidance_scale * (
+            noise_pred_cond - noise_pred_uncond
+        )
+
+    def _predict_noise_cfg_parallel(
+        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,
+        cfg_rank: int,
+        noisy_frame_mask: torch.Tensor | None = None,
+    ) -> 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"``).
+        """
+        if cfg_rank == 0:
+            noise_pred = 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",
+                noisy_frame_mask=noisy_frame_mask,
+            )
+            partial = guidance_scale * noise_pred
+        else:
+            noise_pred = 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",
+                noisy_frame_mask=noisy_frame_mask,
+            )
+            partial = (1.0 - guidance_scale) * noise_pred
+
+        return cfg_model_parallel_all_reduce(partial)
+
+
+class Cosmos3DecodingStage(PipelineStage):
+    """
+    VAE decoding stage for Cosmos3.
+
+    Decodes latents to pixel space using the VAE.
+    Returns OutputBatch instead of Req to signal pipeline completion.
+    """
+
+    parallelism_type = StageParallelismType.REPLICATED
+
+    def __init__(self, vae, guardrails: bool = False):
+        super().__init__()
+        self.vae = vae
+        self._latents_mean = None
+        self._latents_std = None
+        self._guardrails = guardrails
+        # Use VideoProcessor for postprocessing (same as other video pipelines)
+        from diffusers.video_processor import VideoProcessor
+
+        vae_scale_factor = getattr(vae.config, "scale_factor_spatial", 16)
+        self.video_processor = VideoProcessor(vae_scale_factor=vae_scale_factor)
+        if guardrails:
+            from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_guardrails import (
+                _init_guardrails,
+            )
+
+            _init_guardrails()
+
+    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+        result = VerificationResult()
+        result.add_check("latents", batch.latents, V.is_tensor)
+        return result
+
+    def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        """Decode latents to video frames. Returns tensor in [B, C, T, H, W] format."""
+        device = latents.device
+        # Get VAE dtype from its parameters
+        vae_dtype = next(self.vae.parameters()).dtype
+        latents = latents.to(vae_dtype)
+
+        # Apply latent normalization if configured
+        if hasattr(self.vae.config, "latents_mean") and hasattr(
+            self.vae.config, "latents_std"
+        ):
+            if self._latents_mean is None:
+                self._latents_mean = (
+                    torch.tensor(self.vae.config.latents_mean)
+                    .view(1, -1, 1, 1, 1)
+                    .to(device, vae_dtype)
+                )
+                self._latents_std = (
+                    torch.tensor(self.vae.config.latents_std)
+                    .view(1, -1, 1, 1, 1)
+                    .to(device, vae_dtype)
+                )
+            latents = (latents * self._latents_std) + self._latents_mean
+        else:
+            scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
+            latents = latents / scaling_factor
+
+        # Decode - returns [B, C, T, H, W]
+        video = self.vae.decode(latents)
+        # Handle both dict return and direct tensor return
+        if hasattr(video, "sample"):
+            video = video.sample
+        elif isinstance(video, tuple):
+            video = video[0]
+
+        return video
+
+    def forward(self, batch: Req, server_args: ServerArgs):
+        """Decode latents to video, or to a single image for T2I."""
+        from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
+            OutputBatch,
+        )
+
+        is_image_gen = batch.data_type == DataType.IMAGE
+        self.log_info(
+            "Decoding latents to image..."
+            if is_image_gen
+            else "Decoding latents to video..."
+        )
+
+        device = batch.latents.device
+        if server_args.vae_cpu_offload:
+            self.vae.to(device)
+
+        with torch.no_grad():
+            decoded = self._decode_latents(batch.latents)
+
+        if server_args.vae_cpu_offload and not getattr(batch, "is_warmup", False):
+            self.vae.to("cpu", non_blocking=True)
+
+        self.log_info(f"Decoded tensor shape: {decoded.shape}")
+
+        if is_image_gen:
+            output = self.video_processor.postprocess(
+                decoded.squeeze(2), output_type="np"
+            )
+        else:
+            output = self.video_processor.postprocess_video(decoded, output_type="np")
+            self.log_info(f"Postprocessed video shape: {output.shape}")
+
+        if self._guardrails:
+            from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_guardrails import (
+                check_video_safety,
+            )
+
+            if is_image_gen:
+                # check_video_safety expects [B, T, H, W, C]; wrap then unwrap.
+                output = check_video_safety(output[:, np.newaxis, ...])[:, 0, ...]
+            else:
+                output = check_video_safety(output)
+
+        return OutputBatch(
+            output=output,
+            metrics=batch.metrics if hasattr(batch, "metrics") else None,
+        )

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

@@ -0,0 +1,461 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 guardrail stages.
+
+Text:
+  1. Blocklist — ``better_profanity`` + nvidia/Cosmos-Guardrail1 word lists.
+  2. Qwen3Guard — 0.6B LLM classifier (Qwen/Qwen3Guard-Gen-0.6B).
+
+Video:
+  1. SigLIP content-safety filter — 7-class frame classifier; blocks if
+     more than 10% of frames are unsafe.
+  2. RetinaFace face blur — detects faces and pixelates them.
+
+Enabled by default; opt out with ``SGLANG_DISABLE_COSMOS3_GUARDRAILS=1``.
+"""
+
+from __future__ import annotations
+
+import os
+import warnings
+from typing import Callable
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+
+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.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+GUARDRAIL_HF_REPO = "nvidia/Cosmos-Guardrail1"
+GUARDRAIL_HF_REVISION = "d6d4bfa899a71454a700907664f3e88f503950cf"
+CUTOFF_UNSAFE_FRAMES_PERCENT = 10
+
+TextGuardrailFn = Callable[[str], None]
+VideoGuardrailFn = Callable[[np.ndarray], np.ndarray]
+
+
+# ---------------------------------------------------------------------------
+# Video safety classifier (SigLIP so400m + 3-layer head)
+# ---------------------------------------------------------------------------
+class SafetyClassifier(nn.Module):
+    """3-layer classifier with BatchNorm (1152 -> 512 -> 256 -> 7)."""
+
+    def __init__(self, input_size: int = 1152, num_classes: int = 7):
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(input_size, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(),
+            nn.Linear(256, num_classes),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.layers(x)
+
+
+CLASS_IDX_TO_NAME = {
+    0: "Safe",
+    1: "Sexual_Content",
+    3: "Drugs",
+    4: "Child_Abuse",
+    5: "Hate_and_Harassment",
+    6: "Self-Harm",
+}
+
+
+# ---------------------------------------------------------------------------
+# Face pixelation utility
+# ---------------------------------------------------------------------------
+def _pixelate_face(face_img: np.ndarray, blocks: int = 5) -> np.ndarray:
+    h, w = face_img.shape[:2]
+    if h == 0 or w == 0:
+        return face_img
+    temp = cv2.resize(face_img, (blocks, blocks), interpolation=cv2.INTER_LINEAR)
+    return cv2.resize(temp, (w, h), interpolation=cv2.INTER_NEAREST)
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint download helper
+# ---------------------------------------------------------------------------
+def _download_checkpoint() -> str:
+    from huggingface_hub import snapshot_download
+
+    return snapshot_download(GUARDRAIL_HF_REPO, revision=GUARDRAIL_HF_REVISION)
+
+
+# ---------------------------------------------------------------------------
+# Text guardrail builder
+# ---------------------------------------------------------------------------
+def _build_text_guardrail(offload_to_cpu: bool) -> TextGuardrailFn:
+    checkers: list[Callable[[str], tuple[bool, str]]] = []
+
+    # 1. Blocklist
+    try:
+        import nltk
+        from better_profanity import profanity as profanity_filter
+
+        ckpt_dir = _download_checkpoint()
+        blocklist_dir = os.path.join(ckpt_dir, "blocklist")
+        nltk.data.path.append(os.path.join(blocklist_dir, "nltk_data"))
+
+        def _read_keywords(dirpath: str) -> list[str]:
+            words: list[str] = []
+            if not os.path.isdir(dirpath):
+                return words
+            for fname in sorted(os.listdir(dirpath)):
+                fpath = os.path.join(dirpath, fname)
+                if os.path.isfile(fpath):
+                    with open(fpath) as f:
+                        words.extend(line.strip() for line in f if line.strip())
+            return words
+
+        blocklist_words = _read_keywords(os.path.join(blocklist_dir, "custom"))
+        whitelist_words = _read_keywords(os.path.join(blocklist_dir, "whitelist"))
+        profanity_filter.load_censor_words(
+            custom_words=blocklist_words, whitelist_words=whitelist_words
+        )
+
+        def _blocklist_check(prompt: str) -> tuple[bool, str]:
+            if profanity_filter.contains_profanity(prompt):
+                return False, "Blocked by keyword filter"
+            return True, ""
+
+        checkers.append(_blocklist_check)
+        logger.info("Blocklist guardrail loaded (%d keywords)", len(blocklist_words))
+    except ImportError:
+        logger.warning(
+            "better-profanity or nltk not installed; skipping blocklist guardrail"
+        )
+
+    # 2. Qwen3Guard
+    try:
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        model_id = "Qwen/Qwen3Guard-Gen-0.6B"
+        qwen_tokenizer = AutoTokenizer.from_pretrained(model_id)
+        device = "cpu" if offload_to_cpu else "cuda"
+        qwen_model = (
+            AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+            .to(device)
+            .eval()
+        )
+
+        def _qwen_check(prompt: str) -> tuple[bool, str]:
+            conversations = [{"role": "user", "content": prompt}]
+            inputs = qwen_tokenizer.apply_chat_template(
+                conversations,
+                tokenize=True,
+                return_tensors="pt",
+                add_generation_prompt=True,
+                return_dict=True,
+            ).to(device)
+            input_len = inputs["input_ids"].shape[1]
+            with torch.no_grad():
+                output_ids = qwen_model.generate(**inputs, max_new_tokens=128)
+            response = qwen_tokenizer.decode(
+                output_ids[0][input_len:], skip_special_tokens=True
+            )
+            if "unsafe" in response.lower():
+                return False, f"Qwen3Guard: {response.strip()}"
+            return True, ""
+
+        checkers.append(_qwen_check)
+        logger.info("Qwen3Guard guardrail loaded")
+    except ImportError:
+        logger.warning("transformers not installed; skipping Qwen3Guard")
+
+    def text_guardrail(prompt: str) -> None:
+        for checker in checkers:
+            is_safe, msg = checker(prompt)
+            if not is_safe:
+                raise ValueError(f"Guardrail blocked prompt: {msg}")
+
+    return text_guardrail
+
+
+# ---------------------------------------------------------------------------
+# Video guardrail builder
+# ---------------------------------------------------------------------------
+def _build_video_guardrail(offload_to_cpu: bool) -> VideoGuardrailFn:
+    ckpt_dir = _download_checkpoint()
+    safety_checker: Callable[[np.ndarray], tuple[bool, str]] | None = None
+    face_blurrer: Callable[[np.ndarray], np.ndarray] | None = None
+
+    # 1. Video content safety filter: SigLIP so400m + SafetyClassifier
+    try:
+        from PIL import Image
+        from transformers import SiglipModel, SiglipProcessor
+
+        device = "cpu" if offload_to_cpu else "cuda"
+        siglip_id = "google/siglip-so400m-patch14-384"
+        siglip_model = (
+            SiglipModel.from_pretrained(siglip_id)
+            .to(device, dtype=torch.float32)
+            .eval()
+        )
+        siglip_processor = SiglipProcessor.from_pretrained(siglip_id)
+
+        classifier = SafetyClassifier(input_size=1152, num_classes=7)
+        ckpt_path = os.path.join(
+            ckpt_dir, "video_content_safety_filter", "safety_filter.pt"
+        )
+        checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+        # Checkpoint keys have "network." prefix from the VideoSafetyModel wrapper.
+        state = {k.removeprefix("network."): v for k, v in checkpoint["model"].items()}
+        classifier.load_state_dict(state)
+        classifier = classifier.to(device, dtype=torch.float32).eval()
+
+        def _safety_check(frames: np.ndarray) -> tuple[bool, str]:
+            nonlocal siglip_model, classifier
+            if offload_to_cpu:
+                siglip_model = siglip_model.to("cuda")
+                classifier = classifier.to("cuda")
+
+            unsafe_count = 0
+            total = len(frames)
+            for frame in frames:
+                if frame.dtype != np.uint8:
+                    frame = (np.clip(frame, 0.0, 1.0) * 255.0).astype(np.uint8)
+                img = Image.fromarray(frame)
+                inputs = siglip_processor(images=img, return_tensors="pt").to(
+                    "cuda", dtype=torch.float32
+                )
+                with torch.no_grad():
+                    features = siglip_model.get_image_features(**inputs)
+                    if hasattr(features, "pooler_output"):
+                        features = features.pooler_output
+                    features = features / features.norm(dim=-1, keepdim=True)
+                    logits = classifier(features)
+                    pred = logits.argmax(dim=-1).item()
+                class_name = CLASS_IDX_TO_NAME.get(pred, "Unknown")
+                if class_name != "Safe":
+                    unsafe_count += 1
+
+            if offload_to_cpu:
+                siglip_model = siglip_model.to("cpu")
+                classifier = classifier.to("cpu")
+
+            if unsafe_count / total > CUTOFF_UNSAFE_FRAMES_PERCENT / 100:
+                return (
+                    False,
+                    f"Video content safety: {unsafe_count}/{total} frames unsafe",
+                )
+            return True, ""
+
+        safety_checker = _safety_check
+        logger.info("Video content safety filter loaded (SigLIP so400m + classifier)")
+    except (ImportError, FileNotFoundError) as e:
+        logger.warning("Could not load video safety filter: %s", e)
+
+    # 2. Face blur: RetinaFace + pixelation
+    try:
+        from retinaface.data import cfg_re50
+        from retinaface.layers.functions.prior_box import PriorBox
+        from retinaface.models.retinaface import RetinaFace
+        from retinaface.utils.nms.py_cpu_nms import py_cpu_nms
+
+        face_ckpt = os.path.join(ckpt_dir, "face_blur_filter", "Resnet50_Final.pth")
+        if not os.path.exists(face_ckpt):
+            raise FileNotFoundError(face_ckpt)
+
+        cfg = dict(cfg_re50)
+        cfg["pretrain"] = False
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            retinaface_net = RetinaFace(cfg=cfg, phase="test")
+
+        pretrained_dict = torch.load(face_ckpt, map_location="cpu", weights_only=True)
+        if "state_dict" in pretrained_dict:
+            pretrained_dict = pretrained_dict["state_dict"]
+        pretrained_dict = {
+            k.replace("module.", "", 1) if k.startswith("module.") else k: v
+            for k, v in pretrained_dict.items()
+        }
+        retinaface_net.load_state_dict(pretrained_dict, strict=False)
+        retinaface_device = "cpu" if offload_to_cpu else "cuda"
+        retinaface_net = retinaface_net.to(
+            retinaface_device, dtype=torch.float32
+        ).eval()
+
+        CONF_THRESH = 0.7
+        NMS_THRESH = 0.4
+        TOP_K = 5000
+        KEEP_TOP_K = 750
+
+        def _decode_batch(loc, priors, variances):
+            batch_size = loc.size(0)
+            p = priors.unsqueeze(0).expand(batch_size, -1, -1)
+            boxes = torch.cat(
+                (
+                    p[:, :, :2] + loc[:, :, :2] * variances[0] * p[:, :, 2:],
+                    p[:, :, 2:] * torch.exp(loc[:, :, 2:] * variances[1]),
+                ),
+                dim=2,
+            )
+            boxes[:, :, :2] -= boxes[:, :, 2:] / 2
+            boxes[:, :, 2:] += boxes[:, :, :2]
+            return boxes
+
+        def _face_blur(frames: np.ndarray) -> np.ndarray:
+            nonlocal retinaface_net
+            if offload_to_cpu:
+                retinaface_net = retinaface_net.to("cuda")
+
+            prior_data = None
+            scale = None
+            result_frames = []
+
+            for frame in frames:
+                frame_t = torch.from_numpy(frame).to("cuda", dtype=torch.float32)
+                if frame.dtype != np.uint8:
+                    frame_t = frame_t * 255.0
+                frame_t = frame_t.permute(2, 0, 1).unsqueeze(0)  # [1, C, H, W]
+                frame_t = frame_t[:, [2, 1, 0], :, :]  # RGB -> BGR
+                means = torch.tensor(
+                    [104.0, 117.0, 123.0], device="cuda", dtype=torch.float32
+                ).view(1, 3, 1, 1)
+                frame_t = frame_t - means
+
+                h, w = frame_t.shape[2], frame_t.shape[3]
+                if prior_data is None:
+                    priorbox = PriorBox(cfg, image_size=(h, w))
+                    prior_data = priorbox.forward().to("cuda", dtype=torch.float32)
+                if scale is None:
+                    scale = torch.tensor(
+                        [w, h, w, h], device="cuda", dtype=torch.float32
+                    )
+
+                with torch.no_grad():
+                    loc, conf, _ = retinaface_net(frame_t)
+
+                boxes = _decode_batch(loc, prior_data, cfg["variance"])
+                boxes = (boxes * scale).squeeze(0).cpu().numpy()
+                scores = conf.squeeze(0)[:, 1].cpu().numpy()
+
+                # Filter by confidence
+                inds = np.where(scores > CONF_THRESH)[0]
+                boxes_f = boxes[inds]
+                scores_f = scores[inds]
+                order = scores_f.argsort()[::-1][:TOP_K]
+                boxes_f = boxes_f[order]
+                scores_f = scores_f[order]
+
+                # NMS
+                dets = np.hstack((boxes_f, scores_f[:, np.newaxis])).astype(np.float32)
+                keep = py_cpu_nms(dets, NMS_THRESH)
+                dets = dets[keep][:KEEP_TOP_K]
+
+                out_frame = frame.copy()
+                for det in dets:
+                    x1, y1, x2, y2 = map(int, det[:4])
+                    if x2 - x1 < 20 or y2 - y1 < 20:
+                        continue
+                    max_h, max_w = out_frame.shape[:2]
+                    y1c, y2c = max(y1, 0), min(y2, max_h)
+                    x1c, x2c = max(x1, 0), min(x2, max_w)
+                    out_frame[y1c:y2c, x1c:x2c] = _pixelate_face(
+                        out_frame[y1c:y2c, x1c:x2c]
+                    )
+
+                result_frames.append(out_frame)
+
+            if offload_to_cpu:
+                retinaface_net = retinaface_net.to("cpu")
+
+            return np.array(result_frames)
+
+        face_blurrer = _face_blur
+        logger.info("Face blur filter loaded (RetinaFace Resnet50)")
+    except (ImportError, FileNotFoundError) as e:
+        logger.warning("Could not load face blur filter: %s", e)
+
+    def video_guardrail(frames: np.ndarray) -> np.ndarray:
+        if safety_checker is not None:
+            is_safe, msg = safety_checker(frames)
+            if not is_safe:
+                raise ValueError(f"Guardrail blocked video: {msg}")
+        if face_blurrer is not None:
+            frames = face_blurrer(frames)
+        return frames
+
+    return video_guardrail
+
+
+# ---------------------------------------------------------------------------
+# Singleton initialization
+# ---------------------------------------------------------------------------
+_text_guardrail: TextGuardrailFn | None = None
+_video_guardrail: VideoGuardrailFn | None = None
+_initialized = False
+
+
+def _init_guardrails(offload_to_cpu: bool = False) -> None:
+    global _text_guardrail, _video_guardrail, _initialized
+    if _initialized:
+        return
+    logger.info(
+        "Initializing Cosmos3 guardrails (offload_to_cpu=%s)...", offload_to_cpu
+    )
+    _text_guardrail = _build_text_guardrail(offload_to_cpu)
+    _video_guardrail = _build_video_guardrail(offload_to_cpu)
+    _initialized = True
+    logger.info("Cosmos3 guardrails initialized.")
+
+
+# ---------------------------------------------------------------------------
+# Public API — video guardrail function for use inside Cosmos3DecodingStage
+# ---------------------------------------------------------------------------
+def check_video_safety(video: np.ndarray) -> np.ndarray:
+    """Run video guardrails on decoded frames (numpy [B, T, H, W, C] or [T, H, W, C]).
+
+    Raises ``ValueError`` if content is blocked.
+    Returns (potentially face-blurred) frames.
+    """
+    if _video_guardrail is None:
+        return video
+    frames = video[0] if video.ndim == 5 else video
+    frames = _video_guardrail(frames)
+    if video.ndim == 5:
+        frames = frames[np.newaxis]
+    return frames
+
+
+# ---------------------------------------------------------------------------
+# Pipeline stage — text guardrail (runs before generation)
+# ---------------------------------------------------------------------------
+class Cosmos3TextGuardrailStage(PipelineStage):
+    """Check prompt text against safety policies before generation.
+
+    Runs blocklist keyword matching and Qwen3Guard LLM classifier.
+    Raises ``ValueError`` if the prompt is blocked.
+    """
+
+    parallelism_type = StageParallelismType.MAIN_RANK_ONLY
+
+    def __init__(self, offload_to_cpu: bool = False):
+        super().__init__()
+        _init_guardrails(offload_to_cpu)
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        if _text_guardrail is not None and batch.prompt is not None:
+            prompt = batch.prompt
+            if isinstance(prompt, list):
+                for p in prompt:
+                    _text_guardrail(p)
+            else:
+                _text_guardrail(prompt)
+        return batch

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

@@ -339,15 +339,24 @@ def prepare_diffusers_component_path_for_loading(component_path: str) -> str:
             return local_component_path
 
         config["quantization_config"] = normalized_quant_config
-        with open(config_path, "w", encoding="utf-8") as f:
-            json.dump(config, f, indent=2, sort_keys=True)
-            f.write("\n")
-        logger.warning(
-            "Patched legacy flat ModelOpt quantization_config at %s with quant_type=%s "
-            "for diffusers compatibility.",
-            config_path,
-            normalized_quant_config.get("quant_type"),
-        )
+        try:
+            with open(config_path, "w", encoding="utf-8") as f:
+                json.dump(config, f, indent=2, sort_keys=True)
+                f.write("\n")
+        except OSError as exc:
+            logger.warning(
+                "Could not persist normalized ModelOpt config at %s (%s); "
+                "normalization will be applied in memory at load time.",
+                config_path,
+                exc,
+            )
+        else:
+            logger.warning(
+                "Patched legacy flat ModelOpt quantization_config at %s with quant_type=%s "
+                "for diffusers compatibility.",
+                config_path,
+                normalized_quant_config.get("quant_type"),
+            )
 
     return local_component_path
 
@@ -374,6 +383,12 @@ def get_diffusers_component_config(
         lambda acc, path: acc | load_dict(path), config_file_paths, {}
     )
 
+    quant_config = combined_config.get("quantization_config")
+    if quant_config is not None:
+        combined_config["quantization_config"] = normalize_flat_modelopt_quant_config(
+            quant_config
+        )
+
     _clean_hf_config_inplace(combined_config)
 
     logger.debug("HF model config: %s", combined_config)

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

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for Cosmos3 config, weight mapping, and sampling params."""
+
+import unittest
+
+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.sampling_params import DataType
+from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+
+
+def _apply(mapping_fn, key):
+    """Return (target_key, merge_index, total_splits) for a diffusers weight key."""
+    return mapping_fn(key)
+
+
+class TestCosmos3ParamNamesMapping(unittest.TestCase):
+    """Verify diffusers → sglang weight key translations."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.fn = staticmethod(
+            get_param_names_mapping(_build_cosmos3_param_names_mapping())
+        )
+
+    # --- skipped / dropped weights ---
+
+    def test_lm_head_dropped(self):
+        key, idx, total = _apply(self.fn, "lm_head.weight")
+        self.assertEqual(key, "")
+
+    def test_model_norm_dropped(self):
+        key, idx, total = _apply(self.fn, "model.norm.weight")
+        self.assertEqual(key, "")
+
+    # --- top-level pass-through ---
+
+    def test_embed_tokens(self):
+        key, *_ = _apply(self.fn, "model.embed_tokens.weight")
+        self.assertEqual(key, "language_model.embed_tokens.weight")
+
+    def test_norm_moe_gen(self):
+        key, *_ = _apply(self.fn, "model.norm_moe_gen.weight")
+        self.assertEqual(key, "norm_moe_gen.weight")
+
+    # --- time embedder ---
+
+    def test_time_embedder_mlp_0(self):
+        key, *_ = _apply(self.fn, "time_embedder.mlp.0.weight")
+        self.assertEqual(key, "time_embedder.linear_1.weight")
+
+    def test_time_embedder_mlp_2(self):
+        key, *_ = _apply(self.fn, "time_embedder.mlp.2.bias")
+        self.assertEqual(key, "time_embedder.linear_2.bias")
+
+    # --- GEN pathway: Q/K/V merge (must not be claimed by UND catch-all) ---
+
+    def test_gen_q_proj_key_and_merge_index(self):
+        key, idx, total = _apply(
+            self.fn, "model.layers.3.self_attn.q_proj_moe_gen.weight"
+        )
+        self.assertEqual(key, "gen_layers.3.cross_attention.to_qkv.weight")
+        self.assertEqual(idx, 0)
+        self.assertEqual(total, 3)
+
+    def test_gen_k_proj_merge_index(self):
+        _, idx, total = _apply(
+            self.fn, "model.layers.0.self_attn.k_proj_moe_gen.weight"
+        )
+        self.assertEqual(idx, 1)
+        self.assertEqual(total, 3)
+
+    def test_gen_v_proj_merge_index(self):
+        _, idx, total = _apply(
+            self.fn, "model.layers.0.self_attn.v_proj_moe_gen.weight"
+        )
+        self.assertEqual(idx, 2)
+        self.assertEqual(total, 3)
+
+    def test_gen_o_proj(self):
+        key, idx, total = _apply(
+            self.fn, "model.layers.5.self_attn.o_proj_moe_gen.weight"
+        )
+        self.assertEqual(key, "gen_layers.5.cross_attention.to_out.weight")
+        self.assertIsNone(idx)
+
+    def test_gen_mlp_gate_proj(self):
+        key, idx, total = _apply(self.fn, "model.layers.2.mlp_moe_gen.gate_proj.weight")
+        self.assertEqual(key, "gen_layers.2.mlp.gate_up_proj.weight")
+        self.assertEqual(idx, 0)
+        self.assertEqual(total, 2)
+
+    def test_gen_mlp_up_proj(self):
+        key, idx, total = _apply(self.fn, "model.layers.2.mlp_moe_gen.up_proj.weight")
+        self.assertEqual(key, "gen_layers.2.mlp.gate_up_proj.weight")
+        self.assertEqual(idx, 1)
+        self.assertEqual(total, 2)
+
+    def test_gen_mlp_down_proj_passthrough(self):
+        key, idx, _ = _apply(self.fn, "model.layers.2.mlp_moe_gen.down_proj.weight")
+        self.assertEqual(key, "gen_layers.2.mlp.down_proj.weight")
+        self.assertIsNone(idx)
+
+    # --- UND pathway: Q/K/V merge ---
+
+    def test_und_q_proj_key_and_merge_index(self):
+        key, idx, total = _apply(self.fn, "model.layers.7.self_attn.q_proj.weight")
+        self.assertEqual(key, "language_model.layers.7.self_attn.to_qkv.weight")
+        self.assertEqual(idx, 0)
+        self.assertEqual(total, 3)
+
+    def test_und_k_proj_merge_index(self):
+        _, idx, total = _apply(self.fn, "model.layers.0.self_attn.k_proj.weight")
+        self.assertEqual(idx, 1)
+        self.assertEqual(total, 3)
+
+    def test_und_v_proj_merge_index(self):
+        _, idx, total = _apply(self.fn, "model.layers.0.self_attn.v_proj.weight")
+        self.assertEqual(idx, 2)
+        self.assertEqual(total, 3)
+
+    def test_und_mlp_gate_proj(self):
+        key, idx, total = _apply(self.fn, "model.layers.1.mlp.gate_proj.weight")
+        self.assertEqual(key, "language_model.layers.1.mlp.gate_up_proj.weight")
+        self.assertEqual(idx, 0)
+        self.assertEqual(total, 2)
+
+    def test_und_mlp_up_proj(self):
+        _, idx, total = _apply(self.fn, "model.layers.1.mlp.up_proj.weight")
+        self.assertEqual(idx, 1)
+        self.assertEqual(total, 2)
+
+    def test_und_layernorm_catch_all(self):
+        key, idx, _ = _apply(self.fn, "model.layers.0.input_layernorm.weight")
+        self.assertEqual(key, "language_model.layers.0.input_layernorm.weight")
+        self.assertIsNone(idx)
+
+    # --- ordering: GEN patterns must not be swallowed by UND catch-all ---
+
+    def test_gen_layernorm_not_mapped_to_und(self):
+        key, *_ = _apply(self.fn, "model.layers.0.input_layernorm_moe_gen.weight")
+        self.assertIn("gen_layers", key)
+        self.assertNotIn("language_model", key)
+
+    def test_gen_post_attention_layernorm_not_mapped_to_und(self):
+        key, *_ = _apply(
+            self.fn, "model.layers.4.post_attention_layernorm_moe_gen.weight"
+        )
+        self.assertIn("gen_layers", key)
+        self.assertNotIn("language_model", key)
+
+
+class TestCosmos3AdjustNumFrames(unittest.TestCase):
+    """Verify VAE-aligned frame rounding in Cosmos3Config."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.cfg = Cosmos3Config()
+
+    def test_single_frame_t2i_bypass(self):
+        self.assertEqual(self.cfg.adjust_num_frames(1), 1)
+
+    def test_already_aligned(self):
+        # (81 - 1) = 80, 80 % 4 == 0
+        self.assertEqual(self.cfg.adjust_num_frames(81), 81)
+
+    def test_rounds_down_to_nearest_aligned(self):
+        # (83 - 1) = 82 → floor(82/4)*4 + 1 = 81
+        self.assertEqual(self.cfg.adjust_num_frames(83), 81)
+        # (6 - 1) = 5 → floor(5/4)*4 + 1 = 5
+        self.assertEqual(self.cfg.adjust_num_frames(6), 5)
+
+    def test_minimum_video_frame_count(self):
+        # 2 frames: (2-1)=1, 1//4=0, 0*4+1=1 → rounds to 1, but 1 is T2I — still valid
+        self.assertEqual(self.cfg.adjust_num_frames(2), 1)
+
+
+class TestCosmos3SamplingParamsDataType(unittest.TestCase):
+    """Verify num_frames==1 flips data_type to IMAGE before file name derivation."""
+
+    def test_single_frame_sets_image_data_type(self):
+        params = Cosmos3SamplingParams(prompt="test", num_frames=1)
+        params._set_output_file_name()
+        self.assertEqual(params.data_type, DataType.IMAGE)
+        self.assertTrue(
+            params.output_file_name.endswith((".png", ".jpg", ".jpeg", ".webp")),
+            f"Expected image extension, got: {params.output_file_name}",
+        )
+
+    def test_multi_frame_keeps_video_data_type(self):
+        params = Cosmos3SamplingParams(prompt="test", num_frames=81)
+        params._set_output_file_name()
+        self.assertEqual(params.data_type, DataType.VIDEO)
+
+    def test_default_num_frames_is_video(self):
+        params = Cosmos3SamplingParams(prompt="test")
+        params._set_output_file_name()
+        self.assertEqual(params.data_type, DataType.VIDEO)
+
+
+if __name__ == "__main__":
+    unittest.main()