### Download Example Code

Source: https://github.com/ascend/op-plugin/blob/master/examples/cpp_extension_structured/README.md

Commands to clone the repository and navigate to the structured extension example directory.

```bash
# 下载样例代码
git clone https://gitcode.com/Ascend/op-plugin
# 进入代码目录
cd examples/cpp_extension_structured
```

--------------------------------

### Build and Install Commands

Source: https://github.com/ascend/op-plugin/blob/master/examples/framwork_cpp_extension/README.md

Commands to compile the adapter layer and install the generated wheel package.

```bash
python3 setup.py build bdist_wheel
```

```bash
cd ${BASE_DIR}
pip3 install dist/*.whl
```

--------------------------------

### Build and Install Commands

Source: https://github.com/ascend/op-plugin/blob/master/examples/cpp_extension/README.md

Commands to compile the extension and install the resulting wheel package.

```bash
python setup.py bdist_wheel
```

```bash
cd dist
pip install *.whl
```

```bash
cd test
python test.py
```

--------------------------------

### NPU AlltoAllv Quant GMM Execution Example

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu/torch_npu-npu_alltoallv_quant_gmm.md

Demonstrates the setup and execution of the npu_alltoallv_quant_gmm operation within a distributed process group.

```python
def run_npu_alltoallv_quant_gmm(rank, world_size, master_ip, master_port):
    torch_npu.npu.set_device(rank)
    init_method = f"tcp://{master_ip}:{master_port}"
    dist.init_process_group(backend="hccl", rank=rank, world_size=world_size, init_method=init_method)
    from torch.distributed.distributed_c10d import _get_default_group
    default_pg = _get_default_group()
    if torch.__version__ > "2.0.1":
        hcom_info = default_pg._get_backend(torch.device("npu")).get_hccl_comm_name(rank)
    else:
        hcom_info = default_pg.get_hccl_comm_name(rank)

    BS = 128
    K = 2
    e = 2
    H1, N1 = 256, 256
    H2, N2 = 256, 128
    total_tokens = BS * K
    send_counts, recv_counts = generate_counts(world_size, e, total_tokens, seed=rank)
    gmm_x = torch.ones(total_tokens, H1, dtype=torch.int8).to(torch.float8_e4m3fn).npu()
    gmm_weight = torch.ones(e, H1, N1, dtype=torch.int8).to(torch.float8_e5m2).npu()
    gmm_x_scale = torch.ones(total_tokens, math.ceil(H1 / 64), 2, dtype=torch.int8).npu()
    gmm_weight_scale = torch.ones(e, math.ceil(H1 / 64), N1, 2, dtype=torch.int8).npu()
    mm_x = torch.ones(BS, H2, dtype=torch.int8).to(torch.float8_e4m3fn).npu()
    mm_weight = torch.ones(H2, N2, dtype=torch.int8).to(torch.float8_e5m2).npu()
    mm_x_scale = torch.ones(BS, math.ceil(H2 / 64), 2, dtype=torch.int8).npu()
    mm_weight_scale = torch.ones(math.ceil(H2 / 64), N2, 2, dtype=torch.int8).npu()
    quant_mode = 6
    out_dtype = torch.float16

    gmm_y, mm_y, permute_out = torch_npu.npu_alltoallv_quant_gmm(
        gmm_x=gmm_x,
        gmm_weight=gmm_weight,
        gmm_x_scale=gmm_x_scale,
        gmm_weight_scale=gmm_weight_scale,
        hcom=hcom_info,
        ep_world_size=world_size,
        send_counts=send_counts,
        recv_counts=recv_counts,
        gmm_y_dtype=out_dtype,
        mm_x=mm_x,
        mm_weight=mm_weight,
        mm_x_scale=mm_x_scale,
        mm_weight_scale=mm_weight_scale,
        gmm_x_quant_mode=quant_mode,
        gmm_weight_quant_mode=quant_mode,
        mm_x_quant_mode=quant_mode,
        mm_weight_quant_mode=quant_mode,
        permute_out_flag=True,
        gmm_x_dtype=None,
        gmm_weight_dtype=None,
        gmm_x_scale_dtype=torch_npu.float8_e8m0fnu,
        gmm_weight_scale_dtype=torch_npu.float8_e8m0fnu,
        mm_x_dtype=None,
        mm_weight_dtype=None,
        mm_x_scale_dtype=torch_npu.float8_e8m0fnu,
        mm_weight_scale_dtype=torch_npu.float8_e8m0fnu,
        mm_y_dtype=out_dtype,
        send_counts_tensor=None,
        recv_counts_tensor=None,
        group_size=None
    )

if __name__ == "__main__":
    world_size = 2
    master_ip = "127.0.0.1"
    master_port = "50001"
    mp.spawn(run_npu_alltoallv_quant_gmm, args=(world_size, master_ip, master_port), nprocs=world_size, join=True)
```

--------------------------------

### ChannelShuffle Usage Example

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-contrib/（beta）torch_npu-contrib-module-ChannelShuffle.md

Basic usage example demonstrating initialization and forward pass with NPU tensors.

