Is this page helpful?

core.inference.moe.sigmoid_topk#

Sigmoid top-k routing kernel for inference.

Replaces the compiled topk_routing_with_score_function sequence for sigmoid scores with an expert bias (an inductor sigmoid+bias kernel, aten topk, and two inductor gather/normalize kernels) with one launch. Expert selection and ordering match torch.topk(..., sorted=False). Probabilities match the compiled left-fold normalization bitwise for top-k below eight; for larger top-k the normalization sum is reduced in a different order, so a probability can differ in its last bit. With batch_invariant=True the launch configuration does not depend on the number of tokens, so every row’s result is independent of the batch; batch-invariant routing uses the kernel this way in training and inference.

Module Contents#

Functions#

is_sigmoid_topk_supported

Whether the fused kernel supports this routing shape.

_sigmoid_topk_kernel

Top-k of sigmoid(logits) + bias, probabilities from sigmoid(logits).

sigmoid_topk_routing

Dense sigmoid top-k routing with an expert bias.

API#

core.inference.moe.sigmoid_topk.is_sigmoid_topk_supported(num_experts: int, topk: int) → bool#

Whether the fused kernel supports this routing shape.

core.inference.moe.sigmoid_topk._sigmoid_topk_kernel(
logits_ptr,
expert_bias_ptr,
probs_ptr,
indices_ptr,
num_tokens,
scaling_factor,
NUM_EXPERTS: triton.language.constexpr,
BIAS_STRIDE: triton.language.constexpr,
EXPERTS_POW2: triton.language.constexpr,
TOPK: triton.language.constexpr,
TOPK_POW2: triton.language.constexpr,
USE_TL_TOPK: triton.language.constexpr,
APPLY_SCALING: triton.language.constexpr,
BLOCK_T: triton.language.constexpr,
)#

Top-k of sigmoid(logits) + bias, probabilities from sigmoid(logits).

Matches the compiled reference:

  • selection and order follow aten topk(sorted=False): a row’s k-th largest key under radix (total float) order, then every column above it in ascending column order, then columns equal to it in ascending order;

  • probabilities are (s_j / (s_0 + … + s_{k-1} + 1e-20)) * scaling_factor. For top-k below eight the sum is folded left in output order, as inductor generates it, so probabilities match bit for bit. Larger top-k uses a tree reduction, which can differ from the compiled sum in the last bit.

core.inference.moe.sigmoid_topk.sigmoid_topk_routing(
logits: torch.Tensor,
expert_bias: torch.Tensor,
topk: int,
scaling_factor: Optional[float] = None,
block_tokens: Optional[int] = None,
batch_invariant: bool = False,
) → tuple[torch.Tensor, torch.Tensor]#

Dense sigmoid top-k routing with an expert bias.

Equivalent to topk_routing_with_score_function(logits, topk, score_function="sigmoid", expert_bias=expert_bias, scaling_factor=..., dense_output=True) evaluated by torch.compile with gradients disabled (which selects with torch.topk(sorted=False)) for finite fp32 logits and bias in the shape range covered by is_sigmoid_topk_supported. Indices always match bitwise. Probabilities match bitwise for top-k below eight; for larger top-k the compiled reduction order depends on autotuning and prior batch shapes, so a probability can differ in its last bit.

Parameters:

  • logits – [num_tokens, num_experts] fp32 router logits.

  • expert_bias – [num_experts] fp32 routing bias on the same device as logits (selection only). Strided views are supported.

  • topk – experts per token, in [2, 32] and no greater than num_experts.

  • scaling_factor – optional multiplier applied to the normalized probabilities.

  • block_tokens – optional rows per CTA; defaults to one for top-k of eight or more or up to 128 experts, and four otherwise.

  • batch_invariant – keep the launch configuration independent of the number of tokens, so each row’s probabilities and indices do not depend on the batch.

Returns:

probs [num_tokens, topk] fp32 and top_indices [num_tokens, topk] int64.