EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Fix topk inference performance reduce (#6474)

Browse Source
This commit is contained in:
Li Hui
2025-05-23 17:58:31 +08:00
committed by GitHub
parent d8189660a9
commit 2f42749184
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -264,6 +264,8 @@ def biased_grouped_topk(
# TODO merge into kernel for this branch # TODO merge into kernel for this branch
topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
# TODO will fuse this into kernel, thus use slow manual operation now # TODO will fuse this into kernel, thus use slow manual operation now
if num_token_non_padded is None:
return topk_weights, topk_ids
torch.compile( torch.compile(
_mask_topk_ids_padded_region, dynamic=True, backend=get_compiler_backend() _mask_topk_ids_padded_region, dynamic=True, backend=get_compiler_backend()
)(topk_ids, num_token_non_padded) )(topk_ids, num_token_non_padded)