Feature: add expert map support for shared experts & EPLB

[TianDi101]

Motivation

TODO

Technical Details

To map expert index, caller should pass a tensor of shape [total_num_expert] to dispatch. No changes need to be made to combine.

        total_num_expert = self.config.num_experts_per_rank * self.config.world_size
        expert_map = torch.randperm(total_num_expert).to(torch.int32).to(self.device)
        op.dispatch(
                all_rank_input[self.config.rank],
                all_rank_weights[self.config.rank],
                all_rank_scales[self.config.rank],
                all_rank_indices[self.config.rank],
                index_to_exp_id=expert_map,
        )

Test Plan

TODO