cosmos-framework/packages/diffusers-cosmos3/diffusers_cosmos3/transformer.py

# Copyright 2025 The NVIDIA Team and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
from typing import Optional, Tuple

import torch
import torch.nn as nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.attention_dispatch import dispatch_attention_fn
from diffusers.models.modeling_utils import ModelMixin
from transformers.activations import ACT2FN
from transformers.integrations import use_kernel_forward_from_hub
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
from transformers.models.qwen3_vl.modeling_qwen3_vl import apply_rotary_pos_emb

from diffusers_cosmos3.sequence_packing import (
    FactoredSequencePack,
    from_joint,
    from_mode_splits,
    from_und_gen_splits,
    get_all_seq,
    get_causal_seq,
    get_device_and_dtype,
    get_full_only_seq,
    get_gen_seq,
    get_und_seq,
    set_gen_seq,
    set_und_seq,
    zeros_like,
)


def _pack_to_batch(tokens: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int) -> torch.Tensor:
    """Unpack (total_tokens, heads, dim) → (batch, max_seqlen, heads, dim)."""
    batch = cu_seqlens.shape[0] - 1
    cu = cu_seqlens.tolist()
    out = tokens.new_zeros(batch, max_seqlen, *tokens.shape[1:])
    for i in range(batch):
        n = cu[i + 1] - cu[i]
        out[i, :n] = tokens[cu[i] : cu[i + 1]]
    return out


def _batch_to_pack(batched: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
    """Repack (batch, max_seqlen, heads, dim) → (total_tokens, heads, dim)."""
    cu = cu_seqlens.tolist()
    return torch.cat([batched[i, : cu[i + 1] - cu[i]] for i in range(len(cu) - 1)], dim=0)


def _kv_padding_mask(cu_seqlens: torch.Tensor, max_seqlen: int, dtype: torch.dtype, device: torch.device):
    """Float mask (batch, 1, 1, max_seqlen) with -inf at padding positions, or None if uniform."""
    batch = cu_seqlens.shape[0] - 1
    cu = cu_seqlens.tolist()
    mask = torch.zeros(batch, 1, 1, max_seqlen, dtype=dtype, device=device)
    for i in range(batch):
        kl = cu[i + 1] - cu[i]
        if kl < max_seqlen:
            mask[i, 0, 0, kl:] = float("-inf")
    return None if (mask == 0).all() else mask


class CosmosAttnProcessor3_0:
    """
    Packed two-way attention processor for Cosmos3. Implements separate causal
    (understanding) and full (generation) attention pathways via dispatch_attention_fn.
    """

    def __call__(
        self,
        packed_query_states: FactoredSequencePack,
        packed_key_states: FactoredSequencePack,
        packed_value_states: FactoredSequencePack,
    ) -> FactoredSequencePack:
        causal_q, causal_offsets = get_causal_seq(packed_query_states)
        causal_k, _ = get_causal_seq(packed_key_states)
        causal_v, _ = get_causal_seq(packed_value_states)
        full_q, full_offsets = get_full_only_seq(packed_query_states)
        sample_offsets = packed_query_states["sample_offsets"]
        max_causal = packed_query_states["max_causal_len"]
        max_full = packed_query_states["max_full_len"]
        max_sample = packed_query_states["max_sample_len"]

        # Causal (understanding) self-attention
        causal_out = dispatch_attention_fn(
            _pack_to_batch(causal_q, causal_offsets, max_causal),
            _pack_to_batch(causal_k, causal_offsets, max_causal),
            _pack_to_batch(causal_v, causal_offsets, max_causal),
            is_causal=True,
            enable_gqa=True,
        )
        causal_out = _batch_to_pack(causal_out, causal_offsets).flatten(-2, -1)

        # Full (generation) cross-attention: Q = gen tokens, K/V = all tokens
        all_k = get_all_seq(packed_key_states)
        all_v = get_all_seq(packed_value_states)
        full_out = dispatch_attention_fn(
            _pack_to_batch(full_q, full_offsets, max_full),
            _pack_to_batch(all_k, sample_offsets, max_sample),
            _pack_to_batch(all_v, sample_offsets, max_sample),
            attn_mask=_kv_padding_mask(sample_offsets, max_sample, causal_q.dtype, causal_q.device),
            is_causal=False,
            enable_gqa=True,
        )
        full_out = _batch_to_pack(full_out, full_offsets).flatten(-2, -1)

        return from_mode_splits(causal_out, full_out, packed_query_states)


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

    def __init__(self, hidden_size, frequency_embedding_size=256):
        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

    def _init_weights(self):
        std = 1.0 / math.sqrt(self.frequency_embedding_size)
        torch.nn.init.trunc_normal_(self.mlp[0].weight, std=std, a=-3 * std, b=3 * std)
        torch.nn.init.zeros_(self.mlp[0].bias)

        std = 1.0 / math.sqrt(self.hidden_size)
        torch.nn.init.trunc_normal_(self.mlp[2].weight, std=std, a=-3 * std, b=3 * std)
        torch.nn.init.zeros_(self.mlp[2].bias)

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        half = dim // 2
        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
            device=t.device
        )
        args = t[:, None].float() * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
        return embedding

    def forward(self, t):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        return self.linear_2(self.act(self.linear_1(t_freq)))


