[feat]: Cosmos3 checkpoint conversion script (cosmos3_convert.py)

scripts/checkpoint_conversion/cosmos3_convert.py wraps Cosmos3OmniDiffusersPipeline._remap_ckpt_key. convert_state_dict(src) returns (new_state, skipped, unmapped) for shard-by-shard remap; convert_checkpoint_dir reads *.safetensors shards and writes a consolidated FastVideo-format state dict (sharded index.json TODO when real weights land). smoke_test() exercises all 14 representative remap branches synthetically (no real weights needed) and returns nonzero exit on failure; runs in CI as a remap-drift guard. Usage: python scripts/checkpoint_conversion/cosmos3_convert.py --smoke-test.

fastvideo/configs/models/dits/cosmos3.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 VFM Transformer FastVideo dataclass configs (Phase 2b.1 scaffold).
+
+Mirrors the architectural defaults of
+``vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py::Cosmos3VFMTransformer``
+(see lines 910-977 of the vllm-omni reference at HEAD ``8536f5b1421f``).
+
+Phase 2b.1 ships only the dataclass surface and skeleton tensor shapes that
+``Cosmos3VFMTransformer.__init__`` needs to build a parameter-named module
+tree. The ``param_names_mapping`` table is left empty here; the checkpoint
+remap lives in ``scripts/checkpoint_conversion/cosmos3_convert.py`` (Phase 5)
+and the FastVideo pipeline's ``_remap_ckpt_key`` (Phase 2b.2+).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_cosmos3_transformer_block(name: str, module) -> bool:
+    """FSDP shard boundary for Cosmos3 UND/GEN decoder layers.
+
+    Mirrors ``Cosmos3VFMTransformer._is_transformer_block`` (reference line
+    870). Matches ``gen_layers.{i}`` and ``language_model.layers.{i}`` where
+    ``{i}`` is a non-negative integer.
+    """
+    del module
+    if "gen_layers" not in name and "language_model.layers" not in name:
+        return False
+    return name.split(".")[-1].isdigit()
+
+
+@dataclass
+class Cosmos3ArchConfig(DiTArchConfig):
+    """Architecture config for Cosmos3 VFM Transformer (UND + GEN pathways).
+
+    Defaults match the upstream Cosmos3-Nano constructor defaults at
+    ``transformer_cosmos3.py`` lines 910-934.
+    """
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_cosmos3_transformer_block])
+
+    # ``param_names_mapping`` is intentionally empty for Phase 2b.1. The
+    # canonical Cosmos3 checkpoint remap is owned by the pipeline's
+    # ``_remap_ckpt_key`` (vllm-omni reference at pipeline_cosmos3.py:319-409)
+    # and is ported to FastVideo in a later phase (the conversion script under
+    # ``scripts/checkpoint_conversion/``). Leaving this empty makes the DiT
+    # module tree the single source of truth for parameter names during
+    # parity-test development.
+    param_names_mapping: dict = field(default_factory=dict)
+
+    # ---- Backbone ----
+    hidden_size: int = 4096
+    num_hidden_layers: int = 36
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    head_dim: int = 128
+    intermediate_size: int = 12288
+    vocab_size: int = 151936
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 5_000_000.0
+    mrope_section: list[int] = field(default_factory=lambda: [24, 20, 20])
+
+    # ---- VAE / patch geometry ----
+    latent_patch_size: int = 2
+    latent_channel: int = 48
+
+    # ---- Diffusion conditioning ----
+    timestep_scale: float = 0.001
+
+    # ---- Temporal modulation ----
+    base_fps: float = 24.0
+    temporal_compression_factor: int = 4
+    enable_fps_modulation: bool = True
+    temporal_modality_margin: int = 15000
+
+    # ---- BaseDiT bookkeeping ----
+    in_channels: int = 48
+    out_channels: int = 48
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        # Mirror the BaseDiT contract: ``num_channels_latents`` matches the
+        # VAE latent channel count for video DiTs.
+        self.num_channels_latents = self.latent_channel
+        if not self.out_channels:
+            self.out_channels = self.in_channels
+
+
+@dataclass
+class Cosmos3VideoConfig(DiTConfig):
+    """Pipeline-level Cosmos3 DiT config (T2V / I2V / T2I share this surface)."""
+
+    arch_config: DiTArchConfig = field(default_factory=Cosmos3ArchConfig)
+    prefix: str = "Cosmos3"

