MoE

moe

MoEGate

MoEGate(hidden_size: int, num_experts: int, top_k: int)

Bases: Module

Top-k gating router for Mixture of Experts.

Computes softmax scores over experts and returns top-k indices and weights.

Reference: Hunyuan3D-2.1 moe_layers.py MoEGate

Parameters:

Name	Type	Description	Default
hidden_size	int	Input feature dimension.	required
num_experts	int	Number of expert networks.	required
top_k	int	Number of experts to route each token to.	required

MoELayer

MoELayer(hidden_size: int, num_experts: int, top_k: int, expert_fn: Callable[[], Module], shared_expert: Module | None = None)

Bases: Module

Mixture of Experts layer with optional shared expert.

Routes tokens to top-k experts via a gating network, computes weighted sum of expert outputs, and optionally adds a shared expert's output.

Reference: Hunyuan3D-2.1 moe_layers.py MoEBlock

Parameters:

Name	Type	Description	Default
hidden_size	int	Input feature dimension.	required
num_experts	int	Number of routed expert networks.	required
top_k	int	Number of experts per token.	required
expert_fn	Callable[[], Module]	Callable that returns an nn.Module expert.	required
shared_expert	Module | None	Optional always-active expert module.	None

moe

MoEGate

MoEGate(hidden_size: int, num_experts: int, top_k: int)

Bases: Module

Top-k gating router for Mixture of Experts.

Computes softmax scores over experts and returns top-k indices and weights.

Reference: Hunyuan3D-2.1 moe_layers.py MoEGate

Parameters:

Name	Type	Description	Default
hidden_size	int	Input feature dimension.	required
num_experts	int	Number of expert networks.	required
top_k	int	Number of experts to route each token to.	required

MoELayer

MoELayer(hidden_size: int, num_experts: int, top_k: int, expert_fn: Callable[[], Module], shared_expert: Module | None = None)

Bases: Module

Mixture of Experts layer with optional shared expert.

Routes tokens to top-k experts via a gating network, computes weighted sum of expert outputs, and optionally adds a shared expert's output.

Reference: Hunyuan3D-2.1 moe_layers.py MoEBlock

Parameters:

Name	Type	Description	Default
hidden_size	int	Input feature dimension.	required
num_experts	int	Number of routed expert networks.	required
top_k	int	Number of experts per token.	required
expert_fn	Callable[[], Module]	Callable that returns an nn.Module expert.	required
shared_expert	Module | None	Optional always-active expert module.	None