```python
>>> import torch, torch_npu
>>> from torch_npu.contrib.module import ChannelShuffle
>>> x1 = torch.randn(2, 32, 7, 7).npu()
>>> x2 = torch.randn(2, 32, 7, 7).npu()
>>> m = ChannelShuffle(64, split_shuffle=True)
>>> out1, out2 = m(x1, x2)
>>> print(out1.shape)
torch.Size([2, 32, 7, 7])
>>> print(out2.shape)
torch.Size([2, 32, 7, 7])
```

--------------------------------

### Usage Example

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu/（beta）torch_npu-npu_softmax_cross_entropy_with_logits.md

Demonstrates the basic usage of the API with NPU-based tensors.

```python
>>> import torch, torch_npu
>>> batch_size = 4
>>> num_classes = 12
>>> features = torch.rand(1, batch_size * num_classes).npu()
>>> labels = torch.rand(1, batch_size * num_classes).npu()
>>> output = torch_npu.npu_softmax_cross_entropy_with_logits(features, labels)
>>> print(output)
tensor([97.9450], device='npu:0')
```

--------------------------------

### Usage Example

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-contrib/（beta）torch_npu-contrib-function-npu_single_level_responsible_flags.md

Example demonstrating how to iterate through multiple feature map levels to generate responsible flags on an NPU device.

```python
>>> import torch
>>> from torch_npu.contrib.function import npu_single_level_responsible_flags
>>> featmap_sizes = [[10, 10], [20, 20], [40, 40]]
>>> stride = [[32, 32], [16, 16], [8, 8]]
>>> gt_bboxes = torch.randint(0, 512, size=(128, 4))
>>> num_base_anchors = 3
>>> featmap_level = len(featmap_sizes)
>>> for i in range(featmap_level):
...     gt_bboxes = gt_bboxes.npu()
...     out = npu_single_level_responsible_flags(featmap_sizes[i],gt_bboxes,stride[i],num_base_anchors)
...     print(out.shape, out.max(), out.min())
torch.Size([300]) tensor(1, device='npu:0', dtype=torch.uint8) tensor(0, device='npu:0', dtype=torch.uint8)
torch.Size([1200]) tensor(1, device='npu:0', dtype=torch.uint8) tensor(0, device='npu:0', dtype=torch.uint8)
torch.Size([4800]) tensor(1, device='npu:0', dtype=torch.uint8) tensor(0, device='npu:0', dtype=torch.uint8)
```

--------------------------------

### Usage Example of torch_npu.empty_with_format

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu/（beta）torch_npu-empty_with_format.md

Basic usage example showing how to initialize a tensor on an NPU device. Note that the returned values are uninitialized and will vary.

```python
>>> torch_npu.empty_with_format((2, 3), dtype=torch.float32, device="npu")
tensor([[3.1415e-45, 0.0000e+00, 1.2345e+20],
        [9.8765e-10, 0.0000e+00, 1.1111e+30]], device='npu:0')
# 注：empty_with_format返回未初始化张量，示例中的数值仅用于说明，实际值不固定。
```

--------------------------------

### Usage Example for MSTX Profiling

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-npu/torch_npu-npu-mstx-mark.md

Example configuration for enabling MSTX profiling within a torch_npu profiler context.

```python
import torch
import torch_npu

experimental_config = torch_npu.profiler._ExperimentalConfig(
    profiler_level=torch_npu.profiler.ProfilerLevel.Level_none,
    mstx=True,
    export_type=[
        torch_npu.profiler.ExportType.Db
        ])
with torch_npu.profiler.profile(
    schedule=torch_npu.profiler.schedule(wait=1, warmup=1, active=2, repeat=1, skip_first=1),
    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result"),
    experimental_config=experimental_config) as prof:
       
    for step in range(steps):
        train_one_step()    # 用户代码，包含调用mstx接口
        prof.step()
```

--------------------------------

### Usage Example for npu_multi_head_attention

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu/（beta）torch_npu-npu_multi_head_attention.md

A complete example demonstrating how to initialize input tensors and invoke the npu_multi_head_attention function on an NPU device.

