moe¶

moe ¶

Moe: a mixture of experts, an Mlp whose routing and experts are standard values.

Nearly every mixture in transformers computes the same thing::

logits, w, idx = router(x)          # logits [tokens, num_experts]; w, idx [tokens, top_k]; tokens = batch * seq
routed = experts(x, idx, w)         # [tokens, hidden_size]
out = routed (+ shared(x))

and the experts module's call receives (hidden, top_k_index, top_k_weights). So the values are read where the model consumes them: the logits at the operation inside the router's forward that produces them (before any scoring), the weights and indices as the experts' arguments, the per-slot outputs inside transformers' shared experts implementation, the routed sum as the experts' output. Every write lands on the tensor the model uses.

The routing tensors are flat over tokens; each value is served as this invoke's [batch, seq, ...] (a TokenEProperty), a view, so in-place edits land, and an assignment is spliced back into the flat tensor.

RouterLogits `module-attribute` ¶

RouterLogits = Float[Tensor, f'{BATCH} {SEQ} {EXPERTS}']

ExpertWeights `module-attribute` ¶

ExpertWeights = Float[Tensor, f'{BATCH} {SEQ} {TOP_K}']

ExpertIndices `module-attribute` ¶

ExpertIndices = Int[Tensor, f'{BATCH} {SEQ} {TOP_K}']

ExpertOutputs `module-attribute` ¶

ExpertOutputs = Float[Tensor, f'{BATCH} {SEQ} {TOP_K} {HIDDEN}']

LOGITS `module-attribute` ¶

LOGITS = 'F_linear_0'

DISPATCH `module-attribute` ¶

DISPATCH = 'experts_forward_1'

PER_SLOT `module-attribute` ¶

PER_SLOT = 'weighted_out_view_0'

SHARED_NAMES `module-attribute` ¶

SHARED_NAMES = ('shared_experts', 'shared_expert', 'shared_mlp')

ROUTED `module-attribute` ¶

ROUTED = 'routed experts; read per-expert via the Moe component'

Moe ¶

Moe(*args: Any, **kwargs: Any)

Bases: Mlp

A mixture of experts: its routing, each routed expert's output, and the shared expert's.

mlp_output is what the block adds, as on any Mlp. The rest, read where the model consumes them (router, experts and shared_experts are the standard child names; a family aliases its own):

router_logits [batch, seq, num_experts]: before the scoring, at the operation in the router's forward that produces them (LOGITS), so a write changes the routing.
expert_weights / expert_indices [batch, seq, top_k]: the experts module's arguments. Zero a weight to ablate that slot; write an index to reroute (the weight is not recomputed).
expert_outputs [batch, seq, top_k, hidden_size]: each slot's weighted output, inside transformers' grouped_mm / batched_mm experts forward; unavailable under experts_implementation="eager".
routed_output: the experts' sum, expert_outputs.sum(2).
shared_expert_output: the shared expert's contribution, where there is one; routed_output + shared_expert_output is the mixture's output.

Its sizes, num_experts and top_k, are read off the router or the experts module. SCORING says what the logits mean. A family whose module runs the mixture only on some checkpoints (Gemma-4's dense MLP, which hosts the block's experts) says why not in no_mixture, the one place every value's availability starts from.

SCORING `class-attribute` `instance-attribute` ¶

SCORING = 'softmax'

num_experts `property` ¶

num_experts: int

Routed experts in this mixture: the router's or experts module's num_experts / num_local_experts / n_routed_experts.

top_k `property` ¶

top_k: int

Experts each token is routed to: the module's or router's top_k / num_experts_per_tok.

mlp_activation ¶

mlp_activation(value: Tensor) -> Neurons

mlp_neurons ¶

mlp_neurons(value: Tensor) -> Neurons

no_mixture ¶

no_mixture() -> str | None

Why this module runs no mixture of experts, or None (every Moe module does, unless its family says otherwise).

router_logits ¶

router_logits(value: Tensor) -> RouterLogits

The router's logits, [batch, seq, num_experts], where the router's forward produces them.

Before the scoring (SCORING) and any selection bias, so softmax(router_logits) is the routing distribution on a softmax router. Assign or edit in place: the router scores what is written, so the weights and indices follow.

expert_weights ¶

expert_weights(value: Tensor) -> ExpertWeights

The weights the experts receive, [batch, seq, top_k], one per routing slot.

After the router's normalization and scaling: what multiplies each slot's expert output. Zero [b, t, j] to remove slot j of token t; w.masked_fill(expert_indices == e, 0) removes expert e.

expert_indices ¶

expert_indices(value: Tensor) -> ExpertIndices

The experts each token is routed to, [batch, seq, top_k], int64.

Write an index to reroute that slot to another expert; its weight stays what the router gave the expert it chose.

expert_outputs ¶

expert_outputs(value: Tensor) -> ExpertOutputs

Each slot's expert output times its weight, [batch, seq, top_k, hidden_size], before the slots are summed.

expert_outputs.sum(2) == routed_output. The model's own tensor in transformers' grouped_mm (the default) and batched_mm experts forwards; unavailable under eager and hub kernels.

routed_output ¶

routed_output(value: Tensor) -> Residual

The routed experts' weighted sum, [batch, seq, hidden_size], without the shared expert.

shared_expert_output ¶

shared_expert_output(value: Tensor) -> Residual

What the shared expert adds beside the routed ones, [batch, seq, hidden_size]; unavailable on a mixture without one.

mixture_reason ¶

mixture_reason(envoy: Envoy) -> str | None

Why this module runs no mixture of experts, or None: the one place a family decides it (Moe.no_mixture).

routed ¶

routed(name: str)

Why a dense-MLP value at the child name is unavailable on a Moe: ROUTED where it runs a mixture, else absent.

needs_grouped_experts ¶

needs_grouped_experts(envoy: Envoy) -> str | None

Why the per-slot expert outputs are unavailable, or None: they exist only in the grouped_mm / batched_mm forwards.

no_shared_expert ¶

no_shared_expert(envoy: Envoy) -> str | None

Why shared_expert_output is unavailable on this mixture, or None: it has no shared expert.