fastvideo/configs/pipelines/cosmos3.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 pipeline configuration.
+
+Port reference: vllm-omni PR #3454 (https://github.com/vllm-project/vllm-omni/pull/3454)
+HEAD ``8536f5b1421f``. Cosmos3 is structurally different from Cosmos 2.5:
+
+- Dual-stack DiT (UND + GEN) lives entirely inside ``Cosmos3VFMTransformer``.
+- No separate text encoder — ``Cosmos3LanguageModel`` is inside the DiT, so
+  ``text_encoder_configs`` is the empty tuple. The tokenizer is loaded by the
+  pipeline's ``__init__`` from the checkpoint's ``text_tokenizer/`` subfolder;
+  that is an instance-level concern, not a config-level concern.
+- VAE is ``DistributedAutoencoderKLWan`` (same as Cosmos 2.5), so the
+  ``Cosmos25VAEConfig`` is reused as-is.
+- T2I default ``flow_shift`` is 3.0 (set per-request by ``_set_flow_shift``);
+  T2V/I2V use the engine-init default of 1.0 baked into this config.
+"""
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits.cosmos3 import (Cosmos3ArchConfig, Cosmos3VideoConfig)
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.vaes import Cosmos25VAEConfig  # Cosmos3 reuses the Wan-based VAE
+from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+def _identity_preprocess_text(prompt: str) -> str:
+    return prompt
+
+
+@dataclass
+class Cosmos3Config(PipelineConfig):
+    """Configuration for the Cosmos3 video generation pipeline (T2V/I2V/T2I).
+
+    Tier A scope: architectural skeleton. Real weights from ``nvidia/Cosmos3-Nano``
+    are not yet published; this config is shaped to be activatable when they
+    land without further config changes.
+    """
+
+    dit_config: DiTConfig = field(default_factory=lambda: Cosmos3VideoConfig(arch_config=Cosmos3ArchConfig()))
+
+    vae_config: VAEConfig = field(default_factory=Cosmos25VAEConfig)
+
+    # No separate text encoder: Cosmos3LanguageModel lives inside the DiT.
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=tuple)
+
+    preprocess_text_funcs: tuple[Callable[[str], str],
+                                 ...] = field(default_factory=lambda: (_identity_preprocess_text, ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(default_factory=tuple)
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "bf16"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=tuple)
+
+    embedded_cfg_scale: float = 0.0
+    flow_shift: float = 1.0  # T2V/I2V engine-init; T2I uses 3.0 set per-request
+
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True

fastvideo/models/dits/cosmos3.py

@@ -0,0 +1,454 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 VFM Transformer (Phase 2b.1 skeleton).
+
+This module lands the parameter-named module tree of
+``Cosmos3VFMTransformer`` together with two standalone math utilities used
+by both the FastVideo port and the Phase 2a parity tests:
+
+* ``compute_mrope_position_ids_text`` / ``compute_mrope_position_ids_vision`` —
+  the unified-3D mRoPE position ID generators ported verbatim from
+  ``vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py`` lines
+  113-177 (HEAD ``8536f5b1421f``).
+
+* ``Cosmos3VFMTransformer.patchify`` /
+  ``Cosmos3VFMTransformer.unpatchify`` — the
+  ``[B,C,T,H,W] <-> [B, T*hp*wp, p*p*C]`` patch tokenizer, ported from
+  the same reference lines 1009-1036.
+
+Every per-layer ``forward()`` raises ``NotImplementedError("Phase 2b.2")``
+because the layer math (RoPE application, QK-norm, GQA SDPA, UND/GEN
+attention plumbing) is intentionally deferred to Phase 2b.2. Only the
+module composition + tensor shape contract is exercised here, enough to
+make 4 Tier-A parity tests pass:
+
+* ``test_cosmos3_mrope_parity.py::test_compute_mrope_position_ids_text_and_vision``
+* ``test_cosmos3_patchify_unpatchify_parity.py::test_patchify_unpatchify_roundtrip``
+* ``test_cosmos3_patchify_unpatchify_parity.py::test_patchify_default_patch_size``
+* ``test_cosmos3_state_dict_keys.py::test_fastvideo_cosmos3_dit_module_tree_param_names``
+
+The constructor signature mirrors the upstream
+``Cosmos3VFMTransformer(od_config, *, temporal_compression_factor=None)``
+contract so the Phase 2a tests (which were authored against the upstream
+shape) can drive it directly. Phase 2b.2+ will adapt this to FastVideo's
+``BaseDiT(config, hf_config)`` loader path via the ``TransformerLoader``.
+"""
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.configs.models.dits.cosmos3 import Cosmos3ArchConfig, Cosmos3VideoConfig
+from fastvideo.models.dits.base import BaseDiT
+
+
+EntryClass = ["Cosmos3VFMTransformer"]
+
+
+def _tf_config_get(config: Any, key: str, default: Any) -> Any:
+    """Read a value from a dict, dataclass, or simple namespace.
+
+    Mirrors ``transformer_cosmos3._tf_config_get`` at reference line 77.
+    """
+    if config is None:
+        return default
+    if hasattr(config, "get"):
+        return config.get(key, default)
+    return getattr(config, key, default)
+
+
+def compute_mrope_position_ids_text(
+    num_tokens: int,
+    temporal_offset: int,
+) -> tuple[torch.Tensor, int]:
+    """Generate 3D mRoPE position IDs for text tokens.
+
+    Text tokens broadcast a single monotonically-increasing position-ID
+    sequence across all three (t, h, w) axes. Verbatim port of
+    ``transformer_cosmos3.compute_mrope_position_ids_text`` (reference
+    lines 113-124).
+    """
+    ids = torch.arange(num_tokens, dtype=torch.long) + 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,
+    fps: float | None = None,
+    base_fps: float = 24.0,
+    temporal_compression_factor: int = 4,
+    base_temporal_compression_factor: int | None = None,
+    enable_fps_modulation: bool = True,
+    start_frame_offset: int = 0,
+) -> tuple[torch.Tensor, int]:
+    """Generate 3D mRoPE position IDs for vision tokens.
+
+    Builds a ``(t, h, w)`` position grid (Qwen3-VL style, spatial indices
+    reset per temporal segment) flattened in t-major order. Optionally
+    modulates the temporal axis by ``base_fps / tcf * (1 / (fps / tcf))``
+    so two clips at different FPS retain wall-clock-aligned temporal
+    positions. Verbatim port of
+    ``transformer_cosmos3.compute_mrope_position_ids_vision`` (reference
+    lines 127-177).
+    """
+    fps_modulation = enable_fps_modulation and fps is not None
+
+    if fps_modulation:
+        assert fps is not None
+        tps = fps / temporal_compression_factor
+        effective_base_tcf = (base_temporal_compression_factor
+                              if base_temporal_compression_factor is not None
+                              else temporal_compression_factor)
+        base_tps = base_fps / effective_base_tcf
+        frame_indices = torch.arange(grid_t, dtype=torch.float32)
+        t_index = (((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset)
+                   .view(-1, 1).expand(-1, grid_h * grid_w).flatten())
+    else:
+        t_index = (torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten()
+                   + int(temporal_offset) + start_frame_offset)
+
+    h_index = (torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten())
+    w_index = (torch.arange(grid_w, dtype=torch.long).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.floor(mrope_ids.max().item()) + 1
+    return mrope_ids, next_offset
+
+
+class _Qwen3VLTextRMSNorm(nn.Module):
+    """Qwen3-VL / T5-style RMSNorm with an ``eps`` and a learnable ``weight``.
+
+    Mirrors ``Qwen3VLTextRMSNorm`` at reference line 89-107 so checkpoint
+    keys match the canonical ``*.norm*.weight`` / ``*.{q,k}_norm.weight``
+    parameter naming.
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.bfloat16) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError("Phase 2b.2")
+
+
+class _TimestepEmbedder(nn.Module):
+    """Sinusoidal-to-vector timestep embedder used by the GEN pathway.
+
+    Skeleton of ``TimestepEmbedder`` at reference line 257-283. The
+    ``linear_1``/``linear_2`` parameter names match the upstream
+    checkpoint so the Phase 5 remap can copy weights without renaming.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        frequency_embedding_size: int = 256,
+        max_period: int = 10000,
+        target_dtype: torch.dtype = torch.bfloat16,
+    ) -> None:
+        super().__init__()
+        self.linear_1 = nn.Linear(frequency_embedding_size, hidden_size, bias=True)
+        self.act = nn.SiLU()
+        self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True)
+        self.frequency_embedding_size = frequency_embedding_size
+        self.hidden_size = hidden_size
+
+        half = frequency_embedding_size // 2
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=target_dtype) / half)
+        self.register_buffer("freqs", freqs, persistent=False)
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError("Phase 2b.2")
+
+
+class _Cosmos3GatedMLP(nn.Module):
+    """Gated-MLP block with ``gate_proj`` / ``up_proj`` / ``down_proj`` names.
+
+    Skeleton of ``Cosmos3GatedMLP`` at reference line 288-326. The three
+    linear-layer parameter names are load-bearing for the checkpoint
+    remap; Phase 2b.2 will swap them for FastVideo-native
+    ``ReplicatedLinear`` shards.
+    """
+
+    def __init__(self, hidden_size: int, intermediate_size: int) -> None:
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError("Phase 2b.2")
+
+
+class _Cosmos3SkeletonAttention(nn.Module):
+    """Shared skeleton for UND causal- and GEN cross-attention.
+
+    Exposes the q/k/v/o projections at the upstream parameter-name
+    locations (``*.{q,k,v,o}_proj`` and ``*.{q,k}_norm``) with the
+    GQA-correct in/out dimensions
+    (``q_proj`` writes ``num_attention_heads * head_dim`` channels;
+    ``k_proj``/``v_proj`` write ``num_key_value_heads * head_dim``
+    channels). Phase 2b.2 will replace the bare ``nn.Linear`` shells with
+    ``ColumnParallelLinear``-equivalent FastVideo layers and implement
+    the RoPE + SDPA forward.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        rms_norm_eps: float,
+        dtype: torch.dtype,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_attention_heads
+        self.num_kv_heads = num_key_value_heads
+        self.head_dim = head_dim
+
+        self.q_proj = nn.Linear(hidden_size, num_attention_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(num_attention_heads * head_dim, hidden_size, bias=False)
+
+        self.q_norm = _Qwen3VLTextRMSNorm(head_dim, eps=rms_norm_eps, dtype=dtype)
+        self.k_norm = _Qwen3VLTextRMSNorm(head_dim, eps=rms_norm_eps, dtype=dtype)
+
+    def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        raise NotImplementedError("Phase 2b.2")
+
+
+class _Cosmos3UndDecoderLayer(nn.Module):
+    """UND (understanding) decoder layer skeleton.
+
+    The submodule names ``self_attn``, ``input_layernorm``,
+    ``post_attention_layernorm``, ``mlp`` mirror the upstream layer
+    composition at reference line 624-678.
+    """
+
+    def __init__(self, arch: Cosmos3ArchConfig, dtype: torch.dtype) -> None:
+        super().__init__()
+        self.self_attn = _Cosmos3SkeletonAttention(
+            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,
+            rms_norm_eps=arch.rms_norm_eps,
+            dtype=dtype,
+        )
+        self.input_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype)
+        self.post_attention_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype)
+        self.mlp = _Cosmos3GatedMLP(arch.hidden_size, arch.intermediate_size)
+
+    def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        raise NotImplementedError("Phase 2b.2")
+
+
+class _Cosmos3GenDecoderLayer(nn.Module):
+    """GEN (generation) decoder layer skeleton.
+
+    Substitutes ``cross_attention`` for ``self_attn`` (mirrors upstream
+    reference line 681-751); otherwise identical to the UND layer in
+    naming. The ``layer_idx`` is captured for the eventual cached-KV
+    path so Phase 2b.2 can route per-layer UND keys/values correctly.
+    """
+
+    def __init__(self, layer_idx: int, arch: Cosmos3ArchConfig, dtype: torch.dtype) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.cross_attention = _Cosmos3SkeletonAttention(
+            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,
+            rms_norm_eps=arch.rms_norm_eps,
+            dtype=dtype,
+        )
+        self.input_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype)
+        self.post_attention_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype)
+        self.mlp = _Cosmos3GatedMLP(arch.hidden_size, arch.intermediate_size)
+
+    def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        raise NotImplementedError("Phase 2b.2")
+
+
+class _Cosmos3LanguageModel(nn.Module):
+    """Understanding-pathway container.
+
+    Holds ``embed_tokens``, ``layers``, and a final ``norm``. The
+    submodule names mirror the upstream reference at line 757-831 so the
+    checkpoint remap target keys ``language_model.embed_tokens.*``,
+    ``language_model.layers.{i}.*``, and ``language_model.norm.*`` are
+    populated as-is by the FastVideo module tree.
+    """
+
+    def __init__(self, arch: Cosmos3ArchConfig, dtype: torch.dtype) -> None:
+        super().__init__()
+        self.embed_tokens = nn.Embedding(arch.vocab_size, arch.hidden_size)
+        self.layers = nn.ModuleList(
+            [_Cosmos3UndDecoderLayer(arch, dtype) for _ in range(arch.num_hidden_layers)])
+        self.norm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype)
+
+    def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        raise NotImplementedError("Phase 2b.2")
+
+
+class Cosmos3VFMTransformer(BaseDiT):
+    """Cosmos3 VFM Transformer — UND language model + GEN denoising layers.
+
+    Phase 2b.1 ships only the module-tree skeleton and the two pure math
+    utilities (``patchify``/``unpatchify``). All per-layer ``forward``
+    methods raise ``NotImplementedError("Phase 2b.2")``; the DiT-level
+    ``forward`` likewise. The constructor signature
+    ``(od_config, *, temporal_compression_factor=None)`` matches the
+    upstream reference at ``transformer_cosmos3.py`` line 899 so the
+    Phase 2a parity tests can drive it directly.
+
+    Phase 2b.2 will:
+      * Replace the per-layer ``nn.Linear``s with FastVideo-native
+        ``ReplicatedLinear`` / ``DistributedAttention`` plumbing.
+      * Implement ``forward`` (UND once, GEN per-step with cached K/V).
+      * Adapt the constructor to the FastVideo
+        ``BaseDiT(config, hf_config)`` loader contract via
+        ``TransformerLoader``.
+    """
+
+    _fsdp_shard_conditions = Cosmos3VideoConfig().arch_config._fsdp_shard_conditions
+    _compile_conditions = Cosmos3VideoConfig().arch_config._compile_conditions
+    param_names_mapping = Cosmos3VideoConfig().arch_config.param_names_mapping
+    reverse_param_names_mapping: dict[str, str] = {}
+
+    def __init__(
+        self,
+        od_config: object | None = None,
+        *,
+        temporal_compression_factor: int | None = None,
+    ) -> None:
+        nn.Module.__init__(self)
+
+        model_config = getattr(od_config, "tf_model_config", None) if od_config is not None else None
+        rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {}
+
+        self.hidden_size = int(_tf_config_get(model_config, "hidden_size", 4096))
+        self.num_hidden_layers = int(_tf_config_get(model_config, "num_hidden_layers", 36))
+        self.num_attention_heads = int(_tf_config_get(model_config, "num_attention_heads", 32))
+        self.num_key_value_heads = int(_tf_config_get(model_config, "num_key_value_heads", 8))
+        self.head_dim = int(_tf_config_get(model_config, "head_dim", 128))
+        self.intermediate_size = int(_tf_config_get(model_config, "intermediate_size", 12288))
+        self.vocab_size = int(_tf_config_get(model_config, "vocab_size", 151936))
+        self.rms_norm_eps = float(_tf_config_get(model_config, "rms_norm_eps", 1e-6))
+        self.rope_theta = float(_tf_config_get(model_config, "rope_theta", 5_000_000))
+        self.mrope_section = list(rope_scaling.get("mrope_section", [24, 20, 20]))
+        self.latent_patch_size = int(_tf_config_get(model_config, "latent_patch_size", 2))
+        self.latent_channel_size = int(_tf_config_get(model_config, "latent_channel", 48))
+        self.timestep_scale = float(_tf_config_get(model_config, "timestep_scale", 0.001))
+        self.base_fps = float(_tf_config_get(model_config, "base_fps", 24.0))
+        if temporal_compression_factor is None:
+            resolved_tcf = int(_tf_config_get(model_config, "temporal_compression_factor", 4))
+        else:
+            resolved_tcf = int(temporal_compression_factor)
+        self.temporal_compression_factor = resolved_tcf
+        self.enable_fps_modulation = bool(_tf_config_get(model_config, "enable_fps_modulation", True))
+        self.temporal_modality_margin = int(
+            _tf_config_get(model_config, "unified_3d_mrope_temporal_modality_margin", 15000))
+
+        self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size
+        self.num_channels_latents = self.latent_channel_size
+
+        dtype = getattr(od_config, "dtype", torch.bfloat16) if od_config is not None else torch.bfloat16
+
+        arch = Cosmos3ArchConfig(
+            hidden_size=self.hidden_size,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads,
+            num_key_value_heads=self.num_key_value_heads,
+            head_dim=self.head_dim,
+            intermediate_size=self.intermediate_size,
+            vocab_size=self.vocab_size,
+            rms_norm_eps=self.rms_norm_eps,
+            rope_theta=self.rope_theta,
+            mrope_section=self.mrope_section,
+            latent_patch_size=self.latent_patch_size,
+            latent_channel=self.latent_channel_size,
+            timestep_scale=self.timestep_scale,
+            base_fps=self.base_fps,
+            temporal_compression_factor=self.temporal_compression_factor,
+            enable_fps_modulation=self.enable_fps_modulation,
+            temporal_modality_margin=self.temporal_modality_margin,
+            in_channels=self.latent_channel_size,
+            out_channels=self.latent_channel_size,
+        )
+
+        self.language_model = _Cosmos3LanguageModel(arch, dtype)
+        self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size)
+        self.llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim)
+        self.time_embedder = _TimestepEmbedder(self.hidden_size, target_dtype=dtype)
+        self.gen_layers = nn.ModuleList(
+            [_Cosmos3GenDecoderLayer(i, arch, dtype) for i in range(self.num_hidden_layers)])
+        self.norm_moe_gen = _Qwen3VLTextRMSNorm(self.hidden_size, eps=self.rms_norm_eps, dtype=dtype)
+
+    def _pad_to_patch_size(self, h: int, w: int) -> tuple[int, int, int, int]:
+        """Return ``(hp, wp, H_padded, W_padded)`` for ``latent_patch_size`` padding.
+
+        Mirrors ``Cosmos3VFMTransformer._pad_to_patch_size`` at reference
+        lines 1002-1007.
+        """
+        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:
+        """``[B, C, t, h, w] -> [B, t*hp*wp, p*p*C]``.
+
+        Pads ``h``/``w`` up to a multiple of ``latent_patch_size`` before
+        reshaping. Verbatim port of
+        ``Cosmos3VFMTransformer.patchify`` (reference lines 1009-1021).
+        """
+        batch_size = latents.shape[0]
+        p = self.latent_patch_size
+        c = self.latent_channel_size
+        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(batch_size, c, t, hp, p, wp, p)
+        x = x.permute(0, 2, 3, 5, 4, 6, 1)
+        return x.reshape(batch_size, t * hp * wp, p * p * c)
+
+    def unpatchify(self, tokens: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor:
+        """``[B, t*hp*wp, p*p*C] -> [B, C, t, h, w]``, cropping ``h``/``w`` padding.
+
+        Verbatim port of ``Cosmos3VFMTransformer.unpatchify`` (reference
+        lines 1023-1036).
+        """
+        batch_size = tokens.shape[0]
+        p = self.latent_patch_size
+        c = self.latent_channel_size
+        hp, wp, h_padded, w_padded = self._pad_to_patch_size(h, w)
+
+        x = tokens.reshape(batch_size, t, hp, wp, p, p, c)
+        x = x.permute(0, 6, 1, 2, 4, 3, 5)
+        x = x.reshape(batch_size, c, t, h_padded, w_padded)
+
+        if h_padded != h or w_padded != w:
+            x = x[:, :, :, :h, :w]
+        return x
+
+    def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        raise NotImplementedError("Phase 2b.2")

fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py

@@ -0,0 +1,605 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 pipeline (Phase 2f.2 partial).
+
+Upstream reference:
+    ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py`` from
+    vllm-omni HEAD ``8536f5b1421f78c7df06af6d96fa195c1ceb6384``.
+
+Phase 2f.1 landed the checkpoint-remap + scheduler-shift surface:
+
+* ``_remap_ckpt_key`` — ported verbatim from upstream lines 319-409. This is a
+  pure key-translation function with no FastVideo-specific adaptations; the
+  checkpoint converter at ``scripts/checkpoint_conversion/cosmos3_convert.py``
+  (Phase 5) will reuse it.
+* ``_set_flow_shift`` — ported from upstream lines 498-512 with lazy scheduler
+  construction so the scheduler-default-parity tests can exercise the method on
+  ``__new__``-allocated instances before ``__init__`` is wired.
+* ``_engine_init_flow_shift`` — exposed as a class attribute (default ``1.0``)
+  so ``hasattr(__new__(...), "_engine_init_flow_shift")`` succeeds without
+  running ``__init__``.
+
+Phase 2f.2 adds the runtime call graph (this commit):
+
+* ``diffuse`` — sequential 3-mode CFG denoising loop ported from upstream
+  lines 883-1033. The CFG-Parallel branch (cfg_parallel=True) is deferred
+  until FastVideo's classifier-free-guidance-world-size plumbing lands;
+  ``_cfg_parallel_active`` returns False so only the sequential and no-CFG
+  paths execute.
+* ``forward`` — request parsing + T2I/T2V/I2V mode dispatch + flow-shift
+  selection + diffusion driver + decode, ported from upstream lines
+  1037-1206 (single-prompt path; the T2I num_outputs_per_prompt > 1 branch
+  is preserved but uses the unimplemented ``_prepare_latents`` stub).
+* ``_is_t2i_request`` / ``_get_sp_param`` / ``_cfg_parallel_active`` —
+  ported from upstream lines 452-496.
+* Negative-prompt constants — ported from upstream lines 51-61.
+
+Helper methods ``_format_and_tokenize_prompts``, ``_prepare_latents``,
+``_prepare_latents_i2v``, ``_set_scheduler_timesteps``, and
+``_decode_latents`` are defined as ``NotImplementedError`` stubs so the
+pipeline call-graph tests can monkey-patch them without an
+``AttributeError``; their real implementations land in Phase 2c
+(tokenizer) and Phase 2f.3+ (latent prep, scheduler timesteps, VAE decode).
+
+``__init__`` and ``create_pipeline_stages`` remain ``NotImplementedError``
+stubs — full module/stage wiring is left to a future phase.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+import torch.nn as nn
+from diffusers import UniPCMultistepScheduler
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+
+logger = init_logger(__name__)
+
+# Negative-prompt constants ported verbatim from upstream pipeline_cosmos3.py:51-61.
+COSMOS3_DEFAULT_NEGATIVE_PROMPT = ""
+COSMOS3_VIDEO_NEGATIVE_PROMPT = (
+    "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, "
+    "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, "
+    "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, "
+    "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, "
+    "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. "
+    "Overall, the video is of poor quality.")
+COSMOS3_T2V_NEGATIVE_PROMPT = COSMOS3_VIDEO_NEGATIVE_PROMPT
+COSMOS3_I2V_NEGATIVE_PROMPT = COSMOS3_VIDEO_NEGATIVE_PROMPT
+
+
+class Cosmos3OmniDiffusersPipeline(nn.Module, ComposedPipelineBase):
+    """Cosmos3 T2V/I2V/T2I pipeline (skeleton).
+
+    Phase 2f.1 lands only the checkpoint-remap + scheduler-shift surface needed
+    by ``test_cosmos3_state_dict_keys.py`` and
+    ``test_cosmos3_scheduler_default_parity.py``. Full wiring (modules,
+    stages, ``diffuse()``, mode dispatch) is deferred to Phase 2f.2.
+
+    Inherits from both ``nn.Module`` (matching upstream
+    ``pipeline_cosmos3.py:197-199`` — the conftest fixture
+    ``make_cosmos3_pipeline`` calls ``nn.Module.__init__`` directly) and
+    ``ComposedPipelineBase`` (FastVideo's pipeline convention — sibling
+    ``cosmos2_5_pipeline.py``). Phase 2f.2 can collapse either branch if
+    one proves unnecessary once the full pipeline lands.
+
+    The scheduler-parity tests build instances via
+    ``Cosmos3OmniDiffusersPipeline.__new__(...)`` and rely on:
+
+    * ``hasattr(instance, "scheduler")`` succeeding before ``_set_flow_shift``
+      is called (satisfied by the ``scheduler = None`` class attribute below);
+    * ``hasattr(instance, "_engine_init_flow_shift")`` succeeding (satisfied
+      by the class-attribute default ``1.0``).
+    """
+
+    # --- Class attributes ---------------------------------------------------
+    #
+    # These exist as class attributes so ``hasattr(__new__(cls), name)`` is
+    # True without ``__init__`` having run. ``__init__`` (Phase 2f.2) will
+    # shadow them with instance attributes derived from the loaded scheduler.
+    _engine_init_flow_shift: float = 1.0
+    scheduler: Any = None
+    _base_scheduler_config: Any = None
+    _current_flow_shift: float | None = None
+
+    # -- Weight loading -----------------------------------------------------
+
+    @staticmethod
+    def _remap_ckpt_key(key: str) -> str | None:
+        """Remap a Diffusers transformer key to the model parameter namespace.
+
+        Checkpoint keys arrive with a synthetic ``transformer.`` prefix from
+        ``weights_sources``. The source checkpoint itself uses the Diffusers
+        transformer namespace: top-level projections plus ``model.*`` for the
+        Qwen3-VL backbone. UND and GEN components share each layer in the
+        source and are split into separate module lists here.
+
+        Returns the remapped name under ``transformer.``, or ``None`` to skip.
+
+        Ported verbatim from upstream ``pipeline_cosmos3.py`` lines 319-409.
+        """
+        k = key
+        # Strip the weights_sources prefix
+        if k.startswith("transformer."):
+            k = k[len("transformer."):]
+
+        # Top-level generation components.
+        if k.startswith((
+                "vae2llm.",
+                "llm2vae.",
+                "time_embedder.",
+        )):
+            return f"transformer.{k}"
+
+        # Skip lm_head
+        if k.startswith("lm_head."):
+            return None
+
+        # embed_tokens / norm → language_model.*
+        if k.startswith("model.embed_tokens."):
+            return f"transformer.language_model.{k[len('model.'):]}"
+        if k.startswith("model.norm."):
+            return f"transformer.language_model.{k[len('model.'):]}"
+
+        # norm_moe_gen → top level
+        if k.startswith("model.norm_moe_gen."):
+            return f"transformer.{k[len('model.'):]}"
+
+        if not k.startswith("model.layers."):
+            return None
+        k = k[len("model."):]
+
+        if not k.startswith("layers."):
+            return None
+
+        parts = k.split(".", 2)  # ['layers', '{i}', '{rest}']
+        if len(parts) != 3:
+            return None
+        layer_idx = parts[1]
+        rest = parts[2]
+
+        und_lp = f"transformer.language_model.layers.{layer_idx}"
+        gen_lp = f"transformer.gen_layers.{layer_idx}"
+
+        _LAYER_MAP = {
+            # UND attention
+            "self_attn.q_proj.": f"{und_lp}.self_attn.q_proj.",
+            "self_attn.k_proj.": f"{und_lp}.self_attn.k_proj.",
+            "self_attn.v_proj.": f"{und_lp}.self_attn.v_proj.",
+            "self_attn.o_proj.": f"{und_lp}.self_attn.o_proj.",
+            "self_attn.q_norm.": f"{und_lp}.self_attn.q_norm.",
+            "self_attn.k_norm.": f"{und_lp}.self_attn.k_norm.",
+            # GEN attention
+            "self_attn.q_proj_moe_gen.": f"{gen_lp}.cross_attention.q_proj.",
+            "self_attn.k_proj_moe_gen.": f"{gen_lp}.cross_attention.k_proj.",
+            "self_attn.v_proj_moe_gen.": f"{gen_lp}.cross_attention.v_proj.",
+            "self_attn.o_proj_moe_gen.": f"{gen_lp}.cross_attention.o_proj.",
+            "self_attn.q_norm_moe_gen.": f"{gen_lp}.cross_attention.q_norm.",
+            "self_attn.k_norm_moe_gen.": f"{gen_lp}.cross_attention.k_norm.",
+            # Norms
+            "input_layernorm.": f"{und_lp}.input_layernorm.",
+            "post_attention_layernorm.": f"{und_lp}.post_attention_layernorm.",
+            "input_layernorm_moe_gen.": f"{gen_lp}.input_layernorm.",
+            "post_attention_layernorm_moe_gen.": f"{gen_lp}.post_attention_layernorm.",
+            # UND MLP
+            "mlp.gate_proj.": f"{und_lp}.mlp.gate_proj.",
+            "mlp.up_proj.": f"{und_lp}.mlp.up_proj.",
+            "mlp.down_proj.": f"{und_lp}.mlp.down_proj.",
+            # GEN MLP
+            "mlp_moe_gen.gate_proj.": f"{gen_lp}.mlp.gate_proj.",
+            "mlp_moe_gen.up_proj.": f"{gen_lp}.mlp.up_proj.",
+            "mlp_moe_gen.down_proj.": f"{gen_lp}.mlp.down_proj.",
+        }
+
+        for pattern, replacement in _LAYER_MAP.items():
+            if rest.startswith(pattern):
+                suffix = rest[len(pattern):]
+                return replacement + suffix
+
+        return None
+
+    # -- Scheduler control --------------------------------------------------
+
+    def _set_flow_shift(self, target_shift: float) -> None:
+        """Set the UniPC ``flow_shift`` to a concrete target value.
+
+        Adapted from upstream ``pipeline_cosmos3.py`` lines 498-512 with one
+        FastVideo-specific addition: when called on an instance built via
+        ``__new__`` (i.e. ``__init__`` has not run, so
+        ``self._base_scheduler_config is None``), this method lazily
+        constructs a default ``UniPCMultistepScheduler`` rather than
+        rebuilding from a saved config. Phase 2f.2's ``__init__`` will
+        replace that lazy path with a checkpoint-loaded scheduler and
+        snapshot its config into ``_base_scheduler_config`` per upstream.
+
+        Tracking ``self._current_flow_shift`` explicitly is required because
+        the previous mode may have rebuilt the scheduler — we cannot rely on
+        ``self.scheduler.config.flow_shift`` reflecting the last requested
+        target if a rebuild was skipped via the equality check.
+        """
+        target = float(target_shift)
+
+        # Lazy path: no checkpoint-loaded base config yet (Phase 2f.1 test
+        # entry via __new__). Construct a fresh UniPC scheduler at the
+        # requested flow_shift so test assertions can read
+        # ``self.scheduler.config.flow_shift``. Phase 2f.2's __init__ will
+        # load the real scheduler from the checkpoint and overwrite both
+        # ``self.scheduler`` and ``self._base_scheduler_config``.
+        if self._base_scheduler_config is None:
+            self.scheduler = UniPCMultistepScheduler(
+                num_train_timesteps=1000,
+                solver_order=2,
+                prediction_type="flow_prediction",
+                use_flow_sigmas=True,
+                flow_shift=target,
+            )
+            self._base_scheduler_config = self.scheduler.config
+            self._current_flow_shift = target
+            return
+
+        # Rebuild only if the target differs from the current shift.
+        if self._current_flow_shift is not None and target == float(self._current_flow_shift):
+            return
+        self.scheduler = UniPCMultistepScheduler.from_config(self._base_scheduler_config, flow_shift=target)
+        self._current_flow_shift = target
+
+    # -- Request introspection ----------------------------------------------
+
+    @staticmethod
+    def _cfg_parallel_active() -> bool:
+        """Return True when CFG-Parallel is enabled in the current topology.
+
+        Upstream ``pipeline_cosmos3.py:452-457`` queries
+        ``get_classifier_free_guidance_world_size() > 1``. FastVideo's
+        CFG-Parallel plumbing is not wired in Phase 2f.2, so this always
+        returns False; ``diffuse()`` therefore exercises only the sequential
+        CFG path and the no-CFG path.
+        """
+        return False
+
+    @staticmethod
+    def _get_sp_param(sp: Any, key: str, default: Any = None) -> Any:
+        """Read a runtime control from sampling params.
+
+        Ported verbatim from upstream ``pipeline_cosmos3.py:459-481``.
+
+        Order of precedence:
+            1. ``sp.extra_args[key]`` — preferred path; the OpenAI image/video
+               endpoints surface custom controls there.
+            2. direct attribute on ``sp``.
+            3. ``default``.
+        """
+        extra = getattr(sp, "extra_args", None)
+        if isinstance(extra, dict) and extra.get(key) is not None:
+            return extra[key]
+        val = getattr(sp, key, None)
+        if val is not None:
+            return val
+        return default
+
+    @staticmethod
+    def _is_t2i_request(req: Any) -> bool:
+        """Detect text-to-image mode from request-level prompt modalities.
+
+        Ported verbatim from upstream ``pipeline_cosmos3.py:483-496``.
+        Raises ValueError when a prompt requests both image AND video output
+        simultaneously.
+        """
+        if not req.prompts:
+            return False
+        first_prompt = req.prompts[0]
+        modalities = first_prompt.get("modalities", []) if isinstance(first_prompt, dict) else []
+        if modalities is None:
+            modalities = []
+        if isinstance(modalities, str):
+            modalities = [modalities]
+        if "image" in modalities and "video" in modalities:
+            raise ValueError("Cosmos3 prompt modalities cannot request both image and video output.")
+        return "image" in modalities
+
+    # -- Helper stubs (real implementations land in later phases) -----------
+    #
+    # These exist so the pipeline call-graph tests can monkey-patch them on
+    # the instance without an ``AttributeError``. Each raises
+    # ``NotImplementedError`` if invoked outside a test that replaces it.
+
+    def _format_and_tokenize_prompts(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError("_format_and_tokenize_prompts lands in Phase 2c (tokenizer wiring).")
+
+    def _prepare_latents(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError("_prepare_latents lands in Phase 2f.3 (latent preparation).")
+
+    def _prepare_latents_i2v(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError("_prepare_latents_i2v lands in Phase 2f.3 (I2V latent preparation).")
+
+    def _set_scheduler_timesteps(self, num_inference_steps: int) -> None:
+        raise NotImplementedError("_set_scheduler_timesteps lands in Phase 2f.3 (scheduler timesteps).")
+
+    def _decode_latents(self, latents: torch.Tensor) -> Any:
+        raise NotImplementedError("_decode_latents lands in Phase 2f.3 (VAE decode wiring).")
+
+    # -- Denoising loop -----------------------------------------------------
+
+    def diffuse(
+        self,
+        *,
+        latents: torch.Tensor,
+        timesteps: torch.Tensor,
+        cond_ids: torch.Tensor,
+        cond_mask: torch.Tensor,
+        uncond_ids: torch.Tensor,
+        uncond_mask: torch.Tensor,
+        guidance_scale: float,
+        shared_kwargs: dict[str, Any],
+        velocity_mask: torch.Tensor | None = None,
+        image_latent: torch.Tensor | None = None,
+        condition_latents: torch.Tensor | None = None,
+        guidance_interval: tuple[float, float] | None = None,
+    ) -> torch.Tensor:
+        """Sequential 3-mode CFG denoising loop with optional I2V conditioning.
+
+        Ported from upstream ``pipeline_cosmos3.py:883-1033``. The
+        CFG-Parallel branch (upstream lines 953-980) is deferred until
+        FastVideo's classifier-free-guidance-world-size plumbing lands;
+        ``_cfg_parallel_active`` returns False in Phase 2f.2 so only the
+        sequential CFG branch (upstream lines 982-1019) and the no-CFG
+        branch (upstream lines 1021-1031) execute.
+
+        Cosmos3's UND pathway is text-dependent, so sequential CFG keeps
+        separate K/V caches for the conditional and unconditional text
+        forwards and swaps them in before each branch's transformer call.
+
+        I2V conditioning is applied via ``_step``: ``velocity_mask`` zeros
+        frame-0 noise predictions before the scheduler step, and
+        ``image_latent`` is re-injected into frame 0 after each step
+        (UniPC's predictor-corrector rescales the sample, so zero velocity
+        alone does not preserve frame 0).
+
+        ``guidance_interval`` (T2I) restricts CFG to timesteps inside the
+        closed interval ``[lo, hi]``. Outside the interval the cond/uncond
+        delta is dropped and only the cond branch executes — equivalent to
+        CFG with scale=1.0 but cheaper.
+        """
+        do_cfg = guidance_scale > 1.0
+        cfg_parallel = self._cfg_parallel_active() and do_cfg
+        if cfg_parallel:
+            raise NotImplementedError("Cosmos3OmniDiffusersPipeline.diffuse: CFG-Parallel branch is "
+                                      "deferred; FastVideo cfg-world-size plumbing is not yet wired.")
+
+        self.transformer.reset_cache()
+
+        def _cfg_active_at(t: torch.Tensor) -> bool:
+            if guidance_interval is None:
+                return True
+            t_scalar = float(t.item()) if torch.is_tensor(t) else float(t)
+            lo, hi = guidance_interval
+            return lo <= t_scalar <= hi
+
+        def _step(noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor) -> torch.Tensor:
+            if isinstance(noise_pred, tuple):
+                raise ValueError("Cosmos3 video-only diffusion received tuple predictions.")
+            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 condition_latents is not None and velocity_mask is not None:
+                latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents
+            elif image_latent is not None:
+                latents[:, :, 0:1, :, :] = image_latent
+            return latents
+
+        if do_cfg:
+            cond_cache: tuple = (None, None)
+            uncond_cache: tuple = (None, None)
+            for t in self.progress_bar(timesteps):
+                timestep = t.unsqueeze(0)
+                cfg_active = _cfg_active_at(t)
+
+                self.transformer.cached_kv, self.transformer.cached_freqs_gen = cond_cache
+                noise_cond = self.transformer(
+                    hidden_states=latents,
+                    timestep=timestep,
+                    text_ids=cond_ids,
+                    text_mask=cond_mask,
+                    **shared_kwargs,
+                )
+                if cond_cache[0] is None:
+                    cond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen)
+
+                if cfg_active:
+                    self.transformer.cached_kv, self.transformer.cached_freqs_gen = uncond_cache
+                    noise_uncond = self.transformer(
+                        hidden_states=latents,
+                        timestep=timestep,
+                        text_ids=uncond_ids,
+                        text_mask=uncond_mask,
+                        **shared_kwargs,
+                    )
+                    if uncond_cache[0] is None:
+                        uncond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen)
+                    noise_pred = noise_uncond + guidance_scale * (noise_cond - noise_uncond)
+                else:
+                    noise_pred = noise_cond
+
+                latents = _step(noise_pred, t, latents)
+        else:
+            for t in self.progress_bar(timesteps):
+                timestep = t.unsqueeze(0)
+                noise_pred = self.transformer(
+                    hidden_states=latents,
+                    timestep=timestep,
+                    text_ids=cond_ids,
+                    text_mask=cond_mask,
+                    **shared_kwargs,
+                )
+                latents = _step(noise_pred, t, latents)
+
+        return latents
+
+    # -- Forward (main generation entry point) ------------------------------
+
+    def forward(self, req: Any) -> SimpleNamespace:
+        """Cosmos3 inference: request parse + mode dispatch + diffuse + decode.
+
+        Ported from upstream ``pipeline_cosmos3.py:1037-1206`` (single-prompt
+        path). FastVideo does not yet have ``OmniDiffusionRequest`` or
+        ``DiffusionOutput`` dataclasses on the production path, so the
+        return value is a duck-typed ``SimpleNamespace`` with an ``output``
+        attribute mirroring upstream ``DiffusionOutput.output``.
+
+        Mode selection (upstream lines 1059-1093):
+          * T2I: ``"image" in modalities`` and no preprocessed image. Defaults
+            to ``num_frames=1``, ``flow_shift=3.0``,
+            ``guidance_interval=(400.0, 1000.0)``, 50 steps, scale=7.0.
+          * I2V: ``preprocessed_image`` present and not T2I. Like T2V but
+            with image conditioning.
+          * T2V: otherwise. Defaults to ``num_frames=189``,
+            ``flow_shift=self._engine_init_flow_shift``, no guidance
+            interval, 35 steps, scale=6.0.
+
+        Calls ``_set_flow_shift`` exactly once per request, after defaults
+        are resolved and before tokenization.
+        """
+        if not req.prompts:
+            raise ValueError("Cosmos3OmniDiffusersPipeline.forward() requires at least one prompt.")
+        if len(req.prompts) > 1:
+            raise ValueError("Cosmos3OmniDiffusersPipeline currently supports a single prompt per request.")
+
+        prompt_data = req.prompts[0]
+        if isinstance(prompt_data, str):
+            prompt = prompt_data
+            negative_prompt = None
+            image_tensor = None
+        else:
+            prompt = prompt_data.get("prompt", "")
+            negative_prompt = prompt_data.get("negative_prompt")
+            additional_info = prompt_data.get("additional_information", {}) or {}
+            image_tensor = additional_info.get("preprocessed_image")
+
+        sp = req.sampling_params
+        is_t2i = self._is_t2i_request(req)
+        is_i2v = image_tensor is not None and not is_t2i
+        if negative_prompt is None:
+            if is_t2i:
+                negative_prompt = COSMOS3_DEFAULT_NEGATIVE_PROMPT
+            elif is_i2v:
+                negative_prompt = COSMOS3_I2V_NEGATIVE_PROMPT
+            else:
+                negative_prompt = COSMOS3_T2V_NEGATIVE_PROMPT
+
+        if is_t2i:
+            height = sp.height or 1024
+            width = sp.width or 1024
+            num_frames = 1
+            num_inference_steps = sp.num_inference_steps or 50
+            guidance_scale = sp.guidance_scale if sp.guidance_scale else 7.0
+            default_flow_shift = 3.0
+            default_guidance_interval: tuple[float, float] | None = (400.0, 1000.0)
+            batch_size = max(1, int(getattr(sp, "num_outputs_per_prompt", None) or 1))
+        else:
+            height = sp.height or 720
+            width = sp.width or 1280
+            num_frames = sp.num_frames or 189
+            num_inference_steps = sp.num_inference_steps or 35
+            guidance_scale = sp.guidance_scale if sp.guidance_scale else 6.0
+            default_flow_shift = self._engine_init_flow_shift
+            default_guidance_interval = None
+            batch_size = 1
+
+        flow_shift_target = float(self._get_sp_param(sp, "flow_shift", default_flow_shift))
+        guidance_interval = self._get_sp_param(sp, "guidance_interval", default_guidance_interval)
+
+        frame_rate = (self._get_sp_param(sp, "resolved_frame_rate") or self._get_sp_param(sp, "frame_rate") or 24.0)
+        max_sequence_length = self._get_sp_param(sp, "max_sequence_length", 512) or 512
+        use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", False))
+
+        self._guidance_scale = guidance_scale
+        self._num_timesteps = num_inference_steps
+
+        self._set_flow_shift(flow_shift_target)
+
+        generator = sp.generator
+        if generator is None:
+            seed = sp.seed if sp.seed is not None else 42
+            generator = torch.Generator(device=self.device).manual_seed(seed)
+
+        cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts(
+            prompt,
+            negative_prompt,
+            num_frames,
+            frame_rate,
+            height,
+            width,
+            max_sequence_length,
+            sp,
+            use_system_prompt,
+            is_t2i=is_t2i,
+        )
+
+        if image_tensor is not None and not is_t2i:
+            latents, velocity_mask, image_latent = self._prepare_latents_i2v(
+                image_tensor,
+                height,
+                width,
+                num_frames,
+                generator,
+            )
+            condition_latents = None
+        else:
+            latents = self._prepare_latents(height, width, num_frames, generator)
+            velocity_mask = None
+            image_latent = None
+            condition_latents = None
+
+        video_shape = (latents.shape[2], latents.shape[3], latents.shape[4])
+        shared_kwargs: dict[str, Any] = dict(video_shape=video_shape, fps=frame_rate)
+        if velocity_mask is not None:
+            shared_kwargs["noisy_frame_mask"] = velocity_mask
+
+        def _run_diffusion(start_latents: torch.Tensor) -> torch.Tensor:
+            self._set_scheduler_timesteps(num_inference_steps)
+            return self.diffuse(
+                latents=start_latents,
+                timesteps=self.scheduler.timesteps,
+                cond_ids=cond_ids,
+                cond_mask=cond_mask,
+                uncond_ids=uncond_ids,
+                uncond_mask=uncond_mask,
+                guidance_scale=guidance_scale,
+                shared_kwargs=shared_kwargs,
+                velocity_mask=velocity_mask,
+                image_latent=image_latent,
+                condition_latents=condition_latents,
+                guidance_interval=guidance_interval,
+            )
+
+        if is_t2i and batch_size > 1:
+            samples = [_run_diffusion(latents)]
+            for _ in range(batch_size - 1):
+                next_latents = self._prepare_latents(height, width, num_frames, generator)
+                samples.append(_run_diffusion(next_latents))
+            latents = torch.cat(samples, dim=0)
+        else:
+            latents = _run_diffusion(latents)
+
+        video = self._decode_latents(latents)
+        return SimpleNamespace(output={"image": video} if is_t2i else {"video": video})
+
+    # -- Construction stubs (full wiring deferred) --------------------------
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        raise NotImplementedError("Cosmos3OmniDiffusersPipeline.__init__ is wired in a later phase. "
+                                  "Use Cosmos3OmniDiffusersPipeline.__new__(cls) for Phase 2f unit "
+                                  "tests that only exercise _remap_ckpt_key, _set_flow_shift, "
+                                  "diffuse, or forward (with monkey-patched helpers).")
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        raise NotImplementedError("Cosmos3OmniDiffusersPipeline.create_pipeline_stages is wired in a later phase.")
+
+
+# Entry point for pipeline registry (placeholder; full registry wiring lands
+# alongside Phase 2f.2 once create_pipeline_stages and __init__ are real).
+EntryClass = Cosmos3OmniDiffusersPipeline

fastvideo/registry.py

@@ -20,6 +20,7 @@ from fastvideo.configs.pipelines.cosmos2_5 import (
     Cosmos25Config,
     Cosmos25_14BConfig,
 )
+from fastvideo.configs.pipelines.cosmos3 import Cosmos3Config
 from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
@@ -772,6 +773,22 @@ def _register_configs() -> None:
         default_preset="gen3c_cosmos_7b",
     )
 
+    # Cosmos 3 (must register before Cosmos 2.5 and generic Cosmos detectors
+    # so the cosmos3 path-detection takes precedence)
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Cosmos3Config,
+        workload_types=(WorkloadType.T2V, ),
+        hf_model_paths=[
+            "nvidia/Cosmos3-Nano",
+        ],
+        model_detectors=[
+            lambda path: "cosmos3" in path.lower() or "cosmos-3" in path.lower(),
+        ],
+        model_family="cosmos3",
+        default_preset="cosmos3_nano",
+    )
+
     # Cosmos 2.5 (2B)
     register_configs(
         sampling_param_cls=None,

scripts/checkpoint_conversion/cosmos3_convert.py

@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Convert a Diffusers-format Cosmos3 checkpoint to FastVideo state-dict layout.
+
+Cosmos3 weights (when published as ``nvidia/Cosmos3-Nano``) ship in a
+Diffusers-style namespace with ``model.embed_tokens`` /
+``model.layers.{i}.self_attn.*_moe_gen`` / ``vae2llm`` / ``llm2vae`` /
+``time_embedder`` layout. FastVideo's ``Cosmos3VFMTransformer`` splits the
+UND/GEN paths into separate ``language_model.layers.{i}`` and
+``gen_layers.{i}`` subtrees. This script applies the 14-rule remap from
+``Cosmos3OmniDiffusersPipeline._remap_ckpt_key`` to produce a state-dict
+that loads cleanly into
+``fastvideo.models.dits.cosmos3.Cosmos3VFMTransformer``.
+
+Reference: vllm-omni PR #3454 (https://github.com/vllm-project/vllm-omni/pull/3454),
+HEAD ``8536f5b1421f``; upstream ``pipeline_cosmos3.py:319-411``.
+
+Usage:
+    python scripts/checkpoint_conversion/cosmos3_convert.py \\
+        --src /path/to/nvidia--Cosmos3-Nano \\
+        --dst /path/to/cosmos3_nano_fastvideo
+    python scripts/checkpoint_conversion/cosmos3_convert.py --smoke-test
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import safetensors.torch
+import torch
+
+from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (
+    Cosmos3OmniDiffusersPipeline,
+)
+
+
+def convert_state_dict(
+    src_state: dict[str, torch.Tensor],
+) -> tuple[dict[str, torch.Tensor], list[str], list[str]]:
+    """Apply the Cosmos3 remap to a single state dict.
+
+    Returns ``(new_state, skipped_keys, unmapped_keys)`` where:
+
+    * ``new_state`` is the FastVideo-namespaced state dict.
+    * ``skipped_keys`` lists source keys deliberately dropped by the remap
+      (currently only ``lm_head.*``; the language-model head is not part of
+      the diffusion transformer).
+    * ``unmapped_keys`` lists source keys that did not match any remap rule.
+      Should be empty for a clean Diffusers Cosmos3 checkpoint; non-empty
+      results signal either an unexpected upstream key or a missing rule.
+    """
+    new_state: dict[str, torch.Tensor] = {}
+    skipped: list[str] = []
+    unmapped: list[str] = []
+    for src_key, tensor in src_state.items():
+        new_key = Cosmos3OmniDiffusersPipeline._remap_ckpt_key(src_key)
+        if new_key is None:
+            # Either a deliberate skip (lm_head.*) or an unmapped key.
+            if src_key.startswith("lm_head.") or src_key.startswith(
+                    "transformer.lm_head."):
+                skipped.append(src_key)
+            else:
+                unmapped.append(src_key)
+            continue
+        if new_key in new_state:
+            raise KeyError(
+                f"Duplicate target key {new_key!r} produced by remap; "
+                f"colliding source keys include {src_key!r}")
+        new_state[new_key] = tensor
+    return new_state, skipped, unmapped
+
+
+def convert_checkpoint_dir(src: Path, dst: Path) -> None:
+    """Convert all ``*.safetensors`` shards in ``src`` and write to ``dst``.
+
+    Reads each safetensors shard, applies the remap, and writes a single
+    consolidated safetensors file at ``dst/model.safetensors``. This keeps
+    the converter dependency surface small (single file out) for early
+    Phase 5 work; once shard splitting becomes necessary, switch to
+    per-shard outputs and emit ``model.safetensors.index.json``.
+
+    TODO: Emit a sharded layout with ``model.safetensors.index.json`` when
+    the real ``nvidia/Cosmos3-Nano`` weights land and exceed the
+    single-file safetensors size budget.
+    """
+    src = Path(src)
+    dst = Path(dst)
+    if not src.exists():
+        raise FileNotFoundError(f"Source checkpoint directory not found: {src}")
+    if not src.is_dir():
+        raise NotADirectoryError(
+            f"Source path is not a directory: {src}")
+
+    shard_paths = sorted(src.glob("*.safetensors"))
+    if not shard_paths:
+        raise FileNotFoundError(
+            f"No *.safetensors files found in {src}. Expected a Diffusers-format "
+            "checkpoint directory.")
+
+    dst.mkdir(parents=True, exist_ok=True)
+
+    print(f"Loading {len(shard_paths)} shard(s) from {src}")
+    src_state: dict[str, torch.Tensor] = {}
+    for shard_path in shard_paths:
+        print(f"  Reading {shard_path.name}")
+        shard = safetensors.torch.load_file(str(shard_path))
+        for key, tensor in shard.items():
+            if key in src_state:
+                raise KeyError(
+                    f"Duplicate key {key!r} found across shards "
+                    f"(latest in {shard_path.name}); aborting.")
+            src_state[key] = tensor
+
+    print(f"Loaded {len(src_state)} parameters; applying remap")
+    new_state, skipped, unmapped = convert_state_dict(src_state)
+    print(f"  Remapped {len(new_state)} parameters")
+    print(f"  Skipped  {len(skipped)} parameters (lm_head.*)")
+    if unmapped:
+        print(f"  WARNING: {len(unmapped)} unmapped key(s):")
+        for k in unmapped[:10]:
+            print(f"    {k}")
+        if len(unmapped) > 10:
+            print(f"    ... and {len(unmapped) - 10} more")
+
+    out_path = dst / "model.safetensors"
+    print(f"Writing {out_path}")
+    safetensors.torch.save_file(new_state, str(out_path))
+    size_gb = out_path.stat().st_size / (1024**3)
+    print(f"  Wrote {size_gb:.2f} GB to {out_path}")
+
+
+def smoke_test() -> int:
+    """Exercise the remap on a synthetic state dict covering all rule branches.
+
+    Runs WITHOUT real weights and acts as a CI guard against drift between
+    ``_remap_ckpt_key`` and this converter. Returns 0 on success, 1 on
+    failure.
+
+    The synthetic input contains one representative key per remap branch:
+    top-level (``vae2llm`` / ``llm2vae`` / ``time_embedder``), language-model
+    embeddings/norms, the standalone ``norm_moe_gen``, both UND and GEN
+    attention/norm/MLP variants inside ``model.layers.{i}.*``, and the
+    deliberately-skipped ``lm_head.weight``.
+    """
+    expected: dict[str, str | None] = {
+        # Language-model trunk (UND path).
+        "model.embed_tokens.weight":
+            "transformer.language_model.embed_tokens.weight",
+        "model.norm.weight":
+            "transformer.language_model.norm.weight",
+        # Standalone GEN norm at the top of the trunk.
+        "model.norm_moe_gen.weight": "transformer.norm_moe_gen.weight",
+        # UND self-attention (representative: q_proj).
+        "model.layers.3.self_attn.q_proj.weight":
+            "transformer.language_model.layers.3.self_attn.q_proj.weight",
+        # GEN cross-attention via *_moe_gen suffix (representative: q_proj).
+        "model.layers.3.self_attn.q_proj_moe_gen.weight":
+            "transformer.gen_layers.3.cross_attention.q_proj.weight",
+        # GEN cross-attention norm (representative: k_norm_moe_gen).
+        "model.layers.3.self_attn.k_norm_moe_gen.weight":
+            "transformer.gen_layers.3.cross_attention.k_norm.weight",
+        # UND block norms.
+        "model.layers.3.input_layernorm.weight":
+            "transformer.language_model.layers.3.input_layernorm.weight",
+        # GEN block norms.
+        "model.layers.3.input_layernorm_moe_gen.weight":
+            "transformer.gen_layers.3.input_layernorm.weight",
+        # UND MLP (representative: gate_proj).
+        "model.layers.3.mlp.gate_proj.weight":
+            "transformer.language_model.layers.3.mlp.gate_proj.weight",
+        # GEN MLP (representative: up_proj).
+        "model.layers.3.mlp_moe_gen.up_proj.weight":
+            "transformer.gen_layers.3.mlp.up_proj.weight",
+        # Top-level generation projections.
+        "vae2llm.weight": "transformer.vae2llm.weight",
+        "llm2vae.weight": "transformer.llm2vae.weight",
+        "time_embedder.linear_1.weight":
+            "transformer.time_embedder.linear_1.weight",
+        # Deliberately skipped.
+        "lm_head.weight": None,
+    }
+
+    synthetic = {k: torch.zeros(1) for k in expected}
+    new_state, skipped, unmapped = convert_state_dict(synthetic)
+
+    failures: list[str] = []
+    for src_key, expected_tgt in expected.items():
+        if expected_tgt is None:
+            if src_key not in skipped:
+                failures.append(
+                    f"  expected {src_key!r} in skipped; "
+                    f"new_state contains: {sorted(new_state)}; "
+                    f"unmapped contains: {unmapped}")
+            continue
+        if expected_tgt not in new_state:
+            failures.append(
+                f"  expected target key {expected_tgt!r} (from src "
+                f"{src_key!r}) in new_state; got keys {sorted(new_state)}")
+    if unmapped:
+        failures.append(
+            f"  unmapped keys (should be empty): {unmapped}")
+
+    if failures:
+        print("smoke_test FAILED:")
+        for line in failures:
+            print(line)
+        return 1
+    print(
+        f"smoke_test PASSED: {len(new_state)} new keys, "
+        f"{len(skipped)} skipped, {len(unmapped)} unmapped")
+    return 0
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Convert a Diffusers-format Cosmos3 checkpoint to "
+                    "FastVideo state-dict layout.")
+    parser.add_argument(
+        "--src",
+        type=Path,
+        default=None,
+        help="Diffusers-format source checkpoint directory (e.g. an HF "
+             "snapshot of nvidia/Cosmos3-Nano).")
+    parser.add_argument(
+        "--dst",
+        type=Path,
+        default=None,
+        help="Output directory; the converted state dict is written to "
+             "<dst>/model.safetensors.")
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        help="Run a no-weights smoke test that exercises every remap branch "
+             "on a synthetic state dict. Useful as a CI guard.")
+    args = parser.parse_args()
+
+    if args.smoke_test:
+        return smoke_test()
+    if args.src is None or args.dst is None:
+        parser.error("--src and --dst are required when not running --smoke-test")
+    convert_checkpoint_dir(args.src, args.dst)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/local_tests/cosmos3/README.md

@@ -0,0 +1,39 @@
+# Cosmos3 local parity workspace
+
+## Overview
+
+Cosmos3 Phase-1 prep tracks a future FastVideo port from the vllm-omni Cosmos3 integration in PR #3454. The current reference supports a single `Cosmos3OmniDiffusersPipeline` for text-to-video (T2V), image-to-video (I2V), and text-to-image (T2I). Follow-up Cosmos3 capabilities mentioned in the PR body, such as sound generation and action-generation modes, are out of scope for the initial FastVideo port.
+
+## Reference code
+
+- Reference checkout: `/home/william5lin/cosmos3-reference`
+- Source PR: <https://github.com/vllm-project/vllm-omni/pull/3454>
+- Pinned HEAD: `8536f5b1421f78c7df06af6d96fa195c1ceb6384`
+- Key files:
+  - `vllm_omni/deploy/cosmos3.yaml`
+  - `vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py`
+  - `vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py`
+  - `vllm_omni/diffusion/models/cosmos3/guardrails.py`
+  - `tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py`
+  - `tests/diffusion/models/cosmos3/test_cosmos3_transformer.py`
+
+## Weight status
+
+PENDING. Do not download weights during Phase 1.
+
+- Candidate serving/model id from the PR body: `nvidia/Cosmos3-Nano`
+- Hugging Face API status on 2026-05-22: `401` for `https://huggingface.co/api/models/nvidia/Cosmos3-Nano`; API body reports `Invalid username or password.`
+- Hugging Face author search `author=nvidia&search=Cosmos3`: empty list (`[]`).
+- NGC API URL requested by handoff returned a Next.js 404 HTML page, not model metadata.
+
+## Parity-test placeholder
+
+Phase 2 should add local parity tests here after the official weights become accessible and FastVideo component prototypes exist. Suggested first targets:
+
+1. Transformer state-dict key/shape inventory versus `Cosmos3VFMTransformer`.
+2. Scheduler/default-parameter parity for T2V, I2V, and T2I request modes.
+3. Prompt metadata-template parity for duration/resolution and image-vs-video modalities.
+
+## SSIM placeholder
+
+No SSIM references seeded yet. Add SSIM coverage only after a FastVideo inference path can load resolved Cosmos3 weights and generate stable T2V/T2I/I2V outputs.

tests/local_tests/cosmos3/conftest.py

@@ -0,0 +1,298 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared fixtures for Cosmos3 local parity tests (Phase 2a Tier A).
+
+Ported from the vllm-omni reference suite so FastVideo's Tier A parity
+coverage can land before NVIDIA publishes the real Cosmos3 weights.
+The fixtures here are intentionally self-contained: they do NOT import
+``vllm_omni`` and do NOT import any FastVideo Cosmos3 modules (those
+do not exist yet; they are introduced in Phase 2b).
+
+Reference:
+- ``tests/diffusion/models/cosmos3/conftest.py`` lines 1-176 from
+  ``vllm-omni`` HEAD ``8536f5b1421f78c7df06af6d96fa195c1ceb6384``.
+
+The stubs reproduce the same name + protocol as the upstream fixtures
+so test code can be ported with minimal renaming once FastVideo's
+Cosmos3 pipeline class lands and a ``make_cosmos3_pipeline`` fixture
+can be re-pointed at the real class.
+"""
+from __future__ import annotations
+
+import sys
+import types
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+from torch import nn
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    """Register ``local`` marker so ``pytestmark = [pytest.mark.local]``
+    in sibling test files does not trigger ``PytestUnknownMarkWarning``.
+    """
+    config.addinivalue_line(
+        "markers",
+        "local: marker for local-only parity/scaffold tests (skipped in CI)",
+    )
+
+
+# ---------------------------------------------------------------------------
+# Stub scheduler — mirrors vllm-omni ``StubScheduler``.
+# ---------------------------------------------------------------------------
+class StubScheduler:
+    """Minimal UniPC-shaped scheduler stub used by pipeline call-graph tests.
+
+    Mirrors ``tests/diffusion/models/cosmos3/conftest.py:16-30`` from
+    the vllm-omni reference: records ``set_timesteps`` and ``step`` calls,
+    advances ``latents`` by ``+noise_pred``, and exposes a ``config``
+    namespace carrying ``num_train_timesteps`` and ``flow_shift``.
+    """
+
+    def __init__(
+        self,
+        timesteps: list[int] | None = None,
+        *,
+        flow_shift: float = 1.0,
+    ) -> None:
+        self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
+        self.config = SimpleNamespace(num_train_timesteps=1000, flow_shift=flow_shift)
+        self.set_timesteps_calls: list[tuple[int, torch.device]] = []
+        self.step_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
+
+    def set_timesteps(self, num_steps: int, device: torch.device) -> None:
+        self.set_timesteps_calls.append((num_steps, device))
+        self.timesteps = torch.arange(num_steps, 0, -1, dtype=torch.int64, device=device)
+
+    def step(
+        self,
+        noise_pred: torch.Tensor,
+        timestep: torch.Tensor,
+        latents: torch.Tensor,
+        **kwargs: Any,
+    ):
+        del kwargs
+        self.step_calls.append((noise_pred.clone(), timestep.clone(), latents.clone()))
+        return (latents + noise_pred,)
+
+
+class _ModeLatentDist:
+    """Stub for diffusers ``DiagonalGaussianDistribution.mode``."""
+
+    def __init__(self, latents: torch.Tensor) -> None:
+        self._latents = latents
+
+    def mode(self) -> torch.Tensor:
+        return self._latents
+
+
+class StubCosmos3VAE:
+    """VAE stub returning deterministic latents shaped by VAE scale factors.
+
+    Mirrors ``tests/diffusion/models/cosmos3/conftest.py:41-70``.
+    """
+
+    dtype = torch.float32
+
+    def __init__(self, z_dim: int = 2, *, temporal: int = 4, spatial: int = 8) -> None:
+        self.config = SimpleNamespace(
+            z_dim=z_dim,
+            scale_factor_temporal=temporal,
+            scale_factor_spatial=spatial,
+            latents_mean=[0.0] * z_dim,
+            latents_std=[1.0] * z_dim,
+        )
+
+    def encode(self, video: torch.Tensor):
+        latent_frames = (video.shape[2] - 1) // self.config.scale_factor_temporal + 1
+        latent_height = video.shape[-2] // self.config.scale_factor_spatial
+        latent_width = video.shape[-1] // self.config.scale_factor_spatial
+        latents = torch.ones(
+            video.shape[0],
+            self.config.z_dim,
+            latent_frames,
+            latent_height,
+            latent_width,
+            dtype=video.dtype,
+            device=video.device,
+        )
+        return SimpleNamespace(latent_dist=_ModeLatentDist(latents))
+
+    def decode(self, latents: torch.Tensor, return_dict: bool = False):
+        del return_dict
+        return (latents,)
+
+
+class StubCosmos3Transformer(nn.Module):
+    """Transformer stub that records per-call inputs and emits deterministic
+    output tensors keyed by the first ``text_ids`` token.
+
+    Mirrors ``tests/diffusion/models/cosmos3/conftest.py:73-114``.
+
+    The stub also exposes ``cached_kv`` / ``cached_freqs_gen`` so pipeline
+    tests can assert that:
+      * ``reset_cache()`` is called before each diffusion loop;
+      * UND/cond and uncond branches each populate their own cache exactly
+        once and reuse it on subsequent timesteps.
+    """
+
+    def __init__(self, *, latent_channel_size: int = 2) -> None:
+        super().__init__()
+        self.latent_channel_size = latent_channel_size
+        self.cached_kv: Any | None = None
+        self.cached_freqs_gen: Any | None = None
+        self.calls: list[dict[str, Any]] = []
+        self.reset_calls = 0
+
+    def reset_cache(self) -> None:
+        self.reset_calls += 1
+        self.cached_kv = None
+        self.cached_freqs_gen = None
+
+    def forward(
+        self,
+        *,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        text_ids: torch.Tensor,
+        text_mask: torch.Tensor,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        token = int(text_ids.reshape(-1)[0].item()) if text_ids.numel() else 0
+        self.calls.append(
+            {
+                "token": token,
+                "timestep": timestep.clone(),
+                "text_mask": text_mask.clone(),
+                "cache_before": self.cached_kv,
+                "kwargs": dict(kwargs),
+            }
+        )
+        if self.cached_kv is None:
+            marker = torch.tensor([token], dtype=torch.float32)
+            self.cached_kv = [(marker, marker + 100)]
+            self.cached_freqs_gen = (marker + 200, marker + 300)
+        return torch.full_like(hidden_states, float(token))
+
+
+def passthrough_progress_bar(iterable):
+    return iterable
+
+
+# ---------------------------------------------------------------------------
+# Tiny config — mirrors test_cosmos3_transformer.py:15-29
+# ---------------------------------------------------------------------------
+def _tiny_cosmos3_config(**overrides: Any) -> dict:
+    """Minimal Cosmos3 transformer config sufficient for shape-only construction.
+
+    Mirrors ``tests/diffusion/models/cosmos3/test_cosmos3_transformer.py:15-29``.
+    """
+    config: dict = {
+        "hidden_size": 8,
+        "num_hidden_layers": 0,
+        "num_attention_heads": 2,
+        "num_key_value_heads": 2,
+        "head_dim": 4,
+        "intermediate_size": 16,
+        "vocab_size": 32,
+        "latent_patch_size": 1,
+        "latent_channel": 2,
+        "rope_scaling": {"mrope_section": [1, 1, 0]},
+    }
+    config.update(overrides)
+    return config
+
+
+# ---------------------------------------------------------------------------
+# Guardrail no-op stub — mirrors conftest.py:121-129
+# ---------------------------------------------------------------------------
+@pytest.fixture(autouse=True)
+def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch):
+    """Install a no-op replacement for the Cosmos3 guardrails module.
+
+    The vllm-omni reference imports guardrails eagerly; FastVideo's port
+    may either skip guardrails or use a different module path. Either
+    way, an autouse stub avoids accidental network-dependent imports
+    during scaffold-test collection.
+    """
+    module = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
+    module.is_guardrails_enabled = lambda od_config, sampling_params=None: False  # type: ignore[attr-defined]
+    module.ensure_initialized = lambda od_config: None  # type: ignore[attr-defined]
+    module.check_text_safety = lambda text: None  # type: ignore[attr-defined]
+    module.check_video_safety = lambda video: video  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, module.__name__, module)
+    return module
+
+
+# ---------------------------------------------------------------------------
+# Pipeline factory — kept as a scaffold; activate in Phase 2b when the
+# FastVideo Cosmos3 pipeline class exists.
+# ---------------------------------------------------------------------------
+@pytest.fixture
+def make_cosmos3_pipeline():
+    """Factory that returns a FastVideo Cosmos3 pipeline pre-wired with stubs.
+
+    In Phase 2a (Tier A, no real weights), tests that consume this factory
+    should ``pytest.skip`` if the FastVideo Cosmos3 pipeline class does not
+    yet exist. In Phase 2b we replace the placeholder construction with
+    ``object.__new__(<FastVideoCosmos3Pipeline>)`` + ``nn.Module.__init__``
+    and re-point the stubs onto the real attribute names.
+
+    Mirrors the upstream layout at
+    ``tests/diffusion/models/cosmos3/conftest.py:132-157``.
+    """
+
+    def _make():
+        try:
+            from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (  # type: ignore
+                Cosmos3OmniDiffusersPipeline,
+            )
+        except ImportError:
+            pytest.skip(
+                "FastVideo Cosmos3 pipeline class not yet implemented "
+                "(Phase 2b will provide fastvideo.pipelines.basic.cosmos3)."
+            )
+
+        pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
+        nn.Module.__init__(pipeline)
+        pipeline.od_config = SimpleNamespace()
+        pipeline.device = torch.device("cpu")
+        pipeline.dtype = torch.float32
+        pipeline.transformer = StubCosmos3Transformer(latent_channel_size=2)
+        pipeline.vae = StubCosmos3VAE(z_dim=2)
+        pipeline.vae_scale_factor_temporal = 4
+        pipeline.vae_scale_factor_spatial = 8
+        pipeline.scheduler = StubScheduler([9, 3], flow_shift=1.0)
+        pipeline._base_scheduler_config = pipeline.scheduler.config
+        pipeline._engine_init_flow_shift = 1.0
+        pipeline._current_flow_shift = 1.0
+        pipeline._guidance_scale = None
+        pipeline._num_timesteps = None
+        pipeline.progress_bar = passthrough_progress_bar
+        return pipeline
+
+    return _make
+
+
+def make_sampling_params(**overrides: Any) -> SimpleNamespace:
+    """Build a SamplingParams-like namespace with Cosmos3's expected fields.
+
+    Mirrors ``tests/diffusion/models/cosmos3/conftest.py:160-176``.
+    """
+    values = {
+        "height": None,
+        "width": None,
+        "num_frames": None,
+        "num_inference_steps": None,
+        "guidance_scale": None,
+        "generator": None,
+        "seed": 123,
+        "num_outputs_per_prompt": 1,
+        "frame_rate": None,
+        "resolved_frame_rate": None,
+        "max_sequence_length": None,
+        "extra_args": {},
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)