```python
import torch
import torch_npu
import numpy as np
 
batch = 8
attn_head_num = 16
attn_dim_per_head = 64
src_len = 64
tgt_len = 64
dropout_prob = 0.0
softmax_use_float = True
 
weight_col = attn_head_num * attn_dim_per_head
query = torch.from_numpy(np.random.uniform(-1, 1, (batch * tgt_len, weight_col)).astype("float16")).npu()
key = torch.from_numpy(np.random.uniform(-1, 1, (batch * src_len, weight_col)).astype("float16")).npu()
value = torch.from_numpy(np.random.uniform(-1, 1, (batch * tgt_len, weight_col)).astype("float16")).npu()
query_weight = torch.from_numpy(np.random.uniform(-1, 1, (weight_col, weight_col)).astype("float16")).npu()
key_weight = torch.from_numpy(np.random.uniform(-1, 1, (weight_col, weight_col)).astype("float16")).npu()
value_weight = torch.from_numpy(np.random.uniform(-1, 1, (weight_col, weight_col)).astype("float16")).npu()
out_proj_weight = torch.from_numpy(np.random.uniform(-1, 1, (weight_col, weight_col)).astype("float16")).npu()
attn_mask = torch.from_numpy(np.random.uniform(-1, 1, (batch, attn_head_num, tgt_len, src_len)).astype("float16")).npu()
query_bias = torch.from_numpy(np.random.uniform(-1, 1, (weight_col,)).astype("float16")).npu()
key_bias = torch.from_numpy(np.random.uniform(-1, 1, (weight_col,)).astype("float16")).npu()
value_bias = torch.from_numpy(np.random.uniform(-1, 1, (weight_col,)).astype("float16")).npu()
out_proj_bias = torch.from_numpy(np.random.uniform(-1, 1, (weight_col,)).astype("float16")).npu()
dropout_mask = torch.from_numpy(np.random.uniform(-1, 1, (weight_col,)).astype("float16")).npu()
            
npu_result, npu_dropout_mask, npu_query_res, npu_key_res, npu_value_res, npu_attn_scores, npu_attn_res, npu_context = torch_npu.npu_multi_head_attention (query, key, value, query_weight, key_weight, value_weight, attn_mask, out_proj_weight, query_bias, key_bias, value_bias, out_proj_bias,  dropout_mask, attn_head_num, attn_dim_per_head, src_len, tgt_len, dropout_prob, softmax_use_float)
 
print(npu_result)
 
 
 
tensor([[ 623.5000,   75.5000,  307.0000,  ...,   25.3125, -418.7500,
           35.9688],
        [-254.2500, -165.6250,  176.2500,  ...,   87.3750,   78.0000,
           65.2500],
        [ 233.2500,  207.3750,  324.7500,  ...,   38.6250, -264.2500,
          153.7500],
        ...,
        [-110.2500,  -92.5000,  -74.0625,  ...,  -68.0625,  195.6250,
         -157.6250],
        [ 300.0000, -184.6250,   -6.0039,  ...,  -15.7969, -299.0000,
          -93.1875],
        [  -2.5996,   36.8750,  100.0625,  ...,  112.7500,  202.0000,
         -166.3750]], device='npu:0', dtype=torch.float16)
```

--------------------------------

### Call Custom Operator in Python

Source: https://github.com/ascend/op-plugin/blob/master/examples/cpp_extension/README.md

Example of invoking the custom operator from Python after installation.

```python
import torch
x = torch.randint(low=1, high=100, size=length, device='cpu', dtype=torch.int)
y = torch.randint(low=1, high=100, size=length, device='cpu', dtype=torch.int)

x_npu = x.npu()
y_npu = y.npu()
output = op_extension.ops.custom_add(x_npu, y_npu)
```

--------------------------------

### Download and Navigate to Example Code

Source: https://github.com/ascend/op-plugin/blob/master/examples/cpp_extension_asc/README.md

Use these commands to clone the repository and enter the base extension directory.

```bash
# 下载样例代码
git clone https://gitcode.com/Ascend/op-plugin
# 进入代码目录
cd examples/cpp_extension_base
```

--------------------------------

### Usage Example for _KinetoProfile

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-profiler/torch_npu-profiler-_KinetoProfile.md

Demonstrates the lifecycle of the profiler including initialization, starting, stopping, and exporting trace data.

```python
import torch
import torch_npu

...

prof = torch_npu.profiler._KinetoProfile(activities=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, with_modules=False, experimental_config=None)
for epoch in range(epochs):
        train_model_step()
        if epoch == 0:
                prof.start()
        if epoch == 1:
                prof.stop()
prof.export_chrome_trace("result_dir/trace.json")
```

--------------------------------

### Build and Run Sample

Source: https://github.com/ascend/op-plugin/blob/master/examples/kernel_extension_aclgraph/torch_library/README.md

Commands to build the wheel package, install it, and execute the test script.

```bash
python setup.py bdist_wheel
pip install dist/*.whl --force-reinstall
cd test
python ./add_aclgraph_test.py
```

--------------------------------

### Usage Example for FlopsCounter

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-utils/（beta）torch_npu-utils-FlopsCounter.md

Demonstrates the lifecycle of FLOPs tracking including starting, pausing, resuming, and stopping the counter during matrix multiplication operations.

```python
import torch
import torch_npu
 
def matmul():
    x = torch.randn(3, 4).npu()
    y = torch.randn(4, 3).npu()
    torch.matmul(x,y)
 
FlopsCounter = torch_npu.utils.FlopsCounter()
 
# 1.开启统计后进行统计
FlopsCounter.start()
matmul() # 算子计算
print(f"FlopsCounter.start():{FlopsCounter.get_flops()}") # 打印统计结果，含重计算的Flops和不含重计算的Flops累计
 
# 2. 暂停Flops不含重计算统计后进行统计
FlopsCounter.pause()
matmul() # 这里视作重计算操作
print(f"FlopsCounter.pause():{FlopsCounter.get_flops()}") # 仅含重计算Flops累计
 
# 3. 恢复Flops不含重计算统计后进行统计
FlopsCounter.resume()
matmul()
print(f"FlopsCounter.resume():{FlopsCounter.get_flops()}") # 含重计算Flops和不含重计算Flops均累计
 
# 4.关闭Flops统计
FlopsCounter.stop()
matmul()
print(f"FlopsCounter.stop():{FlopsCounter.get_flops()}") # 含重计算Flops和不含重计算Flops清0且均不累计
```

