Mixture of Experts (MoE)

About

The d9d.module.block.moe package implements sparse Mixture-of-Experts layers. MoELayer combines a router, a token dispatcher and the experts, with an optional shared expert.

Expert Parallelism

To set up expert parallelism, see Horizontal Parallelism. Expert parallelism requires the d9d[backend-expert-parallel-deepep] extra. You must build and install DeepEP manually first.

Components

Router

TopKRouter is a learned router that selects the top-k experts for each token. It computes the routing probabilities in fp32 for numerical stability. An optional expert bias changes which experts are selected but not their probabilities. You can use it for loss-free load balancing.

Token Dispatcher

ExpertCommunicationHandler is the interface that moves tokens to their experts and back.

NoCommunicationHandler is the default. It is used when all experts are local, such as on a single GPU or with tensor parallelism.
With expert parallelism, parallelize_expert_parallel switches the layer to a DeepEP handler. It runs the all-to-all communication over NVLink and RDMA with DeepEP.

Experts

GroupedSwiGLU holds a set of SwiGLU experts. It does not loop over experts. It runs all experts in one grouped GEMM per projection, for any number of tokens per expert. It computes SiLU(gate(x)) * up(x) with a fused Triton kernel.

Kernel Benchmarks (bf16, H100)

Grouped GEMM Backends

GroupedSwiGLU runs its projections with GroupedLinear, which passes the grouped GEMM to a backend. The available backends are:

`kind`	Library	GPUs	Extra
`"quack"`	QuACK	sm90+	`d9d[backend-grouped-mm-quack]`
`"grouped_gemm"`	grouped_gemm	sm80+	`d9d[backend-grouped-mm-grouped-gemm]`
`"torch"`	PyTorch `grouped_mm`	Any	-
`"eager"`	Plain PyTorch ops	Any	-

The "quack" backend compiles its kernels on first use and does not support fp32 on sm90. The "torch" backend is fast only for bf16 on sm90 and sm100. The "eager" backend has no extra dependencies. Use it as a fallback and as a correctness reference for the other backends.

Backend Selection

A layer picks its backend in this order:

The grouped_mm_backend configuration passed to MoELayer, GroupedSwiGLU or GroupedLinear.
The D9D_BACKEND_AUTO_GROUPED_MM environment variable.
Auto-detection: "quack" on sm90+ if it is installed, "grouped_gemm" on sm80 if it is installed, "torch" on sm90 and sm100, and "eager" otherwise.

D9D_BACKEND_AUTO_GROUPED_MM holds a backend configuration as JSON, keyed by the kind field:

# Force the grouped_gemm backend.
export D9D_BACKEND_AUTO_GROUPED_MM='{"kind": "grouped_gemm"}'

Shared Experts

A shared expert processes every token, whatever the router chooses. Configure it with SharedExpertParameters. It can scale its output with a learned sigmoid gate (enable_gate).

Usage

import torch

from d9d.module.block.moe import MoELayer, SharedExpertParameters

moe = MoELayer(
    hidden_dim=2048,
    intermediate_dim_grouped=768,
    num_grouped_experts=128,
    top_k=8,
    router_renormalize_probabilities=True,
    shared_expert=SharedExpertParameters(intermediate_size=4096, enable_gate=True),
).to(device="cuda", dtype=torch.bfloat16)
moe.reset_parameters()

hidden_states = torch.randn(2, 16, 2048, device="cuda", dtype=torch.bfloat16)
output = moe(hidden_states)  # (2, 16, 2048)

# Number of tokens routed to each expert since the last reset.
print(moe.tokens_per_expert)
moe.reset_stats()

API Reference

`d9d.module.block.moe`

Building blocks for Mixture-of-Experts (MoE) layers.

`GroupedLinear`

Bases: Module, ModuleLateInit

Linear layer with a separate transformation for each group of tokens, computed with grouped GEMM.

Each group (expert) has its own weight and processes a variable number of tokens. This is the compute core of the Mixture-of-Experts layer.

`init(n_groups, in_features, out_features, device=None, dtype=None, grouped_mm_backend=None)`

Constructs the GroupedLinear object.

Parameters:

