low_latency_dispatch

Applicable Products

• Ascend 950PR/Ascend 950DT: Not supported
• Atlas A3 training products/Atlas A3 inference products: Supported
• Atlas A2 training products/Atlas A2 inference products: Supported
• Atlas 200I/500 A2 inference products: Not supported
• Atlas inference products: Not supported
• Atlas training products: Not supported

Function

• Description: use it together with low_latency_combine to dispatch and combine tokens in parallel MoE deployment.
  • Supports the dynamic quantization use case. Quantize the token data first (optional), then perform alltoallv communication in the EP (Expert Parallelism) domain.
  • Supports the special expert use case.
• Formula:

  • Dynamic quantization scenario:

    If quant_mode is not 2, that is, the non-dynamic quantization scenario:

    $$\ quant\_out= \begin{cases} \ x, & \quad \text{if}\ quant\_mode = 0 \\ \ CastToInt8(\ CastToFp32(x) \times \ scales ), & \quad \text{if } quant\_mode ≠ 0 \\ \end{cases}$$

    $$\ alltoall\_x\_out= \ alltoallv(\ quant\_out)$$

    $$\ expand\_x= \ alltoall\_x\_out$$

    If quant_mode is 2, that is, the dynamic quantization scenario:

    $$\ x\_fp32= \ CastToFp32(x) \times \ scales$$

    $$\ dynamic\_scales\_value = 127.0/Max(Abs(x\_fp32))$$

    $$\ quant\_out=CastToInt8(\ x\_fp32 \times \ dynamic\_scales\_value )$$

    $$\ alltoall\_x\_out= \ alltoallv(\ quant\_out)$$

    $$\ alltoall\_dynamic\_scales\_out = alltoall(1.0/dynamic\_scales)$$

    $$\ expand\_x=\ alltoall\_x\_out$$

    $$\ dynamic\_scales=\ alltoall\_dynamic\_scales\_out$$

  • Special expert scenario:

    Zero-expert scenario, that is, zero_Expert_Num is not 0:

    $$Moe(ori\_x)=0$$

    Copy-expert scenario, that is, copy_Expert_Num is not 0:

    $$Moe(ori\_x)=ori\_x$$

    For the constant expert case, that is, when const_expert_num is not 0:

    $$Moe(ori\_x)=const\_expert\_alpha\_1*ori\_x+const\_expert\_alpha\_2*const\_expert\_v$$

    For parameters ori_x, const_expert_alpha_1, const_expert_alpha_2, and const_expert_v, see low_latency_combine.

Prototype

MoeDistributeBuffer.low_latency_dispatch(x, topk_idx, num_experts, *, quant_mode=0, comm_alg="", x_smooth_scale=None, x_active_mask=None, topk_weights=None, zero_expert_num=0, copy_expert_num=0, const_expert_num=0, elastic_info=None, expert_shard_type=0, shared_expert_num=1, shared_expert_rank_num=0, expert_token_nums_type=1, num_max_dispatch_tokens_per_rank=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)

Parameters

Return Value

Constraints

• This API supports use in inference tasks.

• This API supports single-operator mode and torchair graph mode. low_latency_dispatch and low_latency_combine must be used together.

• The constraints on the list[Tensor] length, whether to transpose, and whether non-contiguous tensors are supported for each input tensor are as follows:

  Parameter	list[Tensor] Length	Whether to Transpose	Whether Non-contiguous Tensor Is Supported
  x	-	No	Supported
  topk_idx	-	No	Supported
  x_smooth_scale	-	No	Not supported
  x_active_mask	-	No	Supported

• The variables used in the Shape parameter are as follows:

  • A: indicates the maximum number of tokens received by this device. The value range is as follows:

    • For shared experts, A must satisfy A=BS*shared_expert_num/shared_expert_rank_num.
    • For MoE experts, when num_max_dispatch_tokens_per_rank is 0, A must satisfy A >= BS * ep_world_size * min(local_expert_num, K); when num_max_dispatch_tokens_per_rank is not 0, A must satisfy A >= num_max_dispatch_tokens_per_rank * ep_world_size * min(local_expert_num, K).

  • H: indicates the hidden size. The value is [1024, 8192].

  • BS: indicates the batch sequence size, that is, the number of tokens finally output by the local device. The value range is 0 < BS ≤ 512.

  • K: indicates the number of top-K experts selected. The value range is 0 < K ≤ 16, and 0 < K ≤ num_experts + zero_expert_num + copy_expert_num + const_expert_num.

  • local_expert_num: indicates the number of experts on the local device.

    • For a shared expert device, local_expert_num is 1.
    • For an MoE expert device, local_expert_num=num_experts/(ep_world_size-shared_expert_rank_num), and 0 < local_expert_num * ep_world_size ≤ 2048 must be satisfied.