--------------------------------

### Download and Navigate to Operator Example

Source: https://github.com/ascend/op-plugin/blob/master/examples/cpp_extension_full/torch_lib_impl/README.md

Clone the repository and enter the specific directory for the C++ extension implementation.

```bash
# 下载样例代码
git clone https://gitcode.com/Ascend/op-plugin
# 进入代码目录
cd examples/cpp_extension_full/torch_library_impl
```

--------------------------------

### MX Quantization Scenario Example (MXFP8)

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu/torch_npu-npu_alltoallv_quant_gmm.md

Provides the setup for an MXFP8 quantization scenario, including the necessary imports and helper function for generating counts.

```python
import torch
import torch_npu
import torch.distributed as dist
import torch.multiprocessing as mp
import numpy as np
import math

def generate_counts(ep_world_size, e, total_tokens, seed=None):
    np.random.seed(seed if seed is not None else 42)
    per_rank_total = total_tokens
    base = per_rank_total // (ep_world_size * e)
    remainder = per_rank_total % (ep_world_size * e)
    send_counts = [base] * (ep_world_size * e)
    for i in range(remainder):
        send_counts[-1 - i] += 1
    recv_counts = send_counts.copy()
    return send_counts, recv_counts
```

--------------------------------

### Download and Execute Operator Example

Source: https://github.com/ascend/op-plugin/blob/master/examples/cpp_extension_base/README.md

Commands to clone the repository and run the build and test script.

```bash
# 下载样例代码
git clone https://gitcode.com/Ascend/op-plugin
# 进入代码目录
cd examples/cpp_extension_base
```

```bash
bash build_and_run.sh
```

```bash
Ran xx tests in xx s
OK
```

--------------------------------

### Single Operator Mode Call Example

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu/torch_npu-npu_attention_to_ffn.md

Demonstrates the setup and execution of the npu_attention_to_ffn interface in a distributed environment, including process group initialization and tensor preparation.

```python
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
import time

# 控制模式
quant_mode = 2  # 2为动态量化
sync_flag = 1   # 1为异步
is_mask = True  # 是否剪枝
is_attn2ffn_scales = True  # 动态量化可选择是否传scales
input_dtype = torch.bfloat16  # 输出dtype
server_num = 1
server_index = 0
port = 50001
master_ip = '127.0.0.1'
dev_num = 4
world_size = server_num * dev_num
rank_per_dev = int(world_size / server_num)  # 每个host有几个die
micro_batch_num = 1
X = 1
L = 1
bs = 8  # token数量
h = 7168  # 每个token的长度
k = 4
hs = h + 128
random_seed = 0
shared_expert_num = 1  # 共享专家数
rank_num_per_shared_expert = 1
shared_ffn_rank_num = shared_expert_num * rank_num_per_shared_expert
moe_expert_per_rank = 2  # 各FFN卡上的MoE专家数
moe_ffn_rank_num = 2  # MoE FFN卡数
moe_expert_num = moe_ffn_rank_num * moe_expert_per_rank
ffn_worker_num = moe_ffn_rank_num + shared_ffn_rank_num
attention_worker_num = world_size - ffn_worker_num
expert_num_per_token = k + shared_expert_num
is_quant = (quant_mode > 0)
ffn_token_info_table_shape = [attention_worker_num, micro_batch_num, 2 + bs * expert_num_per_token]
ffn_token_data_shape = [attention_worker_num, micro_batch_num, bs, expert_num_per_token, hs if is_quant else h]
attn_token_info_table_shape = [micro_batch_num, bs, expert_num_per_token]


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 set_windows(rank, comm_group, hcomm_info):
    if rank >= ffn_worker_num:
        # 当前 rank 属于后半部分，需与前 ffn_rank_num 个 rank 通信
        target_ranks = list(range(ffn_worker_num))
    else:
        # 当前 rank 属于前半部分，需与后面 rank 通信
        target_ranks = list(range(ffn_worker_num, world_size))
    window_size = 120 * 1024 * 200
    comm_group._get_backend(torch.device('npu'))._window_register_and_exchange(window_size, target_ranks)


def run_npu_process(rank):
    torch_npu.npu.set_device(rank)
    rank = rank + dev_num * server_index
    dist.init_process_group(backend='hccl', rank=rank, world_size=world_size, init_method=f'tcp://{master_ip}:{port}')
    rank_list = list(range(world_size))
    comm_group = dist.new_group(backend="hccl", ranks=rank_list)
    hcomm_info = get_hcomm_info(rank, comm_group)
    set_windows(rank, comm_group, hcomm_info)

    # 创建输入tensor
    x = torch.randn(X, bs, h, dtype=input_dtype).npu()
    session_id = torch.tensor([rank - ffn_worker_num], dtype=torch.int32).npu()
    micro_batch_id = torch.tensor([0], dtype=torch.int32).npu()
    layer_id = torch.tensor([0], dtype=torch.int32).npu()
    expert_ids = torch.tensor([[[0, 1, 2, 3],
                            [1, 2, 3, 0],
                            [2, 3, 0, 1],
                            [3, 0, 1, 2],
                            [0, 2, 1, 3],
                            [1, 3, 2, 0],
                            [2, 0, 3, 1],
                            [3, 1, 0, 2]]], dtype=torch.int32).npu()
    expert_rank_table = torch.tensor([[[1, 0, 0],
                                    [1, 0, 1],
                                    [1, 1, 2],
                                    [1, 1, 3],
                                    [1, 2, 4]]], dtype=torch.int32).npu()

    scales_shape = (L, shared_expert_num + moe_expert_num, h)
    if is_attn2ffn_scales:
        scales = torch.randn(scales_shape, dtype=torch.float32).npu()
    else:
        scales = None
    
    if is_mask:
        active_mask = torch.tensor([[1, 1, 1, 1, 0, 0, 0, 0]], dtype=torch.bool).npu()
    else:
        active_mask = None
```

