flashinfer.fused_moe.trtllm_gen_routing

flashinfer.fused_moe.trtllm_gen_routing(routing_logits: Tensor, routing_bias: Tensor | None, routing_method: RoutingMethodType, top_k: int, *, num_fused_shared_experts: int = 0, n_group: int = 0, topk_group: int = 0, local_expert_offset: int = 0, local_num_experts: int | None = None, routed_scaling_factor: float = 1.0, tile_tokens_dim: int = 8, norm_topk_prob: bool = True, enable_pdl: bool | None = None) → TrtllmGenRoutingResult

Standalone trtllm-gen MoE routing (expert selection + permutation).

Runs the same routing kernels the trtllm-gen fused MoE launchers execute before their GEMMs (Routing::Runner::run), and returns the routing outputs directly instead of feeding them into a GEMM. This makes the routing stage unit-testable in isolation from quantization/GEMM axes.

Parameters:

  • routing_logits (torch.Tensor) – Router logits of shape (num_tokens, num_experts), float32 or bfloat16.

  • routing_bias (Optional[torch.Tensor]) – Per-expert routing bias of shape (num_experts,), float32 or bfloat16 (used by DeepSeekV3/MiniMax2-style methods).

  • routing_method (RoutingMethodType) – The routing method to run (all methods except Unspecified).

  • top_k (int) – Number of experts selected per token.

  • num_fused_shared_experts (int) – Extra fused shared-expert slots appended per token (DeepSeekV3 path).

  • n_group (int) – Expert-group parameters for grouped routing (DeepSeekV3). 0 disables grouping.

  • topk_group (int) – Expert-group parameters for grouped routing (DeepSeekV3). 0 disables grouping.

  • local_expert_offset (int) – Expert-parallel shard description. local_num_experts defaults to num_experts.

  • local_num_experts (int) – Expert-parallel shard description. local_num_experts defaults to num_experts.

  • routed_scaling_factor (float) – Output weight scale (DeepSeekV3/MiniMax2-style methods).

  • tile_tokens_dim (int) – Token-tile size the downstream grouped GEMM would use; must be a power of two. The permutation/padding outputs depend on it.

  • norm_topk_prob (bool) – Whether SigmoidRenorm renormalizes the selected probabilities. Only consulted for RoutingMethodType.SigmoidRenorm.

  • enable_pdl (Optional[bool]) – Whether to launch with programmatic dependent launch. Defaults to auto-detection.

Returns:

Named tuple with expert selection (topk_ids, topk_weights) and permutation/bookkeeping tensors (total_num_padded_tokens, expanded_idx_to_permuted_idx, permuted_idx_to_token_idx, cta_idx_xy_to_batch_idx, cta_idx_xy_to_mn_limit, num_non_exiting_ctas).

Return type:

TrtllmGenRoutingResult

Notes

topk_weights is always bfloat16: the routing dispatcher hard-codes its output dtype regardless of the logits dtype.

topk_ids is reconstructed from the permutation (the kernels emit no direct id output in from-logits mode): permuted slot p lies in CTA tile p // tile_tokens_dim, whose expert is cta_idx_xy_to_batch_idx[p // tile_tokens_dim] + local_expert_offset.