sgl-project--sglang/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py

# SPDX-License-Identifier: Apache-2.0
"""Cosmos3 Omni transformer.

Dual-pathway DiT: an Understanding (UND) pathway runs causal self-attention
over the text tokens once and caches its K/V; a Generation (GEN) pathway
cross-attends from noisy visual tokens to that cache at every denoising step.
"""

import math
from collections.abc import Iterable, Iterator
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F

from sglang.multimodal_gen.configs.models.dits.cosmos3video import Cosmos3VideoConfig
from sglang.multimodal_gen.runtime.distributed import (
    get_sp_group,
    get_sp_world_size,
    get_tp_world_size,
    sequence_model_parallel_all_gather,
)
from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.layers.layernorm import (
    RMSNorm,
    apply_qk_norm,
    apply_qk_norm_rope,
)
from sglang.multimodal_gen.runtime.layers.linear import (
    MergedColumnParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
    UnquantizedLinearMethod,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
    QuantizationConfig,
)
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
    Qwen3VLTextRotaryEmbedding,
)
from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding
from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
    VocabParallelEmbedding,
)
from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.srt.utils import add_prefix

logger = init_logger(__name__)


# -----------------------------------------------------------------------------
# mRoPE position ID computation (Qwen3VL-style)
# -----------------------------------------------------------------------------


def compute_mrope_position_ids_text(
    num_tokens: int,
    temporal_offset: int,
    device: torch.device,
) -> tuple[torch.Tensor, int]:
    """Generate 3D mRoPE position IDs for text tokens.

    Text tokens: all three axes (T, H, W) share the same monotonically
    increasing position IDs: (0,0,0), (1,1,1), (2,2,2), ...

    Returns:
        (position_ids [3, num_tokens], next_temporal_offset)
    """
    ids = torch.arange(num_tokens, dtype=torch.long, device=device) + temporal_offset
    mrope_ids = ids.unsqueeze(0).expand(3, -1).contiguous()
    return mrope_ids, temporal_offset + num_tokens


def compute_mrope_position_ids_vision(
    grid_t: int,
    grid_h: int,
    grid_w: int,
    temporal_offset: int | float,
    device: torch.device,
    fps: float | None = None,
    base_fps: float = 24.0,
    temporal_compression_factor: int = 4,
    base_temporal_compression_factor: int | None = None,
    start_frame_offset: int = 0,
) -> tuple[torch.Tensor, int | float]:
    """Generate 3D mRoPE position IDs for vision tokens.

    Creates a (T, H, W) position grid. Spatial indices reset to 0
    per vision segment (Qwen3VL-style).
    Flattened in T-major order.

    When the token rate (``temporal_compression_factor``) differs from the
    base-fps rate (``base_temporal_compression_factor``), the two no longer
    cancel: action tokens run at frame rate (factor 1) while their temporal
    positions are scaled by the video factor. ``base_temporal_compression_factor``
    defaults to ``temporal_compression_factor`` so vision/sound are unchanged.

    Returns:
        (position_ids [3, grid_t * grid_h * grid_w], next_temporal_offset)
    """
    fps_modulation = fps is not None and grid_t > 1

    if fps_modulation:
        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, device=device)
        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, device=device)
            .view(-1, 1)
            .expand(-1, grid_h * grid_w)
            .flatten()
            + int(temporal_offset)
            + start_frame_offset
        )

    h_index = (
        torch.arange(grid_h, dtype=torch.long, device=device)
        .view(1, -1, 1)
        .expand(grid_t, -1, grid_w)
        .flatten()
    )
    w_index = (
        torch.arange(grid_w, dtype=torch.long, device=device)
        .view(1, 1, -1)
        .expand(grid_t, grid_h, -1)
        .flatten()
    )

    if fps_modulation:
        mrope_ids = torch.stack(
            [t_index, h_index.to(torch.float32), w_index.to(torch.float32)], dim=0
        )
    else:
        mrope_ids = torch.stack([t_index, h_index, w_index], dim=0)

    next_offset = math.ceil(mrope_ids.max().item()) + 1
    return mrope_ids, next_offset


def compute_mrope_position_ids_sound(
    grid_t: int,
    temporal_offset: int | float,
    sound_latent_fps: float,
    device: torch.device,
    base_fps: float = 24.0,
    temporal_compression_factor_sound: int = 1,
) -> tuple[torch.Tensor, int | float]:
    """mRoPE position IDs for sound tokens: a (T, 1, 1) grid."""
    return compute_mrope_position_ids_vision(
        grid_t=grid_t,
        grid_h=1,
        grid_w=1,
        temporal_offset=temporal_offset,
        device=device,
        fps=sound_latent_fps,
        base_fps=base_fps,
        temporal_compression_factor=temporal_compression_factor_sound,
    )


def compute_mrope_position_ids_action(
    grid_t: int,
    temporal_offset: int | float,
    action_fps: float | None,
    device: torch.device,
    base_fps: float = 24.0,
    base_temporal_compression_factor: int = 4,
    start_frame_offset: int = 1,
) -> tuple[torch.Tensor, int | float]:
    """mRoPE position IDs for action tokens: a (T, 1, 1) grid.

    Action tokens run at frame rate (``temporal_compression_factor=1``) but
    their positions are scaled by the video's ``base_temporal_compression_factor``
    so they share the video's temporal coordinate frame. ``start_frame_offset=1``
    (default) shifts them one frame ahead so they align with the video frame
    they condition on.
    """
    return compute_mrope_position_ids_vision(
        grid_t=grid_t,
        grid_h=1,
        grid_w=1,
        temporal_offset=temporal_offset,
        device=device,
        fps=action_fps,
        base_fps=base_fps,
        temporal_compression_factor=1,
        base_temporal_compression_factor=base_temporal_compression_factor,
        start_frame_offset=start_frame_offset,
    )


# -----------------------------------------------------------------------------
# Qwen3-style RoPE functions
# -----------------------------------------------------------------------------


