Module nemotron_moe

Module nemotron_moe 

Source
Expand description

2026-09-25: Nemotron-H standalone MoE FFN layer, in two shapes:

  • direct (moe_latent_size == 0): the routed experts work on the hidden vector;
  • LatentMoE (moe_latent_size > 0): the routed experts work in a latent space of moe_latent_size, with the fc1 / fc2 latent projections between hidden and latent.

Forward: RMS norm, gate, sigmoid top-k routing, (fc1 if latent), routed up, relu² and down, weighted sum, (fc2 if latent), plus the shared expert (up, relu², down), summed and added to hidden. Experts are selected on the device through pointer tables; routing reads nothing back to the host.

Owner: model-arch (Nemotron-H). Invariants:

  • new refuses a top_k outside 1..=num_experts or over the routing kernels’ bounds (MOE_TOPK_SIGMOID_MAX_TOP_K, MOE_TOPK_SIGMOID_MAX_EXPERTS).

Structs§

NemotronMoeLayer
2026-09-25: Nemotron-H standalone MoE FFN layer.

Functions§

set_shared_expert_e4m3_from_cli
2026-09-29: Publish the command line’s value; returns the value in force, which differs when something read it first.
shared_expert_e4m3
2026-09-29: Whether the Nemotron-H shared expert’s prefill runs the E4M3 tile GEMM. Off unless the serve published otherwise: the checkpoint’s precision (NVFP4 weights, 16-bit activations on the BF16-MMA w4a16_gemm) is the default.