core.transformer.moe.megakernel.mok.route_adapter

Convert authoritative MCore routes into MOK scheduler inputs.

Module Contents

Classes

RoutingMapToIndices

Compact a dense routing map into fixed-width expert indices and probabilities. Unlike torch.topk(probs), this operation treats routing_map as the authoritative route set. Each row is compacted in increasing expert-index order. Rows with fewer than topk routes are padded with -1 indices and zero probabilities. The backward pass scatters router-weight gradients into the corresponding positions of the dense probability tensor.

Functions

_routing_map_to_indices_kernel	Compact one dense routing-map row and gather its router probabilities.
_routing_map_to_indices_backward_kernel	Scatter compact router-weight gradients back into the dense probability row.
routing_map_to_mok_inputs	Compact an authoritative routing map into fixed-width MOK inputs.

API

core.transformer.moe.megakernel.mok.route_adapter._routing_map_to_indices_kernel(

probs_ptr,

routing_map_ptr,

probs_indices_ptr,

indices_ptr,

num_experts: triton.language.constexpr,

topk: triton.language.constexpr,

num_experts_block_size: triton.language.constexpr,

topk_block_size: triton.language.constexpr,

)

Compact one dense routing-map row and gather its router probabilities.

core.transformer.moe.megakernel.mok.route_adapter._routing_map_to_indices_backward_kernel(

grad_probs_indices_ptr,

indices_ptr,

grad_probs_ptr,

num_experts: triton.language.constexpr,

topk: triton.language.constexpr,

num_experts_block_size: triton.language.constexpr,

topk_block_size: triton.language.constexpr,

)

Scatter compact router-weight gradients back into the dense probability row.

class core.transformer.moe.megakernel.mok.route_adapter.RoutingMapToIndices

Bases: torch.autograd.Function

Compact a dense routing map into fixed-width expert indices and probabilities. Unlike torch.topk(probs), this operation treats routing_map as the authoritative route set. Each row is compacted in increasing expert-index order. Rows with fewer than topk routes are padded with -1 indices and zero probabilities. The backward pass scatters router-weight gradients into the corresponding positions of the dense probability tensor.

static forward(ctx, probs, routing_map, topk)

Compact authoritative dense routes into fixed-width MOK inputs.

static backward(ctx, grad_probs_indices, grad_indices)

Scatter compact routing-weight gradients back to the dense tensor.

core.transformer.moe.megakernel.mok.route_adapter.routing_map_to_mok_inputs(

probs: torch.Tensor,

routing_map: torch.Tensor,

topk: int,

) → tuple[torch.Tensor, torch.Tensor]

Compact an authoritative routing map into fixed-width MOK inputs.

Parameters:

• probs – Dense router probabilities with shape [num_tokens, num_experts].

• routing_map – Dense boolean/integer route map with the same shape as probs.

• topk – Maximum number of valid routes emitted per token.

Returns:

A pair (probs_indices, indices) with shapes [num_tokens, topk]. probs_indices is float32 and indices is int32 with -1 padding.