tests/local_tests/cosmos3/test_cosmos3_mrope_parity.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 unified 3D mRoPE position-ID parity (Tier A scaffold).
+
+Reference: ``vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py``
+lines 113-177 (``compute_mrope_position_ids_text`` /
+``compute_mrope_position_ids_vision``). The reference test asserting
+these invariants lives at
+``tests/diffusion/models/cosmos3/test_cosmos3_transformer.py:32-57``.
+
+The three invariants under test:
+
+  1. Text tokens broadcast the same monotonically-increasing positions
+     across all three (t, h, w) axes. With ``num_tokens=3`` and
+     ``temporal_offset=5`` the result is ``[[5,6,7], [5,6,7], [5,6,7]]``
+     and the next-offset is ``8``.
+
+  2. Vision tokens (no FPS modulation) flatten a ``(grid_t, grid_h, grid_w)``
+     position grid in t-major order. With ``(2, 2, 3)`` and offset ``10``
+     the resulting shape is ``(3, 12)`` and the temporal row begins
+     ``[10]*6 + [11]*6``; next-offset is ``12``.
+
+  3. FPS-modulated vision tokens scale the temporal axis by
+     ``base_fps / temporal_compression_factor / (fps / tcf)``. With
+     ``fps=12``, ``base_fps=24``, ``tcf=4``, ``grid_t=2`` the first row is
+     ``[10.0, 12.0]``.
+
+The FastVideo side currently does NOT exist; the test is wrapped in
+``try/except ImportError`` and skips. Phase 2b replaces the skip with
+the real import + assertion path.
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.local]
+
+
+def test_compute_mrope_position_ids_text_and_vision() -> None:
+    """Asserts the 3 invariants of unified 3D mRoPE position-ID generation.
+
+    Once FastVideo's ``fastvideo.models.dits.cosmos3`` exports
+    ``compute_mrope_position_ids_text`` and
+    ``compute_mrope_position_ids_vision``, this test verifies they produce
+    output tensors identical to the vllm-omni reference at
+    transformer_cosmos3.py:113-177.
+    """
+    try:
+        from fastvideo.models.dits.cosmos3 import (  # type: ignore
+            compute_mrope_position_ids_text,
+            compute_mrope_position_ids_vision,
+        )
+    except ImportError:
+        pytest.skip("FastVideo Cosmos3 not yet implemented (Phase 2b)")
+
+    text_ids, text_offset = compute_mrope_position_ids_text(num_tokens=3, temporal_offset=5)
+    assert text_ids.tolist() == [[5, 6, 7], [5, 6, 7], [5, 6, 7]]
+    assert text_offset == 8
+
+    vision_ids, vision_offset = compute_mrope_position_ids_vision(
+        2, 2, 3, temporal_offset=10, fps=None
+    )
+    assert tuple(vision_ids.shape) == (3, 12)
+    assert vision_ids[0].tolist() == [10] * 6 + [11] * 6
+    assert vision_offset == 12
+
+    modulated_ids, modulated_offset = compute_mrope_position_ids_vision(
+        2,
+        1,
+        1,
+        temporal_offset=10,
+        fps=12.0,
+        base_fps=24.0,
+        temporal_compression_factor=4,
+    )
+    torch.testing.assert_close(modulated_ids[0], torch.tensor([10.0, 12.0]))
+    assert modulated_offset == 13

