mstar.model.components.moe

mstar.model.components.moe#

Mixture-of-Experts blocks.

Top-K router + dispatch path for the standard fused-expert checkpoint layout (experts.gate_up_proj and experts.down_proj packed as (num_experts, ...) parameters). Two block flavors:

  • SparseMoeBlock — Top-K MoE with no shared expert (e.g. the Qwen3-Omni Thinker text backbone).

  • SparseMoeBlockWithSharedExpert — Top-K MoE plus a shared expert with sigmoid gating (e.g. the Qwen3-Omni Talker text backbone). The shared expert is passed in as an nn.Module so callers can pick any MLP shape they need.

Parallel (TP-aware) variants:

When triton is installed and inputs are on CUDA, dispatch goes through the Triton fused-MoE kernel in mstar.utils.fused_moe; otherwise it falls back to the naive per-expert loop in dispatch_experts_fused().

Router contract#

Each block computes its expert assignment with a router. The block keeps the router in self.gate. The default router is TopKRouter.

A router is an nn.Module. A router does not extend a base class. The forward method of the router must obey this contract:

router(hidden_states, router_states=None)
    -> (routing_weights, selected_experts, router_states_next)

The router reads these inputs:

  • hidden_states is the input tensor. The shape is (tokens, hidden_size).

  • router_states is the state from the previous call. A stateless router ignores this input.

The router returns these outputs:

  • routing_weights is the tensor of expert weights. The shape is (tokens, top_k).

  • selected_experts is the tensor of expert indices. The shape is (tokens, top_k). The type is int64.

  • router_states_next is the state for the next call. A stateless router returns None.

A stateless router ignores router_states. A stateless router returns None for router_states_next.

A stateful router reads router_states. A stateful router returns a new state. To thread the state through a block, set the return_router_states flag on the block.

Functions

dispatch_experts_fused(hidden_states, ...)

Naive per-expert dispatch using the fused HF checkpoint layout.

Classes

ParallelSparseMoeBlock(hidden_size, ...[, ...])

TP-aware Top-K sparse MoE.

ParallelSparseMoeBlockWithSharedExpert(...)

TP-aware Top-K sparse MoE with a shared expert + sigmoid gating.

SparseMoeBlock(hidden_size, num_experts, ...)

Top-K sparse MoE with fused expert weights, no shared expert.

SparseMoeBlockWithSharedExpert(hidden_size, ...)

Top-K sparse MoE with a shared expert + sigmoid gating.

TopKRouter(hidden_size, num_experts, ...[, ...])

Softmax top-k router.
class mstar.model.components.moe.ParallelSparseMoeBlock(hidden_size, num_experts, num_experts_per_tok, moe_intermediate_size, norm_topk_prob=True, router=None, comm_group=None)[source]#

Bases: Module

TP-aware Top-K sparse MoE.

When tp_size == 1, the forward is identical to SparseMoeBlock (full fused kernel, no communication). When tp_size > 1, expert weights are sharded along the intermediate dimension and an all-reduce is inserted between the down-projection GEMM and the top-k sum-reduce.

Parameters:

  • hidden_size (int)

  • num_experts (int)

  • num_experts_per_tok (int)

  • moe_intermediate_size (int)

  • norm_topk_prob (bool)

  • router (nn.Module | None)

  • comm_group (CommGroup | None)

forward(hidden_states, router_states=None, *, return_router_states=False)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

  • hidden_states (Tensor)

  • router_states (Tensor | None)

  • return_router_states (bool)

class mstar.model.components.moe.ParallelSparseMoeBlockWithSharedExpert(hidden_size, num_experts, num_experts_per_tok, moe_intermediate_size, shared_expert, norm_topk_prob=False, router=None, comm_group=None)[source]#

Bases: Module

TP-aware Top-K sparse MoE with a shared expert + sigmoid gating.

The shared expert should be a ParallelGatedMLP constructed with the same comm_group so its all-reduce is handled internally.

