[Refactor] Optimize select_experts (#28069)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user