Name	Type	Description	Default
`n_groups`	`int`	Number of groups (experts).	required
`in_features`	`int`	Size of each input sample.	required
`out_features`	`int`	Size of each output sample.	required
`device`	`device \| str \| None`	Target device.	`None`
`dtype`	`dtype \| None`	Target data type.	`None`
`grouped_mm_backend`	`AnyGroupedMmBackendConfig \| None`	Grouped GEMM backend. If `None`, the backend is auto-detected.	`None`

`forward(x, x_groups)`

Performs the grouped matrix multiplication.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tokens of all groups, with the tokens of each group stored together. Shape: `(num_tokens, in_features)`.	required
`x_groups`	`ExpertTokenGroups`	How the tokens are split between groups.	required

Returns:

Type	Description
`Tensor`	The output tensor. Shape: `(num_tokens, out_features)`.

`reset_parameters()`

Initializes the weights from U(-1 / sqrt(in_features), 1 / sqrt(in_features)).

`with_features(in_features, out_features)`

Creates a new GroupedLinear like this one, but with other feature sizes.

The new layer has the same number of groups, grouped GEMM backend configuration, device and dtype, and freshly initialized weights. Layers attached to this one, such as LoRA adapters, use it to run on the same backend without knowing which one it is.

Parameters:

Name	Type	Description	Default
`in_features`	`int`	Input hidden size of the new layer.	required
`out_features`	`int`	Output hidden size of the new layer.	required

Returns:

Type	Description
`GroupedLinear`	The new layer.

`GroupedSwiGLU`

Bases: Module, ModuleLateInit

Set of SwiGLU experts that run with grouped GEMM.

Each expert computes down_proj(SiLU(gate_proj(x)) * up_proj(x)). All experts run in one grouped GEMM per projection, without padding or masking.

`init(hidden_dim, intermediate_dim, num_experts, grouped_mm_backend=None)`

Constructs the GroupedSwiGLU object.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Hidden size of the input and output.	required
`intermediate_dim`	`int`	Intermediate size of each expert.	required
`num_experts`	`int`	Number of experts held by this module.	required
`grouped_mm_backend`	`AnyGroupedMmBackendConfig \| None`	Grouped GEMM backend for the projections. If `None`, the backend is auto-detected.	`None`

`forward(permuted_x, permuted_probs, expert_groups)`

Computes expert outputs for sorted input tokens.

Parameters:

Name	Type	Description	Default
`permuted_x`	`Tensor`	Input tokens sorted by their assigned expert. Shape: `(num_tokens, hidden_size)`.	required
`permuted_probs`	`Tensor`	Routing probabilities of the sorted tokens. Shape: `(num_tokens,)`.	required
`expert_groups`	`ExpertTokenGroups`	How the tokens are split between consecutive experts.	required

Returns:

Name	Type	Description
	`Tensor`	The expert outputs weighted by the routing probabilities, in the sorted order of `permuted_x`.
`Shape`	`Tensor`	`(num_tokens, hidden_size)`.

`reset_parameters()`

Resets module parameters.

`MoELayer`

Bases: Module, ModuleLateInit

Mixture-of-Experts (MoE) block with routing, communication and expert computation.

The layer combines:

Router: Selects experts for each token.
Communicator: Dispatches tokens to local or remote experts (expert parallelism).
Experts: Compute the outputs with grouped SwiGLU.

`init(hidden_dim, intermediate_dim_grouped, num_grouped_experts, top_k, router_renormalize_probabilities, shared_expert=None, grouped_mm_backend=None)`

Constructs the MoELayer object.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Hidden size.	required
`intermediate_dim_grouped`	`int`	Intermediate size of each expert FFN.	required
`num_grouped_experts`	`int`	Total number of experts.	required
`top_k`	`int`	Number of experts to route each token to.	required
`router_renormalize_probabilities`	`bool`	If `True`, the router renormalizes the probabilities of the selected experts to sum to 1.	required
`shared_expert`	`SharedExpertParameters \| None`	Optional configuration for a shared expert.	`None`
`grouped_mm_backend`	`AnyGroupedMmBackendConfig \| None`	Grouped GEMM backend for the experts. If `None`, the backend is auto-detected.	`None`

`enable_distributed_communicator(group)`

Switches from local communication to distributed DeepEP communication.

Call it during model initialization if the model runs with expert parallelism.

Parameters:

Name	Type	Description	Default
`group`	`ProcessGroup`	Process group that spans the expert-parallel ranks.	required

`forward(hidden_states)`