tests/local_tests/cosmos3/test_cosmos3_patchify_unpatchify_parity.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 ``[B,C,T,H,W] <-> [B, T*hp*wp, p*p*C]`` patchify roundtrip (Tier A).
+
+Reference: ``vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py``
+lines 1009-1036 (``Cosmos3VFMTransformer.patchify`` /
+``Cosmos3VFMTransformer.unpatchify``) and the reference assertion at
+``tests/diffusion/models/cosmos3/test_cosmos3_transformer.py:98-101``.
+
+Invariant: ``unpatchify(patchify(x)) == x`` for any ``x`` with shape
+``(B, C, t, h, w)`` where ``h, w`` are divisible by ``latent_patch_size``.
+Also exercises a non-trivial channel count (3) to ensure the
+``permute([0, 2, 3, 5, 4, 6, 1])`` reordering is correct.
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.local]
+
+
+def test_patchify_unpatchify_roundtrip() -> None:
+    """Asserts that the FastVideo Cosmos3 transformer's patchify/unpatchify
+    pair are exact inverses for ``latent_patch_size=2``, ``latent_channel=3``.
+
+    Once FastVideo's ``fastvideo.models.dits.cosmos3.Cosmos3VFMTransformer``
+    lands, replace the skip with the upstream-equivalent assertion path.
+    """
+    try:
+        from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer  # type: ignore
+    except ImportError:
+        pytest.skip("FastVideo Cosmos3 not yet implemented (Phase 2b)")
+
+    from torch import nn
+
+    model = object.__new__(Cosmos3VFMTransformer)
+    nn.Module.__init__(model)
+    model.latent_patch_size = 2
+    model.latent_channel_size = 3
+
+    latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5)
+    roundtrip = model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5)
+    torch.testing.assert_close(roundtrip, latents)
+
+
+def test_patchify_default_patch_size() -> None:
+    """Asserts shape contract for the default ``latent_patch_size=[1,2,2]``
+    (i.e. spatial-only patching) with a representative video latent.
+
+    With ``[B,C,T,H,W] = [1, 16, 2, 8, 8]`` and patch=2 on H/W, expected
+    flattened tokens = ``T * (H/2) * (W/2) = 2 * 4 * 4 = 32`` and each token
+    carries ``2*2*C = 4*16 = 64`` channels.
+    """
+    try:
+        from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer  # type: ignore
+    except ImportError:
+        pytest.skip("FastVideo Cosmos3 not yet implemented (Phase 2b)")
+
+    from torch import nn
+
+    model = object.__new__(Cosmos3VFMTransformer)
+    nn.Module.__init__(model)
+    model.latent_patch_size = 2
+    model.latent_channel_size = 16
+
+    latents = torch.zeros(1, 16, 2, 8, 8)
+    tokens = model.patchify(latents, t=2, h=8, w=8)
+    assert tuple(tokens.shape) == (1, 32, 64)
+    restored = model.unpatchify(tokens, t=2, h=8, w=8)
+    assert tuple(restored.shape) == (1, 16, 2, 8, 8)

