class_docstr: Basic MoE layer.
class_docstr_tok: ["Basic","MoE","layer","."]
docstr: Dispatch tokens to local experts. It's composed of two stages:
  (1) Permute the tokens across the expert parallel devices. After this stage,
  each device receives all of the tokens assigned to its local set of experts
  in its local HBM.
  (2) Permute the tokens locally so that they are grouped by their expert
  assignment. After the stage (1), the tokens are grouped by which device
  they came from. We re-order them locally for subsequent efficient computation.
  
  Args:
      hidden_states: input tokens of shape [SeqLen/TP, MBS, HiddenSize]
  
  Returns:
      permuted_local_hidden_states: Permutation of tokens to local experts group.
      tokens_per_expert: the number of tokens each local expert to process.
      indices: The indices of `local_indices` (which holds the un-sorted expert
      indices of tokens that local expert can process) that give its sorted order along dim 0.
      global_local_map (optional): 2D tensor. A mask of mapping between global and local tokens where each
      element is True if it's between the local_expert_indices. Only useful
      when cross device token permutation is enabled and **AllGahter** is performed.
docstr_tok: ["Dispatch","tokens","to","local","experts",".","It","'s","composed","of","two","stages",":","(","1",")","Permute","the","tokens","across","the","expert","parallel","devices",".","After","this","stage",",","each","device","receives","all","of","the","tokens","assigned","to","its","local","set","of","experts","in","its","local","HBM",".","(","2",")","Permute","the","tokens","locally","so","that","they","are","grouped","by","their","expert","assignment",".","After","the","stage","(","1",")",",","the","tokens","are","grouped","by","which","device","they","came","from",".","We","re-order","them","locally","for","subsequent","efficient","computation",".","Args",":","hidden_states",":","input","tokens","of","shape","[","SeqLen\/TP",",","MBS",",","HiddenSize","]","Returns",":","permuted_local_hidden_states",":","Permutation","of","tokens","to","local","experts","group",".","tokens_per_expert",":","the","number","of","tokens","each","local","expert","to","process",".","indices",":","The","indices","of","`","local_indices","`","(","which","holds","the","un-sorted","expert","indices","of","tokens","that","local","expert","can","process",")","that","give","its","sorted","order","along","dim","0.","global_local_map","(","optional",")",":","2D","tensor",".","A","mask","of","mapping","between","global","and","local","tokens","where","each","element","is","True","if","it","'s","between","the","local_expert_indices",".","Only","useful","when","cross","device","token","permutation","is","enabled","and","*","*","AllGahter","*","*","is","performed","."]
code: def token_permutation(self, hidden_states):
      """Dispatch tokens to local experts. It's composed of two stages:
      (1) Permute the tokens across the expert parallel devices. After this stage,
      each device receives all of the tokens assigned to its local set of experts
      in its local HBM.
      (2) Permute the tokens locally so that they are grouped by their expert
      assignment. After the stage (1), the tokens are grouped by which device
      they came from. We re-order them locally for subsequent efficient computation.
  
      Args:
          hidden_states: input tokens of shape [SeqLen/TP, MBS, HiddenSize]
  
      Returns:
          permuted_local_hidden_states: Permutation of tokens to local experts group.
          tokens_per_expert: the number of tokens each local expert to process.
          indices: The indices of `local_indices` (which holds the un-sorted expert
          indices of tokens that local expert can process) that give its sorted order along dim 0.
          global_local_map (optional): 2D tensor. A mask of mapping between global and local tokens where each
          element is True if it's between the local_expert_indices. Only useful
          when cross device token permutation is enabled and **AllGahter** is performed.
      """
      self.hidden_shape = hidden_states.shape
      route = self.router(hidden_states)
      route = route.view(-1, self.config.num_moe_experts)
  
      if self.training:
          with torch.no_grad():
              norm_route = self.route_algo(
                  route.detach().to(dtype=torch.float32)
              )  # explicit fp32 conversion for stability
              _, max_ind = torch.topk(norm_route, k=self.k, dim=1)
          route = self.router_activation(route)
          # max_ind = max_ind.view(-1)
          max_prob = torch.gather(route, 1, max_ind)
      else:
          route = self.router_activation(route)
          max_prob, max_ind = torch.topk(route, k=self.k, dim=1)
      # [S/TP, B, H] -> [S*B/TP, H]
      hidden_states = hidden_states.view(-1, self.hidden_shape[-1])
  
      # Stage1: permute the tokens across the expert parallel devices.
      if self.sequence_parallel or (self.expert_parallel_size > 1):
          # [S*B/TP, H] -> [S*B, H]
          global_hidden_states = tensor_parallel.gather_from_sequence_parallel_region_to_moe(
              hidden_states
          )
          with torch.no_grad():
              global_indices = self.gather_indices(max_ind)
              # Create a mask of mapping between global and local tokens where each
              # element is True if it's between the local_expert_indices
              global_local_map = (
                  global_indices >= self.local_expert_indices[0]
              ) & (global_indices <= self.local_expert_indices[-1])
              local_indices = global_indices.masked_select(
                  global_local_map
              )
              if self.k > 1:  # k > 1
                  global_probs = self.gather_indices(max_prob)
                  local_probs = global_probs.masked_select(
                      global_local_map
                  )
              else:
                  local_probs = max_prob
              # Reshape global_local_map to be compatible with Tensor.gather
              global_local_map = global_local_map.nonzero()[:, 0]
              global_local_map = global_local_map.view(-1, 1).expand(
                  -1, hidden_states.shape[-1]
              )
          local_hidden_states = torch.gather(
              global_hidden_states, 0, global_local_map
          )
      else:
          local_indices = max_ind
          local_probs = max_prob
          local_hidden_states = hidden_states
          global_local_map = None
      self.max_prob = local_probs
  
      with torch.no_grad():
          # The indices of local_indices that give its sorted order along dim 0.
          indices = torch.argsort(local_indices, dim=0)
          tokens_per_expert = torch.histc(
              local_indices,
              bins=self.num_local_experts,
              min=self.local_expert_indices[0],
              max=self.local_expert_indices[-1],
          )
          tokens_per_expert = tokens_per_expert.cpu().to(torch.long)
  
      # Stage2: permute the tokens locally so that they are grouped by their expert assignment
      # Reshape indices to be compatible with Tensor.gather
      indices = indices.view(-1, 1).expand(-1, hidden_states.shape[-1])
      permuted_local_hidden_states = torch.gather(
          local_hidden_states, 0, indices
      )
  
      return (
          permuted_local_hidden_states,
          tokens_per_expert,
          indices,
          global_local_map,
      )
  
file_path: megatron-lm/megatron/core/transformer/moe/base_moe_layer.py
imports: from abc import ABC, abstractmethod
  import torch
  from megatron.core import parallel_state, tensor_parallel
  from megatron.core.parallel_state import get_tensor_and_expert_parallel_group
  from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_data_parallel_rng_tracker_name
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.transformer_config import TransformerConfig
label_desc: Use image node_id 3 for calling the BaseMoELayer obj's underlying member method code with example usage: obj.token_permutation(hidden_states) and returns: permuted_local_hidden_states,  tokens_per_expert,  indices,  global_local_map