• Across different product models, communication algorithms, or versions, the element values in the Tensor outputs assist_info_for_combine, ep_recv_counts, tp_recv_counts, and expand_scales of low_latency_dispatch may differ. Pass these Tensors directly to the corresponding parameters of low_latency_combine. Other business logic of the model must not depend on them.

• The values of num_experts, expert_shard_type, shared_expert_num, shared_expert_rank_num, and num_max_dispatch_tokens_per_rank used during the API call must be consistent across all devices. expert_shard_type and num_max_dispatch_tokens_per_rank must also be consistent across different layers in the network, and consistent with the corresponding parameters of low_latency_combine.

• In this scenario, a single device contains two DIEs (or "bare dies"), so "this device" in the parameter description refers to a single DIE.

• num_experts + zero_expert_num + copy_expert_num + const_expert_num < MAX_INT32.

• HCCL communicator cache size: Before calling this API, check whether the value of the HCCL_BUFFSIZE environment variable is reasonable. This environment variable indicates the memory size occupied by a single communicator, in MB. If not configured, it defaults to 200 MB.

  • This use case supports configuration only through the environment variable HCCL_BUFFSIZE. The environment variable is managed at communicator granularity. Each communicator exclusively occupies a group of “2*HCCL_BUFFSIZE”-sized memory.
  • In the EP communicator, when comm_alg is set to "fullmesh_v1" or "": the size requirement is >= 2 * (local_expert_num * max_bs * ep_world_size * Align512(Align32(2 * H) + 64) + (K + shared_expert_num) * max_bs * Align512(2 * H)).
  • In the EP communicator, when comm_alg is set to "fullmesh_v2": the size requirement is >= 2 * (local_expert_num * max_bs * ep_world_size * 480Align512(Align32(2 * H) + 64) + (K + shared_expert_num) * max_bs * Align512(2 * H)).
  • Where 480Align512(x) = ((x+480-1)/480)*512,Align512(x) = ((x+512-1)/512)*512,Align32(x) = ((x+32-1)/32)*32.
  • The communicator size can be calculated by calling the MoeDistributeBuffer.get_low_latency_ccl_buffer_size API.

• The "/" in the formulas in this document indicates integer division.

• Communicator usage constraints:

  • The low_latency_dispatch and low_latency_combine operators in a model support only the same EP communicator. No other operators are allowed in this communicator.

  • The low_latency_dispatch and low_latency_combine operators in a model support only the same TP communicator, or neither supports a TP communicator. When a TP communicator exists, no other operators are allowed in it.

  • Nodes in a communicator must be within one SuperPoD. Cross-SuperPoD is not supported.

• Version compatibility constraints:

  In static graph mode, starting from TorchNPU 8.0.0, the TorchNPU framework performs strict validation of the Meta inference and inferShape inference results for the output of the last node in the static graph. When the graph contains only one Dispatch operator and the CANN version is older than the TorchNPU version, a shape mismatch error occurs. Upgrade the CANN version. For detailed version compatibility, see "Related Product Version Compatibility" in TorchNPU Release Notes.

Deterministic Computation

Defaults to a deterministic computation.

Example

