Module nemotron_moe

Module nemotron_moe 

Source
Expand description

Nemotron-H standalone MoE FFN layer.

Supports two variants:

  • Nano 30B: Direct MoE — experts operate on full hidden_size.
  • Super 120B: LatentMoE — routed experts operate in latent space [moe_latent_size], with fc1/fc2 latent projections bridging hidden↔latent.

Forward: RMS norm → gate → sigmoid topK routing → (fc1_latent if latent) → batched up GEMV → fused relu²+down → weighted_sum → (fc2_latent if latent) → shared expert up+relu²+down → sum routed+shared → residual add.

All expert dispatch is device-side (pointer tables) — zero D2H sync.

Structs§

NemotronMoeLayer
Nemotron-H standalone MoE FFN layer.