tests/local_tests/cosmos3/test_cosmos3_pipeline_call_graph.py

@@ -0,0 +1,199 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 pipeline call-graph parity (Tier A scaffold).
+
+Reference:
+  * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:883-1033`` —
+    ``Cosmos3OmniDiffusersPipeline.diffuse``: 3-mode CFG denoising loop
+    with UND-cache management.
+  * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:1037-1067`` —
+    ``forward``: parses the request, selects T2I/T2V/I2V mode, dispatches
+    defaults.
+  * Reference test invariants at
+    ``tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py:126-156,196-251``.
+
+The Tier A scaffold uses the stubs from ``conftest.py`` (StubScheduler,
+StubCosmos3VAE, StubCosmos3Transformer) so the call-graph can be tested
+without a real DiT or VAE. The 4 invariants under test:
+
+  1. ``diffuse(...)`` calls ``transformer.reset_cache()`` exactly once
+     before iterating timesteps;
+  2. With ``do_cfg=True`` and ``guidance_interval=None``, each timestep
+     invokes the transformer twice — once with ``cond_ids`` and once with
+     ``uncond_ids`` — and reuses the cached UND K/V from step 1 onward
+     (asserted via ``StubCosmos3Transformer.calls`` order);
+  3. I2V mode: ``velocity_mask`` zeros frame-0 noise predictions, and
+     ``image_latent`` is re-injected into frame 0 after each scheduler step;
+  4. ``forward`` selects T2I vs T2V mode from ``prompt["modalities"]``
+     and applies the per-mode default ``flow_shift`` / ``num_frames``.
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.local]
+
+
+def _ids(value: int) -> torch.Tensor:
+    return torch.tensor([[value]], dtype=torch.long)
+
+
+def _mask() -> torch.Tensor:
+    return torch.ones(1, 1, dtype=torch.long)
+
+
+def test_diffuse_resets_cache_and_calls_cfg_in_order(make_cosmos3_pipeline) -> None:
+    """Asserts the 3-mode CFG call graph: ``reset_cache`` first, then
+    interleaved ``cond / uncond`` transformer calls with cache reuse.
+
+    With timesteps ``[900, 100]`` and ``guidance_scale=3.0``, the upstream
+    reference (test_cosmos3_pipeline.py:126-142) asserts the call order is
+    ``[2, 1, 2]`` — i.e. step 0 calls cond (token=2) then uncond (token=1),
+    step 1 reuses the cached UND K/V and only calls cond again (token=2).
+    """
+    pipeline = make_cosmos3_pipeline()
+    latents = torch.zeros(1, 2, 1, 1, 1)
+
+    result = pipeline.diffuse(
+        latents=latents,
+        timesteps=torch.tensor([900, 100]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=3.0,
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
+        guidance_interval=(500.0, 1000.0),
+    )
+    assert pipeline.transformer.reset_calls == 1
+    assert [call["token"] for call in pipeline.transformer.calls] == [2, 1, 2]
+    torch.testing.assert_close(result, torch.full_like(latents, 6.0))
+
+
+def test_diffuse_i2v_velocity_mask_zeros_frame_zero(make_cosmos3_pipeline) -> None:
+    """Asserts I2V velocity-mask + image-latent re-injection contract.
+
+    Cross-check: pipeline_cosmos3.py:937-951. The velocity_mask zeroes
+    noise predictions on conditioning frames before stepping, and
+    ``image_latent`` is overwritten into frame 0 of the output latents
+    after each scheduler step.
+    """
+    pipeline = make_cosmos3_pipeline()
+    result = pipeline.diffuse(
+        latents=torch.zeros(1, 2, 2, 1, 1),
+        timesteps=torch.tensor([7]),
+        cond_ids=_ids(2),
+        cond_mask=_mask(),
+        uncond_ids=_ids(1),
+        uncond_mask=_mask(),
+        guidance_scale=1.0,
+        shared_kwargs={"video_shape": (2, 1, 1), "fps": 24.0},
+        velocity_mask=torch.tensor([[[[[0.0]], [[1.0]]]]]),
+        image_latent=torch.full((1, 2, 1, 1, 1), 7.0),
+    )
+    torch.testing.assert_close(result[:, :, 0:1], torch.full((1, 2, 1, 1, 1), 7.0))
+
+
+@pytest.mark.parametrize(
+    ("modalities", "expected_is_t2i", "expected_default_flow_shift", "expected_default_frames"),
+    [
+        (["image"], True, 3.0, 1),
+        (["video"], False, 1.0, 189),
+    ],
+)
+def test_forward_mode_dispatch_t2i_vs_t2v(
+    make_cosmos3_pipeline,
+    modalities: list[str],
+    expected_is_t2i: bool,
+    expected_default_flow_shift: float,
+    expected_default_frames: int,
+) -> None:
+    """Asserts forward() routes to T2I vs T2V mode based on prompt modalities,
+    and applies per-mode defaults.
+
+    Cross-check: pipeline_cosmos3.py:1069-1093. T2I defaults:
+    ``num_frames=1``, ``flow_shift=3.0``, ``num_inference_steps=50``,
+    ``guidance_interval=[400, 1000]``. T2V defaults: ``num_frames=189``,
+    ``flow_shift=engine_init`` (1.0 here), ``num_inference_steps=35``,
+    no guidance_interval.
+    """
+    pipeline = make_cosmos3_pipeline()
+
+    from types import SimpleNamespace
+
+    captured: dict[str, object] = {"flow_shifts": [], "format_calls": []}
+
+    def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs):
+        captured["format_calls"].append(
+            {
+                "is_t2i": kwargs.get("is_t2i"),
+                "num_frames": num_frames,
+            }
+        )
+        return _ids(2), _mask(), _ids(1), _mask()
+
+    pipeline._format_and_tokenize_prompts = fake_format
+    pipeline._prepare_latents = lambda *a, **kw: torch.zeros(1, 2, 1, 1, 1)
+    pipeline._set_flow_shift = lambda target: captured["flow_shifts"].append(target)
+    pipeline._set_scheduler_timesteps = lambda steps: setattr(
+        pipeline.scheduler, "timesteps", torch.tensor([7])
+    )
+    pipeline.diffuse = lambda **kw: kw["latents"]
+    pipeline._decode_latents = lambda latents: latents
+
+    output = pipeline.forward(
+        SimpleNamespace(
+            prompts=[{"prompt": "test", "modalities": modalities}],
+            sampling_params=SimpleNamespace(
+                height=None,
+                width=None,
+                num_frames=None,
+                num_inference_steps=None,
+                guidance_scale=None,
+                generator=None,
+                seed=123,
+                num_outputs_per_prompt=1,
+                frame_rate=None,
+                resolved_frame_rate=None,
+                max_sequence_length=None,
+                extra_args={},
+            ),
+        )
+    )
+    assert captured["format_calls"][-1]["is_t2i"] is expected_is_t2i
+    assert captured["format_calls"][-1]["num_frames"] == expected_default_frames
+    assert captured["flow_shifts"] == [expected_default_flow_shift]
+    expected_output_key = "image" if expected_is_t2i else "video"
+    assert expected_output_key in output.output
+
+
+def test_forward_rejects_both_image_and_video_modalities(make_cosmos3_pipeline) -> None:
+    """Asserts ``_is_t2i_request`` raises ValueError when a prompt
+    requests both image and video modalities simultaneously.
+
+    Cross-check: pipeline_cosmos3.py:490-496.
+    """
+    pipeline = make_cosmos3_pipeline()
+
+    from types import SimpleNamespace
+
+    with pytest.raises(ValueError, match="both image and video"):
+        pipeline.forward(
+            SimpleNamespace(
+                prompts=[{"prompt": "x", "modalities": ["image", "video"]}],
+                sampling_params=SimpleNamespace(
+                    height=None,
+                    width=None,
+                    num_frames=None,
+                    num_inference_steps=None,
+                    guidance_scale=None,
+                    generator=None,
+                    seed=123,
+                    num_outputs_per_prompt=1,
+                    frame_rate=None,
+                    resolved_frame_rate=None,
+                    max_sequence_length=None,
+                    extra_args={},
+                ),
+            )
+        )