Parameters:

  • hidden_size (int)

  • num_experts (int)

  • num_experts_per_tok (int)

  • moe_intermediate_size (int)

  • shared_expert (nn.Module)

  • norm_topk_prob (bool)

  • router (nn.Module | None)

  • comm_group (CommGroup | None)

forward(hidden_states, router_states=None, *, return_router_states=False)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

  • hidden_states (Tensor)

  • router_states (Tensor | None)

  • return_router_states (bool)

class mstar.model.components.moe.SparseMoeBlock(hidden_size, num_experts, num_experts_per_tok, moe_intermediate_size, norm_topk_prob=True, router=None)[source]#

Bases: Module

Top-K sparse MoE with fused expert weights, no shared expert.

Expert weights match the HF fused checkpoint layout:

  • experts.gate_up_proj: (num_experts, 2 * moe_intermediate_size, hidden_size)

  • experts.down_proj: (num_experts, hidden_size, moe_intermediate_size)

Parameters:

  • hidden_size (int)

  • num_experts (int)

  • num_experts_per_tok (int)

  • moe_intermediate_size (int)

  • norm_topk_prob (bool)

  • router (nn.Module | None)

forward(hidden_states, router_states=None, *, return_router_states=False)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

  • hidden_states (Tensor)

  • router_states (Tensor | None)

  • return_router_states (bool)

class mstar.model.components.moe.SparseMoeBlockWithSharedExpert(hidden_size, num_experts, num_experts_per_tok, moe_intermediate_size, shared_expert, norm_topk_prob=False, router=None)[source]#

Bases: Module

Top-K sparse MoE with a shared expert + sigmoid gating.

Final output is:

out = routed(x) + sigmoid(shared_gate(x)) * shared_expert(x)

The shared expert is supplied by the caller (any nn.Module matching the hidden_size → hidden_size interface). The routed path uses the same fused checkpoint layout as SparseMoeBlock.

Parameters:

  • hidden_size (int)

  • num_experts (int)

  • num_experts_per_tok (int)

  • moe_intermediate_size (int)

  • shared_expert (nn.Module)

  • norm_topk_prob (bool)

  • router (nn.Module | None)

forward(hidden_states, router_states=None, *, return_router_states=False)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

  • hidden_states (Tensor)

  • router_states (Tensor | None)

  • return_router_states (bool)

class mstar.model.components.moe.TopKRouter(hidden_size, num_experts, num_experts_per_tok, norm_topk_prob=True)[source]#

Bases: Module

Softmax top-k router. This is the default router for all MoE blocks.

See the module docstring for the router contract. This router is stateless. It ignores router_states. It returns None for router_states_next.

Parameters:

  • hidden_size (int) – input hidden dimension.

  • num_experts (int) – total number of routed experts.

  • num_experts_per_tok (int) – number of experts each token is dispatched to (top-k).

  • norm_topk_prob (bool) – if True, renormalize the top-k probabilities so they sum to 1.

forward(hidden_states, router_states=None)[source]#
Parameters:

  • hidden_states (Tensor) – (tokens, hidden_size) router input.

  • router_states (Tensor | None) – unused (stateless router); accepted for interface compatibility with stateful routers.

Returns:

(tokens, top_k) top-k probabilities

(optionally renormalized).

selected_experts: (tokens, top_k) int64 indices. router_states_next: always None (stateless).

Return type:

routing_weights

mstar.model.components.moe.dispatch_experts_fused(hidden_states, gate_up_proj, down_proj, num_experts, selected_experts, routing_weights)[source]#

Naive per-expert dispatch using the fused HF checkpoint layout.

Used as a fallback when the Triton fused-MoE kernel isn’t available. Loops over the experts that received any tokens and runs SwiGLU per expert.

Parameters:

  • hidden_states (Tensor) – (tokens, hidden_size).

  • gate_up_proj (Tensor) – (num_experts, 2 * moe_intermediate_size, hidden_size).

  • down_proj (Tensor) – (num_experts, hidden_size, moe_intermediate_size).

  • selected_experts (Tensor) – (tokens, top_k) int64.

  • routing_weights (Tensor) – (tokens, top_k) float.

  • num_experts (int)

Return type:

Tensor