flashinfer.fused_moe.hash_topk

flashinfer.fused_moe.hash_topk(router_logits: Tensor, input_ids: Tensor, tid2eid: Tensor, num_fused_shared_experts: int = 0, routed_scaling_factor: float = 1.0, launch_with_pdl: bool = True) → Tuple[Tensor, Tensor]

Hash-based MoE expert routing for DeepSeek-V4.

DSv4-Pro hash-MoE layers select experts from a precomputed token-to-expert table (tid2eid) instead of running a dynamic top-k. Routing is therefore an \(O(1)\) table lookup followed by a sqrt(softplus(.)) score normalization. One warp processes one token.

The routed weight for a selected expert is sqrt(softplus(router_logits[token, expert])) / sum_over_routed. When num_fused_shared_experts == 1, an extra shared-expert slot is appended with id num_routed_experts and weight 1 / routed_scaling_factor.

Parameters:

  • router_logits (torch.Tensor) – Router logits of shape (num_tokens, num_routed_experts), float32.

  • input_ids (torch.Tensor) – Vocabulary token ids of shape (num_tokens,), int64. Indexes the tid2eid table.

  • tid2eid (torch.Tensor) – Precomputed token-to-expert table of shape (vocab, topk), int32.

  • num_fused_shared_experts (int) – Number of fused shared experts to append (0 or 1). Default 0.

  • routed_scaling_factor (float) – Scaling factor for the shared-expert weight. Default 1.0.

  • launch_with_pdl (bool) – Whether to launch with programmatic dependent launch (SM90+). Default True.

Returns:

  • topk_weights (torch.Tensor) – Routing weights of shape (num_tokens, topk + num_fused_shared_experts), float32.

  • topk_ids (torch.Tensor) – Selected expert ids of shape (num_tokens, topk + num_fused_shared_experts), int32.

Notes

The signature matches SGLang’s sglang.jit_kernel.deepseek_v4.hash_topk so this can be used as a drop-in replacement. Implements MOE-01-HASH from the DSv4 tracker.