For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

sink_gate_router_kernel

def sink_gate_router_kernel[scores_type: DType, bias_type: DType, ExpertIndicesLayoutType: TensorLayout, ExpertWeightsLayoutType: TensorLayout, SinkWeightsLayoutType: TensorLayout, LogitsLayoutType: TensorLayout, ExpertBiasLayoutType: TensorLayout, GlobalScaleLayoutType: TensorLayout, n_routed_experts: Int, n_experts_per_tok: Int, n_shared_experts: Int, num_threads: Int](expert_indices: TileTensor[.int32, ExpertIndicesLayoutType, MutAnyOrigin], expert_weights: TileTensor[scores_type, ExpertWeightsLayoutType, MutAnyOrigin], sink_weights: TileTensor[scores_type, SinkWeightsLayoutType, MutAnyOrigin], logits: TileTensor[scores_type, LogitsLayoutType, ImmutAnyOrigin], expert_bias: TileTensor[bias_type, ExpertBiasLayoutType, ImmutAnyOrigin], global_scale: TileTensor[scores_type, GlobalScaleLayoutType, ImmutAnyOrigin], route_scale: Float32)

Fused sigmoid-gate MoE router with always-on sink (shared-expert) lanes.

Sink lanes are gated shared experts, not attention sinks.

One block per token, one thread per routed expert. Fuses: sigmoid(logit) + bias -> top-k selection (_block_top_k) -> softmax over the log-sigmoid of the selected experts' raw (unbiased) logits concatenated with n_shared_experts always-selected sink logits -> scale by route_scale * global_scale.

Softmax over log-sigmoids equals sigmoid(z_i) / sum_j sigmoid(z_j), computed in log space so it stays finite where the sigmoids themselves would underflow.

Expert bucketing stays in moe_create_indices: it needs every token's assignment before it can build the per-expert CSR, which this per-token-block kernel cannot provide without a grid-wide sync.

Parameters:

• ​scores_type (DType): DType of the logits and the output weights.
• ​bias_type (DType): DType of the per-routed-expert selection bias.
• ​ExpertIndicesLayoutType (TensorLayout): TensorLayout of the expert_indices output tensor.
• ​ExpertWeightsLayoutType (TensorLayout): TensorLayout of the expert_weights output tensor.
• ​SinkWeightsLayoutType (TensorLayout): TensorLayout of the sink_weights output tensor.
• ​LogitsLayoutType (TensorLayout): TensorLayout of the logits input tensor.
• ​ExpertBiasLayoutType (TensorLayout): TensorLayout of the expert_bias input tensor.
• ​GlobalScaleLayoutType (TensorLayout): TensorLayout of the global_scale input tensor.
• ​n_routed_experts (Int): Total number of routed experts scored per token. Also equals the thread count per block.
• ​n_experts_per_tok (Int): Number of routed experts selected per token.
• ​n_shared_experts (Int): Number of always-selected sink experts. Together with n_experts_per_tok, must sum to a power of two no greater than the warp size (the two are jointly softmax-normalized by a single warp-level reduction).
• ​num_threads (Int): Threads per block; must equal n_routed_experts.

Args:

• ​expert_indices (TileTensor[.int32, ExpertIndicesLayoutType, MutAnyOrigin]): Output selected routed-expert index per token. Shape [num_tokens, n_experts_per_tok].
• ​expert_weights (TileTensor[scores_type, ExpertWeightsLayoutType, MutAnyOrigin]): Output routing weight per selected routed expert. Shape [num_tokens, n_experts_per_tok].
• ​sink_weights (TileTensor[scores_type, SinkWeightsLayoutType, MutAnyOrigin]): Output routing weight per sink expert. Shape [num_tokens, n_shared_experts].
• ​logits (TileTensor[scores_type, LogitsLayoutType, ImmutAnyOrigin]): Input raw (pre-sigmoid) gate logits, routed experts followed by sink experts. Shape [num_tokens, at least n_routed_experts + n_shared_experts]; a wider row's tail is not read.
• ​expert_bias (TileTensor[bias_type, ExpertBiasLayoutType, ImmutAnyOrigin]): Per-routed-expert bias added during selection only. Shape [n_routed_experts].
• ​global_scale (TileTensor[scores_type, GlobalScaleLayoutType, ImmutAnyOrigin]): Single scalar multiplied into every weight. Shape [1].
• ​route_scale (Float32): Compile-time-known-per-model scalar multiplied into every weight alongside global_scale.