--------------------------------

### Prefetcher Usage Example

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-contrib/（beta）torch_npu-contrib-Prefetcher.md

Demonstrates initializing the Prefetcher for single-use and multi-epoch training scenarios with stream management.

```python
>>> import torch
>>> import torch_npu
>>> from torch_npu.contrib import Prefetcher
>>> # 创建DataLoader
>>> dataset = torch.utils.data.TensorDataset(torch.randn(100, 3, 224, 224), torch.randint(0, 10, (100,)))
>>> loader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)
>>> # 初始化Prefetcher（仅初始化一次，无需指定stream）
>>> prefetcher = Prefetcher(loader)
>>> # 迭代获取数据
>>> input, target = prefetcher.next()
>>> while input is not None:
...     # 对input和target进行训练操作
...     input, target = prefetcher.next()
>>> # 重复初始化Prefetcher时（如多epoch训练），需指定stream防止内存泄漏
>>> stream = torch.npu.Stream()
>>> for epoch in range(10):
...     prefetcher = Prefetcher(loader, stream=stream)
...     input, target = prefetcher.next()
...     while input is not None:
...         # 对input和target进行训练操作
...         input, target = prefetcher.next()
```

--------------------------------

### FFNToAttention Single-Operator Mode Invocation

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu/torch_npu-npu_ffn_to_attention.md

This example demonstrates the setup of the NPU environment, distributed process group initialization, and helper functions for preparing input arguments and communication windows for the FFNToAttention operator.

```python
import os
import math
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
import time

# 控制模式
input_dtype = torch.bfloat16  # 输出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)  # 每个host有几个die
micro_batch_num = 1
bs = 8  # token数量
h = 7168  # 每个token的长度
scale = 128
hs = h + scale 
k = 4
random_seed = 0
shared_expert_num = 1  # 共享专家数
rank_num_per_shared_expert = 1
shared_ffn_rank_num = shared_expert_num * rank_num_per_shared_expert
moe_expert_per_rank = 2  # 各ffn卡moe专家数
moe_ffn_rank_num = 4    #FFN卡数
moe_expert_num = moe_ffn_rank_num * moe_expert_per_rank
ffn_worker_num = moe_ffn_rank_num + shared_ffn_rank_num
attention_worker_num = world_size - ffn_worker_num
expert_num_per_token = k + shared_expert_num
token_info_table_shape = [micro_batch_num, bs, expert_num_per_token]
token_data_table_shape = [micro_batch_num, bs, expert_num_per_token, hs]
Y = int(math.ceil(micro_batch_num * bs * attention_worker_num * expert_num_per_token / ffn_worker_num))


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 ffn2attn_get_kwargs(
    x, session_ids, micro_batch_ids, token_ids, expert_offsets, actual_token_num,
    attn_rank_table, group, world_size, token_info_table_shape, token_data_shape
):
    x = x.to(input_dtype).npu()
    session_ids = session_ids.to(torch.int32).npu()
    micro_batch_ids = micro_batch_ids.to(torch.int32).npu()
    token_ids = token_ids.to(torch.int32).npu()
    expert_offsets = expert_offsets.to(torch.int32).npu()
    actual_token_num = actual_token_num.to(torch.int64).npu()
    attn_rank_table = attn_rank_table.to(torch.int32).npu()
    
    return {
        'x':x,
        'session_ids':session_ids,
        'micro_batch_ids':micro_batch_ids,
        'token_ids':token_ids,
        'expert_offsets':expert_offsets,
        'actual_token_num':actual_token_num,
        'attn_rank_table':attn_rank_table,
        'group':group,
        'world_size':world_size,
        'token_info_table_shape':token_info_table_shape,
        'token_data_shape':token_data_shape
    }

def set_windows(rank, comm_group, hcomm_info):
    if attention_worker_num <= rank < world_size:
        target_ranks = list(range(attention_worker_num))
    else:
        target_ranks = list(range(attention_worker_num, world_size))
    window_size = 1024 * 1024 * 200
    comm_group._get_backend(torch.device('npu'))._window_register_and_exchange(window_size,target_ranks)

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}')
    rank_list = list(range(world_size))
    comm_group = dist.new_group(backend="hccl", ranks=rank_list)
    hcomm_info = get_hcomm_info(rank, comm_group)
    set_windows(rank, comm_group, hcomm_info)
```