class DomainAwareLinear(nn.Module):
    """Linear projection with one weight/bias pair per action embodiment domain."""

    def __init__(self, input_size: int, output_size: int, num_domains: int) -> None:
        super().__init__()
        self.input_size = int(input_size)
        self.output_size = int(output_size)
        self.num_domains = int(num_domains)
        self.fc = nn.Embedding(self.num_domains, self.output_size * self.input_size)
        self.bias = nn.Embedding(self.num_domains, self.output_size)
        nn.init.xavier_uniform_(self.fc.weight)
        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)
        if x.shape[0] != domain_id.shape[0]:
            raise ValueError(
                "Cosmos3 action domain_id batch size must match action tokens: "
                f"tokens={x.shape[0]}, domain_id={domain_id.shape[0]}."
            )
        if torch.any((domain_id < 0) | (domain_id >= self.num_domains)):
            raise ValueError(f"Cosmos3 action domain_id must be in [0, {self.num_domains}), got {domain_id.tolist()}.")

        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
        if x.ndim == 3:
            return torch.bmm(x, weight) + bias.unsqueeze(1)
        raise ValueError(f"Cosmos3 DomainAwareLinear expected rank-2 or rank-3 input, got {tuple(x.shape)}.")


class LayerTypes:
    def __init__(self, is_moe: bool):
        self.is_moe = is_moe
        if is_moe:  # TODO: moe is not yet tested
            self.mlp = Qwen3VLMoeTextMLP
            self.rms_norm = Qwen3VLMoeTextRMSNorm
            self.rotary_embedding = Qwen3VLMoeTextRotaryEmbedding
        else:
            self.mlp = Cosmos3VLTextMLP
            self.rms_norm = Cosmos3VLTextRMSNorm
            self.rotary_embedding = Cosmos3VLTextRotaryEmbedding


class Cosmos3VLTextRotaryEmbedding(nn.Module):
    def __init__(self, config):
        super().__init__()
        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
            self.rope_type = config.rope_scaling.get("rope_type", "default")
        else:
            self.rope_type = "default"
        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        self.mrope_section = (
            config.rope_scaling.get("mrope_section", [24, 20, 20]) if config.rope_scaling is not None else [24, 20, 20]
        )

    def init_weights(self, buffer_device: torch.device | None = None) -> None:
        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, buffer_device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def apply_interleaved_mrope(self, freqs, mrope_section):
        """Apply interleaved MRoPE to 3D rotary embeddings.
        Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
        interleaved [THTHWHTHW...TT], preserving frequency continuity.
        args:
            x: (3, bs, seq_len, head_dim // 2)
            mrope_section: (3,)
        returns:
            x_t: (bs, seq_len, head_dim // 2)
        """
        freqs_t = freqs[0]  # just overwrite the first dimension T
        for dim, offset in enumerate((1, 2), start=1):  # H, W
            length = mrope_section[dim] * 3
            idx = slice(offset, length, 3)
            freqs_t[..., idx] = freqs[dim, ..., idx]
        return freqs_t

    @torch.no_grad()
    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
    def forward(self, x, position_ids):
        assert self.inv_freq.dtype == torch.float32, f"inv_freq must be float32, but got {self.inv_freq.dtype}"

        # In contrast to other models, Cosmos3Omni has different position ids for the grids
        # So we expand the inv_freq to shape (3, ...)
        if position_ids.ndim == 2:
            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)  # [3,B,N]
        inv_freq_expanded = (
            self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1).to(x.device)
        )  # [3,B,head_dim//2,1]
        position_ids_expanded = position_ids[:, :, None, :].float()  # [3,B,1,N]

        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)  # [3,B,N,head_dim//2]
        freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)  # [B,N,head_dim//2]
        emb = torch.cat((freqs, freqs), dim=-1)  # [B,N,head_dim]
        cos = emb.cos() * self.attention_scaling  # [B,N,head_dim]
        sin = emb.sin() * self.attention_scaling  # [B,N,head_dim]

        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)  # each: [B,N,head_dim]