• Single-operator mode call:

  import os
  import torch
  import random
  import torch_npu
  import numpy as np
  from torch.multiprocessing import Process
  import torch.distributed as dist
  from torch.distributed import ReduceOp
  from cann_ops_transformer.ops import MoeDistributeBuffer
  
  # Control mode
  quant_mode = 2  # 2 indicates dynamic quantization.
  is_dispatch_scales = True  # For dynamic quantization, optionally pass scales.
  input_dtype = torch.bfloat16  # Output dtype
  server_num = 1
  server_index = 0
  port = 50001
  master_ip = '127.0.0.1'
  dev_num = 16
  world_size = server_num * dev_num
  rank_per_dev = int(world_size / server_num)  # Number of dies per host
  shared_expert_rank_num = 0  # Number of shared experts
  num_experts = 32  # Number of MoE experts
  bs = 8  # Number of tokens
  h = 7168  # Length of each token.
  k = 8
  random_seed = 0
  tp_world_size = 1
  ep_world_size = int(world_size / tp_world_size)
  moe_rank_num = ep_world_size - shared_expert_rank_num
  local_moe_expert_num = num_experts // moe_rank_num
  globalBS = bs * ep_world_size
  is_shared = (shared_expert_rank_num > 0)
  is_quant = (quant_mode > 0)
  zero_expert_num = 1
  copy_expert_num = 1
  const_expert_num = 0
  
  
  def gen_const_expert_alpha_1():
      const_expert_alpha_1 = torch.empty(size=[const_expert_num, h], dtype=input_dtype).uniform_(-1, 1)
      return const_expert_alpha_1
  
  
  def gen_const_expert_alpha_2():
      const_expert_alpha_2 = torch.empty(size=[const_expert_num, h], dtype=input_dtype).uniform_(-1, 1)
      return const_expert_alpha_2
  
  
  def gen_const_expert_v():
      const_expert_v = torch.empty(size=[const_expert_num, h], dtype=input_dtype).uniform_(-1, 1)
      return const_expert_v
  
  
  def get_new_group(rank):
      for i in range(tp_world_size):
          # If tp_world_size = 2 and ep_world_size = 8, the result is [[0, 2, 4, 6, 8, 10, 12, 14], [1, 3, 5, 7, 9, 11, 13, 15]].
          ep_ranks = [x * tp_world_size + i for x in range(ep_world_size)]
          ep_group = dist.new_group(backend="hccl", ranks=ep_ranks)
          if rank in ep_ranks:
              ep_group_t = ep_group
              print(f"rank:{rank} ep_ranks:{ep_ranks}")
      for i in range(ep_world_size):
          # If tp_world_size = 2 and ep_world_size = 8, the result is [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13], [14, 15]].
          tp_ranks = [x + tp_world_size * i for x in range(tp_world_size)]
          tp_group = dist.new_group(backend="hccl", ranks=tp_ranks)
          if rank in tp_ranks:
              tp_group_t = tp_group
              print(f"rank:{rank} tp_ranks:{tp_ranks}")
      return ep_group_t, tp_group_t
  
  
  def get_hcomm_info(rank, comm_group):
      if torch.__version__ > '2.0.1':
          hcomm_info = comm_group._get_backend(torch.device("npu")).get_hccl_comm_name(rank)
      else:
          hcomm_info = comm_group.get_hccl_comm_name(rank)
      return hcomm_info
  
  
  def get_dispatch_kwargs_warmup(
      x_warm_up, topk_idx_warm_up, group_ep, group_tp, ep_rank_id, tp_rank_id,
  ):
      x_warm_up = x_warm_up.to(input_dtype).npu()
      topk_idx_warm_up = topk_idx_warm_up.to(torch.int32).npu()
      return {
          'x': x_warm_up,
          'topk_idx': topk_idx_warm_up,
          'x_active_mask': None,
          'shared_expert_num': 0,
          'shared_expert_rank_num': shared_expert_rank_num,
          'num_experts': num_experts,
          'quant_mode': 2,
          'num_max_dispatch_tokens_per_rank': 16,
      }
  
  
  def run_npu_process(rank):
      torch_npu.npu.set_device(rank)
      rank = rank + 16 * server_index
      # Initialize the distributed process group.
      dist.init_process_group(backend='hccl', rank=rank, world_size=world_size, init_method=f'tcp://{master_ip}:{port}')
      ep_group, tp_group = get_new_group(rank)
      ep_hcomm_info = get_hcomm_info(rank, ep_group)
      tp_hcomm_info = get_hcomm_info(rank, tp_group)
  
      # Create the input tensor.
      x = torch.randn(bs, h, dtype=input_dtype).npu()
      topk_idx = torch.tensor([[5, 7, 17, 4, 2, 6, 11, 16],
                              [10, 12, 13, 15, 19, 4, 18, 1],
                              [19, 33, 1, 17, 9, 5, 0, 32],
                              [19, 11, 17, 0, 10, 5, 7, 9],
                              [10, 16, 11, 17, 33, 8, 9, 3],
                              [12, 19, 5, 7, 1, 3, 18, 16],
                              [11, 9, 13, 16, 12, 33, 17, 14],
                              [16, 4, 9, 5, 0, 10, 11, 17]], dtype=torch.int32).npu()
      topk_weights = torch.randn(bs, k, dtype=torch.float32).npu()
  
      scales_shape = (1 + num_experts, h) if shared_expert_rank_num else (num_experts, h)
      if is_dispatch_scales:
          scales = torch.randn(scales_shape, dtype=torch.float32).npu()
      else:
          scales = None
  
      const_expert_alpha_1 = None
      const_expert_alpha_2 = None
      const_expert_v = None
  
      # Construct the communicator buffer.
      distribute_buffer = MoeDistributeBuffer(ep_group)
  
      # Call dispatch to quantize tokens and distribute them through alltoallv in the EP domain.
      expand_x, dynamic_scales, assist_info_for_combine, expert_token_nums, ep_recv_counts, _ = distribute_buffer.low_latency_dispatch(
          x=x,
          topk_idx=topk_idx,
          shared_expert_num=0,
          shared_expert_rank_num=shared_expert_rank_num,
          num_experts=num_experts,
          quant_mode=quant_mode,
          num_max_dispatch_tokens_per_rank=globalBS,
          zero_expert_num=zero_expert_num,
          copy_expert_num=copy_expert_num,
          const_expert_num=const_expert_num)
  
      if is_quant:
          expand_x = expand_x.to(input_dtype)
  
      # Pass the dispatch output to combine to collect tokens.
      x = distribute_buffer.low_latency_combine(x=expand_x,
                                              topk_idx=topk_idx,
                                              assist_info_for_combine=assist_info_for_combine,
                                              ep_send_counts=ep_recv_counts,
                                              topk_weights=topk_weights,
                                              shared_expert_num=0,
                                              shared_expert_rank_num=shared_expert_rank_num,
                                              num_experts=num_experts,
                                              num_max_dispatch_tokens_per_rank=globalBS,
                                              ori_x=x,
                                              const_expert_alpha_1=const_expert_alpha_1,
                                              const_expert_alpha_2=const_expert_alpha_2,
                                              const_expert_v=const_expert_v,
                                              zero_expert_num=zero_expert_num,
                                              copy_expert_num=copy_expert_num,
                                              const_expert_num=const_expert_num)
      print(f'rank {rank} epid {rank // tp_world_size} tpid {rank % tp_world_size} npu finished! \n')
  
  
  if __name__ == "__main__":
      print(f"bs={bs}")
      print(f"num_max_dispatch_tokens_per_rank={globalBS}")
      print(f"shared_expert_rank_num={shared_expert_rank_num}")
      print(f"num_experts={num_experts}")
      print(f"k={k}")
      print(f"quant_mode={quant_mode}", flush=True)
      print(f"local_moe_expert_num={local_moe_expert_num}", flush=True)
      print(f"tp_world_size={tp_world_size}", flush=True)
      print(f"ep_world_size={ep_world_size}", flush=True)
  
      if tp_world_size != 1 and local_moe_expert_num > 1:
          print("unSupported tp = 2 and local moe > 1")
          exit(0)
      if shared_expert_rank_num > ep_world_size:
          print("shared_expert_rank_num cannot be greater than ep_world_size.")
          exit(0)
      if shared_expert_rank_num > 0 and ep_world_size % shared_expert_rank_num != 0:
          print("ep_world_size must be an integer multiple of shared_expert_rank_num.")
          exit(0)
      if num_experts % moe_rank_num != 0:
          print("num_experts must be an integer multiple of moe_rank_num")
          exit(0)
      p_list = []
      for rank in range(rank_per_dev):
          p = Process(target=run_npu_process, args=(rank,))
          p_list.append(p)
  
      for p in p_list:
          p.start()
      for p in p_list:
          p.join()
  
      print("run npu success.")

• Graph mode call:

  import os
  import torch
  import random
  import torch_npu
  import torchair
  import numpy as np
  from torch.multiprocessing import Process
  import torch.distributed as dist
  from torch.distributed import ReduceOp
  import time
  from cann_ops_transformer.ops import MoeDistributeBuffer
  
  # Control mode
  quant_mode = 2  # 2 indicates dynamic quantization
  is_dispatch_scales = True  # For dynamic quantization, choose whether to pass scales
  input_dtype = torch.bfloat16  # Output dtype
  server_num = 1
  server_index = 0
  port = 50001
  master_ip = '127.0.0.1'
  dev_num = 16
  world_size = server_num * dev_num
  rank_per_dev = int(world_size / server_num)  # Number of dies per host
  shared_expert_rank_num = 0  # Number of shared experts
  num_experts = 32  # Number of MoE experts
  bs = 8  # Number of tokens
  h = 7168  # Length of each token.
  k = 8
  random_seed = 0
  tp_world_size = 1
  ep_world_size = int(world_size / tp_world_size)
  moe_rank_num = ep_world_size - shared_expert_rank_num
  local_moe_expert_num = num_experts // moe_rank_num
  globalBS = bs * ep_world_size
  is_shared = (shared_expert_rank_num > 0)
  is_quant = (quant_mode > 0)
  
  zero_expert_num = 1
  copy_expert_num = 1
  const_expert_num = 0
  
  class MOE_DISTRIBUTE_GRAPH_Model(torch.nn.Module):
      def __init__(self, group_ep):
          super().__init__()
          self.group = group_ep
          self.distribute_buffer = MoeDistributeBuffer(group_ep)
  
      def forward(self, x, topk_idx, group_ep, group_tp, ep_world_size, tp_world_size,
                  ep_rank_id, tp_rank_id, expert_shard_type, shared_expert_rank_num, num_experts,
                  scales, quant_mode, num_max_dispatch_tokens_per_rank, topk_weights, elastic_info, const_expert_alpha_1, const_expert_alpha_2, const_expert_v, zero_expert_num, copy_expert_num, const_expert_num):
          output_dispatch_npu = self.distribute_buffer.low_latency_dispatch(
              x=x,
              topk_idx=topk_idx,
              shared_expert_num=0,
              shared_expert_rank_num=shared_expert_rank_num,
              num_experts=num_experts,
              num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank,
              elastic_info=elastic_info,
              zero_expert_num=zero_expert_num,
              copy_expert_num=copy_expert_num,
              const_expert_num=const_expert_num
          )
          expand_x_npu, _, assist_info_for_combine_npu, _, ep_recv_counts_npu, expand_scales = output_dispatch_npu
          if expand_x_npu.dtype == torch.int8:
              expand_x_npu = expand_x_npu.to(input_dtype)
          output_combine_npu = self.distribute_buffer.low_latency_combine(
              x=expand_x_npu,
              topk_idx=topk_idx,
              topk_weights=topk_weights,
              assist_info_for_combine=assist_info_for_combine_npu,
              ep_send_counts=ep_recv_counts_npu,
              shared_expert_num=0,
              shared_expert_rank_num=shared_expert_rank_num,
              num_experts=num_experts,
              num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank,
              elastic_info=elastic_info,
              ori_x=x,
              const_expert_alpha_1=const_expert_alpha_1,
              const_expert_alpha_2=const_expert_alpha_2,
              const_expert_v=const_expert_v,
              zero_expert_num=zero_expert_num,
              copy_expert_num=copy_expert_num,
              const_expert_num=const_expert_num
          )
          x = output_combine_npu
          x_combine_res = output_combine_npu
          return [x_combine_res, output_combine_npu]
  
  def gen_const_expert_alpha_1():
      const_expert_alpha_1 = torch.empty(size=[const_expert_num, h], dtype=input_dtype).uniform_(-1, 1)
      return const_expert_alpha_1
  
  def gen_const_expert_alpha_2():
      const_expert_alpha_2 = torch.empty(size=[const_expert_num, h], dtype=input_dtype).uniform_(-1, 1)
      return const_expert_alpha_2
  
  def gen_const_expert_v():
      const_expert_v = torch.empty(size=[const_expert_num, h], dtype=input_dtype).uniform_(-1, 1)
      return const_expert_v
  
  def get_new_group(rank):
      for i in range(tp_world_size):
          ep_ranks = [x * tp_world_size + i for x in range(ep_world_size)]
          ep_group = dist.new_group(backend="hccl", ranks=ep_ranks)
          if rank in ep_ranks:
              ep_group_t = ep_group
              print(f"rank:{rank} ep_ranks:{ep_ranks}")
      for i in range(ep_world_size):
          tp_ranks = [x + tp_world_size * i for x in range(tp_world_size)]
          tp_group = dist.new_group(backend="hccl", ranks=tp_ranks)
          if rank in tp_ranks:
              tp_group_t = tp_group
              print(f"rank:{rank} tp_ranks:{tp_ranks}")
      return ep_group_t, tp_group_t
  
  def get_hcomm_info(rank, comm_group):
      if torch.__version__ > '2.0.1':
          hcomm_info = comm_group._get_backend(torch.device("npu")).get_hccl_comm_name(rank)
      else:
          hcomm_info = comm_group.get_hccl_comm_name(rank)
      return hcomm_info
  
  def get_dispatch_kwargs_warmup(
      x_warm_up, topk_idx_warm_up, group_ep, group_tp, ep_rank_id, tp_rank_id,
  ):
      x_warm_up = x_warm_up.to(input_dtype).npu()
      topk_idx_warm_up = topk_idx_warm_up.to(torch.int32).npu()
      return {
          'x': x_warm_up,
          'topk_idx': topk_idx_warm_up,
          'x_active_mask': None,
          'expert_shard_type': 0,
          'shared_expert_num': 0,
          'shared_expert_rank_num': shared_expert_rank_num,
          'num_experts': num_experts,
          'scales': None,
          'quant_mode': 2,
          'num_max_dispatch_tokens_per_rank': 1,
      }
  
  def run_npu_process(rank):
      torch_npu.npu.set_device(rank)
      rank = rank + 16 * server_index
      dist.init_process_group(
          backend='hccl',
          rank=rank,
          world_size=world_size,
          init_method=f'tcp://{master_ip}:{port}'
      )
      ep_group, tp_group = get_new_group(rank)
      ep_hcomm_info = get_hcomm_info(rank, ep_group)
      tp_hcomm_info = get_hcomm_info(rank, tp_group)
      # Create the input tensor.
      x = torch.randn(bs, h, dtype=input_dtype).npu()
      topk_idx = torch.tensor([
          [0, 8, 4, 1, 6, 12, 14, 17],
          [14, 10, 7, 3, 0, 12, 11, 17],
          [12, 0, 5, 11, 19, 4, 6, 18],
          [17, 3, 4, 10, 18, 0, 1, 2],
          [13, 16, 9, 10, 15, 6, 7, 14],
          [17, 15, 14, 8, 16, 18, 3, 12],
          [4, 12, 2, 17, 15, 3, 9, 10],
          [16, 7, 12, 9, 18, 3, 19, 17]
      ], dtype=torch.int32).npu()
      topk_weights = torch.randn(bs, k, dtype=torch.float32).npu()
      scales_shape = (1 + num_experts, h) if shared_expert_rank_num else (num_experts, h)
      if is_dispatch_scales:
          scales = torch.randn(scales_shape, dtype=torch.float32).npu()
      else:
          scales = None
      elastic_info = None
      const_expert_alpha_1 = None
      const_expert_alpha_2 = None
      const_expert_v = None
      model = MOE_DISTRIBUTE_GRAPH_Model(ep_group)
      model = model.npu()
      npu_backend = torchair.get_npu_backend()
      model = torch.compile(model, backend=npu_backend, dynamic=False)
      output = model.forward(
          x, topk_idx, ep_hcomm_info, tp_hcomm_info, ep_world_size, tp_world_size,
          rank // tp_world_size, rank % tp_world_size, 0, shared_expert_rank_num, num_experts, scales,
          quant_mode, globalBS, topk_weights, elastic_info, const_expert_alpha_1, const_expert_alpha_2, const_expert_v,
          zero_expert_num, copy_expert_num, const_expert_num
      )
      torch.npu.synchronize()
      print(f'rank {rank} epid {rank // tp_world_size} tpid {rank % tp_world_size} npu finished! \n')
  
      time.sleep(10)
  
  
  if __name__ == "__main__":
      print(f"bs={bs}")
      print(f"num_max_dispatch_tokens_per_rank={globalBS}")
      print(f"shared_expert_rank_num={shared_expert_rank_num}")
      print(f"num_experts={num_experts}")
      print(f"k={k}")
      print(f"quant_mode={quant_mode}", flush=True)
      print(f"local_moe_expert_num={local_moe_expert_num}", flush=True)
      print(f"tp_world_size={tp_world_size}", flush=True)
      print(f"ep_world_size={ep_world_size}", flush=True)
  
      if tp_world_size != 1 and local_moe_expert_num > 1:
          print("unSupported tp = 2 and local moe > 1")
          exit(0)
  
      if shared_expert_rank_num > ep_world_size:
          print("shared_expert_rank_num cannot be greater than ep_world_size.")
          exit(0)
  
      if shared_expert_rank_num > 0 and ep_world_size % shared_expert_rank_num != 0:
          print("ep_world_size must be an integer multiple of shared_expert_rank_num.")
          exit(0)
  
      if num_experts % moe_rank_num != 0:
          print("num_experts must be an integer multiple of moe_rank_num.")
          exit(0)
  
      p_list = []
      for rank in range(rank_per_dev):
          p = Process(target=run_npu_process, args=(rank,))
          p_list.append(p)
      for p in p_list:
          p.start()
      for p in p_list:
          p.join()
  
      print("run npu success.")