--------------------------------

### MoE Routing Initialization Examples

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu/torch_npu-npu_moe_init_routing_v2.md

Demonstrates how to execute the MoE routing initialization operator using different quantization configurations.

```python
def demo_no_quant():
    """无量化模式"""
    print("=" * 50)
    print("Demo: 无量化模式 (quant_mode=-1)")
    print("=" * 50)
    
    bs, h, k = 32, 200, 5
    expert_range = [0, 16]
    
    # 生成输入
    x, expert_idx, scale, offset, expert_tokens_num_type, row_idx_type, active_num, expert_capacity = \
        MoeInitRoutingV2CPU.generate_inputs(bs, h, k, np.float16, (bs,), False, True, 0)
    
    # 执行计算
    expanded_x, expanded_row_idx, expert_tokens_count, expanded_scale = \
        MoeInitRoutingV2CPU.cpu_op_exec(
            x, expert_idx, scale, offset,
            expert_range=expert_range,
            quant_mode=-1,
            row_idx_type=row_idx_type,
            expert_tokens_num_flag=True,
            expert_tokens_num_type=expert_tokens_num_type,
            drop_pad_mode=0,
            active_num=active_num,
            expert_capacity=expert_capacity
        )
    
    print(f"Input x shape: {x.shape}, dtype: {x.dtype}")
    print(f"Input expert_idx shape: {expert_idx.shape}")
    print(f"Output expanded_x shape: {expanded_x.shape}, dtype: {expanded_x.dtype}")
    print(f"Output expanded_row_idx shape: {expanded_row_idx.shape}")
    print(f"Output expert_tokens_count: {expert_tokens_count}")
    print(f"Output expanded_scale: {expanded_scale}")
    print()
```

```python
def demo_static_quant():
    """静态量化模式"""
    print("=" * 50)
    print("Demo: 静态量化模式 (quant_mode=0)")
    print("=" * 50)
    
    bs, h, k = 32, 200, 5
    expert_range = [0, 16]
    
    # 生成输入 (需要scale和offset)
    x, expert_idx, scale, offset, expert_tokens_num_type, row_idx_type, active_num, expert_capacity = \
        MoeInitRoutingV2CPU.generate_inputs(bs, h, k, np.float32, (1,), False, False, 0)
    
    # 执行计算
    expanded_x, expanded_row_idx, expert_tokens_count, expanded_scale = \
        MoeInitRoutingV2CPU.cpu_op_exec(
            x, expert_idx, scale, offset,
            expert_range=expert_range,
            quant_mode=0,
            row_idx_type=row_idx_type,
            expert_tokens_num_flag=True,
            expert_tokens_num_type=1,
            drop_pad_mode=0,
            active_num=active_num,
            expert_capacity=expert_capacity
        )
    
    print(f"Input x shape: {x.shape}, dtype: {x.dtype}")
    print(f"Input scale shape: {scale.shape}, offset shape: {offset.shape}")
    print(f"Output expanded_x shape: {expanded_x.shape}, dtype: {expanded_x.dtype}")
    print(f"Output expanded_row_idx shape: {expanded_row_idx.shape}")
    print(f"Output expert_tokens_count: {expert_tokens_count}")
    print()
```

```python
def demo_dynamic_quant():
    """动态量化模式"""
    print("=" * 50)
    print("Demo: 动态量化模式 (quant_mode=1)")
    print("=" * 50)
    
    bs, h, k = 32, 200, 8
    expert_range = [0, 16]
    expert_range_length = expert_range[1] - expert_range[0]
    
    # 生成输入 (可选scale)
    x, expert_idx, scale, offset, expert_tokens_num_type, row_idx_type, active_num, expert_capacity = \
        MoeInitRoutingV2CPU.generate_inputs(bs, h, k, np.float32, (expert_range_length, h), False, True, 0)
    
    # 执行计算
    expanded_x, expanded_row_idx, expert_tokens_count, expanded_scale = \
        MoeInitRoutingV2CPU.cpu_op_exec(
            x, expert_idx, scale, offset,
            expert_range=expert_range,
            quant_mode=1,
            row_idx_type=row_idx_type,
            expert_tokens_num_flag=True,
            expert_tokens_num_type=1,
            drop_pad_mode=0,
            active_num=active_num,
            expert_capacity=expert_capacity
        )
    
    print(f"Input x shape: {x.shape}, dtype: {x.dtype}")
    print(f"Input scale shape: {scale.shape}")
    print(f"Output expanded_x shape: {expanded_x.shape}, dtype: {expanded_x.dtype}")
    print(f"Output expanded_scale shape: {expanded_scale.shape if expanded_scale is not None else None}")
    print()
```

