Cosmos 3 · OmniMoT ▸ MoT 主干 ▸ 单层展开

MoTDecoderLayer — 算子级数据流

主干里第 i 层的内部。两条模态通路——理解 _und（蓝）与 生成 moe_gen（绿）——各自持有 norm / projection / MLP 权重， 并行流过整层；唯一发生跨模态交互的地方是共享注意力——而且是单向的：gen 读 und，und 读不到 gen。 norm 是 pre-norm 结构，两处残差（attention 后、MLP 后）分别在各通路内相加。

理解通路 _und 生成通路 moe_gen 共享算子 RMSNorm pre-norm ⊕ 残差相加
MoTDecoderLayer [ i ] · × N MoTDecoderLayer .forward
und tokens text / VLM gen tokens noisy vision · sound · action
打包联合序列 · 入层
┈ residual #1：input ─────────────────┈
① pre-attention RMSNorm（分叉）
input_layernormRMSNorm(hidden)
input_layernorm_moe_genRMSNorm(hidden)
▼
🔗 共享 Self-Attention PackedAttentionMoT PackedAttentionMoT .forward two_way_attention
und 投影 q_proj
  • q_proj / k_proj / v_proj
  • q_norm · k_norm (qk_norm_for_text)
gen 投影 q_proj_moe_gen
  • q/k/v_proj_moe_gen
  • q_norm_moe_gen · k_norm_moe_gen (for_diffusion)
und Q/K/V  ⇘  ⇙  gen Q/K/V
RoPE · GQA(num_key_value_heads) · two_way 注意力（非对称）
▸ gen query 看 und ⊕ gen 拼接的 K/V（全注意力，读得到文本条件）；und query 只做 und 内部因果自注意力，看不到带噪 gen —— 跨模态混合是单向的 gen ← und
o_proj → und out
o_proj_moe_gen → gen out
⊕ und_in + attn_und
⊕ gen_in + attn_gen
▼
┈ residual #2：post-attn ─────────────┈
② post-attention RMSNorm（分叉）
post_attention_layernormRMSNorm(hidden)
post_attention_layernorm_moe_genRMSNorm(hidden)
▼
③ 前馈 —— 每通路独立权重
mlpSwiGLU MLP（und tower）
mlp_moe_genSwiGLU MLP（gen tower）
MoE 层（每 decoder_sparse_step 层）：两侧换成 Qwen3VLMoeTextSparseMoeBlock —— router → top-k experts → 加权合并，附 load-balance 辅助损失（LBL metadata）；gen tower 可开 noisy gating。MLP 前先 unpad，避免 padding token 造成 expert 失衡。 Qwen3VLMoeTextSparseMoeBlock
⊕ resid_und + mlp_und
⊕ resid_gen + mlp_gen
▼
und tokens → 下一层 / Reasoner 头 gen tokens → 下一层 / Diffusion 头
出层 · 打包联合序列
为什么这是 “Mixture-of-Transformers” 而非普通 MoE：两套完整的 transformer 权重（norm + q/k/v/o + FFN）按 token 的模态静态路由， 只共享一次注意力。相比之下 MoE 只在 FFN 处按 token 动态选 expert。二者可叠加——MoT 的每条通路内部还能再是 MoE。 另有 gen_only 快路径：自回归推理时跳过 und 通路，und 的 K/V 从 cache 读取。 gen_only