diff --git a/docs/diffusion/compatibility_matrix.md b/docs/diffusion/compatibility_matrix.md
index cd7f2d5b5..7f760575a 100644
--- a/docs/diffusion/compatibility_matrix.md
+++ b/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
1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| LTX-2.3 (one/two-stage/TI2V/HQ) | `Lightricks/LTX-2.3` | 768×512
1536×1024
1920×1088 (HQ default) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| Cosmos3-Nano (T2V / I2V / T2I) | `nvidia/Cosmos3-Nano` | 720p · 480p
1024×1024 (T2I) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| Cosmos3-Super (T2V / I2V / T2I) | `nvidia/Cosmos3-Super` | 720p · 480p
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
diff --git a/docs/diffusion/index.md b/docs/diffusion/index.md
index e0790d9e7..91ec58d39 100644
--- a/docs/diffusion/index.md
+++ b/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
diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py
index 63ac893dd..c3ee938b6 100644
--- a/python/sglang/multimodal_gen/configs/models/dits/__init__.py
+++ b/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",
diff --git a/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py
new file mode 100644
index 000000000..6efb44987
--- /dev/null
+++ b/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"
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
index 005bc862c..831f6bfcf 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
+++ b/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",
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py
new file mode 100644
index 000000000..160bad35d
--- /dev/null
+++ b/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
diff --git a/python/sglang/multimodal_gen/configs/sample/cosmos3.py b/python/sglang/multimodal_gen/configs/sample/cosmos3.py
new file mode 100644
index 000000000..c67b7e5b9
--- /dev/null
+++ b/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()
diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py
index 0eef02b01..84fdf29c7 100644
--- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py
+++ b/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")
diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py
index 9e7612162..9f310a8da 100644
--- a/python/sglang/multimodal_gen/registry.py
+++ b/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,
diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py
index 71aa1c4f1..5eb38bb3f 100644
--- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py
+++ b/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,
diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
index 2e9cf9e4e..281f4d163 100644
--- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
+++ b/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(
diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
index 14e9a64b9..5d1cee0f2 100644
--- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
+++ b/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,
diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py
new file mode 100644
index 000000000..8d6192b3f
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py
@@ -0,0 +1,1359 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 Omni transformer.
+
+Dual-pathway DiT: an Understanding (UND) pathway runs causal self-attention
+over the text tokens once and caches its K/V; a Generation (GEN) pathway
+cross-attends from noisy visual tokens to that cache at every denoising step.
+"""
+
+import math
+from collections.abc import Iterable, Iterator
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.configs.models.dits.cosmos3video import Cosmos3VideoConfig
+from sglang.multimodal_gen.runtime.distributed import (
+ get_sp_group,
+ get_sp_world_size,
+ sequence_model_parallel_all_gather,
+)
+from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
+from sglang.multimodal_gen.runtime.layers.attention import USPAttention
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
+from sglang.multimodal_gen.runtime.layers.linear import (
+ MergedColumnParallelLinear,
+ ReplicatedLinear,
+ RowParallelLinear,
+ UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+ QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding
+from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
+ VocabParallelEmbedding,
+)
+from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.utils import add_prefix
+
+logger = init_logger(__name__)
+
+
+# -----------------------------------------------------------------------------
+# mRoPE position ID computation (Qwen3VL-style)
+# -----------------------------------------------------------------------------
+
+
+def compute_mrope_position_ids_text(
+ num_tokens: int,
+ temporal_offset: int,
+ device: torch.device,
+) -> tuple[torch.Tensor, int]:
+ """Generate 3D mRoPE position IDs for text tokens.
+
+ Text tokens: all three axes (T, H, W) share the same monotonically
+ increasing position IDs: (0,0,0), (1,1,1), (2,2,2), ...
+
+ Returns:
+ (position_ids [3, num_tokens], next_temporal_offset)
+ """
+ ids = torch.arange(num_tokens, dtype=torch.long, device=device) + temporal_offset
+ mrope_ids = ids.unsqueeze(0).expand(3, -1).contiguous()
+ return mrope_ids, temporal_offset + num_tokens
+
+
+def compute_mrope_position_ids_vision(
+ grid_t: int,
+ grid_h: int,
+ grid_w: int,
+ temporal_offset: int | float,
+ device: torch.device,
+ fps: float | None = None,
+ base_fps: float = 24.0,
+ temporal_compression_factor: int = 4,
+) -> tuple[torch.Tensor, int | float]:
+ """Generate 3D mRoPE position IDs for vision tokens.
+
+ Creates a (T, H, W) position grid. Spatial indices reset to 0
+ per vision segment (Qwen3VL-style).
+ Flattened in T-major order.
+
+ Returns:
+ (position_ids [3, grid_t * grid_h * grid_w], next_temporal_offset)
+ """
+ fps_modulation = fps is not None and grid_t > 1
+
+ if fps_modulation:
+ tps = fps / temporal_compression_factor
+ base_tps = base_fps / temporal_compression_factor
+ frame_indices = torch.arange(grid_t, dtype=torch.float32, device=device)
+ t_index = (
+ (frame_indices / tps * base_tps + temporal_offset)
+ .view(-1, 1)
+ .expand(-1, grid_h * grid_w)
+ .flatten()
+ )
+ else:
+ t_index = torch.arange(grid_t, dtype=torch.long, device=device).view(
+ -1, 1
+ ).expand(-1, grid_h * grid_w).flatten() + int(temporal_offset)
+
+ h_index = (
+ torch.arange(grid_h, dtype=torch.long, device=device)
+ .view(1, -1, 1)
+ .expand(grid_t, -1, grid_w)
+ .flatten()
+ )
+ w_index = (
+ torch.arange(grid_w, dtype=torch.long, device=device)
+ .view(1, 1, -1)
+ .expand(grid_t, grid_h, -1)
+ .flatten()
+ )
+
+ if fps_modulation:
+ mrope_ids = torch.stack(
+ [t_index, h_index.to(torch.float32), w_index.to(torch.float32)], dim=0
+ )
+ else:
+ mrope_ids = torch.stack([t_index, h_index, w_index], dim=0)
+
+ next_offset = math.ceil(mrope_ids.max().item()) + 1
+ return mrope_ids, next_offset
+
+
+# -----------------------------------------------------------------------------
+# Qwen3-style RoPE functions
+# -----------------------------------------------------------------------------
+
+
+def qwen3_rotate_half(x: torch.Tensor) -> torch.Tensor:
+ """Qwen3/Llama-style rotate_half: split first/second half of head_dim."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def qwen3_apply_rotary_pos_emb(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Qwen3-style RoPE: (x * cos) + (rotate_half(x) * sin).
+
+ Args:
+ q: [B, S, H, D]
+ k: [B, S, H_kv, D]
+ cos: [1, S, 1, D] or broadcastable
+ sin: [1, S, 1, D] or broadcastable
+ """
+ q_embed = (q * cos) + (qwen3_rotate_half(q) * sin)
+ k_embed = (k * cos) + (qwen3_rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+# -----------------------------------------------------------------------------
+# Qwen3VL-style Rotary Embedding
+# -----------------------------------------------------------------------------
+
+
+class Qwen3VLTextRotaryEmbedding(nn.Module):
+ """Qwen3VL-style multi-dimensional rotary embedding."""
+
+ def __init__(
+ self,
+ head_dim: int = 128,
+ rope_theta: float = 5000000.0,
+ mrope_section: tuple[int, int, int] = (24, 20, 20),
+ ):
+ super().__init__()
+ self.rope_type = "default"
+ self.max_seq_len_cached = 262144
+ self.mrope_section = list(mrope_section)
+ self.head_dim = head_dim
+
+ # Compute inverse frequencies
+ dim = head_dim
+ inv_freq = 1.0 / (
+ rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)
+ )
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.attention_scaling = 1.0
+
+ def apply_interleaved_mrope(
+ self, freqs: torch.Tensor, mrope_section: list[int]
+ ) -> torch.Tensor:
+ """Apply interleaved MRoPE to 3D rotary embeddings.
+
+ Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
+ interleaved [THTHWHTHW...TT], preserving frequency continuity.
+
+ Args:
+ freqs: (3, bs, seq_len, head_dim // 2)
+ mrope_section: (3,) section sizes
+
+ Returns:
+ freqs_t: (bs, seq_len, head_dim // 2)
+ """
+ freqs_t = freqs[0].clone()
+ for dim, offset in enumerate((1, 2), start=1): # H, W
+ length = mrope_section[dim] * 3
+ idx = slice(offset, length, 3)
+ freqs_t[..., idx] = freqs[dim, ..., idx]
+ return freqs_t
+
+ @torch.no_grad()
+ def forward(
+ self, x: torch.Tensor, position_ids: torch.Tensor
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Compute cos and sin for rotary embeddings.
+
+ Args:
+ x: dummy tensor for dtype
+ position_ids: [3, B, S] or [B, S] position IDs
+
+ Returns:
+ (cos, sin) each of shape [B, S, D]
+ """
+ if position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+ # Expand inv_freq: [3, B, D//2, 1]
+ inv_freq_expanded = (
+ self.inv_freq[None, None, :, None]
+ .float()
+ .expand(3, position_ids.shape[1], -1, 1)
+ .to(position_ids.device)
+ )
+ # position_ids_expanded: [3, B, 1, S]
+ position_ids_expanded = position_ids[:, :, None, :].float()
+
+ # freqs: [3, B, D//2, S] -> transpose -> [3, B, S, D//2]
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(
+ 2, 3
+ )
+ freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+# -----------------------------------------------------------------------------
+# Cosmos3 Timestep Embedder
+# -----------------------------------------------------------------------------
+
+
+class Cosmos3TimestepEmbedder(nn.Module):
+ """Embeds scalar timesteps into vector representations.
+
+ Uses ReplicatedLinear for consistency with other SGLang models and
+ to support quantization (though timestep embedders are typically excluded).
+ """
+
+ def __init__(
+ self,
+ hidden_size: int,
+ frequency_embedding_size: int = 256,
+ max_period: int = 10000,
+ timestep_scale: float = 0.001,
+ prefix: str = "",
+ quant_config: QuantizationConfig | None = None,
+ ):
+ super().__init__()
+ self.timestep_scale = timestep_scale
+ self.frequency_embedding_size = frequency_embedding_size
+ self.hidden_size = hidden_size
+ self.max_period = max_period
+
+ # Use ReplicatedLinear for consistency (typically excluded from quantization)
+ self.linear_1 = ReplicatedLinear(
+ frequency_embedding_size,
+ hidden_size,
+ bias=True,
+ quant_config=quant_config,
+ prefix=add_prefix("linear_1", prefix),
+ )
+ self.act = nn.SiLU()
+ self.linear_2 = ReplicatedLinear(
+ hidden_size,
+ hidden_size,
+ bias=True,
+ quant_config=quant_config,
+ prefix=add_prefix("linear_2", prefix),
+ )
+
+ def forward(self, t: torch.Tensor) -> torch.Tensor:
+ """Embed timesteps.
+
+ Args:
+ t: [B] timestep values
+
+ Returns:
+ [B, hidden_size] timestep embeddings
+ """
+ # Scale timestep
+ t_scaled = t * self.timestep_scale
+
+ # Compute sinusoidal embeddings in fp32
+ t_freq = timestep_embedding(
+ t_scaled,
+ self.frequency_embedding_size,
+ self.max_period,
+ dtype=torch.float32,
+ )
+
+ # Project through MLP
+ # When fp8-quantized, weight.dtype is float8_e4m3fn — keep input in
+ # float32 (the quant kernel handles input quantization internally).
+ w_dtype = self.linear_1.weight.dtype
+ if w_dtype.is_floating_point and w_dtype.itemsize >= 2:
+ x = t_freq.to(w_dtype)
+ else:
+ x = t_freq # already float32 from timestep_embedding
+ x, _ = self.linear_1(x)
+ x = self.act(x)
+ x, _ = self.linear_2(x)
+ return x
+
+
+# -----------------------------------------------------------------------------
+# Cosmos3 Gated MLP
+# -----------------------------------------------------------------------------
+
+
+class Cosmos3GatedMLP(nn.Module):
+ """Gated MLP (SwiGLU-style) for Cosmos3."""
+
+ def __init__(
+ self,
+ hidden_size: int,
+ intermediate_size: int,
+ prefix: str = "",
+ quant_config: QuantizationConfig | None = None,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+
+ self.gate_up_proj = MergedColumnParallelLinear(
+ hidden_size,
+ [intermediate_size] * 2,
+ bias=False,
+ gather_output=False,
+ quant_config=quant_config,
+ prefix=add_prefix("gate_up_proj", prefix),
+ )
+ self.down_proj = RowParallelLinear(
+ intermediate_size,
+ hidden_size,
+ bias=False,
+ input_is_parallel=True,
+ quant_config=quant_config,
+ prefix=add_prefix("down_proj", prefix),
+ )
+ self.act_fn = SiluAndMul()
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ gate_up, _ = self.gate_up_proj(x)
+ out, _ = self.down_proj(self.act_fn(gate_up))
+ return out
+
+
+# -----------------------------------------------------------------------------
+# Cosmos3 UND Causal Attention
+# -----------------------------------------------------------------------------
+
+
+class Cosmos3CausalAttention(nn.Module):
+ """Understanding pathway: causal self-attention on text tokens."""
+
+ def __init__(
+ self,
+ hidden_size: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ head_dim: int,
+ prefix: str = "",
+ quant_config: QuantizationConfig | None = None,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads
+ self.head_dim = head_dim
+
+ self.q_size = num_attention_heads * head_dim
+ self.kv_size = num_key_value_heads * head_dim
+ self.to_qkv = MergedColumnParallelLinear(
+ hidden_size,
+ [self.q_size, self.kv_size, self.kv_size],
+ bias=False,
+ gather_output=True,
+ quant_config=quant_config,
+ prefix=add_prefix("to_qkv", prefix),
+ )
+ # Output projection - ReplicatedLinear for quantization support
+ # Input is not parallel (gather_output=True on QKV)
+ self.to_out = ReplicatedLinear(
+ num_attention_heads * head_dim,
+ hidden_size,
+ bias=False,
+ quant_config=quant_config,
+ prefix=add_prefix("to_out", prefix),
+ )
+
+ # Per-head QK norm. Modules hold the weights; F.rms_norm in forward.
+ self.norm_q = RMSNorm(head_dim, eps=1e-6)
+ self.norm_k = RMSNorm(head_dim, eps=1e-6)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ freqs_cos: torch.Tensor,
+ freqs_sin: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Forward with KV cache return.
+
+ Returns:
+ (output, K, V) where K/V are post-norm, post-RoPE
+ """
+ batch_size, seq_len = hidden_states.shape[:2]
+
+ qkv, _ = self.to_qkv(hidden_states)
+ # split returns strided views into qkv; .contiguous() before .view()
+ # because the per-head reshape needs row-major memory.
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+ q = q.contiguous().view(
+ batch_size, seq_len, self.num_attention_heads, self.head_dim
+ )
+ k = k.contiguous().view(
+ batch_size, seq_len, self.num_key_value_heads, self.head_dim
+ )
+ v = v.contiguous().view(
+ batch_size, seq_len, self.num_key_value_heads, self.head_dim
+ )
+
+ q = F.rms_norm(
+ q, (self.head_dim,), self.norm_q.weight, self.norm_q.variance_epsilon
+ )
+ k = F.rms_norm(
+ k, (self.head_dim,), self.norm_k.weight, self.norm_k.variance_epsilon
+ )
+ q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
+
+ out = F.scaled_dot_product_attention(
+ q.transpose(1, 2),
+ k.transpose(1, 2),
+ v.transpose(1, 2),
+ is_causal=True,
+ enable_gqa=True,
+ )
+ out = out.transpose(1, 2).reshape(batch_size, seq_len, -1)
+
+ out, _ = self.to_out(out)
+ return out, k, v
+
+
+# -----------------------------------------------------------------------------
+# Cosmos3 GEN Cross Attention
+# -----------------------------------------------------------------------------
+
+
+class Cosmos3CrossAttention(nn.Module):
+ """Generation pathway: cross-attention where visual Q attends to all K/V."""
+
+ def __init__(
+ self,
+ hidden_size: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ head_dim: int,
+ prefix: str = "",
+ quant_config: QuantizationConfig | None = None,
+ supported_attention_backends: set | None = None,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.num_attention_heads = num_attention_heads
+ self.num_key_value_heads = num_key_value_heads
+ self.head_dim = head_dim
+
+ self.q_size = num_attention_heads * head_dim
+ self.kv_size = num_key_value_heads * head_dim
+ self.to_qkv = MergedColumnParallelLinear(
+ hidden_size,
+ [self.q_size, self.kv_size, self.kv_size],
+ bias=False,
+ gather_output=True,
+ quant_config=quant_config,
+ prefix=add_prefix("to_qkv", prefix),
+ )
+ self.to_out = ReplicatedLinear(
+ num_attention_heads * head_dim,
+ hidden_size,
+ bias=False,
+ quant_config=quant_config,
+ prefix=add_prefix("to_out", prefix),
+ )
+
+ self.norm_q = RMSNorm(head_dim, eps=1e-6)
+ self.norm_k = RMSNorm(head_dim, eps=1e-6)
+
+ self.attn = USPAttention(
+ num_heads=num_attention_heads,
+ head_size=head_dim,
+ num_kv_heads=num_key_value_heads,
+ causal=False,
+ supported_attention_backends=supported_attention_backends,
+ prefix=add_prefix("attn", prefix),
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ k_und: torch.Tensor,
+ v_und: torch.Tensor,
+ freqs_cos: torch.Tensor,
+ freqs_sin: torch.Tensor,
+ ) -> torch.Tensor:
+ """Cross-attention from GEN to cached UND K/V.
+
+ Args:
+ hidden_states: [B, S_gen_local, hidden_size] visual tokens (may be sharded)
+ k_und: [B, S_und, H_kv, D] pre-computed UND keys (always full/replicated)
+ v_und: [B, S_und, H_kv, D] pre-computed UND values (always full/replicated)
+ freqs_cos: [B, S_gen_local, 1, D] cosine part of RoPE (for local shard)
+ freqs_sin: [B, S_gen_local, 1, D] sine part of RoPE (for local shard)
+ """
+ batch_size, seq_len_gen = hidden_states.shape[:2]
+
+ qkv, _ = self.to_qkv(hidden_states)
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+ q = q.contiguous().view(
+ batch_size, seq_len_gen, self.num_attention_heads, self.head_dim
+ )
+ k = k.contiguous().view(
+ batch_size, seq_len_gen, self.num_key_value_heads, self.head_dim
+ )
+ v = v.contiguous().view(
+ batch_size, seq_len_gen, self.num_key_value_heads, self.head_dim
+ )
+
+ q = F.rms_norm(
+ q, (self.head_dim,), self.norm_q.weight, self.norm_q.variance_epsilon
+ )
+ k = F.rms_norm(
+ k, (self.head_dim,), self.norm_k.weight, self.norm_k.variance_epsilon
+ )
+ q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
+
+ # K/V = [text (replicated full on every SP rank) | image (sharded same as Q)].
+ # USPAttention routes through the registered attention backend (FA, sage,
+ # …) and handles the Ulysses all-to-all when SP > 1.
+ num_und = k_und.shape[1]
+ k = torch.cat([k_und, k], dim=1)
+ v = torch.cat([v_und, v], dim=1)
+ out = self.attn(q, k, v, num_replicated_kv_prefix=num_und)
+ out = out.reshape(batch_size, seq_len_gen, -1)
+ out, _ = self.to_out(out)
+ return out
+
+
+# -----------------------------------------------------------------------------
+# Cosmos3 UND Decoder Layer
+# -----------------------------------------------------------------------------
+
+
+class Cosmos3UndDecoderLayer(nn.Module):
+ """Understanding pathway decoder layer: causal self-attention + MLP."""
+
+ def __init__(
+ self,
+ hidden_size: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ head_dim: int,
+ intermediate_size: int,
+ rms_norm_eps: float,
+ layer_idx: int,
+ prefix: str = "",
+ quant_config: QuantizationConfig | None = None,
+ ):
+ super().__init__()
+ self.layer_idx = layer_idx
+
+ self.self_attn = Cosmos3CausalAttention(
+ hidden_size=hidden_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ head_dim=head_dim,
+ prefix=add_prefix("self_attn", prefix),
+ quant_config=quant_config,
+ )
+ self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+ self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+ self.mlp = Cosmos3GatedMLP(
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ prefix=add_prefix("mlp", prefix),
+ quant_config=quant_config,
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ freqs_cos: torch.Tensor,
+ freqs_sin: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Forward pass.
+
+ Returns:
+ (hidden_states, K, V) where K/V are for GEN cross-attention
+ """
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ attn_out, k, v = self.self_attn(hidden_states, freqs_cos, freqs_sin)
+ hidden_states = residual + attn_out
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = residual + self.mlp(hidden_states)
+
+ return hidden_states, k, v
+
+
+# -----------------------------------------------------------------------------
+# Cosmos3 GEN Decoder Layer
+# -----------------------------------------------------------------------------
+
+
+class Cosmos3GenDecoderLayer(nn.Module):
+ """Generation pathway decoder layer: cross-attention + MLP."""
+
+ def __init__(
+ self,
+ hidden_size: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ head_dim: int,
+ intermediate_size: int,
+ rms_norm_eps: float,
+ layer_idx: int,
+ prefix: str = "",
+ quant_config: QuantizationConfig | None = None,
+ supported_attention_backends: set | None = None,
+ ):
+ super().__init__()
+ self.layer_idx = layer_idx
+
+ self.cross_attention = Cosmos3CrossAttention(
+ hidden_size=hidden_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ head_dim=head_dim,
+ prefix=add_prefix("cross_attention", prefix),
+ quant_config=quant_config,
+ supported_attention_backends=supported_attention_backends,
+ )
+ self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+ self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+ self.mlp = Cosmos3GatedMLP(
+ hidden_size=hidden_size,
+ intermediate_size=intermediate_size,
+ prefix=add_prefix("mlp", prefix),
+ quant_config=quant_config,
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ k_und: torch.Tensor,
+ v_und: torch.Tensor,
+ freqs_cos: torch.Tensor,
+ freqs_sin: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ # Fused add+rmsnorm: each `(hidden_states, residual) = norm(...)`
+ # collapses the residual add and RMSNorm into one kernel. The
+ # caller threads `residual` across layers and resolves it before
+ # the post-loop all-gather + `norm_moe_gen`.
+ if residual is None:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ else:
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+ hidden_states = self.cross_attention(
+ hidden_states, k_und, v_und, freqs_cos, freqs_sin
+ )
+
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+ hidden_states = self.mlp(hidden_states)
+
+ return hidden_states, residual
+
+
+# -----------------------------------------------------------------------------
+# Cosmos3 Language Model (UND pathway)
+# -----------------------------------------------------------------------------
+
+
+class Cosmos3LanguageModel(nn.Module):
+ """Understanding pathway: processes text tokens and caches K/V."""
+
+ def __init__(
+ self,
+ hidden_size: int,
+ num_hidden_layers: int,
+ num_attention_heads: int,
+ num_key_value_heads: int,
+ head_dim: int,
+ intermediate_size: int,
+ vocab_size: int,
+ rms_norm_eps: float,
+ rope_theta: float,
+ mrope_section: tuple[int, int, int],
+ quant_config: QuantizationConfig | None = None,
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.num_hidden_layers = num_hidden_layers
+
+ self.embed_tokens = VocabParallelEmbedding(
+ vocab_size,
+ hidden_size,
+ params_dtype=torch.bfloat16,
+ )
+ self.rotary_emb = Qwen3VLTextRotaryEmbedding(
+ head_dim=head_dim,
+ rope_theta=rope_theta,
+ mrope_section=mrope_section,
+ )
+ self.layers = nn.ModuleList(
+ [
+ Cosmos3UndDecoderLayer(
+ hidden_size=hidden_size,
+ num_attention_heads=num_attention_heads,
+ num_key_value_heads=num_key_value_heads,
+ head_dim=head_dim,
+ intermediate_size=intermediate_size,
+ rms_norm_eps=rms_norm_eps,
+ layer_idx=i,
+ prefix=f"layers.{i}",
+ quant_config=quant_config,
+ )
+ for i in range(num_hidden_layers)
+ ]
+ )
+
+ def forward(
+ self,
+ text_ids: torch.Tensor,
+ text_mask: torch.Tensor,
+ freqs_cos: torch.Tensor,
+ freqs_sin: torch.Tensor,
+ ) -> list[tuple[torch.Tensor, torch.Tensor]]:
+ """Process text tokens and return per-layer K/V cache.
+
+ Args:
+ text_ids: [B, S] token IDs
+ text_mask: [B, S] float mask (1=real, 0=pad)
+ freqs_cos: [B, S, D] RoPE cosines
+ freqs_sin: [B, S, D] RoPE sines
+
+ Returns:
+ List of (K, V) per layer for GEN cross-attention
+ """
+ hidden = self.embed_tokens(text_ids)
+ mask_3d = text_mask.unsqueeze(-1)
+
+ # Add dimension for per-head broadcast
+ freqs_cos = freqs_cos.unsqueeze(2) # [B, S, 1, D]
+ freqs_sin = freqs_sin.unsqueeze(2)
+
+ cached_kv: list[tuple[torch.Tensor, torch.Tensor]] = []
+ for layer in self.layers:
+ hidden = hidden * mask_3d
+ hidden, k, v = layer(hidden, freqs_cos, freqs_sin)
+ cached_kv.append((k, v))
+
+ return cached_kv
+
+
+# -----------------------------------------------------------------------------
+# Cosmos3 Omni Transformer
+# -----------------------------------------------------------------------------
+
+
+class Cosmos3OmniTransformer(CachableDiT):
+ """Cosmos3 Omni transformer.
+
+ Dual-pathway architecture:
+ - Understanding (UND): causal LM processing text
+ - Generation (GEN): cross-attention from visual to UND K/V
+ """
+
+ _fsdp_shard_conditions = Cosmos3VideoConfig()._fsdp_shard_conditions
+ _compile_conditions = Cosmos3VideoConfig()._compile_conditions
+ _supported_attention_backends = Cosmos3VideoConfig()._supported_attention_backends
+ param_names_mapping = Cosmos3VideoConfig().arch_config.param_names_mapping
+ reverse_param_names_mapping = (
+ Cosmos3VideoConfig().arch_config.reverse_param_names_mapping
+ )
+ lora_param_names_mapping = Cosmos3VideoConfig().arch_config.lora_param_names_mapping
+
+ def __init__(
+ self,
+ config: Cosmos3VideoConfig,
+ hf_config: dict[str, Any],
+ quant_config: QuantizationConfig | None = None,
+ ) -> None:
+ super().__init__(config=config, hf_config=hf_config)
+
+ arch = config.arch_config
+ self.hidden_size = arch.hidden_size
+ self.num_hidden_layers = arch.num_hidden_layers
+ self.num_attention_heads = arch.num_attention_heads
+ self.num_key_value_heads = arch.num_key_value_heads
+ self.head_dim = arch.head_dim
+ self.intermediate_size = arch.intermediate_size
+ self.latent_patch_size = arch.latent_patch_size
+ self.latent_channel = arch.latent_channel
+ self.num_channels_latents = arch.out_channels
+ self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel
+ self.timestep_scale = arch.timestep_scale
+ self.base_fps = arch.base_fps
+ self.temporal_compression_factor = arch.temporal_compression_factor
+ self.temporal_margin = arch.unified_3d_mrope_temporal_modality_margin
+ self.rms_norm_eps = arch.rms_norm_eps
+
+ # Ulysses sequence parallelism. When CFG-parallel is also enabled
+ # the SP group only spans ranks that share a CFG context (cond or
+ # uncond), so ``sp_size`` here is the per-context shard count.
+ self.sp_size = get_sp_world_size()
+ self.sp_group = get_sp_group() if self.sp_size > 1 else None
+ self.sp_rank = self.sp_group.rank_in_group if self.sp_group else 0
+ if self.sp_size > 1:
+ logger.info(
+ f"Cosmos3 SP enabled: sp_size={self.sp_size}, sp_rank={self.sp_rank}"
+ )
+
+ # Language model (UND pathway)
+ self.language_model = Cosmos3LanguageModel(
+ hidden_size=arch.hidden_size,
+ num_hidden_layers=arch.num_hidden_layers,
+ num_attention_heads=arch.num_attention_heads,
+ num_key_value_heads=arch.num_key_value_heads,
+ head_dim=arch.head_dim,
+ intermediate_size=arch.intermediate_size,
+ vocab_size=arch.vocab_size,
+ rms_norm_eps=arch.rms_norm_eps,
+ rope_theta=arch.rope_theta,
+ mrope_section=arch.mrope_section,
+ quant_config=quant_config,
+ )
+
+ # Latent projection layers - ReplicatedLinear for quantization support
+ self.vae2llm = ReplicatedLinear(
+ self.patch_latent_dim,
+ self.hidden_size,
+ bias=True,
+ quant_config=quant_config,
+ prefix="vae2llm",
+ )
+ self.llm2vae = ReplicatedLinear(
+ self.hidden_size,
+ self.patch_latent_dim,
+ bias=True,
+ quant_config=quant_config,
+ prefix="llm2vae",
+ )
+
+ # Timestep embedder
+ self.time_embedder = Cosmos3TimestepEmbedder(
+ hidden_size=self.hidden_size,
+ frequency_embedding_size=arch.frequency_embedding_size,
+ timestep_scale=arch.timestep_scale,
+ prefix="time_embedder",
+ quant_config=quant_config,
+ )
+
+ # Generation layers (GEN pathway)
+ self.gen_layers = nn.ModuleList(
+ [
+ Cosmos3GenDecoderLayer(
+ hidden_size=arch.hidden_size,
+ num_attention_heads=arch.num_attention_heads,
+ num_key_value_heads=arch.num_key_value_heads,
+ head_dim=arch.head_dim,
+ intermediate_size=arch.intermediate_size,
+ rms_norm_eps=arch.rms_norm_eps,
+ layer_idx=i,
+ prefix=f"gen_layers.{i}",
+ quant_config=quant_config,
+ supported_attention_backends=arch._supported_attention_backends,
+ )
+ for i in range(arch.num_hidden_layers)
+ ]
+ )
+
+ # Output norm
+ self.norm_moe_gen = RMSNorm(self.hidden_size, eps=arch.rms_norm_eps)
+
+ # Cached K/V from UND pathway - dict keyed by cache_key for CFG support
+ # This allows maintaining separate caches for conditional and unconditional
+ # prompts, avoiding recomputation on every denoising step
+ self.cached_kv: dict[str, list[tuple[torch.Tensor, torch.Tensor]]] = {}
+ self.cached_freqs_gen: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+
+ self.__post_init__()
+
+ def _pad_to_patch_size(self, H: int, W: int) -> tuple[int, int, int, int]:
+ """Compute padded spatial dims aligned to patch_size."""
+ p = self.latent_patch_size
+ H_padded = ((H + p - 1) // p) * p
+ W_padded = ((W + p - 1) // p) * p
+ return H_padded // p, W_padded // p, H_padded, W_padded
+
+ def patchify(self, latents: torch.Tensor, T: int, H: int, W: int) -> torch.Tensor:
+ """Convert latents to patches: [B, C, T, H, W] -> [B, T*Hp*Wp, p*p*C]."""
+ B = latents.shape[0]
+ p = self.latent_patch_size
+ C = self.latent_channel
+ Hp, Wp, H_padded, W_padded = self._pad_to_patch_size(H, W)
+
+ if H_padded != H or W_padded != W:
+ latents = F.pad(latents, (0, W_padded - W, 0, H_padded - H))
+
+ x = latents.reshape(B, C, T, Hp, p, Wp, p)
+ x = x.permute(0, 2, 3, 5, 4, 6, 1) # [B, T, Hp, Wp, p, p, C]
+ return x.reshape(B, T * Hp * Wp, p * p * C)
+
+ def unpatchify(self, tokens: torch.Tensor, T: int, H: int, W: int) -> torch.Tensor:
+ """Convert patches back to latents: [B, T*Hp*Wp, p*p*C] -> [B, C, T, H, W]."""
+ B = tokens.shape[0]
+ p = self.latent_patch_size
+ C = self.latent_channel
+ Hp, Wp, H_padded, W_padded = self._pad_to_patch_size(H, W)
+
+ x = tokens.reshape(B, T, Hp, Wp, p, p, C)
+ x = x.permute(0, 6, 1, 2, 4, 3, 5) # [B, C, T, Hp, p, Wp, p]
+ x = x.reshape(B, C, T, H_padded, W_padded)
+
+ if H_padded != H or W_padded != W:
+ x = x[:, :, :, :H, :W]
+ return x
+
+ def _compute_rope_freqs(
+ self,
+ text_mask: torch.Tensor,
+ T: int,
+ Hp: int,
+ Wp: int,
+ fps: float | None,
+ device: torch.device,
+ dtype: torch.dtype,
+ ) -> tuple[
+ tuple[torch.Tensor, torch.Tensor],
+ tuple[torch.Tensor, torch.Tensor],
+ ]:
+ """Compute mRoPE cos/sin for UND (text) and GEN (visual) pathways."""
+ B = text_mask.shape[0]
+ S_text = text_mask.shape[1]
+ text_lengths = text_mask.sum(dim=1).long()
+ effective_fps = fps if fps is not None and T > 1 else None
+
+ text_pos_list = []
+ vis_pos_list = []
+ for b in range(B):
+ real_len = int(text_lengths[b].item())
+ t_pos, t_offset = compute_mrope_position_ids_text(
+ real_len, temporal_offset=0, device=device
+ )
+ v_pos, _ = compute_mrope_position_ids_vision(
+ T,
+ Hp,
+ Wp,
+ temporal_offset=t_offset + self.temporal_margin,
+ device=device,
+ fps=effective_fps,
+ base_fps=self.base_fps,
+ temporal_compression_factor=self.temporal_compression_factor,
+ )
+ if real_len < S_text:
+ t_pos = torch.cat(
+ [
+ t_pos,
+ torch.zeros(
+ 3, S_text - real_len, dtype=t_pos.dtype, device=device
+ ),
+ ],
+ dim=1,
+ )
+ text_pos_list.append(t_pos)
+ vis_pos_list.append(v_pos)
+
+ text_pos_ids = torch.stack(text_pos_list, dim=1).to(device) # [3, B, S_text]
+ vis_pos_ids = torch.stack(vis_pos_list, dim=1).to(device) # [3, B, S_vis]
+
+ rotary_emb = self.language_model.rotary_emb
+ _dummy = torch.tensor([], dtype=dtype, device=device)
+ cos_und, sin_und = rotary_emb(_dummy, position_ids=text_pos_ids)
+ cos_gen, sin_gen = rotary_emb(_dummy, position_ids=vis_pos_ids)
+
+ freqs_und = (cos_und, sin_und)
+ freqs_gen = (cos_gen, sin_gen)
+ return freqs_und, freqs_gen
+
+ def reset_cache(self, cache_key: str | None = None):
+ """Reset cached K/V from UND pathway.
+
+ Args:
+ cache_key: If provided, reset only the specified cache key.
+ If None, reset all caches.
+ """
+ if cache_key is None:
+ # Reset all caches
+ self.cached_kv = {}
+ self.cached_freqs_gen = {}
+ else:
+ # Reset specific cache
+ if cache_key in self.cached_kv:
+ del self.cached_kv[cache_key]
+ if cache_key in self.cached_freqs_gen:
+ del self.cached_freqs_gen[cache_key]
+
+ def _ensure_cache_dicts(self):
+ """Ensure cache dictionaries exist (for backwards compatibility)."""
+ if not isinstance(self.cached_kv, dict):
+ self.cached_kv = {}
+ if not isinstance(self.cached_freqs_gen, dict):
+ self.cached_freqs_gen = {}
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+ timestep: torch.LongTensor,
+ encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+ guidance=None,
+ text_ids: torch.Tensor | None = None,
+ text_mask: torch.Tensor | None = None,
+ fps: float | None = None,
+ cache_key: str = "default",
+ noisy_frame_mask: torch.Tensor | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ """Forward pass for denoising.
+
+ Args:
+ hidden_states: [B, C, T, H, W] noisy latents
+ encoder_hidden_states: Not used (text embedded in transformer)
+ timestep: [B] diffusion timestep per sample
+ text_ids: [B, S_text] tokenized text input
+ text_mask: [B, S_text] attention mask for text (1=real, 0=pad)
+ fps: video frame rate for temporal mRoPE scaling
+ cache_key: Key for the UND K/V cache. Use different keys for
+ conditional ("cond") and unconditional ("uncond") branches
+ in CFG to avoid recomputing the cache every step.
+ noisy_frame_mask: Optional [B, 1, T, 1, 1] mask where 1 marks
+ noisy frames (timestep embedding applied) and 0 marks
+ conditioned frames (clean context, embedding skipped).
+ ``None`` means every frame is noisy (T2V / T2I).
+
+ Returns:
+ [B, C, T, H, W] velocity prediction
+ """
+ if text_ids is None or text_mask is None:
+ raise ValueError("Cosmos3 requires text_ids and text_mask to be passed")
+
+ batch_size, C, T, H, W = hidden_states.shape
+ Hp, Wp, _, _ = self._pad_to_patch_size(H, W)
+ max_real_len = int(text_mask.sum(dim=1).max().item())
+
+ # Check if sequence parallelism is enabled
+ sequence_shard_enabled = self.sp_size > 1
+
+ # Patchify and project to hidden dim
+ hidden_gen, _ = self.vae2llm(self.patchify(hidden_states, T, H, W))
+ seq_len_orig = hidden_gen.shape[1]
+ seq_shard_pad = 0
+
+ # Per-token noisy mask follows the same pad/shard as hidden_gen, so
+ # build it before the SP split.
+ token_noisy_mask: torch.Tensor | None = None
+ if noisy_frame_mask is not None:
+ token_noisy_mask = (
+ noisy_frame_mask[:, 0, :, 0, 0]
+ .unsqueeze(-1)
+ .expand(-1, -1, Hp * Wp)
+ .reshape(batch_size, -1, 1)
+ .to(hidden_gen.dtype)
+ )
+
+ # Shard sequence across GPUs if SP enabled
+ if sequence_shard_enabled:
+ if seq_len_orig % self.sp_size != 0:
+ seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size)
+ pad = torch.zeros(
+ (batch_size, seq_shard_pad, hidden_gen.shape[2]),
+ dtype=hidden_gen.dtype,
+ device=hidden_gen.device,
+ )
+ hidden_gen = torch.cat([hidden_gen, pad], dim=1)
+ if token_noisy_mask is not None:
+ mask_pad = torch.zeros(
+ (batch_size, seq_shard_pad, 1),
+ dtype=token_noisy_mask.dtype,
+ device=token_noisy_mask.device,
+ )
+ token_noisy_mask = torch.cat([token_noisy_mask, mask_pad], dim=1)
+ local_seq_len = hidden_gen.shape[1] // self.sp_size
+ hidden_gen = hidden_gen.view(
+ batch_size, self.sp_size, local_seq_len, hidden_gen.shape[2]
+ )
+ hidden_gen = hidden_gen[:, self.sp_rank, :, :]
+ if token_noisy_mask is not None:
+ token_noisy_mask = token_noisy_mask.view(
+ batch_size, self.sp_size, local_seq_len, 1
+ )[:, self.sp_rank, :, :]
+
+ # Add timestep embedding (computed in float32 for numerical stability, then cast back)
+ time_embed = self.time_embedder(timestep.float())
+ time_embed = time_embed.to(
+ hidden_states.dtype
+ ) # Cast to match hidden_gen dtype
+ if token_noisy_mask is not None:
+ hidden_gen = hidden_gen + time_embed.unsqueeze(1) * token_noisy_mask
+ else:
+ hidden_gen = hidden_gen + time_embed.unsqueeze(1)
+
+ self._ensure_cache_dicts()
+
+ # Compute UND K/V cache for this cache_key if not already cached
+ # This allows reusing the cache across denoising steps for the same text
+ if cache_key not in self.cached_kv:
+ freqs_und, freqs_gen = self._compute_rope_freqs(
+ text_mask, T, Hp, Wp, fps, hidden_states.device, hidden_states.dtype
+ )
+ # UND K/V cache is kept FULL on all ranks (not sharded). Text
+ # sequence is short, so memory impact is minimal, and the GEN
+ # cross-attention needs the full K/V on every SP rank.
+ self.cached_kv[cache_key] = self.language_model(
+ text_ids, text_mask, freqs_und[0], freqs_und[1]
+ )
+ self.cached_freqs_gen[cache_key] = freqs_gen
+
+ freqs_gen = self.cached_freqs_gen[cache_key]
+ cos_gen, sin_gen = freqs_gen
+
+ if sequence_shard_enabled:
+ if seq_shard_pad > 0:
+ pad_cos = cos_gen[:, -1:].expand(-1, seq_shard_pad, -1)
+ pad_sin = sin_gen[:, -1:].expand(-1, seq_shard_pad, -1)
+ cos_gen = torch.cat([cos_gen, pad_cos], dim=1)
+ sin_gen = torch.cat([sin_gen, pad_sin], dim=1)
+ cos_gen = cos_gen.view(batch_size, self.sp_size, local_seq_len, -1)
+ sin_gen = sin_gen.view(batch_size, self.sp_size, local_seq_len, -1)
+ cos_gen = cos_gen[:, self.sp_rank, :, :]
+ sin_gen = sin_gen[:, self.sp_rank, :, :]
+
+ cos_gen = cos_gen.unsqueeze(2) # [B, S, 1, D]
+ sin_gen = sin_gen.unsqueeze(2)
+
+ # Run GEN layers. `residual` is threaded so each layer's
+ # input_layernorm and post_attention_layernorm can use the
+ # fused add+rmsnorm path instead of separate add + norm kernels.
+ cached_kv_for_key = self.cached_kv[cache_key]
+ residual: torch.Tensor | None = None
+ for i, layer in enumerate(self.gen_layers):
+ k_und, v_und = cached_kv_for_key[i]
+ k_und = k_und[:, :max_real_len]
+ v_und = v_und[:, :max_real_len]
+ hidden_gen, residual = layer(
+ hidden_gen,
+ k_und,
+ v_und,
+ cos_gen,
+ sin_gen,
+ residual=residual,
+ )
+
+ # Collapse the trailing residual carry. RMSNorm and the linear
+ # projection that follow are per-token, so we run them on the
+ # local shard and only gather the (much smaller) patch-space
+ # output. With patch_latent_dim ~= hidden_size / 21 for cosmos3,
+ # this cuts the post-loop SP collective bandwidth ~21x.
+ hidden_gen = hidden_gen + residual
+ hidden_gen = self.norm_moe_gen(hidden_gen)
+ output, _ = self.llm2vae(hidden_gen)
+
+ if sequence_shard_enabled:
+ output = sequence_model_parallel_all_gather(output, dim=1)
+ if seq_shard_pad > 0:
+ output = output[:, :seq_len_orig, :]
+
+ return self.unpatchify(output, T, H, W)
+
+ def preprocess_loaded_state_dict(
+ self, iterator: Iterable[tuple[str, torch.Tensor]]
+ ) -> Iterator[tuple[str, torch.Tensor]]:
+ # ModelOpt FP8 emits a 0-d per-tensor scale per source Linear. Where
+ # sources fuse into a single MergedColumnParallelLinear (Q/K/V into
+ # to_qkv, gate/up into gate_up_proj), the FP8 weights of each shard
+ # are quantized against their own scale. Naively concatenating the
+ # FP8 bytes and applying a single fused scale at runtime yields noise
+ # (the K/V tiles get dequant'd with the wrong factor).
+ #
+ # Fix: per fused Linear, dequant each FP8 shard with its own scale,
+ # pick max as the fused scale, requant each shard against the max,
+ # then concat the requantized FP8 bytes. input_scale is shared across
+ # shards (same activation tensor), so just take max — no requant
+ # needed.
+ mapping_fn = get_param_names_mapping(self.param_names_mapping)
+ pending: dict[str, dict[str, dict[int, torch.Tensor]]] = {}
+ expected_count: dict[str, int] = {}
+
+ def _try_emit(linear_target: str):
+ groups = pending.get(linear_target, {})
+ n = expected_count.get(linear_target)
+ if n is None:
+ return
+ weights = groups.get("weight", {})
+ w_scales = groups.get("weight_scale", {})
+ i_scales = groups.get("input_scale", {})
+ if len(weights) != n or len(w_scales) != n:
+ return
+ saw_input_scale = bool(i_scales)
+ if saw_input_scale and len(i_scales) != n:
+ return
+ scales_t = torch.stack([w_scales[i].reshape(()) for i in range(n)])
+ max_w_scale = scales_t.max()
+ rescaled = []
+ for i in range(n):
+ w_fp8 = weights[i]
+ original_scale = w_scales[i].reshape(()).to(torch.float32)
+ w_dequant = w_fp8.to(torch.float32) * original_scale
+ w_requant = (
+ (w_dequant / max_w_scale.to(torch.float32))
+ .clamp(-448.0, 448.0)
+ .to(torch.float8_e4m3fn)
+ )
+ rescaled.append(w_requant)
+ merged_weight = torch.cat(rescaled, dim=0)
+ pending.pop(linear_target, None)
+ expected_count.pop(linear_target, None)
+ yield linear_target + ".weight", merged_weight
+ yield linear_target + ".weight_scale", max_w_scale
+ if saw_input_scale:
+ in_t = torch.stack([i_scales[i].reshape(()) for i in range(n)])
+ yield linear_target + ".input_scale", in_t.max()
+
+ for name, tensor in iterator:
+ target_name, merge_index, num_to_merge = mapping_fn(name)
+ if num_to_merge is None:
+ yield target_name, tensor
+ continue
+ suffix = None
+ for candidate in ("weight_scale", "input_scale", "weight"):
+ if target_name.endswith("." + candidate):
+ suffix = candidate
+ break
+ if suffix is None:
+ yield name, tensor
+ continue
+ if suffix == "weight" and tensor.dtype != torch.float8_e4m3fn:
+ yield name, tensor
+ continue
+ linear_target = target_name[: -(len(suffix) + 1)]
+ pending.setdefault(linear_target, {}).setdefault(suffix, {})[
+ merge_index
+ ] = tensor
+ expected_count[linear_target] = num_to_merge
+ yield from _try_emit(linear_target)
+
+ def post_load_weights(self, target_dtype: torch.dtype = torch.bfloat16) -> None:
+ """Cast non-quantized parameters to their preferred dtypes and rebuild
+ meta-device buffers.
+
+ Time-embedder stays in float32 for numerical stability; embeddings and
+ the VAE/LLM bridge linears go to ``target_dtype``. Quantized modules
+ (e.g. FP8 from a ModelOpt export) are skipped — calling ``.to(dtype)``
+ on them would cast their FP8 weights back to BF16/FP32 and break the
+ quant kernels.
+
+ Also re-materializes the RoPE ``inv_freq`` buffer, which can land on
+ the meta device when the model is constructed under a meta context.
+ """
+ # Get the actual device from a loaded parameter
+ device = next(self.parameters()).device
+
+ # Recompute RoPE inv_freq buffer on the correct device
+ # This is needed because model is created in meta device context
+ rotary_emb = self.language_model.rotary_emb
+ if rotary_emb.inv_freq.is_meta:
+ dim = rotary_emb.head_dim
+ rope_theta = 5000000.0 # From config
+ inv_freq = 1.0 / (
+ rope_theta
+ ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
+ )
+ rotary_emb.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ # Collect all quantized submodules so we can skip them during
+ # dtype casts. Calling .to(dtype) on a quantized layer would
+ # cast FP8 weights back to BF16/FP32, breaking the quant kernels.
+ quantized_modules: set[int] = {
+ id(m)
+ for m in self.modules()
+ if hasattr(m, "quant_method")
+ and not isinstance(m.quant_method, UnquantizedLinearMethod)
+ }
+
+ def _is_quantized(module: torch.nn.Module) -> bool:
+ return id(module) in quantized_modules
+
+ def _cast_direct(module: torch.nn.Module, dtype: torch.dtype) -> None:
+ """Cast only the module's own parameters and buffers, not its
+ children. This avoids the recursive `.to()` which would cast
+ quantized (FP8) weights back to BF16/FP32."""
+ for key, param in module._parameters.items():
+ if param is not None:
+ module._parameters[key] = torch.nn.Parameter(
+ param.data.to(dtype=dtype), requires_grad=False
+ )
+ for key, buf in module._buffers.items():
+ if buf is not None:
+ module._buffers[key] = buf.to(dtype=dtype)
+
+ # Time embedder should stay in float32 for numerical stability.
+ # Cast only non-quantized submodules' own params (non-recursive).
+ for module in self.time_embedder.modules():
+ if not _is_quantized(module):
+ _cast_direct(module, torch.float32)
+
+ # Ensure embeddings and projections are in target dtype
+ self.language_model.embed_tokens.to(target_dtype)
+ for module in self.vae2llm.modules():
+ if not _is_quantized(module):
+ _cast_direct(module, target_dtype)
+ for module in self.llm2vae.modules():
+ if not _is_quantized(module):
+ _cast_direct(module, target_dtype)
+
+ # Convert RMSNorm layers to target dtype
+ for module in self.modules():
+ if isinstance(module, RMSNorm):
+ module.to(target_dtype)
+
+
+EntryClass = Cosmos3OmniTransformer
diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py
index cca8b81b0..7d6c77f1e 100644
--- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py
+++ b/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":
diff --git a/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py
new file mode 100644
index 000000000..db667dd02
--- /dev/null
+++ b/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
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py
new file mode 100644
index 000000000..fed13f075
--- /dev/null
+++ b/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,
+ )
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py
new file mode 100644
index 000000000..a9815aeaf
--- /dev/null
+++ b/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
diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
index 1d821e5e9..1ec4c34a2 100644
--- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
+++ b/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)
diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py
new file mode 100644
index 000000000..db021484f
--- /dev/null
+++ b/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()