class Cosmos3VLTextRMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
        """
        Cosmos3VLTextRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

    def extra_repr(self) -> str:
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"


class Cosmos3VLTextMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj


class Cosmos3VLTextAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = True

        self.to_q = nn.Linear(
            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
        )
        self.to_k = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.to_v = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.to_out = nn.Linear(
            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
        )
        self.norm_q = Cosmos3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)  # unlike olmo, only on the head dim!
        self.norm_k = Cosmos3VLTextRMSNorm(
            self.head_dim, eps=config.rms_norm_eps
        )  # thus post norm_q does not need reshape


class PackedAttentionMoT(Cosmos3VLTextAttention):
    """
    Dual-pathway packed attention for Qwen3VL MoT (Dense version).
    Implements understanding and generation pathways with separate projections.

    Note that this implementation is used for both Qwen3VL and Qwen3VL-MoE variants,
    even though it derives from the dense version of Qwen3VLTextAttention.
    """

    def __init__(self, config, layer_idx: int, layer_types: LayerTypes):
        super().__init__(config, layer_idx)

        # Add missing attributes for MoT compatibility
        self.hidden_size = config.hidden_size
        self.num_attention_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout

        # Generation pathway projections (separate from understanding pathway)
        # Qwen3VL already has query/key norms built in, so we add generation versions
        self.norm_added_q = layer_types.rms_norm(self.head_dim, eps=config.rms_norm_eps)
        self.norm_added_k = layer_types.rms_norm(self.head_dim, eps=config.rms_norm_eps)

        # Generation pathway linear projections
        self.add_q_proj = nn.Linear(
            self.hidden_size, self.num_attention_heads * self.head_dim, bias=config.attention_bias
        )
        self.add_k_proj = nn.Linear(
            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.add_v_proj = nn.Linear(
            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.to_add_out = nn.Linear(
            self.num_attention_heads * self.head_dim, self.hidden_size, bias=config.attention_bias
        )
        self.dispatch_attention_fn = CosmosAttnProcessor3_0()
        self.cp_mesh = None

    def forward(
        self,
        pack: FactoredSequencePack,
        attention_mask,
        packed_position_embeddings: Tuple[FactoredSequencePack, FactoredSequencePack],
        dual_kv_cache=None,
        natten_metadata: dict | None = None,
    ) -> FactoredSequencePack:
        """Forward pass with optional KV cache for autoregressive generation.

        This method is used for frame 0 where we store K/V for both und and gen tokens.
        For frame 1+, forward_with_kv_cache() is used instead (optimized path).

        Args:
            pack: Packed sequence with und/gen tokens
            attention_mask: Attention mask (BlockMask or SplitInfo)
            packed_position_embeddings: RoPE embeddings (cos, sin)
            dual_kv_cache: Optional dual KV cache for AR generation (frame 0).
        """

        q_und_in = self.to_q(get_und_seq(pack))  # [N_und,num_heads*head_dim]
        q_gen_in = self.add_q_proj(get_gen_seq(pack))  # [N_gen,num_heads*head_dim]

        k_und_in = self.to_k(get_und_seq(pack))  # [N_und,num_kv_heads*head_dim]
        k_gen_in = self.add_k_proj(get_gen_seq(pack))  # [N_gen,num_kv_heads*head_dim]

        v_und_in = self.to_v(get_und_seq(pack))  # [N_und,num_kv_heads*head_dim]
        v_gen_in = self.add_v_proj(get_gen_seq(pack))  # [N_gen,num_kv_heads*head_dim]

        q_und = q_und_in.view(-1, self.num_attention_heads, self.head_dim)  # [N_und,num_heads,head_dim]
        k_und = k_und_in.view(-1, self.num_key_value_heads, self.head_dim)  # [N_und,num_kv_heads,head_dim]
        v_und = v_und_in.view(-1, self.num_key_value_heads, self.head_dim)  # [N_und,num_kv_heads,head_dim]

        q_gen = q_gen_in.view(-1, self.num_attention_heads, self.head_dim)  # [N_gen,num_heads,head_dim]
        k_gen = k_gen_in.view(-1, self.num_key_value_heads, self.head_dim)  # [N_gen,num_kv_heads,head_dim]
        v_gen = v_gen_in.view(-1, self.num_key_value_heads, self.head_dim)  # [N_gen,num_kv_heads,head_dim]

        q_und = self.norm_q(q_und)  # [N_und,num_heads,head_dim]
        k_und = self.norm_k(k_und)  # [N_und,num_kv_heads,head_dim]

        q_gen = self.norm_added_q(q_gen)  # [N_gen,num_heads,head_dim]
        k_gen = self.norm_added_k(k_gen)  # [N_gen,num_kv_heads,head_dim]

        if self.config.freeze_und:
            q_und = q_und.detach()
            k_und = k_und.detach()
            v_und = v_und.detach()

        # Attempted port: Apply RoPE (BAGEL qwen-2.5)
        # Note: Position embeddings are now pre-squeezed at model level
        packed_cos = packed_position_embeddings[0]
        packed_sin = packed_position_embeddings[1]

        q_und_, k_und_ = apply_rotary_pos_emb(
            q_und,
            k_und,
            get_und_seq(packed_cos),
            get_und_seq(packed_sin),
            unsqueeze_dim=1,
        )  # q_und_: [N_und,num_heads,head_dim], k_und_: [N_und,num_kv_heads,head_dim]
        q_gen_, k_gen_ = apply_rotary_pos_emb(
            q_gen,
            k_gen,
            get_gen_seq(packed_cos),
            get_gen_seq(packed_sin),
            unsqueeze_dim=1,
        )  # q_gen_: [N_gen,num_heads,head_dim], k_gen_: [N_gen,num_kv_heads,head_dim]

        # === KV CACHE INTEGRATION FOR AUTOREGRESSIVE GENERATION ===
        # Frame 0: Store und and gen K/V (no fetching)
        # Apply cache after RoPE (cached keys already have positional info)
        # CP path: storage happens inside context_parallel_attention() after all-to-all,
        #          so tensors are stored head-sharded [1,S,H/cp,D].
        # Non-CP path: store here as [1,S,H,D] for fetch_kv() dim=1 compat.
        if dual_kv_cache is not None and self.cp_mesh is None:
            und_len = pack["_num_causal_tokens"]
            gen_len = pack["_num_full_tokens"]
            if not dual_kv_cache.und_cache.is_initialized:
                dual_kv_cache.und_cache.store(
                    k_und_[:und_len].unsqueeze(0), v_und[:und_len].unsqueeze(0)
                )  # [1,S_und,H,D]
            dual_kv_cache.gen_cache.store_kv(
                k_gen_[:gen_len].unsqueeze(0), v_gen[:gen_len].unsqueeze(0), frame_idx=0
            )  # [1,S_gen,H,D]

        packed_query_states_ = from_und_gen_splits(q_und_, q_gen_, pack)  # [N_und+N_gen,num_heads,head_dim]
        packed_key_states_ = from_und_gen_splits(k_und_, k_gen_, pack)  # [N_und+N_gen,num_kv_heads,head_dim]
        packed_value_states_ = from_und_gen_splits(v_und, v_gen, pack)  # [N_und+N_gen,num_kv_heads,head_dim]

        # CP: pass dual_kv_cache so context_parallel_attention() stores head-sharded K/V
        dispatch_kwargs: dict = {}
        if self.cp_mesh is not None and dual_kv_cache is not None:
            dispatch_kwargs["dual_kv_cache"] = dual_kv_cache
            dispatch_kwargs["frame_idx"] = 0

        packed_attn_output = self.dispatch_attention_fn(
            packed_query_states_,
            packed_key_states_,
            packed_value_states_,
        )

        # Apply projections directly to get final results
        und_seq = self.to_out(get_und_seq(packed_attn_output))  # [N_und,hidden_size]
        gen_seq = self.to_add_out(get_gen_seq(packed_attn_output))  # [N_gen,hidden_size]
        return from_und_gen_splits(und_seq, gen_seq, pack)  # [N_und+N_gen,hidden_size]


class Cosmos3VLTextMoTDecoderLayer(nn.Module):
    """
    Qwen3VL text MoT (Mixture of Tokens) decoder layer.
    Features dual-pathway attention for understanding vs generation.

    This is used for both Dense and MoE models.
    """

    def __init__(
        self,
        config,
        layer_idx: int,
        layer_types: LayerTypes,
    ):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.freeze_und = config.freeze_und
        self.self_attn = PackedAttentionMoT(config, layer_idx, layer_types)

        # TODO: Qwen3VLMoeTextSparseMoeBlock not supported yet
        self.mlp = layer_types.mlp(config)
        self.mlp_moe_gen = layer_types.mlp(config)

        self.input_layernorm = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps)
        self.input_layernorm_moe_gen = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm_moe_gen = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        input: FactoredSequencePack,
        attention_mask,
        packed_position_embeddings: Tuple[FactoredSequencePack, FactoredSequencePack],
        dual_kv_cache: None = None,
        frame_idx: Optional[int] = None,
        natten_metadata: dict | None = None,
    ) -> FactoredSequencePack:
        """Training forward pass with MoT routing - Attempted port from qwen2_mot

        Args:
            input: Packed sequence with und/gen tokens
            attention_mask: Attention mask
            packed_position_embeddings: RoPE embeddings (cos, sin)
            dual_kv_cache: Optional dual KV cache for AR generation
            frame_idx: Current frame index (default: None, treated as 0)
        """

        # Handle None frame_idx as 0
        if frame_idx is None:
            frame_idx = 0

        # TODO: support gen_only = True and AR generation
        gen_only = False
        # if dual_kv_cache is not None and isinstance(dual_kv_cache, DualKVCache):
        #     gen_only = frame_idx > 0 and dual_kv_cache.und_cache.is_initialized

        # Pre-Attention layernorm
        pack_norm_out = from_und_gen_splits(
            self.input_layernorm(get_und_seq(input)),  # [N_und,hidden_size]
            self.input_layernorm_moe_gen(get_gen_seq(input)),  # [N_gen,hidden_size]
            input,
        )  # [N_und+N_gen,hidden_size]

        # STANDARD PATH: Process both und and gen tokens (frame 0)
        pack_attn_out = self.self_attn(
            pack_norm_out,
            attention_mask,
            packed_position_embeddings,
            dual_kv_cache,
            natten_metadata=natten_metadata,
        )
        residual_und = get_und_seq(input) + get_und_seq(pack_attn_out)  # [N_und,hidden_size]
        residual_gen = get_gen_seq(input) + get_gen_seq(pack_attn_out)  # [N_gen,hidden_size]

        # STANDARD PATH: Process both und and gen tokens
        ln_out_und = self.post_attention_layernorm(residual_und)  # [N_und,hidden_size]
        ln_out_gen = self.post_attention_layernorm_moe_gen(residual_gen)  # [N_gen,hidden_size]

        # UNPAD MLP INPUT ===============
        # NOTE: This is only need for the MoE auxiliary loss computation and to avoid
        #       artificial expert inbalance due to routing padding tokens.
        gen_len = pack_attn_out["_num_full_tokens"]
        und_len = pack_attn_out["_num_causal_tokens"]
        ln_out_und_unpadded = ln_out_und[:und_len]  # [N_und_unpadded,hidden_size]
        ln_out_gen_unpadded = ln_out_gen[:gen_len]  # [N_gen_unpadded,hidden_size]

        mlp_out_und_unpadded = self.mlp(ln_out_und_unpadded)  # [N_und_unpadded,hidden_size]
        mlp_out_gen_unpadded = self.mlp_moe_gen(ln_out_gen_unpadded)  # [N_gen_unpadded,hidden_size]

        # PAD MLP OUTPUT ===============
        mlp_out_und = torch.cat([mlp_out_und_unpadded, ln_out_und[und_len:]], dim=0)  # [N_und,hidden_size]
        mlp_out_gen = torch.cat([mlp_out_gen_unpadded, ln_out_gen[gen_len:]], dim=0)  # [N_gen,hidden_size]

        mlp_out_und_seq = residual_und + mlp_out_und  # [N_und,hidden_size]
        mlp_out_gen_seq = residual_gen + mlp_out_gen  # [N_gen,hidden_size]

        return from_und_gen_splits(mlp_out_und_seq, mlp_out_gen_seq, input)


class Cosmos3OmniTransformer(ModelMixin, ConfigMixin):
    @register_to_config
    def __init__(
        self,
        attention_bias: bool = False,
        attention_dropout: float = 0.0,
        dtype: str = "bfloat16",
        freeze_und: bool = False,
        head_dim: int = 128,
        hidden_act: str = "silu",
        hidden_size: int = 4096,
        initializer_range: float = 0.02,
        intermediate_size: int = 12288,
        base_fps: int = 24,
        enable_fps_modulation: bool = True,
        joint_attn_implementation: str = "two_way",
        latent_channel: int = 48,
        action_dim: int | None = None,
        action_gen: bool = False,
        max_action_dim: int = 32,
        num_embodiment_domains: int = 32,
        position_embedding_type: str = "unified_3d_mrope",
        unified_3d_mrope_reset_spatial_ids: bool = True,
        unified_3d_mrope_temporal_modality_margin: int = 15000,
        video_temporal_causal: bool = False,
        latent_patch_size: int = 2,
        max_position_embeddings: int = 262144,
        model_type: str = "qwen3_vl_text",
        num_attention_heads: int = 32,
        num_hidden_layers: int = 36,
        num_key_value_heads: int = 8,
        patch_latent_dim: int = 192,
        qk_norm: bool = False,
        qk_norm_for_diffusion: bool = True,
        qk_norm_for_text: bool = True,
        rms_norm_eps: float = 1e-6,
        rope_scaling: dict | None = None,
        rope_theta: float = 5000000.0,
        sound_dim: int | None = None,
        sound_gen: bool = False,
        sound_latent_fps: float = 25.0,
        temporal_compression_factor_sound: int = 1,
        timestep_scale: float = 0.001,
        use_cache: bool = True,
        use_moe: bool = True,
        vocab_size: int = 151936,
    ):
        super().__init__()

        if rope_scaling is None:
            rope_scaling = {"mrope_interleaved": True, "mrope_section": [24, 20, 20], "rope_type": "default"}
            self.register_to_config(rope_scaling=rope_scaling)

        layer_types = LayerTypes(is_moe=False)
        self.embed_tokens = nn.Embedding(self.config.vocab_size, self.config.hidden_size)
        self.layers = nn.ModuleList(
            [
                Cosmos3VLTextMoTDecoderLayer(self.config, layer_idx, layer_types)
                for layer_idx in range(self.config.num_hidden_layers)
            ]
        )
        # Understanding pathway final norm
        self.norm = layer_types.rms_norm(self.config.hidden_size, eps=self.config.rms_norm_eps)
        # Generation pathway final norm
        self.norm_moe_gen = layer_types.rms_norm(self.config.hidden_size, eps=self.config.rms_norm_eps)
        self.rotary_emb = Cosmos3VLTextRotaryEmbedding(config=self.config)
        self.vocab_size = vocab_size
        self.action_gen = action_gen
        self.action_dim = int(max_action_dim if action_dim is None else action_dim)
        self.num_embodiment_domains = int(num_embodiment_domains)
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
        self.proj_in = nn.Linear(patch_latent_dim, hidden_size, bias=True)
        self.proj_out = nn.Linear(hidden_size, patch_latent_dim, bias=True)
        self.time_embedder = TimestepEmbedder(hidden_size)
        if action_gen:
            self.action_proj_in = DomainAwareLinear(self.action_dim, hidden_size, self.num_embodiment_domains)
            self.action_proj_out = DomainAwareLinear(hidden_size, self.action_dim, self.num_embodiment_domains)
            self.action_modality_embed = nn.Parameter(torch.zeros(hidden_size))
        if sound_gen:
            if sound_dim is None:
                raise ValueError("`sound_dim` must be provided when `sound_gen=True`.")
            self.audio_proj_in = nn.Linear(sound_dim, hidden_size, bias=True)
            self.audio_proj_out = nn.Linear(hidden_size, sound_dim, bias=True)
            self.audio_modality_embed = nn.Parameter(torch.zeros(hidden_size))

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        model = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
        # inv_freq is a non-persistent buffer absent from the saved state_dict.
        # Initialize it on CPU; it will move to the correct device with .to() / .cuda().
        model.rotary_emb.init_weights(buffer_device=None)
        return model

    def forward(
        self,
        pack: FactoredSequencePack,
        attention_mask,
        position_ids: torch.Tensor,
        dual_kv_cache: None = None,
        frame_idx: Optional[int] = None,
        natten_metadata_list: list | None = None,
    ) -> Tuple[FactoredSequencePack, None]:
        """Training forward pass - simplified to match qwen3_mot.

        Returns:
            (outputs, None) — the None placeholder mirrors the (packed_outputs, lbl_metadata)
            tuple returned by the original language_model so callers can unpack both.
        """
        # Handle None frame_idx as 0
        if frame_idx is None:
            frame_idx = 0

        # Create position embeddings (Qwen3 style) - squeeze once at model level
        # tensor below is only used for its dtype and device
        device, dtype = get_device_and_dtype(pack)
        _meta_tensor = torch.tensor([], dtype=dtype, device=device)  # [0]
        cos, sin = self.rotary_emb(
            _meta_tensor,
            position_ids=position_ids.unsqueeze(0) if position_ids.ndim == 1 else position_ids.unsqueeze(1),
        )  # if ndim == 2, then the mrope position_ids is (3, seq_len), we need to put batch dimension in the middle to make it compatible with the rotary_emb
        # cos, sin: [1,N,head_dim] (1D pos_ids) or [3,1,N,head_dim] (mrope pos_ids)
        cos = cos.squeeze(0)  # [N,head_dim] or [3,N,head_dim]
        sin = sin.squeeze(0)  # [N,head_dim] or [3,N,head_dim]
        position_embeddings = (
            from_joint(cos, pack),
            from_joint(sin, pack),
        )

        # TODO: Add lbl_metadata_all (we don't need it at inference)
        hidden_states = pack

        for i, decoder_layer in enumerate(self.layers):
            hidden_states = decoder_layer(
                hidden_states,
                attention_mask,
                position_embeddings,
                dual_kv_cache[i] if dual_kv_cache is not None else None,
                frame_idx,
                natten_metadata=None if natten_metadata_list is None else natten_metadata_list[i],
            )

        outputs = zeros_like(hidden_states)  # [N_und+N_gen,hidden_size]
        set_und_seq(outputs, self.norm(get_und_seq(hidden_states)))  # [N_und,hidden_size]
        set_gen_seq(outputs, self.norm_moe_gen(get_gen_seq(hidden_states)))  # [N_gen,hidden_size]
        return outputs, None


@use_kernel_forward_from_hub("RMSNorm")
class Qwen3VLMoeTextRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        Qwen3VLMoeTextRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"


class Qwen3VLMoeTextMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj


class Qwen3VLMoeTextRotaryEmbedding(nn.Module):
    def __init__(self, config):
        super().__init__()
        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
            self.rope_type = config.rope_scaling.get("rope_type", "default")
        else:
            self.rope_type = "default"
        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings
        self.mrope_section = config.rope_scaling.get("mrope_section", [24, 20, 20])

        self.config = config
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

    def init_weights(self, buffer_device: torch.device | None = None) -> None:
        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, buffer_device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def apply_interleaved_mrope(self, freqs, mrope_section):
        """Apply interleaved MRoPE to 3D rotary embeddings.
        Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
        interleaved [THTHWHTHW...TT], preserving frequency continuity.
        args:
            x: (3, bs, seq_len, head_dim // 2)
            mrope_section: (3,)
        returns:
            x_t: (bs, seq_len, head_dim // 2)
        """
        freqs_t = freqs[0]  # just overwrite the first dimension T
        for dim, offset in enumerate((1, 2), start=1):  # H, W
            length = mrope_section[dim] * 3
            idx = slice(offset, length, 3)
            freqs_t[..., idx] = freqs[dim, ..., idx]
        return freqs_t

    @torch.no_grad()
    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
    def forward(self, x, position_ids):
        assert self.inv_freq.dtype == torch.float32, f"inv_freq must be float32, but got {self.inv_freq.dtype}"

        # In contrast to other models, Qwen3VLMoe has different position ids for the grids
        # So we expand the inv_freq to shape (3, ...)
        if position_ids.ndim == 2:
            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)  # [3,B,N]
        inv_freq_expanded = (
            self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
        )  # [3,B,head_dim//2,1]
        position_ids_expanded = position_ids[:, :, None, :].float()  # [3,B,1,N]

        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)  # [3,B,N,head_dim//2]
        freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)  # [B,N,head_dim//2]
        emb = torch.cat((freqs, freqs), dim=-1)  # [B,N,head_dim]
        cos = emb.cos() * self.attention_scaling  # [B,N,head_dim]
        sin = emb.sin() * self.attention_scaling  # [B,N,head_dim]

        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)