Skip to main content

MoE

Dense mixture-of-experts: a router picks top-k experts per token, each runs its gated FFN, and the outputs combine under the router weights.

__init__​

__init__(self, num_experts: 'int', hidden_size: 'int', intermediate_size: 'int', top_k: 'int', *, gate_up_bias: 'bool' = False, down_bias: 'bool' = False, dtype: 'DataType' = DataType.Float32, device: 'Device | None' = None, **options: 'Any') -> 'None'

Declare the expert stacks. options forwards the routing kwargs (routing_mode, renormalize, activation, swiglu_fusion, gelu_mode, n_group, topk_group, routed_scaling_factor, sparse_mixer_eps, apply_router_weight_on_input, swiglu_alpha, swiglu_beta, swiglu_limit, expert_output_scale); unset kwargs keep the runtime defaults.

extra_repr​

extra_repr(self) -> 'str'

One line of per-class detail for :meth:__repr__ (a layer prints its geometry here).

forward​

forward(self, x: 'Tensor', router_logits: 'Tensor', shared_output: 'Tensor | None' = None) -> 'Tensor'

Subclasses define the computation here; call the module itself (m(x)), not forward directly.

set_weights​

set_weights(self, gate_up_experts: 'Tensor', down_experts: 'Tensor', *, gate_up_bias: 'Tensor | None' = None, down_bias: 'Tensor | None' = None, gate_experts: 'Tensor | None' = None, gate_bias: 'Tensor | None' = None, e_score_correction_bias: 'Tensor | None' = None) -> 'None'

Bind the declared expert stacks (and the declared biases; the [E] selection bias binds even when not declared).