Skip to content

vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge

Cosmos3 Edge transformer variant with a Nemotron dense UND backbone.

COSMOS3_EDGE_BACKBONE_TYPE module-attribute 

COSMOS3_EDGE_BACKBONE_TYPE = 'cosmos3_edge_nemotron_dense'

Cosmos3EdgeCausalAttention

Bases: Module

Edge UND causal attention.

Reasoner self-attention uses raw UND Q/K with RoPE because qk_norm_for_text=false. The GEN-facing UND key cache is a separate normalized-then-RoPE'd view of the raw K tensor when GEN QK norm is enabled; no-GEN-QK-norm checkpoints reuse the raw RoPE'd UND K.

attn instance-attribute 

attn = FrameworkAttention(
    num_heads=self.num_heads,
    head_size=self.head_dim,
    causal=True,
    softmax_scale=1.0 / self.head_dim**0.5,
    num_kv_heads=self.num_kv_heads,
    skip_sequence_parallel=True,
)

head_dim instance-attribute 

head_dim = head_dim

hidden_size instance-attribute 

hidden_size = hidden_size

k_norm_und_for_gen instance-attribute 

k_norm_und_for_gen = (
    RMSNorm(self.head_dim, eps=rms_norm_eps)
    if use_und_k_norm_for_gen
    else None
)

num_heads instance-attribute 

num_heads = num_attention_heads

num_heads_local instance-attribute 

num_heads_local = self.num_heads // tp_size

num_kv_heads instance-attribute 

num_kv_heads = num_key_value_heads

num_kv_heads_local instance-attribute 

num_kv_heads_local = self.num_kv_heads // tp_size

to_k instance-attribute 

to_k = ColumnParallelLinear(
    hidden_size,
    self.num_kv_heads * self.head_dim,
    bias=False,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.to_k",
)

to_out instance-attribute 

to_out = RowParallelLinear(
    self.num_heads * self.head_dim,
    hidden_size,
    bias=False,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.to_out",
)

to_q instance-attribute 

to_q = ColumnParallelLinear(
    hidden_size,
    self.num_heads * self.head_dim,
    bias=False,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.to_q",
)

to_v instance-attribute 

to_v = ColumnParallelLinear(
    hidden_size,
    self.num_kv_heads * self.head_dim,
    bias=False,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.to_v",
)

forward 

forward(
    hidden_states: Tensor,
    freqs_cos: Tensor,
    freqs_sin: Tensor,
) -> tuple[Tensor, Tensor, Tensor]

Cosmos3EdgeLanguageModel

Bases: Module

Nemotron dense UND tower that returns GEN-facing K plus raw V per layer.

embed_tokens instance-attribute 

embed_tokens = nn.Embedding(vocab_size, hidden_size)

layers instance-attribute 

layers = nn.ModuleList(
    [
        Cosmos3EdgeUndDecoderLayer(
            hidden_size=hidden_size,
            intermediate_size=intermediate_size,
            num_attention_heads=num_attention_heads,
            num_key_value_heads=num_key_value_heads,
            head_dim=head_dim,
            rms_norm_eps=rms_norm_eps,
            use_und_k_norm_for_gen=use_und_k_norm_for_gen,
            quant_config=quant_config,
            prefix=f"{prefix}.layers.{i}",
        )
        for i in range(num_hidden_layers)
    ]
)

norm instance-attribute 

norm = RMSNorm(hidden_size, eps=rms_norm_eps)

rotary_emb instance-attribute 

rotary_emb = Qwen3VLTextRotaryEmbedding(
    head_dim=head_dim,
    rope_theta=rope_theta,
    mrope_section=mrope_section,
)

forward 

forward(
    text_ids: Tensor, freqs: tuple[Tensor, Tensor]
) -> list[tuple[Tensor, Tensor]]

Cosmos3EdgeUndDecoderLayer

Bases: Module

Edge UND decoder layer: causal self-attention + ReLU2 MLP.

input_layernorm instance-attribute 

input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)

mlp instance-attribute 

mlp = Cosmos3Relu2MLP(
    hidden_size=hidden_size,
    intermediate_size=intermediate_size,
    quant_config=quant_config,
    prefix=f"{prefix}.mlp",
)

post_attention_layernorm instance-attribute 

post_attention_layernorm = RMSNorm(
    hidden_size, eps=rms_norm_eps
)

self_attn instance-attribute 

self_attn = Cosmos3EdgeCausalAttention(
    hidden_size=hidden_size,
    num_attention_heads=num_attention_heads,
    num_key_value_heads=num_key_value_heads,
    head_dim=head_dim,
    rms_norm_eps=rms_norm_eps,
    use_und_k_norm_for_gen=use_und_k_norm_for_gen,
    quant_config=quant_config,
    prefix=f"{prefix}.self_attn",
)

forward 

forward(
    hidden_states: Tensor, freqs: tuple[Tensor, Tensor]
) -> tuple[Tensor, Tensor, Tensor]

Cosmos3EdgeVFMTransformer

Bases: Cosmos3VFMTransformer

Cosmos3 Edge variant with Nemotron dense UND and shared GEN diffusion.

validate_loaded_weights 

validate_loaded_weights(loaded: set[str]) -> None

Cosmos3Relu2MLP

Bases: Module

Nemotron dense MLP: down_proj(relu(up_proj(x)) ** 2).

down_proj instance-attribute 

down_proj = RowParallelLinear(
    intermediate_size,
    hidden_size,
    bias=False,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.down_proj",
)

up_proj instance-attribute 

up_proj = ColumnParallelLinear(
    hidden_size,
    intermediate_size,
    bias=False,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.up_proj",
)

forward 

forward(x: Tensor) -> Tensor