Skip to content

⚡️ Speed up function _topk_ids_logical_to_physical_static by 51% #321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

⚡️ Speed up function _topk_ids_logical_to_physical_static by 51% #321

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion python/sglang/srt/eplb/expert_location_dispatch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,14 @@ def topk_ids_logical_to_physical(
def _topk_ids_logical_to_physical_static(
topk_ids: torch.Tensor, info: Optional[ExpertLocationDispatchInfo]
) -> torch.Tensor:
return info.partial_logical_to_rank_dispatch_physical_map[topk_ids]
# Using torch.take is more efficient for advanced indexing on 1D tensors
# and can be up to 2x faster than standard indexing.
# Only applies if partial_logical_to_rank_dispatch_physical_map is 1D.
# Fallback to original if not 1D.
partial_map = info.partial_logical_to_rank_dispatch_physical_map
if partial_map.ndim == 1 and topk_ids.dtype == torch.long:
return torch.take(partial_map, topk_ids)
return partial_map[topk_ids]


def _topk_ids_logical_to_physical_dynamic(
Expand Down