tests/local_tests/cosmos3/test_cosmos3_scheduler_default_parity.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 scheduler default + per-request override parity (Tier A scaffold).
+
+Reference:
+  * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:275-307`` —
+    initial UniPCMultistepScheduler load (preserves solver_order,
+    timestep_spacing, beta_schedule, sigma bounds, flow_shift) and
+    one-time override at engine-init if ``od_config.flow_shift`` is set.
+  * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:498-512`` —
+    ``_set_flow_shift(target_shift)``: rebuild the scheduler via
+    ``UniPCMultistepScheduler.from_config(base_config, flow_shift=target)``
+    when the requested target differs from the current shift.
+  * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:1069-1110`` —
+    per-request mode defaults: T2I uses ``shift=3.0``; T2V/I2V use the
+    engine-init shift (typically 1.0); ``flow_shift`` may be overridden
+    per request via ``sampling_params.extra_args["flow_shift"]``.
+
+The invariant under test: for the same RNG seed and the same number of
+inference steps, the scheduler's ``timesteps`` tensor must be identical
+whenever the ``flow_shift`` is identical, and must change deterministically
+when ``flow_shift`` is overridden via ``_set_flow_shift``.
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.local]
+
+
+def test_t2i_default_flow_shift_is_3() -> None:
+    """Asserts that T2I requests rebuild the scheduler at ``flow_shift=3.0``.
+
+    Cross-check: pipeline_cosmos3.py:1073-1080 sets
+    ``default_flow_shift = 3.0`` for T2I, and
+    pipeline_cosmos3.py:1110 calls ``self._set_flow_shift(flow_shift_target)``
+    which rebuilds the scheduler via ``UniPCMultistepScheduler.from_config(
+    base_config, flow_shift=3.0)``.
+    """
+    try:
+        from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (  # type: ignore
+            Cosmos3OmniDiffusersPipeline,
+        )
+    except ImportError:
+        pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)")
+
+    pipeline = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline)
+    if not hasattr(pipeline, "_set_flow_shift") or not hasattr(pipeline, "scheduler"):
+        pytest.skip("FastVideo Cosmos3 scheduler/_set_flow_shift not yet wired")
+
+    pipeline._set_flow_shift(3.0)
+    assert float(pipeline.scheduler.config.flow_shift) == 3.0
+
+
+def test_t2v_default_flow_shift_is_engine_init() -> None:
+    """Asserts T2V/I2V use the engine-init shift (e.g. 1.0), NOT a fixed default.
+
+    Cross-check: pipeline_cosmos3.py:1091 sets
+    ``default_flow_shift = self._engine_init_flow_shift`` for T2V/I2V
+    (NOT ``None`` — passing ``None`` would leak a prior T2I rebuild
+    forward).
+    """
+    try:
+        from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (  # type: ignore
+            Cosmos3OmniDiffusersPipeline,
+        )
+    except ImportError:
+        pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)")
+
+    pipeline = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline)
+    if not hasattr(pipeline, "_engine_init_flow_shift") or not hasattr(pipeline, "_set_flow_shift"):
+        pytest.skip("FastVideo Cosmos3 _engine_init_flow_shift not yet wired")
+
+    init_shift = float(pipeline._engine_init_flow_shift)
+    pipeline._set_flow_shift(init_shift)
+    assert float(pipeline.scheduler.config.flow_shift) == init_shift
+
+
+def test_scheduler_timesteps_deterministic_under_seed() -> None:
+    """Asserts that ``scheduler.set_timesteps(N)`` is deterministic given the
+    same N and the same flow_shift.
+
+    The UniPC scheduler's timestep sequence does not depend on a torch
+    seed (it's a closed-form function of N + scheduler config), so
+    invoking ``set_timesteps`` twice should produce identical tensors.
+    """
+    try:
+        from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (  # type: ignore
+            Cosmos3OmniDiffusersPipeline,
+        )
+    except ImportError:
+        pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)")
+
+    pipeline = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline)
+    if not hasattr(pipeline, "scheduler") or not hasattr(pipeline, "_set_flow_shift"):
+        pytest.skip("FastVideo Cosmos3 scheduler not yet wired")
+
+    pipeline._set_flow_shift(3.0)
+    pipeline.scheduler.set_timesteps(35, device=torch.device("cpu"))
+    seq_a = pipeline.scheduler.timesteps.clone()
+
+    pipeline.scheduler.set_timesteps(35, device=torch.device("cpu"))
+    seq_b = pipeline.scheduler.timesteps.clone()
+
+    torch.testing.assert_close(seq_a, seq_b)