def _apply_qwen3_qk_norm_rope(
    q: torch.Tensor,
    k: torch.Tensor,
    q_norm: RMSNorm,
    k_norm: RMSNorm,
    head_dim: int,
    cos_sin_cache: torch.Tensor,
    rope_cache_positions: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    return apply_qk_norm_rope(
        q=q.contiguous(),
        k=k.contiguous(),
        q_norm=q_norm,
        k_norm=k_norm,
        head_dim=head_dim,
        cos_sin_cache=cos_sin_cache,
        is_neox=True,
        positions=rope_cache_positions,
    )


def _apply_qwen3_rope_from_cache(
    q: torch.Tensor, k: torch.Tensor, cos_sin_cache: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    batch_size, seq_len = q.shape[:2]
    half = q.shape[-1] // 2
    cos = cos_sin_cache[:, :half].view(batch_size, seq_len, 1, half).to(q.dtype)
    sin = cos_sin_cache[:, half:].view(batch_size, seq_len, 1, half).to(q.dtype)

    q1 = q[..., :half]
    q2 = q[..., half:]
    q_out = torch.empty_like(q)
    q_out[..., :half] = q1 * cos - q2 * sin
    q_out[..., half:] = q2 * cos + q1 * sin

    k1 = k[..., :half]
    k2 = k[..., half:]
    k_out = torch.empty_like(k)
    k_out[..., :half] = k1 * cos - k2 * sin
    k_out[..., half:] = k2 * cos + k1 * sin
    return q_out, k_out


def _apply_qwen3_qk_norm_rope_split(
    q: torch.Tensor,
    k: torch.Tensor,
    q_norm: RMSNorm,
    k_norm: RMSNorm,
    head_dim: int,
    cos_sin_cache: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    q, k = apply_qk_norm(q.contiguous(), k.contiguous(), q_norm, k_norm, head_dim)
    return _apply_qwen3_rope_from_cache(q, k, cos_sin_cache)


# -----------------------------------------------------------------------------
# Action domain-aware projection
# -----------------------------------------------------------------------------


class DomainAwareLinear(nn.Module):
    """Per-domain linear projection for action conditioning.

    Maintains one weight matrix and bias per embodiment domain via embedding
    tables, enabling multi-domain robot action generation from a shared
    backbone.
    """

    def __init__(self, input_size: int, output_size: int, num_domains: int) -> None:
        super().__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.num_domains = num_domains
        self.fc = nn.Embedding(num_domains, output_size * input_size)
        self.bias = nn.Embedding(num_domains, output_size)
        nn.init.xavier_uniform_(
            self.fc.weight.view(num_domains, output_size, input_size)
        )
        nn.init.zeros_(self.bias.weight)

    def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor:
        if domain_id.ndim == 0:
            domain_id = domain_id.unsqueeze(0)
        domain_id = domain_id.to(device=x.device, dtype=torch.long).reshape(-1)
        weight = self.fc(domain_id).view(
            domain_id.shape[0], self.input_size, self.output_size
        )
        bias = self.bias(domain_id).view(domain_id.shape[0], self.output_size)
        if x.ndim == 2:
            return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias
        return torch.bmm(x, weight) + bias.unsqueeze(1)


# -----------------------------------------------------------------------------
# Cosmos3 Timestep Embedder
# -----------------------------------------------------------------------------


class Cosmos3TimestepEmbedder(nn.Module):
    """Embeds scalar timesteps into vector representations.

    Uses ReplicatedLinear for consistency with other SGLang models and
    to support quantization (though timestep embedders are typically excluded).
    """

    def __init__(
        self,
        hidden_size: int,
        frequency_embedding_size: int = 256,
        max_period: int = 10000,
        timestep_scale: float = 0.001,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        self.timestep_scale = timestep_scale
        self.frequency_embedding_size = frequency_embedding_size
        self.hidden_size = hidden_size
        self.max_period = max_period

        # Use ReplicatedLinear for consistency (typically excluded from quantization)
        self.linear_1 = ReplicatedLinear(
            frequency_embedding_size,
            hidden_size,
            bias=True,
            quant_config=quant_config,
            prefix=add_prefix("linear_1", prefix),
        )
        self.act = nn.SiLU()
        self.linear_2 = ReplicatedLinear(
            hidden_size,
            hidden_size,
            bias=True,
            quant_config=quant_config,
            prefix=add_prefix("linear_2", prefix),
        )

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        """Embed timesteps.

        Args:
            t: [B] timestep values

        Returns:
            [B, hidden_size] timestep embeddings
        """
        # Scale timestep
        t_scaled = t * self.timestep_scale

        # Compute sinusoidal embeddings in fp32
        t_freq = timestep_embedding(
            t_scaled,
            self.frequency_embedding_size,
            self.max_period,
            dtype=torch.float32,
        )

        # Project through MLP
        # When fp8-quantized, weight.dtype is float8_e4m3fn — keep input in
        # float32 (the quant kernel handles input quantization internally).
        w_dtype = self.linear_1.weight.dtype
        if w_dtype.is_floating_point and w_dtype.itemsize >= 2:
            x = t_freq.to(w_dtype)
        else:
            x = t_freq  # already float32 from timestep_embedding
        x, _ = self.linear_1(x)
        x = self.act(x)
        x, _ = self.linear_2(x)
        return x


# -----------------------------------------------------------------------------
# Cosmos3 Gated MLP
# -----------------------------------------------------------------------------


class Cosmos3GatedMLP(nn.Module):
    """Gated MLP (SwiGLU-style) for Cosmos3."""

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size

        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size,
            [intermediate_size] * 2,
            bias=False,
            gather_output=False,
            quant_config=quant_config,
            prefix=add_prefix("gate_up_proj", prefix),
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
            input_is_parallel=True,
            quant_config=quant_config,
            prefix=add_prefix("down_proj", prefix),
        )
        self.act_fn = SiluAndMul()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate_up, _ = self.gate_up_proj(x)
        out, _ = self.down_proj(self.act_fn(gate_up))
        return out


# -----------------------------------------------------------------------------
# Cosmos3 UND Causal Attention
# -----------------------------------------------------------------------------


class Cosmos3CausalAttention(nn.Module):
    """Understanding pathway: causal self-attention on text tokens."""

    def __init__(
        self,
        hidden_size: int,
        num_attention_heads: int,
        num_key_value_heads: int,
        head_dim: int,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.tp_size = get_tp_world_size()
        if num_attention_heads % self.tp_size != 0:
            raise ValueError(
                "Cosmos3CausalAttention requires num_attention_heads divisible "
                f"by tp_size, got {num_attention_heads=} {self.tp_size=}."
            )
        if num_key_value_heads % self.tp_size != 0:
            raise ValueError(
                "Cosmos3CausalAttention requires num_key_value_heads divisible "
                f"by tp_size, got {num_key_value_heads=} {self.tp_size=}."
            )
        self.local_num_attention_heads = num_attention_heads // self.tp_size
        self.local_num_key_value_heads = num_key_value_heads // self.tp_size

        self.q_size = num_attention_heads * head_dim
        self.kv_size = num_key_value_heads * head_dim
        self.to_qkv = MergedColumnParallelLinear(
            hidden_size,
            [self.q_size, self.kv_size, self.kv_size],
            bias=False,
            gather_output=False,
            quant_config=quant_config,
            prefix=add_prefix("to_qkv", prefix),
        )
        self.to_out = RowParallelLinear(
            num_attention_heads * head_dim,
            hidden_size,
            bias=False,
            input_is_parallel=True,
            quant_config=quant_config,
            prefix=add_prefix("to_out", prefix),
        )

        # Per-head QK norm.
        self.norm_q = RMSNorm(head_dim, eps=1e-6)
        self.norm_k = RMSNorm(head_dim, eps=1e-6)

    def forward(
        self,
        hidden_states: torch.Tensor,
        cos_sin_cache: torch.Tensor,
        rope_cache_positions: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Forward with KV cache return.

        Returns:
            (output, K, V) where K/V are TP-local, post-norm, post-RoPE
        """
        batch_size, seq_len = hidden_states.shape[:2]

        qkv, _ = self.to_qkv(hidden_states)
        qkv = qkv.view(
            batch_size,
            seq_len,
            self.local_num_attention_heads + 2 * self.local_num_key_value_heads,
            self.head_dim,
        )
        q = qkv[:, :, : self.local_num_attention_heads, :]
        k = qkv[
            :,
            :,
            self.local_num_attention_heads : self.local_num_attention_heads
            + self.local_num_key_value_heads,
            :,
        ]
        v = qkv[
            :,
            :,
            self.local_num_attention_heads + self.local_num_key_value_heads :,
            :,
        ]

        q = F.rms_norm(
            q, (self.head_dim,), self.norm_q.weight, self.norm_q.variance_epsilon
        )
        k = F.rms_norm(
            k, (self.head_dim,), self.norm_k.weight, self.norm_k.variance_epsilon
        )
        q, k = _apply_qwen3_rope_from_cache(q, k, cos_sin_cache)

        out = F.scaled_dot_product_attention(
            q.transpose(1, 2),
            k.transpose(1, 2),
            v.transpose(1, 2),
            is_causal=True,
            enable_gqa=True,
        )
        out = out.transpose(1, 2).reshape(batch_size, seq_len, -1)

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


# -----------------------------------------------------------------------------
# Cosmos3 GEN Cross Attention
# -----------------------------------------------------------------------------


class Cosmos3CrossAttention(nn.Module):
    """Generation pathway: cross-attention where visual Q attends to all K/V."""

    def __init__(
        self,
        hidden_size: int,
        num_attention_heads: int,
        num_key_value_heads: int,
        head_dim: int,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
        supported_attention_backends: set | None = None,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.tp_size = get_tp_world_size()
        if num_attention_heads % self.tp_size != 0:
            raise ValueError(
                "Cosmos3CrossAttention requires num_attention_heads divisible "
                f"by tp_size, got {num_attention_heads=} {self.tp_size=}."
            )
        if num_key_value_heads % self.tp_size != 0:
            raise ValueError(
                "Cosmos3CrossAttention requires num_key_value_heads divisible "
                f"by tp_size, got {num_key_value_heads=} {self.tp_size=}."
            )
        self.local_num_attention_heads = num_attention_heads // self.tp_size
        self.local_num_key_value_heads = num_key_value_heads // self.tp_size

        self.q_size = num_attention_heads * head_dim
        self.kv_size = num_key_value_heads * head_dim
        self.to_qkv = MergedColumnParallelLinear(
            hidden_size,
            [self.q_size, self.kv_size, self.kv_size],
            bias=False,
            gather_output=False,
            quant_config=quant_config,
            prefix=add_prefix("to_qkv", prefix),
        )
        self.to_out = RowParallelLinear(
            num_attention_heads * head_dim,
            hidden_size,
            bias=False,
            input_is_parallel=True,
            quant_config=quant_config,
            prefix=add_prefix("to_out", prefix),
        )

        self.norm_q = RMSNorm(head_dim, eps=1e-6)
        self.norm_k = RMSNorm(head_dim, eps=1e-6)

        self.attn = USPAttention(
            num_heads=self.local_num_attention_heads,
            head_size=head_dim,
            num_kv_heads=self.local_num_key_value_heads,
            causal=False,
            supported_attention_backends=supported_attention_backends,
            prefix=add_prefix("attn", prefix),
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        k_und: torch.Tensor,
        v_und: torch.Tensor,
        cos_sin_cache: torch.Tensor,
        rope_cache_positions: torch.Tensor,
        use_fused_qk_norm_rope: bool,
    ) -> torch.Tensor:
        """Cross-attention from GEN to cached UND K/V.

        Args:
            hidden_states: [B, S_gen_local, hidden_size] visual tokens (may be sharded)
            k_und: [B, S_und, H_kv_local, D] UND keys (replicated over SP)
            v_und: [B, S_und, H_kv_local, D] UND values (replicated over SP)
            cos_sin_cache: [B*S_gen_local, D] local rows of [cos, sin]
            rope_cache_positions: identity row positions into cos_sin_cache
        """
        batch_size, seq_len_gen = hidden_states.shape[:2]

        qkv, _ = self.to_qkv(hidden_states)
        qkv = qkv.view(
            batch_size,
            seq_len_gen,
            self.local_num_attention_heads + 2 * self.local_num_key_value_heads,
            self.head_dim,
        )
        q = qkv[:, :, : self.local_num_attention_heads, :]
        k = qkv[
            :,
            :,
            self.local_num_attention_heads : self.local_num_attention_heads
            + self.local_num_key_value_heads,
            :,
        ]
        v = qkv[
            :,
            :,
            self.local_num_attention_heads + self.local_num_key_value_heads :,
            :,
        ]

        if use_fused_qk_norm_rope:
            q, k = _apply_qwen3_qk_norm_rope(
                q,
                k,
                self.norm_q,
                self.norm_k,
                self.head_dim,
                cos_sin_cache,
                rope_cache_positions,
            )
        else:
            q, k = _apply_qwen3_qk_norm_rope_split(
                q, k, self.norm_q, self.norm_k, self.head_dim, cos_sin_cache
            )

        # K/V = [text (replicated on every SP rank) | image (sharded same as Q)].
        # USPAttention routes through the registered attention backend (FA, sage,
        # …) and handles the Ulysses all-to-all when SP > 1.
        out = self.attn.forward_with_replicated_kv_prefix(q, k_und, v_und, k, v)
        out = out.reshape(batch_size, seq_len_gen, -1)
        out, _ = self.to_out(out)
        return out


# -----------------------------------------------------------------------------
# Cosmos3 UND Decoder Layer
# -----------------------------------------------------------------------------


class Cosmos3UndDecoderLayer(nn.Module):
    """Understanding pathway decoder layer: causal self-attention + MLP."""

    def __init__(
        self,
        hidden_size: int,
        num_attention_heads: int,
        num_key_value_heads: int,
        head_dim: int,
        intermediate_size: int,
        rms_norm_eps: float,
        layer_idx: int,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        self.layer_idx = layer_idx

        self.self_attn = Cosmos3CausalAttention(
            hidden_size=hidden_size,
            num_attention_heads=num_attention_heads,
            num_key_value_heads=num_key_value_heads,
            head_dim=head_dim,
            prefix=add_prefix("self_attn", prefix),
            quant_config=quant_config,
        )
        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
        self.mlp = Cosmos3GatedMLP(
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            prefix=add_prefix("mlp", prefix),
            quant_config=quant_config,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        cos_sin_cache: torch.Tensor,
        rope_cache_positions: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Forward pass.

        Returns:
            (hidden_states, K, V) where K/V are for GEN cross-attention
        """
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)

        attn_out, k, v = self.self_attn(
            hidden_states, cos_sin_cache, rope_cache_positions
        )
        hidden_states = residual + attn_out

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = residual + self.mlp(hidden_states)

        return hidden_states, k, v


# -----------------------------------------------------------------------------
# Cosmos3 GEN Decoder Layer
# -----------------------------------------------------------------------------


class Cosmos3GenDecoderLayer(nn.Module):
    """Generation pathway decoder layer: cross-attention + MLP."""

    def __init__(
        self,
        hidden_size: int,
        num_attention_heads: int,
        num_key_value_heads: int,
        head_dim: int,
        intermediate_size: int,
        rms_norm_eps: float,
        layer_idx: int,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
        supported_attention_backends: set | None = None,
    ):
        super().__init__()
        self.layer_idx = layer_idx

        self.cross_attention = Cosmos3CrossAttention(
            hidden_size=hidden_size,
            num_attention_heads=num_attention_heads,
            num_key_value_heads=num_key_value_heads,
            head_dim=head_dim,
            prefix=add_prefix("cross_attention", prefix),
            quant_config=quant_config,
            supported_attention_backends=supported_attention_backends,
        )
        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
        self.mlp = Cosmos3GatedMLP(
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            prefix=add_prefix("mlp", prefix),
            quant_config=quant_config,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        k_und: torch.Tensor,
        v_und: torch.Tensor,
        cos_sin_cache: torch.Tensor,
        rope_cache_positions: torch.Tensor,
        use_fused_qk_norm_rope: bool,
        residual: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # Fused add+rmsnorm: each `(hidden_states, residual) = norm(...)`
        # collapses the residual add and RMSNorm into one kernel. The
        # caller threads `residual` across layers and resolves it before
        # the post-loop all-gather + `norm_moe_gen`.
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(hidden_states, residual)

        hidden_states = self.cross_attention(
            hidden_states,
            k_und,
            v_und,
            cos_sin_cache,
            rope_cache_positions,
            use_fused_qk_norm_rope,
        )

        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
        hidden_states = self.mlp(hidden_states)

        return hidden_states, residual


# -----------------------------------------------------------------------------
# Cosmos3 Language Model (UND pathway)
# -----------------------------------------------------------------------------


class Cosmos3LanguageModel(nn.Module):
    """Understanding pathway: processes text tokens and caches K/V."""

    def __init__(
        self,
        hidden_size: int,
        num_hidden_layers: int,
        num_attention_heads: int,
        num_key_value_heads: int,
        head_dim: int,
        intermediate_size: int,
        vocab_size: int,
        rms_norm_eps: float,
        rope_theta: float,
        mrope_section: tuple[int, int, int],
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers

        self.embed_tokens = VocabParallelEmbedding(
            vocab_size,
            hidden_size,
            params_dtype=torch.bfloat16,
        )
        self.rotary_emb = Qwen3VLTextRotaryEmbedding(
            head_dim=head_dim,
            rope_theta=rope_theta,
            mrope_section=mrope_section,
        )
        self.layers = nn.ModuleList(
            [
                Cosmos3UndDecoderLayer(
                    hidden_size=hidden_size,
                    num_attention_heads=num_attention_heads,
                    num_key_value_heads=num_key_value_heads,
                    head_dim=head_dim,
                    intermediate_size=intermediate_size,
                    rms_norm_eps=rms_norm_eps,
                    layer_idx=i,
                    prefix=f"layers.{i}",
                    quant_config=quant_config,
                )
                for i in range(num_hidden_layers)
            ]
        )

    def forward(
        self,
        text_ids: torch.Tensor,
        text_mask: torch.Tensor,
        position_ids: torch.Tensor,
    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
        """Process text tokens and return per-layer K/V cache.

        Args:
            text_ids: [B, S] token IDs
            text_mask: [B, S] float mask (1=real, 0=pad)
            position_ids: [3, B, S] mRoPE position IDs

        Returns:
            List of (K, V) per layer for GEN cross-attention
        """
        hidden = self.embed_tokens(text_ids)
        mask_3d = text_mask.unsqueeze(-1)
        cos_sin_cache, rope_cache_positions = self.rotary_emb.build_rope_cache_inputs(
            position_ids, cache_dtype=hidden.dtype
        )

        cached_kv: list[tuple[torch.Tensor, torch.Tensor]] = []
        for layer in self.layers:
            hidden = hidden * mask_3d
            hidden, k, v = layer(hidden, cos_sin_cache, rope_cache_positions)
            cached_kv.append((k, v))

        return cached_kv


# -----------------------------------------------------------------------------
# Cosmos3 Omni Transformer
# -----------------------------------------------------------------------------


class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
    """Cosmos3 Omni transformer.

    Dual-pathway architecture:
    - Understanding (UND): causal LM processing text
    - Generation (GEN): cross-attention from visual to UND K/V
    """

    _fsdp_shard_conditions = Cosmos3VideoConfig()._fsdp_shard_conditions
    _compile_conditions = Cosmos3VideoConfig()._compile_conditions
    _supported_attention_backends = Cosmos3VideoConfig()._supported_attention_backends
    param_names_mapping = Cosmos3VideoConfig().arch_config.param_names_mapping
    reverse_param_names_mapping = (
        Cosmos3VideoConfig().arch_config.reverse_param_names_mapping
    )
    lora_param_names_mapping = Cosmos3VideoConfig().arch_config.lora_param_names_mapping

    def __init__(
        self,
        config: Cosmos3VideoConfig,
        hf_config: dict[str, Any],
        quant_config: QuantizationConfig | None = None,
    ) -> None:
        super().__init__(config=config, hf_config=hf_config)

        arch = config.arch_config
        self.hidden_size = arch.hidden_size
        self.num_hidden_layers = arch.num_hidden_layers
        self.num_attention_heads = arch.num_attention_heads
        self.num_key_value_heads = arch.num_key_value_heads
        self.head_dim = arch.head_dim
        self.intermediate_size = arch.intermediate_size
        self.latent_patch_size = arch.latent_patch_size
        self.latent_channel = arch.latent_channel
        self.num_channels_latents = arch.out_channels
        self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel
        self.timestep_scale = arch.timestep_scale
        self.base_fps = arch.base_fps
        self.temporal_compression_factor = arch.temporal_compression_factor
        self.temporal_margin = arch.unified_3d_mrope_temporal_modality_margin
        self.sound_latent_fps = arch.sound_latent_fps
        self.temporal_compression_factor_sound = arch.temporal_compression_factor_sound
        self.rms_norm_eps = arch.rms_norm_eps

        # Ulysses sequence parallelism. When CFG-parallel is also enabled
        # the SP group only spans ranks that share a CFG context (cond or
        # uncond), so ``sp_size`` here is the per-context shard count.
        self.sp_size = get_sp_world_size()
        self.sp_group = get_sp_group() if self.sp_size > 1 else None
        self.sp_rank = self.sp_group.rank_in_group if self.sp_group else 0
        if self.sp_size > 1:
            logger.info(
                f"Cosmos3 SP enabled: sp_size={self.sp_size}, sp_rank={self.sp_rank}"
            )

        # Language model (UND pathway)
        self.language_model = Cosmos3LanguageModel(
            hidden_size=arch.hidden_size,
            num_hidden_layers=arch.num_hidden_layers,
            num_attention_heads=arch.num_attention_heads,
            num_key_value_heads=arch.num_key_value_heads,
            head_dim=arch.head_dim,
            intermediate_size=arch.intermediate_size,
            vocab_size=arch.vocab_size,
            rms_norm_eps=arch.rms_norm_eps,
            rope_theta=arch.rope_theta,
            mrope_section=arch.mrope_section,
            quant_config=quant_config,
        )

        # Latent projection layers - ReplicatedLinear for quantization support
        self.proj_in = ReplicatedLinear(
            self.patch_latent_dim,
            self.hidden_size,
            bias=True,
            quant_config=quant_config,
            prefix="proj_in",
        )
        self.proj_out = ReplicatedLinear(
            self.hidden_size,
            self.patch_latent_dim,
            bias=True,
            quant_config=quant_config,
            prefix="proj_out",
        )

        self.sound_gen = arch.sound_gen
        self.sound_dim = arch.sound_dim
        if arch.sound_gen:
            self.audio_proj_in = ReplicatedLinear(
                self.sound_dim,
                self.hidden_size,
                bias=True,
                quant_config=quant_config,
                prefix="audio_proj_in",
            )
            self.audio_proj_out = ReplicatedLinear(
                self.hidden_size,
                self.sound_dim,
                bias=True,
                quant_config=quant_config,
                prefix="audio_proj_out",
            )
            self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))

        if arch.action_gen:
            self.action_dim = arch.action_dim
            self.num_embodiment_domains = arch.num_embodiment_domains
            self.action_proj_in = DomainAwareLinear(
                self.action_dim,
                self.hidden_size,
                self.num_embodiment_domains,
            )
            self.action_proj_out = DomainAwareLinear(
                self.hidden_size,
                self.action_dim,
                self.num_embodiment_domains,
            )
            self.action_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))

        # Timestep embedder
        self.time_embedder = Cosmos3TimestepEmbedder(
            hidden_size=self.hidden_size,
            frequency_embedding_size=arch.frequency_embedding_size,
            timestep_scale=arch.timestep_scale,
            prefix="time_embedder",
            quant_config=quant_config,
        )

        # Generation layers (GEN pathway)
        self.gen_layers = nn.ModuleList(
            [
                Cosmos3GenDecoderLayer(
                    hidden_size=arch.hidden_size,
                    num_attention_heads=arch.num_attention_heads,
                    num_key_value_heads=arch.num_key_value_heads,
                    head_dim=arch.head_dim,
                    intermediate_size=arch.intermediate_size,
                    rms_norm_eps=arch.rms_norm_eps,
                    layer_idx=i,
                    prefix=f"gen_layers.{i}",
                    quant_config=quant_config,
                    supported_attention_backends=arch._supported_attention_backends,
                )
                for i in range(arch.num_hidden_layers)
            ]
        )

        # Output norm
        self.norm_moe_gen = RMSNorm(self.hidden_size, eps=arch.rms_norm_eps)

        # Cached K/V from UND pathway - dict keyed by cache_key for CFG support
        # This allows maintaining separate caches for conditional and unconditional
        # prompts, avoiding recomputation on every denoising step
        self.cached_kv: dict[str, list[tuple[torch.Tensor, torch.Tensor]]] = {}
        self.cached_gen_rope_inputs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}

        self.__post_init__()

        self.layer_names = ["gen_layers", "language_model.layers"]

    def _pad_to_patch_size(self, H: int, W: int) -> tuple[int, int, int, int]:
        """Compute padded spatial dims aligned to patch_size."""
        p = self.latent_patch_size
        H_padded = ((H + p - 1) // p) * p
        W_padded = ((W + p - 1) // p) * p
        return H_padded // p, W_padded // p, H_padded, W_padded

    def patchify(self, latents: torch.Tensor, T: int, H: int, W: int) -> torch.Tensor:
        """Convert latents to patches: [B, C, T, H, W] -> [B, T*Hp*Wp, p*p*C]."""
        B = latents.shape[0]
        p = self.latent_patch_size
        C = self.latent_channel
        Hp, Wp, H_padded, W_padded = self._pad_to_patch_size(H, W)

        if H_padded != H or W_padded != W:
            latents = F.pad(latents, (0, W_padded - W, 0, H_padded - H))

        x = latents.reshape(B, C, T, Hp, p, Wp, p)
        x = x.permute(0, 2, 3, 5, 4, 6, 1)  # [B, T, Hp, Wp, p, p, C]
        return x.reshape(B, T * Hp * Wp, p * p * C)

    def unpatchify(self, tokens: torch.Tensor, T: int, H: int, W: int) -> torch.Tensor:
        """Convert patches back to latents: [B, T*Hp*Wp, p*p*C] -> [B, C, T, H, W]."""
        B = tokens.shape[0]
        p = self.latent_patch_size
        C = self.latent_channel
        Hp, Wp, H_padded, W_padded = self._pad_to_patch_size(H, W)

        x = tokens.reshape(B, T, Hp, Wp, p, p, C)
        x = x.permute(0, 6, 1, 2, 4, 3, 5)  # [B, C, T, Hp, p, Wp, p]
        x = x.reshape(B, C, T, H_padded, W_padded)

        if H_padded != H or W_padded != W:
            x = x[:, :, :, :H, :W]
        return x

    def _compute_rope_position_ids(
        self,
        text_mask: torch.Tensor,
        T: int,
        Hp: int,
        Wp: int,
        fps: float | None,
        device: torch.device,
        sound_frames: int = 0,
        action_frames: int = 0,
        action_fps: float | None = None,
        action_start_frame_offset: int = 1,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Compute mRoPE position IDs for UND text and GEN visual + action + sound tokens."""
        B = text_mask.shape[0]
        S_text = text_mask.shape[1]
        text_lengths = text_mask.sum(dim=1).long()
        effective_fps = fps if fps is not None and T > 1 else None

        text_pos_list = []
        vis_pos_list = []
        for b in range(B):
            real_len = int(text_lengths[b].item())
            t_pos, t_offset = compute_mrope_position_ids_text(
                real_len, temporal_offset=0, device=device
            )
            media_offset = t_offset + self.temporal_margin
            v_pos, _ = compute_mrope_position_ids_vision(
                T,
                Hp,
                Wp,
                temporal_offset=media_offset,
                device=device,
                fps=effective_fps,
                base_fps=self.base_fps,
                temporal_compression_factor=self.temporal_compression_factor,
            )
            if action_frames > 0:
                a_pos, _ = compute_mrope_position_ids_action(
                    action_frames,
                    temporal_offset=media_offset,
                    action_fps=action_fps,
                    device=device,
                    base_fps=self.base_fps,
                    base_temporal_compression_factor=self.temporal_compression_factor,
                    start_frame_offset=action_start_frame_offset,
                )
                pos_dtype = torch.promote_types(v_pos.dtype, a_pos.dtype)
                v_pos = torch.cat([v_pos.to(pos_dtype), a_pos.to(pos_dtype)], dim=1)
            if sound_frames > 0:
                s_pos, _ = compute_mrope_position_ids_sound(
                    sound_frames,
                    temporal_offset=media_offset,
                    sound_latent_fps=self.sound_latent_fps,
                    device=device,
                    base_fps=self.base_fps,
                    temporal_compression_factor_sound=self.temporal_compression_factor_sound,
                )
                pos_dtype = torch.promote_types(v_pos.dtype, s_pos.dtype)
                v_pos = torch.cat([v_pos.to(pos_dtype), s_pos.to(pos_dtype)], dim=1)
            if real_len < S_text:
                t_pos = torch.cat(
                    [
                        t_pos,
                        torch.zeros(
                            3, S_text - real_len, dtype=t_pos.dtype, device=device
                        ),
                    ],
                    dim=1,
                )
            text_pos_list.append(t_pos)
            vis_pos_list.append(v_pos)

        text_pos_ids = torch.stack(text_pos_list, dim=1).to(device)  # [3, B, S_text]
        vis_pos_ids = torch.stack(vis_pos_list, dim=1).to(device)  # [3, B, S_gen]

        return text_pos_ids, vis_pos_ids

    def reset_cache(self, cache_key: str | None = None):
        """Reset cached K/V from UND pathway.

        Args:
            cache_key: If provided, reset only the specified cache key.
                      If None, reset all caches.
        """
        if cache_key is None:
            # Reset all caches
            self.cached_kv = {}
            self.cached_gen_rope_inputs = {}
        else:
            # Reset specific cache
            if cache_key in self.cached_kv:
                del self.cached_kv[cache_key]
            if cache_key in self.cached_gen_rope_inputs:
                del self.cached_gen_rope_inputs[cache_key]

    def _ensure_cache_dicts(self):
        """Ensure cache dictionaries exist (for backwards compatibility)."""
        if not isinstance(self.cached_kv, dict):
            self.cached_kv = {}
        if not isinstance(self.cached_gen_rope_inputs, dict):
            self.cached_gen_rope_inputs = {}

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
        timestep: torch.LongTensor,
        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
        guidance=None,
        text_ids: torch.Tensor | None = None,
        text_mask: torch.Tensor | None = None,
        fps: float | None = None,
        cache_key: str = "default",
        noisy_frame_mask: torch.Tensor | None = None,
        max_text_seq_len: int | None = None,
        sound_latents: torch.Tensor | None = None,
        action_latents: torch.Tensor | None = None,
        action_domain_ids: torch.Tensor | None = None,
        action_noisy_mask: torch.Tensor | None = None,
        action_fps: float | None = None,
        action_start_frame_offset: int = 1,
        **kwargs,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """Forward pass for denoising.

        Args:
            hidden_states: [B, C, T, H, W] noisy latents
            encoder_hidden_states: Not used (text embedded in transformer)
            timestep: [B] diffusion timestep per sample
            text_ids: [B, S_text] tokenized text input
            text_mask: [B, S_text] attention mask for text (1=real, 0=pad)
            fps: video frame rate for temporal mRoPE scaling
            cache_key: Key for the UND K/V cache. Use different keys for
                conditional ("cond") and unconditional ("uncond") branches
                in CFG to avoid recomputing the cache every step.
            noisy_frame_mask: Optional [B, 1, T, 1, 1] mask where 1 marks
                noisy frames (timestep embedding applied) and 0 marks
                conditioned frames (clean context, embedding skipped).
                ``None`` means every frame is noisy (T2V / T2I).
            max_text_seq_len: Real text length already computed during
                tokenization. When omitted it is derived from ``text_mask``.
            action_latents: Optional [B, T_action, D_action] noisy action
                latents for action generation.
            action_domain_ids: [B] embodiment domain IDs (0=no-action default).
            action_noisy_mask: [B, T_action, 1] where 1=noisy, 0=conditioned;
                controls which action tokens receive the timestep embedding.
                ``None`` means all tokens are noisy.
            action_fps: Frame rate for action token temporal mRoPE scaling.
                Defaults to the video fps when None.
            action_start_frame_offset: Temporal offset applied to action
                position IDs relative to the video's media_offset (default 1).

        Returns:
            [B, C, T, H, W] velocity prediction, or a tuple
            (video_pred, ...) with extra tensors when action/sound are active.
        """
        if text_ids is None or text_mask is None:
            raise ValueError("Cosmos3 requires text_ids and text_mask to be passed")

        batch_size, C, T, H, W = hidden_states.shape
        Hp, Wp, _, _ = self._pad_to_patch_size(H, W)
        if max_text_seq_len is None:
            max_text_seq_len = int(text_mask.sum(dim=1).max().item())
        if max_text_seq_len < text_ids.shape[1]:
            text_ids = text_ids[:, :max_text_seq_len]
            text_mask = text_mask[:, :max_text_seq_len]

        sound_frames = sound_latents.shape[-1] if sound_latents is not None else 0

        action_frames = 0
        if action_latents is not None:
            if self.sp_size > 1:
                raise NotImplementedError(
                    "Cosmos3 action generation does not support sequence parallelism yet"
                )
            action_frames = action_latents.shape[1]
            if action_domain_ids is None:
                action_domain_ids = torch.zeros(
                    action_latents.shape[0],
                    dtype=torch.long,
                    device=action_latents.device,
                )

        extra_frames = action_frames + sound_frames
        sequence_shard_enabled = self.sp_size > 1

        # Add timestep embedding (computed in float32 for numerical stability, then cast back)
        time_embed = self.time_embedder(timestep.float())
        time_embed = time_embed.to(
            hidden_states.dtype
        )  # Cast to match hidden_gen dtype

        # Patchify and project to hidden dim
        hidden_gen, _ = self.proj_in(self.patchify(hidden_states, T, H, W))
        seq_len_orig = hidden_gen.shape[1]
        seq_shard_pad = 0

        # Per-token noisy mask follows the same pad/shard as hidden_gen, so
        # build it before the SP split.
        token_noisy_mask: torch.Tensor | None = None
        if noisy_frame_mask is not None:
            token_noisy_mask = (
                noisy_frame_mask[:, 0, :, 0, 0]
                .unsqueeze(-1)
                .expand(-1, -1, Hp * Wp)
                .reshape(batch_size, -1, 1)
                .to(hidden_gen.dtype)
            )

        if extra_frames == 0:
            # Video-only: shard the visual tokens, then add the timestep
            # embedding on the local shard.
            if sequence_shard_enabled:
                if seq_len_orig % self.sp_size != 0:
                    seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size)
                    pad = torch.zeros(
                        (batch_size, seq_shard_pad, hidden_gen.shape[2]),
                        dtype=hidden_gen.dtype,
                        device=hidden_gen.device,
                    )
                    hidden_gen = torch.cat([hidden_gen, pad], dim=1)
                    if token_noisy_mask is not None:
                        mask_pad = torch.zeros(
                            (batch_size, seq_shard_pad, 1),
                            dtype=token_noisy_mask.dtype,
                            device=token_noisy_mask.device,
                        )
                        token_noisy_mask = torch.cat(
                            [token_noisy_mask, mask_pad], dim=1
                        )
                local_seq_len = hidden_gen.shape[1] // self.sp_size
                hidden_gen = hidden_gen.view(
                    batch_size, self.sp_size, local_seq_len, hidden_gen.shape[2]
                )
                hidden_gen = hidden_gen[:, self.sp_rank, :, :]
                if token_noisy_mask is not None:
                    token_noisy_mask = token_noisy_mask.view(
                        batch_size, self.sp_size, local_seq_len, 1
                    )[:, self.sp_rank, :, :]
            if token_noisy_mask is not None:
                hidden_gen = hidden_gen + time_embed.unsqueeze(1) * token_noisy_mask
            else:
                hidden_gen = hidden_gen + time_embed.unsqueeze(1)
        else:
            # Multi-modal: assemble the full GEN sequence
            # (video[, action][, sound]) with timestep embeddings, then shard
            # the combined stream so sequence parallelism splits every modality
            # evenly. The per-modality output heads run after the post-loop
            # all-gather reassembles the sequence.
            if token_noisy_mask is not None:
                hidden_gen = hidden_gen + time_embed.unsqueeze(1) * token_noisy_mask
            else:
                hidden_gen = hidden_gen + time_embed.unsqueeze(1)

            if action_latents is not None:
                hidden_action = self.action_proj_in(
                    action_latents.to(hidden_gen.dtype), action_domain_ids
                )
                hidden_action = hidden_action + self.action_modality_embed.to(
                    hidden_action.dtype
                )
                if action_noisy_mask is None:
                    hidden_action = hidden_action + time_embed.unsqueeze(1)
                else:
                    hidden_action = hidden_action + time_embed.unsqueeze(
                        1
                    ) * action_noisy_mask.to(hidden_action.dtype)
                hidden_gen = torch.cat([hidden_gen, hidden_action], dim=1)

            if sound_latents is not None:
                packed_sound = sound_latents.permute(0, 2, 1).to(hidden_gen.dtype)
                hidden_sound, _ = self.audio_proj_in(packed_sound)
                hidden_sound = hidden_sound + self.audio_modality_embed.to(
                    hidden_sound.dtype
                )
                hidden_sound = hidden_sound + time_embed.unsqueeze(1)
                hidden_gen = torch.cat([hidden_gen, hidden_sound], dim=1)

            seq_len_orig = hidden_gen.shape[1]
            if sequence_shard_enabled:
                if seq_len_orig % self.sp_size != 0:
                    seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size)
                    pad = torch.zeros(
                        (batch_size, seq_shard_pad, hidden_gen.shape[2]),
                        dtype=hidden_gen.dtype,
                        device=hidden_gen.device,
                    )
                    hidden_gen = torch.cat([hidden_gen, pad], dim=1)
                local_seq_len = hidden_gen.shape[1] // self.sp_size
                hidden_gen = hidden_gen.view(
                    batch_size, self.sp_size, local_seq_len, hidden_gen.shape[2]
                )
                hidden_gen = hidden_gen[:, self.sp_rank, :, :]

        self._ensure_cache_dicts()

        # Compute UND K/V cache for this cache_key if not already cached
        # This allows reusing the cache across denoising steps for the same text
        if (
            cache_key not in self.cached_kv
            or cache_key not in self.cached_gen_rope_inputs
        ):
            text_pos_ids, vis_pos_ids = self._compute_rope_position_ids(
                text_mask,
                T,
                Hp,
                Wp,
                fps,
                hidden_states.device,
                sound_frames=sound_frames,
                action_frames=action_frames,
                action_fps=action_fps if action_fps is not None else fps,
                action_start_frame_offset=action_start_frame_offset,
            )
            # UND K/V cache is kept FULL on all ranks (not sharded). Text
            # sequence is short, so memory impact is minimal, and the GEN
            # cross-attention needs the full K/V on every SP rank.
            self.cached_kv[cache_key] = self.language_model(
                text_ids, text_mask, text_pos_ids
            )
            if sequence_shard_enabled:
                if seq_shard_pad > 0:
                    pad_pos = vis_pos_ids[:, :, -1:].expand(-1, -1, seq_shard_pad)
                    vis_pos_ids = torch.cat([vis_pos_ids, pad_pos], dim=2)
                vis_pos_ids = vis_pos_ids.view(
                    3, batch_size, self.sp_size, local_seq_len
                )[:, :, self.sp_rank, :]
            self.cached_gen_rope_inputs[cache_key] = (
                self.language_model.rotary_emb.build_rope_cache_inputs(
                    vis_pos_ids, cache_dtype=hidden_gen.dtype
                )
            )

        cos_sin_gen, gen_rope_cache_positions = self.cached_gen_rope_inputs[cache_key]

        # Run GEN layers. `residual` is threaded so each layer's
        # input_layernorm and post_attention_layernorm can use the
        # fused add+rmsnorm path instead of separate add + norm kernels.
        cached_kv_for_key = self.cached_kv[cache_key]
        residual: torch.Tensor | None = None
        use_fused_qk_norm_rope = T > 1
        for i, layer in enumerate(self.gen_layers):
            k_und, v_und = cached_kv_for_key[i]
            hidden_gen, residual = layer(
                hidden_gen,
                k_und,
                v_und,
                cos_sin_gen,
                gen_rope_cache_positions,
                use_fused_qk_norm_rope,
                residual=residual,
            )

        # Collapse the trailing residual carry. RMSNorm and the linear
        # projection that follow are per-token, so we run them on the
        # local shard and only gather the (much smaller) patch-space
        # output. With patch_latent_dim ~= hidden_size / 21 for cosmos3,
        # this cuts the post-loop SP collective bandwidth ~21x.
        hidden_gen = hidden_gen + residual
        hidden_gen = self.norm_moe_gen(hidden_gen)

        if extra_frames == 0:
            # Video-only: project on the local shard and gather the (much
            # smaller) patch-space output. With patch_latent_dim ~=
            # hidden_size / 21 for cosmos3, this cuts the post-loop SP
            # collective bandwidth ~21x.
            output, _ = self.proj_out(hidden_gen)
            if sequence_shard_enabled:
                output = sequence_model_parallel_all_gather(output, dim=1)
                if seq_shard_pad > 0:
                    output = output[:, :seq_len_orig, :]
            return self.unpatchify(output, T, H, W)

        # Multi-modal: gather the full GEN hidden and drop shard padding, then
        # split per modality so each output head sees its own contiguous tokens.
        if sequence_shard_enabled:
            hidden_gen = sequence_model_parallel_all_gather(hidden_gen, dim=1)
            if seq_shard_pad > 0:
                hidden_gen = hidden_gen[:, :seq_len_orig, :]

        s_video = seq_len_orig - extra_frames
        output, _ = self.proj_out(hidden_gen[:, :s_video, :])
        video_pred = self.unpatchify(output, T, H, W)

        extra_outputs: list[torch.Tensor] = []
        idx = s_video
        if action_frames > 0:
            action_hidden = hidden_gen[:, idx : idx + action_frames, :]
            extra_outputs.append(self.action_proj_out(action_hidden, action_domain_ids))
            idx += action_frames
        if sound_frames > 0:
            sound_hidden = hidden_gen[:, idx:, :]
            sound_output, _ = self.audio_proj_out(sound_hidden)
            extra_outputs.append(sound_output.permute(0, 2, 1).contiguous())

        return (video_pred, *extra_outputs)

    def preprocess_loaded_state_dict(
        self, iterator: Iterable[tuple[str, torch.Tensor]]
    ) -> Iterator[tuple[str, torch.Tensor]]:
        # ModelOpt FP8 emits a 0-d per-tensor scale per source Linear. Where
        # sources fuse into a single MergedColumnParallelLinear (Q/K/V into
        # to_qkv, gate/up into gate_up_proj), the FP8 weights of each shard
        # are quantized against their own scale. Naively concatenating the
        # FP8 bytes and applying a single fused scale at runtime yields noise
        # (the K/V tiles get dequant'd with the wrong factor).
        #
        # Fix: per fused Linear, dequant each FP8 shard with its own scale,
        # pick max as the fused scale, requant each shard against the max,
        # then concat the requantized FP8 bytes. input_scale is shared across
        # shards (same activation tensor), so just take max — no requant
        # needed. The emitted fused tensors keep the full Q/K/V layout; the
        # column-parallel loader slices each logical shard for the local TP rank.
        mapping_fn = get_param_names_mapping(self.param_names_mapping)
        pending: dict[str, dict[str, dict[int, torch.Tensor]]] = {}
        expected_count: dict[str, int] = {}

        def _try_emit(linear_target: str):
            groups = pending.get(linear_target, {})
            n = expected_count.get(linear_target)
            if n is None:
                return
            weights = groups.get("weight", {})
            w_scales = groups.get("weight_scale", {})
            i_scales = groups.get("input_scale", {})
            if len(weights) != n or len(w_scales) != n:
                return
            saw_input_scale = bool(i_scales)
            if saw_input_scale and len(i_scales) != n:
                return
            scales_t = torch.stack([w_scales[i].reshape(()) for i in range(n)])
            max_w_scale = scales_t.max()
            rescaled = []
            for i in range(n):
                w_fp8 = weights[i]
                original_scale = w_scales[i].reshape(()).to(torch.float32)
                w_dequant = w_fp8.to(torch.float32) * original_scale
                w_requant = (
                    (w_dequant / max_w_scale.to(torch.float32))
                    .clamp(-448.0, 448.0)
                    .to(torch.float8_e4m3fn)
                )
                rescaled.append(w_requant)
            merged_weight = torch.cat(rescaled, dim=0)
            pending.pop(linear_target, None)
            expected_count.pop(linear_target, None)
            yield linear_target + ".weight", merged_weight
            yield linear_target + ".weight_scale", max_w_scale
            if saw_input_scale:
                in_t = torch.stack([i_scales[i].reshape(()) for i in range(n)])
                yield linear_target + ".input_scale", in_t.max()

        for name, tensor in iterator:
            target_name, merge_index, num_to_merge = mapping_fn(name)
            if num_to_merge is None:
                yield target_name, tensor
                continue
            suffix = None
            for candidate in ("weight_scale", "input_scale", "weight"):
                if target_name.endswith("." + candidate):
                    suffix = candidate
                    break
            if suffix is None:
                yield name, tensor
                continue
            if suffix == "weight" and tensor.dtype != torch.float8_e4m3fn:
                yield name, tensor
                continue
            linear_target = target_name[: -(len(suffix) + 1)]
            pending.setdefault(linear_target, {}).setdefault(suffix, {})[
                merge_index
            ] = tensor
            expected_count[linear_target] = num_to_merge
            yield from _try_emit(linear_target)

    def post_load_weights(self, target_dtype: torch.dtype = torch.bfloat16) -> None:
        """Cast non-quantized parameters to their preferred dtypes and rebuild
        meta-device buffers.

        Time-embedder stays in float32 for numerical stability; embeddings and
        the VAE/LLM bridge linears go to ``target_dtype``. Quantized modules
        (e.g. FP8 from a ModelOpt export) are skipped — calling ``.to(dtype)``
        on them would cast their FP8 weights back to BF16/FP32 and break the
        quant kernels.

        Also re-materializes the RoPE ``inv_freq`` buffer, which can land on
        the meta device when the model is constructed under a meta context.
        """
        # Get the actual device from a loaded parameter
        device = next(self.parameters()).device

        # Recompute RoPE inv_freq buffer on the correct device
        # This is needed because model is created in meta device context
        rotary_emb = self.language_model.rotary_emb
        if rotary_emb.inv_freq.is_meta:
            dim = rotary_emb.head_dim
            rope_theta = 5000000.0  # From config
            inv_freq = 1.0 / (
                rope_theta
                ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
            )
            rotary_emb.register_buffer("inv_freq", inv_freq, persistent=False)

        # Collect all quantized submodules so we can skip them during
        # dtype casts.  Calling .to(dtype) on a quantized layer would
        # cast FP8 weights back to BF16/FP32, breaking the quant kernels.
        quantized_modules: set[int] = {
            id(m)
            for m in self.modules()
            if hasattr(m, "quant_method")
            and not isinstance(m.quant_method, UnquantizedLinearMethod)
        }

        def _is_quantized(module: torch.nn.Module) -> bool:
            return id(module) in quantized_modules

        def _cast_direct(module: torch.nn.Module, dtype: torch.dtype) -> None:
            """Cast only the module's own parameters and buffers, not its
            children.  This avoids the recursive `.to()` which would cast
            quantized (FP8) weights back to BF16/FP32."""
            for key, param in module._parameters.items():
                if param is not None:
                    module._parameters[key] = torch.nn.Parameter(
                        param.data.to(dtype=dtype), requires_grad=False
                    )
            for key, buf in module._buffers.items():
                if buf is not None:
                    module._buffers[key] = buf.to(dtype=dtype)

        # Time embedder should stay in float32 for numerical stability.
        # Cast only non-quantized submodules' own params (non-recursive).
        for module in self.time_embedder.modules():
            if not _is_quantized(module):
                _cast_direct(module, torch.float32)

        # Ensure embeddings and projections are in target dtype
        self.language_model.embed_tokens.to(target_dtype)
        for module in self.proj_in.modules():
            if not _is_quantized(module):
                _cast_direct(module, target_dtype)
        for module in self.proj_out.modules():
            if not _is_quantized(module):
                _cast_direct(module, target_dtype)

        # Convert RMSNorm layers to target dtype
        for module in self.modules():
            if isinstance(module, RMSNorm):
                module.to(target_dtype)


EntryClass = Cosmos3OmniTransformer