Routes tokens to experts, computes the expert outputs and combines them.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Input tensor. Shape: `(batch, seq_len, hidden_size)`.	required

Returns:

Type	Description
`Tensor`	Output tensor. Shape: `(batch, seq_len, hidden_size)`.

`reset_parameters()`

Resets module parameters.

`reset_stats()`

Resets the expert load balancing counters.

`RoutingResult` `dataclass`

The result of expert routing.

Attributes:

Name	Type	Description
`selected_expert_indices`	`Tensor`	Indices of the chosen experts for each token. Shape: `(num_tokens, top_k)`.
`selected_probabilities`	`Tensor`	Probabilities of the chosen experts. Shape: `(num_tokens, top_k)`.

`SharedExpertParameters`

Bases: BaseModel

Configuration for a shared expert.

Attributes:

Name	Type	Description
`intermediate_size`	`int`	Intermediate size of the SwiGLU FFN.
`enable_gate`	`bool`	Whether to scale the output with a learned sigmoid gate.

`SharedSwiGLU`

Bases: Module, ModuleLateInit

Shared expert: a SwiGLU FFN with an optional sigmoid output gate.

`init(hidden_size, params)`

Constructs the SharedSwiGLU object.

Parameters:

Name	Type	Description	Default
`hidden_size`	`int`	Hidden size.	required
`params`	`SharedExpertParameters`	Shared expert configuration.	required

`forward(hidden_states)`

Applies the shared expert computation to the input.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Input tensor. Shape: `(num_tokens, hidden_size)`.	required

Returns:

Type	Description
`Tensor`	Output tensor. Shape: `(num_tokens, hidden_size)`.

`reset_parameters()`

Resets module parameters.

`TopKRouter`

Bases: Module, ModuleLateInit

Router that selects the top-k experts for each token with a learned gate.

The router:

Projects input tokens to one score per expert.
Applies softmax to get expert probabilities.
Selects the top_k experts with the highest probabilities. The optional expert bias affects only this selection, not the returned probabilities.
Optionally renormalizes the selected probabilities to sum to 1.

`init(dim, num_experts, top_k, renormalize_probabilities, enable_expert_bias=False)`

Constructs the TopKRouter object.

Parameters:

Name	Type	Description	Default
`dim`	`int`	Hidden size.	required
`num_experts`	`int`	Total number of experts to choose from.	required
`top_k`	`int`	Number of experts to select for each token.	required
`renormalize_probabilities`	`bool`	If `True`, the probabilities of the selected experts are renormalized to sum to 1.	required
`enable_expert_bias`	`bool`	If `True`, adds a bias to the expert probabilities before top-k selection. It can be used for loss-free load balancing.	`False`

`forward(hidden_states)`

Computes routing decisions for the input tokens.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Input tokens. Shape: `(num_tokens, hidden_size)`.	required

Returns:

Type	Description
`RoutingResult`	The indices of the selected experts and their probabilities.

`reset_parameters()`

Resets module parameters.

`d9d.module.block.moe.communications`

Communication strategies that route tokens between MoE experts.

`ExpertCommunicationHandler`

Bases: ABC

Abstract base class for Mixture-of-Experts communication strategies.

`combine(hidden_states)` `abstractmethod`

Returns processed hidden states to their original order and rank.

Replicated tokens are summed. Must be called after dispatch.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Processed hidden states. Shape: `(num_received_tokens, hidden_size)`.	required

Returns:

Type	Description
`Tensor`	The combined hidden states in the original order. Shape: `(num_tokens, hidden_size)`.

`dispatch(hidden_states, topk_ids, topk_weights)` `abstractmethod`

Routes local hidden states to their target experts, which can be on other ranks.

Tokens assigned to several experts are replicated. The received tokens are sorted by expert index, as grouped GEMM requires.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Input tokens. Shape: `(num_tokens, hidden_size)`.	required
`topk_ids`	`Tensor`	Indices of the experts selected for each token. Shape: `(num_tokens, top_k)`.	required
`topk_weights`	`Tensor`	Routing weights of the selected experts. Shape: `(num_tokens, top_k)`.	required

Returns:

Type	Description
`Tensor`	A tuple of:
`Tensor`	Sorted hidden states received by this rank. Shape: `(num_received_tokens, hidden_size)`.
`ExpertTokenGroups`	Routing weights in the same order. Shape: `(num_received_tokens,)`.
`tuple[Tensor, Tensor, ExpertTokenGroups]`	How the received tokens are split between the local experts.