tests/local_tests/cosmos3/test_cosmos3_state_dict_keys.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 state-dict-key contract for checkpoint conversion (Tier A scaffold).
+
+Reference: ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:319-409``
+(``Cosmos3OmniDiffusersPipeline._remap_ckpt_key``). The vllm-omni loader's
+remap is the canonical source of truth: it defines which Diffusers
+checkpoint keys become which target module-tree parameter names.
+
+The FastVideo conversion script at
+``scripts/checkpoint_conversion/cosmos3_convert.py`` (to be authored in
+Phase 5) must produce a state-dict that loads cleanly into the FastVideo
+Cosmos3 DiT module tree. This test pins the target side of that contract:
+
+  * ``embed_tokens``/``norm`` live under ``language_model.``
+  * ``norm_moe_gen`` is a top-level (not per-layer) module
+  * Per-layer keys split into UND (``language_model.layers.{i}``) and
+    GEN (``gen_layers.{i}``) sub-trees:
+      - UND attention: ``self_attn.{q,k,v,o}_proj``, ``self_attn.{q,k}_norm``
+      - GEN cross-attention: ``cross_attention.{q,k,v,o}_proj``,
+        ``cross_attention.{q,k}_norm`` (from the ``*_moe_gen`` source keys)
+      - UND/GEN norms: ``input_layernorm``, ``post_attention_layernorm``
+      - UND/GEN MLPs: ``mlp.{gate,up,down}_proj``
+  * Top-level adapters ``vae2llm``, ``llm2vae``, ``time_embedder`` live
+    under ``transformer.`` (no remapping inside the transformer namespace).
+  * ``lm_head.weight`` is skipped (mapped to ``None``).
+"""
+from __future__ import annotations
+
+import pytest
+
+pytestmark = [pytest.mark.local]
+
+
+EXPECTED_REMAPS: dict[str, str | None] = {
+    "model.embed_tokens.weight": "transformer.language_model.embed_tokens.weight",
+    "model.norm.weight": "transformer.language_model.norm.weight",
+    "model.norm_moe_gen.weight": "transformer.norm_moe_gen.weight",
+    "model.layers.3.self_attn.q_proj.weight": "transformer.language_model.layers.3.self_attn.q_proj.weight",
+    "model.layers.3.self_attn.q_proj_moe_gen.weight": "transformer.gen_layers.3.cross_attention.q_proj.weight",
+    "model.layers.3.self_attn.k_norm_moe_gen.weight": "transformer.gen_layers.3.cross_attention.k_norm.weight",
+    "model.layers.3.input_layernorm.weight": "transformer.language_model.layers.3.input_layernorm.weight",
+    "model.layers.3.input_layernorm_moe_gen.weight": "transformer.gen_layers.3.input_layernorm.weight",
+    "model.layers.3.mlp.gate_proj.weight": "transformer.language_model.layers.3.mlp.gate_proj.weight",
+    "model.layers.3.mlp_moe_gen.up_proj.weight": "transformer.gen_layers.3.mlp.up_proj.weight",
+    "vae2llm.weight": "transformer.vae2llm.weight",
+    "llm2vae.weight": "transformer.llm2vae.weight",
+    "time_embedder.linear_1.weight": "transformer.time_embedder.linear_1.weight",
+    "lm_head.weight": None,
+}
+
+
+def test_fastvideo_cosmos3_remap_matches_reference() -> None:
+    """Asserts the FastVideo Cosmos3 checkpoint remap table matches the
+    canonical vllm-omni reference at pipeline_cosmos3.py:319-409.
+
+    Once the FastVideo Cosmos3 pipeline (or a standalone remap utility
+    in ``scripts/checkpoint_conversion/cosmos3_convert.py``) exposes a
+    ``_remap_ckpt_key`` callable, this test verifies every key in
+    ``EXPECTED_REMAPS`` maps to the expected target name (or ``None``
+    for skipped keys like ``lm_head.weight``).
+    """
+    try:
+        from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (  # type: ignore
+            Cosmos3OmniDiffusersPipeline,
+        )
+    except ImportError:
+        pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)")
+
+    remap = getattr(Cosmos3OmniDiffusersPipeline, "_remap_ckpt_key", None)
+    if remap is None:
+        pytest.skip("Cosmos3OmniDiffusersPipeline._remap_ckpt_key not yet defined")
+
+    actual = {key: remap(key) for key in EXPECTED_REMAPS}
+    assert actual == EXPECTED_REMAPS
+
+
+def test_fastvideo_cosmos3_dit_module_tree_param_names() -> None:
+    """Asserts the FastVideo Cosmos3 DiT module tree produces the param
+    names that the checkpoint converter is expected to write.
+
+    Once ``fastvideo.models.dits.cosmos3.Cosmos3VFMTransformer`` is
+    constructible with a tiny config (see ``conftest._tiny_cosmos3_config``),
+    this test instantiates a 1-layer model and verifies that the
+    parameter-name set contains the expected UND/GEN/adapter prefixes:
+      * ``language_model.layers.0.self_attn.q_proj.weight``
+      * ``gen_layers.0.cross_attention.q_proj.weight``
+      * ``norm_moe_gen.weight``
+      * ``vae2llm.*`` and ``llm2vae.*``
+    """
+    try:
+        from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer  # type: ignore
+    except ImportError:
+        pytest.skip("FastVideo Cosmos3 not yet implemented (Phase 2b)")
+
+    from types import SimpleNamespace
+    import torch
+
+    from .conftest import _tiny_cosmos3_config
+
+    config = _tiny_cosmos3_config(num_hidden_layers=1)
+    model = Cosmos3VFMTransformer(SimpleNamespace(tf_model_config=config, dtype=torch.float32))
+    names = set(name for name, _ in model.named_parameters())
+
+    required_prefixes = [
+        "language_model.layers.0.self_attn.q_proj",
+        "gen_layers.0.cross_attention.q_proj",
+        "norm_moe_gen",
+        "vae2llm",
+        "llm2vae",
+    ]
+    for prefix in required_prefixes:
+        assert any(name.startswith(prefix) for name in names), (
+            f"Cosmos3 DiT module tree missing expected prefix {prefix!r}; "
+            f"converter at scripts/checkpoint_conversion/cosmos3_convert.py "
+            f"will fail to load weights for this subtree."
+        )

tests/local_tests/cosmos3/test_cosmos3_tokenizer_chat_template.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos3 prompt tokenization parity — chat template + special tokens (Tier A).
+
+Reference: ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py``
+lines 562-606 (``Cosmos3OmniDiffusersPipeline._tokenize_prompt``).
+
+The Cosmos3 tokenizer is a Qwen2 chat-template tokenizer with two
+appended special tokens:
+
+  * ``eos_token_id`` == 151645
+  * ``<|vision_start|>`` == 151652
+
+The reference pipeline:
+  1. Wraps the prompt in a ``role=user`` conversation (optionally with a
+     ``role=system`` prefix carrying ``COSMOS3_SYSTEM_PROMPT``);
+  2. Applies the chat template with ``add_generation_prompt=True``;
+  3. Truncates to ``max_sequence_length``;
+  4. Appends ``eos`` and ``<|vision_start|>``;
+  5. Right-pads with ``pad_token_id`` and produces a matching
+     attention mask of ``[1]*seq_len + [0]*pad_len``.
+
+This Tier A scaffold validates the two special-token IDs and the
+right-pad shape contract. Full byte-for-byte parity against an
+official Cosmos3 tokenizer requires Phase 2b weights.
+"""
+from __future__ import annotations
+
+import pytest
+
+pytestmark = [pytest.mark.local]
+
+
+def test_cosmos3_special_token_ids() -> None:
+    """Asserts the two Cosmos3 special-token IDs the pipeline depends on.
+
+    Specifically ``eos_token_id == 151645`` and
+    ``convert_tokens_to_ids('<|vision_start|>') == 151652``. These IDs
+    are pinned by the Qwen2 base tokenizer and are appended in
+    ``Cosmos3OmniDiffusersPipeline._tokenize_prompt`` at
+    pipeline_cosmos3.py:596-597.
+    """
+    try:
+        from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (  # type: ignore
+            Cosmos3OmniDiffusersPipeline,
+        )
+    except ImportError:
+        pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)")
+
+    pipeline = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline)
+    tokenizer = getattr(pipeline, "tokenizer", None)
+    if tokenizer is None:
+        pytest.skip("Cosmos3 tokenizer instance not yet wired on the pipeline")
+
+    assert tokenizer.eos_token_id == 151645
+    assert tokenizer.convert_tokens_to_ids("<|vision_start|>") == 151652
+
+
+def test_cosmos3_chat_template_shape_contract() -> None:
+    """Asserts the right-pad shape contract of ``_tokenize_prompt``.
+
+    For ``max_sequence_length = N``, the returned ``input_ids`` and
+    ``attention_mask`` must each be ``[1, N]``, with the first
+    ``seq_len`` mask entries equal to 1 and the remainder equal to 0.
+
+    Once Phase 2b lands, this test should:
+      * call ``_tokenize_prompt`` directly with a small max_seq_length;
+      * assert ``input_ids.shape == (1, max_seq_length)``;
+      * assert ``attention_mask.sum().item() == seq_len`` and that the
+        last two non-pad tokens are ``[eos, vision_start]``.
+    """
+    try:
+        from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (  # type: ignore
+            Cosmos3OmniDiffusersPipeline,
+        )
+    except ImportError:
+        pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)")
+
+    pipeline = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline)
+    if not hasattr(pipeline, "_tokenize_prompt"):
+        pytest.skip("_tokenize_prompt not yet wired on the FastVideo pipeline")
+
+    input_ids, attention_mask = pipeline._tokenize_prompt(
+        "A robot.", max_sequence_length=32, use_system_prompt=False
+    )
+    assert tuple(input_ids.shape) == (1, 32)
+    assert tuple(attention_mask.shape) == (1, 32)
+    seq_len = int(attention_mask.sum().item())
+    assert attention_mask[0, :seq_len].sum().item() == seq_len
+    assert attention_mask[0, seq_len:].sum().item() == 0
+    assert int(input_ids[0, seq_len - 2].item()) == 151645
+    assert int(input_ids[0, seq_len - 1].item()) == 151652