Implementation:NVIDIA TransformerEngine Fused Router

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, PyTorch, Optimization
Last Updated	2026-02-07 14:00 GMT

Overview

Fused CUDA implementations of Mixture-of-Experts (MoE) router functions including top-k with score function, MoE auxiliary loss computation, and score computation for aux loss.

Description

This module provides three fused MoE router operations implemented as custom torch.autograd.Function classes:

• FusedTopkScoreFunction / fused_topk_with_score_function() -- Fuses top-k selection with score function (softmax or sigmoid), group top-k, and expert bias into a single kernel. Supports both pre-softmax (softmax then top-k) and post-softmax ordering.

• FusedComputeScoresForMoEAuxLoss / fused_compute_score_for_moe_aux_loss() -- Computes scores and routing maps for auxiliary loss calculation in a fused kernel.

• FusedAuxLoss / fused_moe_aux_loss() -- Computes the MoE auxiliary (load balancing) loss from probability distributions and tokens-per-expert counts in a fused kernel.

All three provide both forward and backward implementations using transformer_engine_torch CUDA kernels.

Usage

Used internally by the MoE router module in Transformer Engine and Megatron-LM for efficient expert selection and load balancing.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/pytorch/router.py

Lines

1--275

Signature

def fused_topk_with_score_function(logits, topk, use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, expert_bias) -> Tuple[torch.Tensor, torch.Tensor]: ...
def fused_compute_score_for_moe_aux_loss(logits, topk, score_function) -> Tuple[torch.Tensor, torch.Tensor]: ...
def fused_moe_aux_loss(probs, tokens_per_expert, total_num_tokens, num_experts, topk, coeff) -> torch.Tensor: ...

Import

from transformer_engine.pytorch.router import fused_topk_with_score_function, fused_compute_score_for_moe_aux_loss, fused_moe_aux_loss

I/O Contract

Inputs

Name	Type	Required	Description
logits	torch.Tensor	Yes	Router logits of shape (*, num_experts)
topk	int	Yes	Number of experts to select per token
score_function	str	Yes	"softmax" or "sigmoid"
use_pre_softmax	bool	Yes	If True, apply softmax before top-k
num_groups	int	No	Number of groups for group top-k
group_topk	int	No	Number of experts per group
scaling_factor	float	No	Scaling factor for scores
expert_bias	torch.Tensor	No	Expert bias for sigmoid routing

Outputs

Name	Type	Description
probs	torch.Tensor	Expert selection probabilities
routing_map	torch.Tensor	Expert assignment map

Usage Examples

from transformer_engine.pytorch.router import fused_topk_with_score_function

probs, routing_map = fused_topk_with_score_function(
    logits=router_logits,
    topk=2,
    use_pre_softmax=True,
    num_groups=1,
    group_topk=0,
    scaling_factor=1.0,
    score_function="softmax",
    expert_bias=torch.zeros(num_experts),
)

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements