low_latency_combine

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:

Must be used together with low_latency_dispatch, equivalent to returning along the same path by which the low_latency_dispatch operator collects data.
- Supports the data integration function, performs alltoallv communication, and finally integrates the received data (multiplies by weights and then adds).
- Supports special expert use cases.
Formula:
- Data integration function:
  
  $ata\_out = AlltoAllv(expand\_x)$
  
  $x = Sum(topk\_weights * ata\_out + topk\_weights * shared\_expert\_x)$
- Special expert use cases:
  
  Zero-expert use case, where zero_expert_num is not 0:
  
  $Moe(ori\_x)=0$
  
  Copy-expert use case, where copy_expert_num is not 0:
  
  $Moe(ori\_x)=ori\_x$
  
  Constant-expert use case, where const_expert_num is not 0:
  
  $Moe(ori\_x)=const\_expert\_alpha\_1*ori\_x+const\_expert\_alpha\_2*const\_expert\_v$

Prototype

MoeDistributeBuffer.low_latency_combine(x, topk_idx, topk_weights, assist_info_for_combine, ep_send_counts, *, num_experts=0, comm_alg="", comm_quant_mode=0, x_active_mask=None, expand_scales=None, shared_expert_x=None, elastic_info=None, ori_x=None, const_expert_alpha_1=None, const_expert_alpha_2=None, const_expert_v=None, zero_expert_num=0, copy_expert_num=0, const_expert_num=0, expert_shared_type=0, shared_expert_num=1, shared_expert_rank_num=0, num_max_dispatch_tokens_per_rank=0) -> Tensor

Parameters

Parameter	Parameter Type	Optional/Mandatory	Description	Data Type	Shape
x	Tensor	Mandatory	Token features expanded based on `topk_idx`. The data format is ND.	bfloat16, float16	(A, H)
topk_idx	Tensor	Mandatory	TopK expert indices of each token, in ND data format. Corresponds to the `topk_idx` input of [low_latency_dispatch](./low_latency_dispatch.md). The value range of each element in the tensor is [0, num_experts), and the K values in the same row cannot be duplicated.	int32	(BS, K)
topk_weights	Tensor	Mandatory	Indicates the weights of the topK experts of each token. Shared experts do not need to be multiplied by a weight coefficient and can be added directly. The data format is ND.	float	(BS, K)
assist_info_for_combine	Tensor	Mandatory	Indicates the number of tokens sent to the same expert. The data format is ND. Corresponds to the `assist_info_for_combine` output of [low_latency_dispatch](./low_latency_dispatch.md).	int32	(A*128, )
ep_send_counts	Tensor	Mandatory	Indicates the number of tokens sent by each expert on this device to each device in the EP (Expert Parallelism) domain. The token count is expressed as a prefix sum. The data format is ND. Corresponds to the `ep_recv_counts` output of [low_latency_dispatch](./low_latency_dispatch.md).	int32	(ep_world_size*local_expert_num, )
num_experts	int	Optional	Number of MoE experts. Value range: [1, 1024]. Must satisfy the following condition: num_experts%(ep_world_size - shared_expert_rank_num)=0. Default value is 0.	int	-
comm_alg	str	Optional	Indicates the communication affinity memory layout algorithm. Currently, only "" is supported. The default value is "".	str	-
comm_quant_mode	int	Optional	Indicates the communication quantization type. Supports 0 and 2. 0 indicates no quantization during communication, and 2 indicates int8 quantization during communication. The default value is 0.	int	-
x_active_mask	Optional[Tensor]	Optional	Indicates whether a token participates in communication. Required to be a 1D or 2D tensor with the ND data format. For 1D input, true indicates that the corresponding token participates in communication, and all true values must precede false values. For example, {true, false, true} is invalid input. For 2D input, true indicates that the `topk_idx` corresponding to the current token participates in communication. If all K bool values of the current token are false, the current token does not participate in communication. By default, all tokens participate in communication. When the BS quantity of each device is inconsistent, all tokens must be valid. The default value is None.	bool	(BS, ) (BS, K)
expand_scales	Optional[Tensor]	Optional	Corresponds to the `expand_scales` output of [low_latency_dispatch](./low_latency_dispatch.md). This parameter is not supported. Use the default value. The default value is None.	float	-
shared_expert_x	Optional[Tensor]	Optional	The data type must keep consistent with `expand_x`. Used only in the scenario of the shared expert device quantity `shared_expert_rank_num` being 0. Indicates the shared expert token, which is the value to be added after combine_v2. When the tensor is 2D, the shape is (BS, H). When the tensor is 3D, the product of the first two dimensions must equal BS, and the third dimension must equal H. The default value is None.	Consistent with x	(BS, H) (, , H)
elastic_info	Optional[Tensor]	Optional	Reserved parameter. Not supported in the current version. Pass the default value None. The default value is None.	int32	-
ori_x	Optional[Tensor]	Optional	Indicates the token data that has not passed through FFN. This input data is required in the scenario where `copy_expert_num` is not 0 or `const_expert_num` is not 0. The data type must keep consistent with `expand_x`. The data format must be ND. You can optionally pass valid data or None. When `copy_expert_num` is not 0 or `const_expert_num` is not 0, you must pass valid data. The default value is None.	Same as x	(BS, H)
const_expert_alpha_1	Optional[Tensor]	Optional	Computation coefficient that requires input in the scenario of `const_expert_num` is not 0. The data type need to keep consistent with `expand_x`, and the data format is required to be ND. You can optionally pass valid data or fill None. When `const_expert_num` is not 0, you must pass valid input. The default value is None.	Consistent with x	(const_expert_num, H)
const_expert_alpha_2	Optional[Tensor]	Optional	Computation coefficient that requires input in the scenario of `const_expert_num` is not 0. The data type need to keep consistent with `expand_x`, and the data format is required to be ND. You can optionally pass valid data or fill None. When `const_expert_num` is not 0, you must pass valid input. The default value is None.	Consistent with x	(const_expert_num, H)
const_expert_v	Optional[Tensor]	Optional	Computation coefficient that requires input when `const_expert_num` is not 0. The data type needs to keep consistent with `expand_x`, and the data format is required to be ND. You can optionally pass valid data or fill None. When `const_expert_num` is not 0, you must pass valid input. The default value is None.	Consistent with x	(const_expert_num, H)
zero_expert_num	int	Optional	Indicates the number of zero experts. Value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. The valid ID values of zero experts are [num_experts, num_experts+zero_expert_num). The default value is 0.	int	-
copy_expert_num	int	Optional	Indicates the number of copy experts. Value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. The valid ID values of copy experts are [num_experts+zero_expert_num, num_experts+zero_expert_num+copy_expert_num). The default value is 0.	int	-
const_expert_num	int	Optional	Indicates the number of constant experts. Value range: [0, MAX_INT32), where MAX_INT32 = 2^31 - 1. The valid constant expert ID value is [num_experts+zero_expert_num+copy_expert_num, num_experts+zero_expert_num+copy_expert_num+const_expert_num). The default value is 0.	int	-
expert_shared_type	int	Optional	Indicates the layout type of shared expert devices. Currently only 0 is supported, which indicates that shared expert devices are placed before MoE expert devices. The default value is 0.	int	-
shared_expert_num	int	Optional	Indicates the number of shared experts. A shared expert can be replicated and deployed on multiple devices. Value range: [0, 4]. 0 indicates no shared expert. The default value is 1.	int	-
shared_expert_rank_num	int	Optional	Indicates the number of shared expert devices. Value range: [0, ep_world_size). A value of 0 indicates no shared experts. A non-zero value must satisfy shared_expert_rank_num%shared_expert_num=0. The default value is 0.	int	-
num_max_dispatch_tokens_per_rank	int	Optional	Indicates the number of tokens on each device. When the BS differs across ranks, this is the maximum BS size. When the BS is the same across ranks, the default value is 0.	int	-

Return Value

Parameter	Parameter Type	Optional/Mandatory	Description	Data Type	Dimension (shape)
x	Tensor	mandatory	Indicates the processed token. The data type is consistent with the input `expand_x`. The data format is ND. Non-contiguous tensors are not supported.	bfloat16, float16	(BS, H)

Constraints

This API supports use in inference tasks.
This API supports single-operator mode and torchair graph mode calls.

The list[Tensor] length, whether to transpose, and whether non-contiguous tensors are supported for each tensor parameter are constrained as follows:

Parameter	list[Tensor] Length	Whether to Transpose	Whether Non-contiguous Tensor Is Supported
x	-	No	Supported
topk_idx	-	No	Supported
topk_weights	-	No	Supported
assist_info_for_combine	-	No	Supported
ep_send_counts	-	No	Supported
x_active_mask	-	No	Supported
expand_scales	-	No	-
shared_expert_x	-	No	-
elastic_info	-	No	-
ori_x	-	No	Supported
const_expert_alpha_1	-	No	Supported
const_expert_alpha_2	-	No	Supported
const_expert_v	-	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, it must satisfy A=BS*ep_world_size*shared_expert_num/shared_expert_rank_num.
  - For MoE experts, when num_max_dispatch_tokens_per_rank is 0, it must satisfy A >= BS*ep_world_size * min(local_expert_num, K); when num_max_dispatch_tokens_per_rank is not 0, it must satisfy A >= num_max_dispatch_tokens_per_rank * ep_world_size * min(local_expert_num, K).
- H: indicates the hidden size. The value range is [1024, 8192].
- BS: indicates the number of tokens to be sent. The value range is 0<BS≤512.
- K: indicates the topK experts to be selected. The value range is 0<K≤16, and 0 < K ≤ num_experts + zero_expert_num + copy_expert_num + const_expert_num must also be satisfied.
- local_expert_num: indicates the number of experts on this device.
  - For shared expert devices, local_expert_num=1.
  - For MoE expert devices, local_expert_num = num_experts/(ep_world_size - shared_expert_rank_num).
  - The following must be satisfied: 0 < local_expert_num * ep_world_size ≤ 2048.
low_latency_dispatch and low_latency_combine must be used together.
Across different product models, communication algorithms, or versions, the element values in the low_latency_dispatch output tensors assist_info_for_combine, ep_recv_counts, tp_recv_counts, and expand_scales may differ. When using them, directly pass the preceding tensors to the corresponding parameters of low_latency_combine. Other model business logic must not depend on them.
During API calls, the values of group_ep, ep_world_size, num_experts, group_tp, tp_world_size, expert_shard_type, shared_expert_num, shared_expert_rank_num, and num_max_dispatch_tokens_per_rank must keep consistent across all devices. group_ep, ep_world_size, group_tp, tp_world_size, expert_shard_type, and num_max_dispatch_tokens_per_rank must also keep consistent across different layers in the network, and keep consistent with the corresponding parameters of low_latency_dispatch.
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 buffer size:

Before calling this API, check whether the communicator buffer size is reasonable. The unit is MB. If not configured, the default value is 200 MB.
- In this use case, the size can only be configured through the environment variable HCCL_BUFFSIZE. This environment variable is managed at the communicator granularity. Each communicator exclusively occupies a group of memory of size “2*HCCL_BUFFSIZE”.
- Within the ep communicator: the configured size must be >= 2 and satisfy >= 2 * (local_expert_num * max_bs * ep_world_size * Align512(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 buffer size opened in the communicator can be calculated by calling the MoeDistributeBuffer.get_low_latency_ccl_buffer_size API.
The "/" in the formulas in this document indicates integer division.
Communication domain usage constraints:
- The low_latency_dispatch and low_latency_combine operators in a model support only the same EP communicator, and no other operators are allowed in this communicator.
- The low_latency_dispatch and low_latency_combine operators in a model do not support the TP communicator.
- The nodes in a communicator must be within one SuperPoD. Cross-SuperPoD is not supported.

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, you can 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 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 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()

    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 first to quantize tokens and distribute them through alltoallv in the EP domain. The assist_info_for_combine and ep_recv_counts required by combine come from the output of this step.
    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)

    # Call combine to return the dispatch output along the original path. Pass assist_info_for_combine and ep_recv_counts together, multiply by topk_weights, and then add.
    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, you can 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_num_experts = 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, expert_scales, 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=expert_scales,
            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()
    expert_scales = 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, expert_scales, 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_num_experts={local_num_experts}", 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_num_experts > 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.")