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
def sink_gate_router[scores_type: DType, bias_type: DType, //, n_routed_experts: Int, n_experts_per_tok: Int, n_shared_experts: Int, target: StringSpan[ImmStaticOrigin]](expert_indices: TileTensor[.int32, Engine=expert_indices.Engine, address_space=expert_indices.address_space, linear_idx_type=expert_indices.linear_idx_type], expert_weights: TileTensor[scores_type, Engine=expert_weights.Engine, address_space=expert_weights.address_space, linear_idx_type=expert_weights.linear_idx_type], sink_weights: TileTensor[scores_type, Engine=sink_weights.Engine, address_space=sink_weights.address_space, linear_idx_type=sink_weights.linear_idx_type], logits: TileTensor[scores_type, Engine=logits.Engine, address_space=logits.address_space, linear_idx_type=logits.linear_idx_type], expert_bias: TileTensor[bias_type, Engine=expert_bias.Engine, address_space=expert_bias.address_space, linear_idx_type=expert_bias.linear_idx_type], global_scale: TileTensor[scores_type, Engine=global_scale.Engine, address_space=global_scale.address_space, linear_idx_type=global_scale.linear_idx_type], route_scale: Float32, context: DeviceContext)
Launch the fused sink-gate MoE router on GPU.
See sink_gate_router_kernel for the fused computation. One block per
token, one thread per routed expert.
Inputs: expert_indices: Output selected expert indices. Shape: [num_tokens, n_experts_per_tok]. expert_weights: Output selected-expert weights. Shape: [num_tokens, n_experts_per_tok]. sink_weights: Output sink-expert weights. Shape: [num_tokens, n_shared_experts]. logits: Input raw gate logits (routed then sink columns). Shape: [num_tokens, at least n_routed_experts + n_shared_experts]; a wider row's tail is not read. expert_bias: Per-routed-expert selection bias. global_scale: Scalar output-scaling weight. route_scale: Scalar output-scaling factor. context: The device context.
Parameters:
- scores_type (
DType): DType of logits and output weights. - bias_type (
DType): DType of the expert selection bias. - n_routed_experts (
Int): Total number of routed experts (e.g. 256 for Inkling-Small). - n_experts_per_tok (
Int): Routed experts selected per token (e.g. 6 for Inkling-Small). - n_shared_experts (
Int): Always-selected sink experts (e.g. 2 for Inkling-Small). - target (
StringSpan[ImmStaticOrigin]): The target device to run the kernel on.