--------------------------------

### Verify Operator Installation

Source: https://github.com/ascend/op-plugin/blob/master/examples/aclnn_extension/README.md

Commands to run the test script after installation.

```bash
cd ..
cd test
python test_npu_fast_gelu_custom.py
```

--------------------------------

### Usage Example for _ExperimentalConfig

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-profiler/torch_npu-profiler-_ExperimentalConfig.md

Demonstrates how to instantiate _ExperimentalConfig and pass it to the torch_npu.profiler.profile context manager.

```python
import torch
import torch_npu

...

experimental_config = torch_npu.profiler._ExperimentalConfig(
    export_type=[
        torch_npu.profiler.ExportType.Text
        ],
    profiler_level=torch_npu.profiler.ProfilerLevel.Level0,
    msprof_tx=False,
    aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
    l2_cache=False,
    op_attr=False,
    data_simplification=False,
    record_op_args=False,
    gc_detect_threshold=None
)

with torch_npu.profiler.profile(
        on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result"),
        experimental_config=experimental_config) as prof:
        for step in range(steps): # 训练函数
                train_one_step() # 训练函数
                prof.step()
```

--------------------------------

### Basic Usage Example

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu/torch_npu-set_stream_limit.md

Demonstrates setting core limits for the current stream and a new stream.

```python
>>> import torch
>>> import torch_npu
>>> torch.npu.set_device(0)
>>> torch.npu.set_stream_limit(torch.npu.current_stream(), 12, 24)
>>> torch.npu.set_stream_limit(torch.npu.Stream(), 13, 23)
```

--------------------------------

### Usage Example of torch_npu.profiler.schedule

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-profiler/torch_npu-profiler-schedule.md

Example demonstrating how to integrate the schedule into a profiler context manager.

```python
import torch
import torch_npu

...

with torch_npu.profiler.profile(
    activities=[
        torch_npu.profiler.ProfilerActivity.CPU,
        torch_npu.profiler.ProfilerActivity.NPU,
    ],
    schedule=torch_npu.profiler.schedule(
        wait=1,                        # 等待阶段，跳过1个step
        warmup=1,                      # 预热阶段，跳过1个step
        active=2,                      # 记录2个step的活动数据，并在之后调用on_trace_ready
        repeat=2,                      # 循环wait+warmup+active过程2遍
        skip_first=1,                  # 跳过1个step
        skip_first_wait=1            # 跳过第一个wait
    ),
    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler('./result')
    ) as prof:
        for _ in range(9):
            train_one_step()
            prof.step()                # 通知profiler完成一个step
```

--------------------------------

### Usage Example for Compatible Implementation

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-npu/torch_npu-npu-use_compatible_impl.md

Demonstrates enabling compatibility mode and applying it to a gelu operation.

```python
import torch
import torch_npu

torch_npu.npu.use_compatible_impl(True)
shape = [100, 400]
mode = "none"
input = torch.rand(shape, dtype=torch.float16).npu()
output = torch.nn.functional.gelu(input, approximate=mode)
```

--------------------------------

### Install CANN Ops Package

Source: https://github.com/ascend/op-plugin/blob/master/examples/kernel_extension_aclgraph/torch_library/README.md

Command to install the CANN ops package with appropriate permissions.

```bash
# 确保安装包具有可执行权限
chmod +x Ascend-cann-${soc_name}-ops_${cann_version}_linux-${arch}.run
# 安装命令
./Ascend-cann-${soc_name}-ops_${cann_version}_linux-${arch}.run  --install --quiet --install-path=${install_path}
```

--------------------------------

### Initialize PMCC Obfuscation Resources

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-npu/（beta）torch_npu-npu-obfuscation_initialize.md

Example demonstrating the initialization of PMCC obfuscation resources using torch_npu.npu.obfuscation_initialize.

```python
import torch
import torch_npu

device = "npu:0"
hidden_size = int(3584)
cmd = 1
data_type = torch.bfloat16
model_obf_seed = 0
data_obf_seed = 0
thread_num = 4
tp_rank = 0
i = 0
hidden_states = torch.randn((1024,3584), dtype=torch.bfloat16, device=device)
obf_cft = 1.0
fd = torch_npu.npu.obfuscation_initialize(hidden_size, tp_rank, cmd, data_type=data_type, thread_num= thread_num, obf_coefficient=obf_cft)
```

--------------------------------

### Usage Example

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-contrib/（beta）torch_npu-contrib-function-npu_batched_multiclass_nms.md

Demonstrates how to import and use the function with dummy tensor inputs on an NPU device.

```python
>>> import torch, torch_npu
>>> from torch_npu.contrib.function import npu_batched_multiclass_nms
>>> boxes = torch.randint(1, 255, size=(4, 200, 80, 4)).npu().half()
>>> scores = torch.randn(4, 200, 81).npu().half()
>>> det_bboxes, det_labels = npu_batched_multiclass_nms(boxes, scores, score_thr=0.3, nms_thr=0.5, max_num=3)
>>> print(det_bboxes.shape)
torch.Size([4, 3, 5])
>>> print(det_labels.shape)
torch.Size([4, 3])
```