`NoCommunicationHandler`

Bases: ExpertCommunicationHandler

MoE communication handler for the case when all experts are on the local device.

This handler does no communication. It only sorts the tokens by expert index.

`init(num_experts)`

Constructs the NoCommunicationHandler object.

Parameters:

Name	Type	Description	Default
`num_experts`	`int`	Number of local experts.	required

`d9d.module.block.moe.grouped_mm`

Provides pluggable grouped GEMM backends for Mixture-of-Experts layers.

`EagerGroupedMmBackendConfig`

Bases: BaseModel

Configuration for the eager backend: one matmul per expert.

A portable, dependency-free reference. Needs the per-expert counts on the host.

Attributes:

Name	Type	Description
`kind`	`Literal['eager']`	Discriminator field. Always `"eager"`.

`ExpertTokenGroups`

How the permuted tokens are split between consecutive experts.

`offsets` `property`

Token offsets of the experts with a leading zero, on the device. Shape: (num_experts + 1,).

`init(counts, host_counts=None)`

Constructs the ExpertTokenGroups object.

Parameters:

Name	Type	Description	Default
`counts`	`Tensor`	Tokens per expert on the device the permuted tokens live on. Shape: `(num_experts,)`.	required
`host_counts`	`Tensor \| None`	The same counts on the host, if the producer already has them. Shape: `(num_experts,)`.	`None`

`host_counts()`

Returns the tokens per expert on the host.

If the producer did not provide them, the first call copies them from the device and synchronizes with it. Later calls return the same tensor.

Returns:

Type	Description
`Tensor`	The tokens per expert. Shape: `(num_experts,)`.

`GroupedGemmBackendConfig`

Bases: BaseModel

Configuration for the grouped_gemm library backend.

Needs the per-expert counts on the host; getting them from the device synchronizes with it. Requires the d9d[backend-grouped-mm-grouped-gemm] extra.

Attributes:

Name	Type	Description
`kind`	`Literal['grouped_gemm']`	Discriminator field. Always `"grouped_gemm"`.

`GroupedMmBackend`

Bases: Protocol

Protocol for grouped matrix multiplication backends used by GroupedLinear.

`call(x, weight, groups)`

Multiplies the rows of every group by the weight of that group.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Rows of all groups, sorted by group. Shape: `(total_rows, in_features)`.	required
`weight`	`Tensor`	Per-group weights. Shape: `(n_groups, in_features, out_features)`.	required
`groups`	`ExpertTokenGroups`	How the rows are split between groups.	required

Returns:

Type	Description
`Tensor`	The result. Shape: `(total_rows, out_features)`.

`GroupedMmParameters` `dataclass`

Structural parameters of a grouped linear layer that a grouped GEMM backend must support.

Attributes:

Name	Type	Description
`in_features`	`int`	Size of the contracted dimension.
`out_features`	`int`	Size of the output dimension.

`QuackGroupedMmBackendConfig`

Bases: BaseModel

Configuration for the QuACK (CuTe-DSL) variable-length GEMM backend.

Reads the expert offsets on the device, so it does not need the per-expert counts on the host. Kernels are JIT-compiled on first use. Requires the d9d[backend-grouped-mm-quack] extra.

Attributes:

Name	Type	Description
`kind`	`Literal['quack']`	Discriminator field. Always `"quack"`.

`TorchGroupedMmBackendConfig`

Bases: BaseModel

Configuration for the PyTorch torch.nn.functional.grouped_mm backend.

Reads the expert offsets on the device, so it does not need the per-expert counts on the host. Fast only for bf16 on sm90 and sm100.

Attributes:

Name	Type	Description
`kind`	`Literal['torch']`	Discriminator field. Always `"torch"`.

`build_grouped_mm_backend(params, backend_config)`

Builds the grouped GEMM backend for the given configuration.

If backend_config is None, the backend config is read from the D9D_BACKEND_AUTO_GROUPED_MM environment variable (a JSON backend config). If the variable is not set, the factory picks QuACK on sm90+ if it is installed, grouped_gemm on sm80 if it is installed, PyTorch grouped_mm on sm90 and sm100, and the eager backend otherwise.

