chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Reorganize topk logic to clean up code and expose logical experts (#16945)

Browse Source
This commit is contained in:
Junrong Lin
2026-02-23 16:30:57 +08:00
committed by GitHub
parent 4f3dc1ef5b
commit 6448de7e96
1 changed files with 62 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

140
python/sglang/srt/layers/moe/topk.py
View File

@@ -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)