--------------------------------

### Usage Example

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu/torch_npu-npu_scaled_masked_softmax.md

Demonstrates how to initialize inputs and call the npu_scaled_masked_softmax function on an NPU device.

```python
>>> import torch
>>> import torch_npu
>>> shape = [4, 4, 2048, 2048]
>>> x = torch.rand(shape).npu()
>>> mask = torch.zeros_like(x).bool()
>>> scale = 1.0
>>> fixed_triu_mask = False
>>> output = torch_npu.npu_scaled_masked_softmax(x, mask, scale, fixed_triu_mask)
>>> print(output.shape)
torch.Size([4, 4, 2048, 2048])
```

--------------------------------

### MultiheadAttention Usage Example

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-contrib/（beta）torch_npu-contrib-module-MultiheadAttention.md

Example demonstrating how to configure, initialize, and execute the MultiheadAttention module on an NPU.

```python
>>> from torch_npu.testing.common_utils import create_common_tensor
>>> from torch_npu.contrib.module import MultiheadAttention
>>> import numpy as np
>>> from torch_npu.contrib.module.multihead_attention import _MHAConfig
>>> _MHAConfig.set_fussion()
>>> model = MultiheadAttention(embed_dim=1024,num_heads=16,dropout=0.1,kdim=1024,vdim=1024,self_attention=True,encoder_decoder_attention=True)
>>> _, query = create_common_tensor([np.float16, 29, (1024,1024)], -1, 1)
>>> _, key = create_common_tensor([np.float16, 29, (1024,1024)], -1, 1)
>>> _, value = create_common_tensor([np.float16, 29, (1024,1024)], -1, 1)
>>> _, key_padding_mask = create_common_tensor([np.float16, 29, (1024,1024)], -1, 1)
>>> bsz = 16
>>> tgt_len = 64
>>> s_len=64
>>> model = model.to("npu")
>>> output = model(query, key, value, bsz, tgt_len, s_len, key_padding_mask)
>>> print(output)
```

--------------------------------

### FusedColorJitter Usage Example

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-contrib/（beta）torch_npu-contrib-module-FusedColorJitter.md

Example demonstrating how to initialize and apply FusedColorJitter to an image on an NPU device.

```python
>>> import torch
>>> from PIL import Image 
>>> from torch_npu.contrib.module import FusedColorJitter
>>> import numpy as np
>>> image = Image.fromarray(torch.randint(0, 256, size=(224, 224, 3)).numpy().astype(np.uint8))
>>> fcj = FusedColorJitter(0.1, 0.1, 0.1, 0.1).npu()
>>> img = fcj(image)
```

--------------------------------

### Usage Example of NpuFusedRMSprop

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu-optim/torch_npu-optim-NpuFusedRMSprop.md

Demonstrate the initialization and step execution of the NpuFusedRMSprop optimizer with sample parameters.

```python
import torch
from torch_npu.npu.amp import GradScaler, autocast
from torch_npu.optim import NpuFusedRMSprop 

def _create_simple_params_and_grads():
    params = [
        torch.arange(6).reshape(2, 3).float().npu(),
        torch.arange(12).reshape(4, 3).float().npu(),
        torch.arange(6).reshape(2, 3).half().npu(),
        torch.arange(12).reshape(4, 3).half().npu(),
        torch.arange(15).reshape(5, 3).float().npu(),
        torch.arange(18).reshape(6, 3).half().npu(),
        torch.arange(6).reshape(2, 3).float().npu()
    ]

    for i, p in enumerate(params):
        if i < len(params) - 1:
            p.requires_grad = True
            p.grad = p.clone().detach() / 100.

    return params

opt_kwargs = dict(eps=0.001, lr=0.01, weight_decay=1e-5)
params = _create_simple_params_and_grads()
fused_opt = NpuFusedRMSprop(params, **opt_kwargs)
with torch.no_grad():
    fused_opt.step()
```

--------------------------------

### Usage Example for torch_npu.npu_iou

Source: https://github.com/ascend/op-plugin/blob/master/docs/zh/custom_APIs/torch_npu/（beta）torch_npu-npu_iou.md

Example demonstrating the calculation of IoU using input tensors on the NPU.

```python
>>> import torch
>>> import torch_npu
>>> bboxes = torch.tensor([[0, 0, 10, 10],
                           [10, 10, 20, 20],
                           [32, 32, 38, 42]], dtype=torch.float16).to("npu")
>>> gtboxes = torch.tensor([[0, 0, 10, 20],
                            [0, 10, 10, 10],
                            [10, 10, 20, 20]], dtype=torch.float16).to("npu")
>>> output_iou = torch_npu.npu_iou(bboxes, gtboxes, 0)
>>> print(output_iou)
tensor([[0.4985, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000], 
       [0.0000, 0.9961, 0.0000]], device='npu:0', dtype=torch.float16)
```