Parameters:

Name	Type	Description	Default
`params`	`GroupedMmParameters`	Structural layer parameters the backend must support.	required
`backend_config`	`AnyGroupedMmBackendConfig \| None`	Explicit grouped GEMM backend configuration, or `None` to auto-detect.	required

Returns:

Type	Description
`GroupedMmBackend`	The grouped GEMM backend.

Raises:

Type	Description
`ValueError`	If the backend configuration type is unknown.

Mixture of Experts (MoE)

About

Expert Parallelism

Components

Router

Token Dispatcher

Experts

Kernel Benchmarks (bf16, H100)

Grouped GEMM Backends

Backend Selection

Shared Experts

Usage

API Reference

d9d.module.block.moe

GroupedLinear

__init__(n_groups, in_features, out_features, device=None, dtype=None, grouped_mm_backend=None)

forward(x, x_groups)

reset_parameters()

with_features(in_features, out_features)

GroupedSwiGLU

__init__(hidden_dim, intermediate_dim, num_experts, grouped_mm_backend=None)

forward(permuted_x, permuted_probs, expert_groups)

reset_parameters()

MoELayer

__init__(hidden_dim, intermediate_dim_grouped, num_grouped_experts, top_k, router_renormalize_probabilities, shared_expert=None, grouped_mm_backend=None)

enable_distributed_communicator(group)

forward(hidden_states)

reset_parameters()

reset_stats()

RoutingResult dataclass

SharedExpertParameters

SharedSwiGLU

__init__(hidden_size, params)

forward(hidden_states)

reset_parameters()

TopKRouter

__init__(dim, num_experts, top_k, renormalize_probabilities, enable_expert_bias=False)

forward(hidden_states)

reset_parameters()

d9d.module.block.moe.communications

ExpertCommunicationHandler

combine(hidden_states) abstractmethod

dispatch(hidden_states, topk_ids, topk_weights) abstractmethod

NoCommunicationHandler

__init__(num_experts)

d9d.module.block.moe.grouped_mm

EagerGroupedMmBackendConfig

ExpertTokenGroups

offsets property

__init__(counts, host_counts=None)

host_counts()

GroupedGemmBackendConfig

GroupedMmBackend

__call__(x, weight, groups)

GroupedMmParameters dataclass

QuackGroupedMmBackendConfig

TorchGroupedMmBackendConfig

build_grouped_mm_backend(params, backend_config)

`d9d.module.block.moe`

`GroupedLinear`

`init(n_groups, in_features, out_features, device=None, dtype=None, grouped_mm_backend=None)`

`forward(x, x_groups)`

`reset_parameters()`

`with_features(in_features, out_features)`

`GroupedSwiGLU`

`init(hidden_dim, intermediate_dim, num_experts, grouped_mm_backend=None)`

`forward(permuted_x, permuted_probs, expert_groups)`

`reset_parameters()`

`MoELayer`

`init(hidden_dim, intermediate_dim_grouped, num_grouped_experts, top_k, router_renormalize_probabilities, shared_expert=None, grouped_mm_backend=None)`

`enable_distributed_communicator(group)`

`forward(hidden_states)`

`reset_parameters()`

`reset_stats()`

`RoutingResult` `dataclass`

`SharedExpertParameters`

`SharedSwiGLU`

`init(hidden_size, params)`

`forward(hidden_states)`

`reset_parameters()`

`TopKRouter`

`init(dim, num_experts, top_k, renormalize_probabilities, enable_expert_bias=False)`

`forward(hidden_states)`

`reset_parameters()`

`d9d.module.block.moe.communications`

`ExpertCommunicationHandler`

`combine(hidden_states)` `abstractmethod`

`dispatch(hidden_states, topk_ids, topk_weights)` `abstractmethod`

`NoCommunicationHandler`

`init(num_experts)`

`d9d.module.block.moe.grouped_mm`

`EagerGroupedMmBackendConfig`

`ExpertTokenGroups`

`offsets` `property`

`init(counts, host_counts=None)`

`host_counts()`

`GroupedGemmBackendConfig`

`GroupedMmBackend`

`call(x, weight, groups)`

`GroupedMmParameters` `dataclass`

`QuackGroupedMmBackendConfig`

`TorchGroupedMmBackendConfig`

`build_grouped_mm_backend(params, backend_config)`