diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index a44fe5dc4..aeb5e793c 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -418,8 +418,6 @@ def fused_topk_cpu( gating_output: torch.Tensor, topk: int, renormalize: bool, - num_token_non_padded: Optional[torch.Tensor] = None, - expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, correction_bias: torch.Tensor = None, scoring_func: str = "softmax", ): @@ -429,8 +427,6 @@ def fused_topk_cpu( topk=topk, renormalize=renormalize, ) - topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) - _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) return topk_weights, topk_ids @@ -453,8 +449,6 @@ def fused_topk( topk: int, renormalize: bool, correction_bias: Optional[torch.Tensor] = None, - num_token_non_padded: Optional[torch.Tensor] = None, - expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, scoring_func: str = "softmax", ): assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" @@ -484,8 +478,6 @@ def fused_topk( else: raise ValueError(f"Invalid scoring function: {scoring_func}") - topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) - _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) return topk_weights, topk_ids @@ -500,8 +492,6 @@ def grouped_topk_gpu( topk_group: Optional[int] = None, num_fused_shared_experts: int = 0, routed_scaling_factor: Optional[float] = None, - num_token_non_padded: Optional[torch.Tensor] = None, - expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, apply_routed_scaling_factor_on_output: Optional[bool] = False, ): assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" @@ -553,8 +543,7 @@ def grouped_topk_gpu( topk_weights *= routed_scaling_factor topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32) - topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) - _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) + return topk_weights, topk_ids @@ -567,12 +556,9 @@ def grouped_topk_cpu( topk_group: Optional[int] = None, num_fused_shared_experts: int = 0, routed_scaling_factor: Optional[float] = None, - num_token_non_padded: Optional[torch.Tensor] = None, - expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, apply_routed_scaling_factor_on_output: Optional[bool] = False, ): assert not apply_routed_scaling_factor_on_output - assert expert_location_dispatch_info is None return torch.ops.sgl_kernel.grouped_topk_cpu( hidden_states, gating_output, @@ -582,7 +568,8 @@ def grouped_topk_cpu( topk_group, num_fused_shared_experts, routed_scaling_factor, - num_token_non_padded, + # num_token_non_padded must be None since it is not supported in kernel + num_token_non_padded=None, ) @@ -594,8 +581,6 @@ def kimi_k2_biased_topk_impl( topk: int, renormalize: bool, routed_scaling_factor: Optional[float] = None, - num_token_non_padded: Optional[torch.Tensor] = None, - expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, apply_routed_scaling_factor_on_output: Optional[bool] = False, ): """ @@ -623,8 +608,6 @@ def kimi_k2_biased_topk_impl( topk_weights *= routed_scaling_factor topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32) - topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) - _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) return topk_weights, topk_ids @@ -639,8 +622,6 @@ def biased_grouped_topk_impl( topk_group: Optional[int] = None, num_fused_shared_experts: int = 0, routed_scaling_factor: Optional[float] = None, - num_token_non_padded: Optional[torch.Tensor] = None, - expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, apply_routed_scaling_factor_on_output: Optional[bool] = False, ): assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" @@ -699,8 +680,7 @@ def biased_grouped_topk_impl( topk_weights *= routed_scaling_factor topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32) - topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) - _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) + return topk_weights, topk_ids @@ -737,8 +717,6 @@ def biased_grouped_topk_gpu( topk_group: Optional[int] = None, num_fused_shared_experts: int = 0, routed_scaling_factor: Optional[float] = None, - num_token_non_padded: Optional[torch.Tensor] = None, - expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, apply_routed_scaling_factor_on_output: Optional[bool] = False, ): @@ -789,12 +767,6 @@ def biased_grouped_topk_gpu( True, ) - if (expert_location_dispatch_info is not None) or ( - num_token_non_padded is not None - ): - topk_ids = _biased_grouped_topk_postprocess( - topk_ids, expert_location_dispatch_info, num_token_non_padded - ) return topk_weights, topk_ids elif ( @@ -813,13 +785,7 @@ def biased_grouped_topk_gpu( routed_scaling_factor if routed_scaling_factor is not None else 1.0, apply_routed_scaling_factor_on_output, ) - # TODO merge into kernel - if (expert_location_dispatch_info is not None) or ( - num_token_non_padded is not None - ): - topk_ids = _biased_grouped_topk_postprocess( - topk_ids, expert_location_dispatch_info, num_token_non_padded - ) + return topk_weights, topk_ids elif _use_aiter: @@ -865,8 +831,6 @@ def biased_grouped_topk_gpu( topk_group, num_fused_shared_experts=num_fused_shared_experts, routed_scaling_factor=routed_scaling_factor, - num_token_non_padded=num_token_non_padded, - expert_location_dispatch_info=expert_location_dispatch_info, apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, ) @@ -882,11 +846,8 @@ def biased_grouped_topk_cpu( compiled: bool = True, num_fused_shared_experts: int = 0, routed_scaling_factor: Optional[float] = None, - num_token_non_padded: Optional[torch.Tensor] = None, - expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, apply_routed_scaling_factor_on_output: Optional[bool] = False, ): - assert expert_location_dispatch_info is None assert not apply_routed_scaling_factor_on_output, "Not implemented" return torch.ops.sgl_kernel.biased_grouped_topk_cpu( hidden_states, @@ -898,7 +859,8 @@ def biased_grouped_topk_cpu( topk_group, num_fused_shared_experts, routed_scaling_factor, - num_token_non_padded, + # num_token_non_padded must be None since it is not supported in kernel + num_token_non_padded=None, ) @@ -913,6 +875,50 @@ else: fused_topk_native = fused_topk_torch_native +def _post_process_topk_ids( + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + topk_config: TopKConfig, + router_logits: torch.Tensor, + layer_id: int, + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, +) -> torch.Tensor: + num_fused_shared_experts = topk_config.num_fused_shared_experts + fused_shared_experts_scaling_factor = ( + topk_config.fused_shared_experts_scaling_factor + ) + get_global_experts_capturer().capture( + layer_id=layer_id, + topk_ids=topk_ids, + ) + if _is_cuda: + topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) + _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) + + if num_fused_shared_experts > 0 and _use_aiter: + M, N = router_logits.shape + scale_factor = ( + 1.0 + if fused_shared_experts_scaling_factor is None + else fused_shared_experts_scaling_factor + ) + + # Lazy import to avoid circular-import issues + from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_kernels import ( + fused_append_shared_experts, + ) + + topk_ids, topk_weights = fused_append_shared_experts( + topk_ids, + topk_weights, + num_fused_shared_experts, + scale_factor, + N, # base id for shared experts + ) + return topk_ids, topk_weights + + def select_experts( hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -936,9 +942,7 @@ def select_experts( apply_routed_scaling_factor_on_output = ( topk_config.apply_routed_scaling_factor_on_output ) - fused_shared_experts_scaling_factor = ( - topk_config.fused_shared_experts_scaling_factor - ) + scoring_func = topk_config.scoring_func router_logits, correction_bias = ( @@ -965,8 +969,6 @@ def select_experts( topk_group=topk_group, num_fused_shared_experts=num_fused_shared_experts, routed_scaling_factor=routed_scaling_factor, - num_token_non_padded=num_token_non_padded, - expert_location_dispatch_info=expert_location_dispatch_info, apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, ) else: @@ -980,8 +982,6 @@ def select_experts( topk_group=topk_group, num_fused_shared_experts=num_fused_shared_experts, routed_scaling_factor=routed_scaling_factor, - num_token_non_padded=num_token_non_padded, - expert_location_dispatch_info=expert_location_dispatch_info, apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, ) elif torch_native and custom_routing_function is None: @@ -1007,8 +1007,6 @@ def select_experts( topk=num_routed_topk if _use_aiter else top_k, renormalize=renormalize, correction_bias=correction_bias, - num_token_non_padded=num_token_non_padded, - expert_location_dispatch_info=expert_location_dispatch_info, scoring_func=scoring_func, ) else: @@ -1024,32 +1022,18 @@ def select_experts( renormalize=renormalize, ) - if num_fused_shared_experts > 0 and _use_aiter: - M, N = router_logits.shape - scale_factor = ( - 1.0 - if fused_shared_experts_scaling_factor is None - else fused_shared_experts_scaling_factor - ) - - # Lazy import to avoid circular-import issues - from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_kernels import ( - fused_append_shared_experts, - ) - - topk_ids, topk_weights = fused_append_shared_experts( - topk_ids, - topk_weights, - num_fused_shared_experts, - scale_factor, - N, # base id for shared experts - ) + topk_ids, topk_weights = _post_process_topk_ids( + topk_ids=topk_ids, + topk_weights=topk_weights, + topk_config=topk_config, + router_logits=router_logits, + num_token_non_padded=num_token_non_padded, + layer_id=layer_id, + expert_location_dispatch_info=expert_location_dispatch_info, + ) get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids) - get_global_experts_capturer().capture( - layer_id=layer_id, - topk_ids=topk_ids, - ) + return StandardTopKOutput(topk_weights, topk_ids, router_logits)