[Refactor] Optimize select_experts (#28069)

Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye
2025-11-19 18:53:15 -05:00
committed by GitHub
parent 3aaa94ac99
commit 5031cd5d55
5 changed files with 7 additions and 15 deletions

View File

@@ -1246,7 +1246,6 @@ def eplb_map_to_physical_and_record(
expert_load_view: torch.Tensor,
logical_to_physical_map: torch.Tensor,
logical_replica_count: torch.Tensor,
indices_type: torch.dtype | None = None,
) -> torch.Tensor:
"""
Map the logical expert ids to physical expert ids
@@ -1260,7 +1259,6 @@ def eplb_map_to_physical_and_record(
expert_load_view: The expert load view.
logical_to_physical_map: The logical to physical map.
logical_replica_count: The logical replica count.
indices_type: The indices type.
Returns:
The physical expert ids.
@@ -1310,9 +1308,6 @@ def eplb_map_to_physical_and_record(
index=topk_ids_flatten.long(),
src=torch.ones_like(topk_ids_flatten).to(expert_load_view),
)
if indices_type is not None:
topk_ids = topk_ids.to(dtype=indices_type)
return topk_ids