Star1122
2723
代码
代码
Issues223
Pull Requests411
Actions
项目讨论
Wiki
分析
Star1122
2723
已合并
dispatchV2 combineV2算子新增<<<>>>直调 #2141
yangzeheng创建于 2月28日
dispatchV2 combineV2算子新增<<<>>>直调 #2141
已合并
yangzeheng创建于 2月28日
从yangzeheng/ops-transformer_1master合入到cann/ops-transformermaster
共 32 个文件变更+6739-3
Mexamples/fast_kernel_launch_example/CMakeLists.txt+2-0
@@ -53,6 +53,8 @@ set(LINK_LIBRARIES
53 register53 register
54 tiling_api54 tiling_api
55 runtime55 runtime
56+ hccl
57+ hcomm
56)58)
57 59 
58# compile options60# compile options
Mexamples/fast_kernel_launch_example/ascend_ops/__init__.py+2-1
@@ -21,4 +21,5 @@ except ImportError as e:
21 "Cannot import _C. Please make sure the `ascend_ops` is properly installed. "21 "Cannot import _C. Please make sure the `ascend_ops` is properly installed. "
22 ) from e22 ) from e
23 23 
24-from . import ops24+from . import ops
25+from .moe_distribute import MoeDistributeBuffer
Aexamples/fast_kernel_launch_example/ascend_ops/moe_distribute.py+129-0
@@ -0,0 +1,129 @@
1+import torch
2+import torch_npu
3+from torch.library import impl
4+from torch_npu.utils._error_code import ErrCode, ops_error
5+ 
6+X_CONTEXT_SIZE=4
7+ 
8+class MoeDistributeBuffer:
9+ def __init__(self, group, ccl_buffer_size: int = 0, comm_alg: int = 0):
10+ self.group = group
11+ self.rank_id = torch.distributed.get_rank(group)
12+ self.world_size = torch.distributed.get_world_size(group)
13+ self.group_name = group._get_backend(torch.device("npu")).get_hccl_comm_name(self.rank_id, init_comm=False)
14+ mb_buffer_size = 200 if ccl_buffer_size == 0 else ccl_buffer_size
15+ self.ccl_buffer_size = mb_buffer_size * 1024 * 1024 # convert from mb
16+ x_context = torch.zeros(X_CONTEXT_SIZE, dtype=torch.int32).npu()
17+ self.context = torch.ops.ascend_ops.updateContext(x_context, self.group_name, self.ccl_buffer_size, self.world_size)
18+ 
19+ @staticmethod
20+ def get_ccl_buffer_size(world_size: int, num_max_dispatch_tokens_per_rank: int, hidden: int,
21+ num_moe_expert: int, topk: int, num_shared_expert: int = 0,
22+ num_shared_expert_ranks: int = 0, comm_alg: str = "",
23+ ) -> int:
24+ def inline_align(value, base):
25+ return (value + base - 1) // base * base
26+ 
27+ max_out_dtype_size = 2 # sizeof(int32)
28+ mb_conversion = 1024 * 1024
29+ ub_align = 32 # 32B
30+ scale_expand_index_buffer = 44 # scale 32B + 3 * 4 expand_idx
31+ full_mesh_data_align = 480
32+ win_addr_align = 512
33+ 
34+ comm_alg_support_list = ["fullmesh_v1", "fullmesh_v2", ""]
35+ torch._check(comm_alg in comm_alg_support_list,
36+ lambda: (f"comm_alg only support {comm_alg_support_list=} "
37+ f"but got {comm_alg=}."),)
38+ 
39+ token_actual_len = inline_align(hidden * max_out_dtype_size, ub_align) + scale_expand_index_buffer
40+ if (comm_alg == "fullmesh_v2"):
41+ token_need_size_dispatch = inline_align(token_actual_len, full_mesh_data_align) // full_mesh_data_align * \
42+ win_addr_align
43+ else:
44+ token_need_size_dispatch = inline_align(token_actual_len, win_addr_align)
45+ token_need_size_combine = inline_align(hidden * max_out_dtype_size, win_addr_align)
46+ 
47+ local_moe_expert_num = num_moe_expert // (world_size - num_shared_expert_ranks)
48+ minimum_buffer_size = 2 * (
49+ (num_max_dispatch_tokens_per_rank * token_need_size_dispatch * world_size * local_moe_expert_num) + \
50+ (num_max_dispatch_tokens_per_rank * token_need_size_combine * (topk + num_shared_expert))) + mb_conversion
51+ return inline_align(minimum_buffer_size, mb_conversion) // mb_conversion
52+ 
53+ def update_ctx(self, new_group):
54+ self.group = new_group
55+ self.rank_id = torch.distributed.get_rank(new_group)
56+ self.group_name = new_group._get_backend(torch.device("npu")).get_hccl_comm_name(self.rank_id, init_comm=False)
57+ x_context = torch.zeros(X_CONTEXT_SIZE, dtype=torch.int32).npu()
58+ self.context = torch.ops.ascend_ops.updateContext(x_context, self.group_name, self.ccl_buffer_size, self.world_size)
59+ return
60+ 
61+ def npu_moe_distribute_dispatch_v2(self, x, expert_ids,
62+ moe_expert_num, *, scales=None, x_active_mask=None,
63+ expert_scales=None, performance_info=None, expert_shard_type=0, shared_expert_num=0,
64+ shared_expert_rank_num=0, quant_mode=0, global_bs=0, expert_token_nums_type=0,
65+ comm_alg="", zero_expert_num=0, copy_expert_num=0, const_expert_num=0):
66+ (expand_x, dynamic_scales, expand_idx, expert_token_nums, ep_recv_counts, expand_scales) \
67+ = torch.ops.ascend_ops.MoeDistributeDispatchV2(
68+ mc2_context=self.context,
69+ x=x,
70+ expert_ids=expert_ids,
71+ group_ep=self.group_name,
72+ ep_world_size=self.world_size,
73+ ep_rank_id=self.rank_id,
74+ moe_expert_num=moe_expert_num,
75+ total_winsize_ep=self.ccl_buffer_size,
76+ scales=scales,
77+ x_active_mask=x_active_mask,
78+ expert_scales=expert_scales,
79+ performance_info=performance_info,
80+ expert_shard_type=expert_shard_type,
81+ shared_expert_num=shared_expert_num,
82+ shared_expert_rank_num=shared_expert_rank_num,
83+ quant_mode=quant_mode,
84+ expert_token_nums_type=expert_token_nums_type,
85+ global_bs=global_bs,
86+ comm_alg=comm_alg,
87+ zero_expert_num=zero_expert_num,
88+ copy_expert_num=copy_expert_num,
89+ const_expert_num=const_expert_num)
90+ return expand_x, dynamic_scales, expand_idx, expert_token_nums, ep_recv_counts, expand_scales
91+ 
92+ 
93+ def npu_moe_distribute_combine_v2(self, expand_x, expert_ids, assist_info_for_combine,
94+ ep_send_counts, expert_scales,
95+ moe_expert_num, *,
96+ x_active_mask=None, shared_expert_x=None, ori_x=None,
97+ const_expert_alpha_1=None, const_expert_alpha_2=None, const_expert_v=None,
98+ performance_info=None, expert_shard_type=0, shared_expert_num=0, shared_expert_rank_num=0,
99+ global_bs=0, comm_quant_mode=0,
100+ comm_alg="", zero_expert_num=0, copy_expert_num=0, const_expert_num=0):
101+ return torch.ops.ascend_ops.MoeDistributeCombineV2(
102+ mc2_context=self.context,
103+ expand_x=expand_x,
104+ expert_ids=expert_ids,
105+ assist_info_for_combine=assist_info_for_combine,
106+ ep_send_counts=ep_send_counts,
107+ expert_scales=expert_scales,
108+ group_ep=self.group_name,
109+ ep_world_size=self.world_size,
110+ ep_rank_id=self.rank_id,
111+ moe_expert_num=moe_expert_num,
112+ total_winsize_ep=self.ccl_buffer_size,
113+ x_active_mask=x_active_mask,
114+ shared_expert_x=shared_expert_x,
115+ ori_x=ori_x,
116+ const_expert_alpha_1=const_expert_alpha_1,
117+ const_expert_alpha_2=const_expert_alpha_2,
118+ const_expert_v=const_expert_v,
119+ performance_info=performance_info,
120+ expert_shard_type=expert_shard_type,
121+ shared_expert_num=shared_expert_num,
122+ shared_expert_rank_num=shared_expert_rank_num,
123+ global_bs=global_bs,
124+ comm_quant_mode=comm_quant_mode,
125+ comm_alg=comm_alg,
126+ zero_expert_num=zero_expert_num,
127+ copy_expert_num=copy_expert_num,
128+ const_expert_num=const_expert_num)
129+
Mexamples/fast_kernel_launch_example/ascend_ops/ops.py+1-1
@@ -60,4 +60,4 @@ def groupedmatmul(
60 group_list, per_token_scale,60 group_list, per_token_scale,
61 split_item, group_type, group_list_type, act_type,61 split_item, group_type, group_list_type, act_type,
62 tuning_config62 tuning_config
63- )63+ )
Aexamples/fast_kernel_launch_example/csrc/common/inc/kernel/mc2_kernel_utils.h+29-0
@@ -0,0 +1,29 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file mc2_kernel_utils.h
13+ * \brief
14+ */
15+ 
16+#ifndef MC2_KERNEL_UTILS_H
17+#define MC2_KERNEL_UTILS_H
18+ 
19+namespace AscendC {
20+ 
21+template<AscendC::HardEvent event>
22+__aicore__ inline void SyncFunc() {
23+ AscendC::TEventID eventID = GetTPipePtr()->FetchEventID(event);
24+ AscendC::SetFlag<event>(eventID);
25+ AscendC::WaitFlag<event>(eventID);
26+}
27+ 
28+}
29+#endif // MC2_KERNEL_UTILS_H
Aexamples/fast_kernel_launch_example/csrc/common/inc/kernel/mc2_moe_context.h+25-0
@@ -0,0 +1,25 @@
1+/**
2+* Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+* CANN Open Software License Agreement Version 2.0 (the "License").
5+* Please refer to the License for details. You may not use this file except in compliance with the License.
6+* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+* See LICENSE in the root of the software repository for the full text of the License.
9+*/
10+ 
11+/*!
12+* \file mc2_moe_context.h
13+* \brief
14+*/
15+ 
16+#ifndef MC2_MOE_CONTEXT_H
17+#define MC2_MOE_CONTEXT_H
18+ 
19+struct Mc2MoeContext {
20+ uint64_t epRankId;
21+ uint64_t kfcContextAddr; // host kfc方案中，需要传递通信API所需的地址
22+ uint64_t epHcclBuffer[1024];
23+};
24+ 
25+#endif //MC2_MOE_CONTEXT_H
Mexamples/fast_kernel_launch_example/csrc/extension.cpp+1-1
@@ -42,4 +42,4 @@ TORCH_LIBRARY(ascend_ops, m)
42 "groupType, int groupListType, int actType,int[]? tuningConfigOptional) -> Tensor");42 "groupType, int groupListType, int actType,int[]? tuningConfigOptional) -> Tensor");
43}43}
44 44 
45-} // namespace ascend_ops45+} // namespace ascend_ops
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_combine_v2/ascend910_93/CMakeLists.txt+12-0
@@ -0,0 +1,12 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# ----------------------------------------------------------------------------
10+ 
11+add_sources("--npu-arch=dav-2201")
12+ 
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_combine_v2/ascend910_93/moe_distribute_combine_v2_entry.cpp+110-0
@@ -0,0 +1,110 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file moe_distribute_combine_v2.cpp
13+ * \brief
14+ */
15+#if ASC_DEVKIT_MAJOR >= 9
16+#include "basic_api/kernel_basic_intf.h"
17+#else
18+#include "kernel_operator.h"
19+#endif
20+#include "op_kernel/moe_distribute_combine_v2.h"
21+ 
22+using namespace MoeDistributeCombineV2Impl;
23+using namespace Mc2Kernel;
24+using namespace AscendC;
25+ 
26+/*
27+* A3 tilingkey说明
28+* 3位的十进制数
29+* 第1位（个位）：quantMode:
30+* 0: 不量化, 2: 使能量化
31+* 第2位（十位）：x输入类型:
32+* 0: float16, 1: bfloat16
33+* 第3位（百位）：无实际含义
34+*/
35+ 
36+template<typename ExpandXType, bool IsInt8Quant>
37+__attribute__((always_inline)) __aicore__ __inline__ void moe_distribute_combine_v2(
38+ GM_ADDR expandX, GM_ADDR expertIds, GM_ADDR expandIdx, GM_ADDR epSendCount, GM_ADDR residualX,
39+ GM_ADDR gamma, GM_ADDR expertScales, GM_ADDR xActiveMask, GM_ADDR sharedExpertX, GM_ADDR oriX,
40+ GM_ADDR constExpertAlpha1, GM_ADDR constExpertAlpha2, GM_ADDR constExpertV, GM_ADDR performanceInfo, GM_ADDR XOut,
41+ GM_ADDR workspaceGM, GM_ADDR mc2Context, MoeDistributeCombineV2Info tilingData)
42+{
43+ TPipe pipe;
44+ MoeDistributeCombineV2<ExpandXType, IsInt8Quant> op;
45+ op.Init(mc2Context, expandX, expertIds, expandIdx, epSendCount, residualX, gamma, expertScales,
46+ xActiveMask, sharedExpertX, oriX, constExpertAlpha1, constExpertAlpha2, constExpertV, performanceInfo,
47+ XOut, workspaceGM, tilingData, &pipe);
48+
49+ op.Process();
50+ return;
51+}
52+ 
53+extern "C" __global__ __aicore__ void moe_distribute_combine_v2_generic(
54+ int32_t tilingKey,
55+ GM_ADDR expandX, GM_ADDR expertIds, GM_ADDR expandIdx, GM_ADDR epSendCount, GM_ADDR residualX,
56+ GM_ADDR gamma, GM_ADDR expertScales, GM_ADDR xActiveMask, GM_ADDR sharedExpertX, GM_ADDR oriX,
57+ GM_ADDR constExpertAlpha1, GM_ADDR constExpertAlpha2, GM_ADDR constExpertV, GM_ADDR performanceInfo, GM_ADDR XOut,
58+ GM_ADDR workspaceGM, GM_ADDR mc2Context, MoeDistributeCombineV2Info tilingData)
59+{
60+ // 根据不同的数据类型调用不同的模板
61+ switch (tilingKey) {
62+ 
63+ case 100:
64+ moe_distribute_combine_v2<float16_t, false>(
65+ expandX, expertIds, expandIdx, epSendCount, residualX, gamma, expertScales, xActiveMask,
66+ sharedExpertX, oriX, constExpertAlpha1, constExpertAlpha2, constExpertV, performanceInfo, XOut,
67+ workspaceGM, mc2Context, tilingData);
68+ break;
69+ 
70+ case 102:
71+ moe_distribute_combine_v2<float16_t, true>(
72+ expandX, expertIds, expandIdx, epSendCount, residualX, gamma, expertScales, xActiveMask,
73+ sharedExpertX, oriX, constExpertAlpha1, constExpertAlpha2, constExpertV, performanceInfo, XOut,
74+ workspaceGM, mc2Context, tilingData);
75+ break;
76+ 
77+ case 110:
78+ moe_distribute_combine_v2<bfloat16_t, false>(
79+ expandX, expertIds, expandIdx, epSendCount, residualX, gamma, expertScales, xActiveMask,
80+ sharedExpertX, oriX, constExpertAlpha1, constExpertAlpha2, constExpertV, performanceInfo, XOut,
81+ workspaceGM, mc2Context, tilingData);
82+ break;
83+ 
84+ case 112:
85+ moe_distribute_combine_v2<bfloat16_t, true>(
86+ expandX, expertIds, expandIdx, epSendCount, residualX, gamma, expertScales, xActiveMask,
87+ sharedExpertX, oriX, constExpertAlpha1, constExpertAlpha2, constExpertV, performanceInfo, XOut,
88+ workspaceGM, mc2Context, tilingData);
89+ break;
90+ 
91+ default:
92+ AscendC::PRINTF("moe_distribute_combine_v2 Error: invalid tilingKey = %d\n", tilingKey);
93+ return;
94+ }
95+ return;
96+}
97+ 
98+// <<<>>>调用函数
99+void moe_distribute_combine_v2_entry(int32_t tilingKey, uint32_t blockDim, void* stream, GM_ADDR expandX, GM_ADDR expertIds,
Y
Yyangzeheng3月16日

为什么要封一个PTA接口

likedislike
100+ GM_ADDR expandIdx, GM_ADDR epSendCount, GM_ADDR residualX,
101+ GM_ADDR gamma, GM_ADDR expertScales, GM_ADDR xActiveMask, GM_ADDR sharedExpertX, GM_ADDR oriX,
102+ GM_ADDR constExpertAlpha1, GM_ADDR constExpertAlpha2, GM_ADDR constExpertV, GM_ADDR performanceInfo, GM_ADDR XOut,
103+ GM_ADDR workspaceGM, GM_ADDR mc2Context, MoeDistributeCombineV2Info tilingData)
104+{
105+ moe_distribute_combine_v2_generic<<<blockDim, nullptr, stream>>>(
106+ tilingKey,
107+ expandX, expertIds, expandIdx, epSendCount, residualX, gamma, expertScales, xActiveMask,
108+ sharedExpertX, oriX, constExpertAlpha1, constExpertAlpha2, constExpertV, performanceInfo, XOut,
109+ workspaceGM, mc2Context, tilingData);
110+}
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_combine_v2/ascend910_93/moe_distribute_combine_v2_entry.h+22-0
@@ -0,0 +1,22 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+#ifndef MOE_DISTRIBUTE_COMBINE_V2_ENTRY_H
11+#define MOE_DISTRIBUTE_COMBINE_V2_ENTRY_H
12+ 
13+#include "op_kernel/moe_distribute_combine_v2_tiling.h"
14+ 
15+// <<<>>>调用函数声明
16+void moe_distribute_combine_v2_entry(int32_t tilingKey, uint32_t blockDim, void* stream, GM_ADDR expandX, GM_ADDR expertIds,
17+ GM_ADDR expandIdx, GM_ADDR epSendCount, GM_ADDR residualX,
18+ GM_ADDR gamma, GM_ADDR expertScales, GM_ADDR xActiveMask, GM_ADDR sharedExpertX, GM_ADDR oriX,
19+ GM_ADDR constExpertAlpha1, GM_ADDR constExpertAlpha2, GM_ADDR constExpertV, GM_ADDR performanceInfo, GM_ADDR yOut,
20+ GM_ADDR workspaceGM, GM_ADDR mc2Context, MoeDistributeCombineV2Info tilingData);
21+ 
22+#endif //MOE_DISTRIBUTE_COMBINE_V2_ENTRY_H
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_combine_v2/ascend910_93/moe_distribute_combine_v2_torch.cpp+375-0
@@ -0,0 +1,375 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify it.
3+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
8+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+ * See LICENSE in the root of the software repository for the full text of the License.
10+ */
11+ 
12+/*!
13+ * \file moe_distribute_combine_v2_torch.cpp
14+ * \brief
15+ */
16+ 
17+#include <ATen/ATen.h>
18+#include <vector>
19+#include <torch/all.h>
20+ 
21+#include "acl/acl.h"
22+#include "tiling/platform/platform_ascendc.h"
23+#include "kernel_operator.h"
24+#include "moe_distribute_combine_v2_entry.h"
25+#include "op_kernel/moe_distribute_combine_v2_tiling.h"
26+#include "../../moe_distribute_dispatch_v2/ascend910_93/moe_distribute_dispatch_v2_torch.h"
27+ 
28+#include "torch_npu/csrc/core/npu/NPUStream.h"
29+#include "torch_npu/csrc/framework/OpCommand.h"
30+ 
31+namespace ascend_ops {
32+ 
33+namespace MoeDistributeCombineV2 {
34+ 
35+TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m)
36+{
37+ m.def("MoeDistributeCombineV2(Tensor expand_x, Tensor expert_ids, Tensor assist_info_for_combine, " \
38+ "Tensor ep_send_counts, Tensor expert_scales, Tensor mc2_context, str group_ep, int ep_world_size, " \
39+ "int ep_rank_id, int moe_expert_num, int total_winsize_ep, *, " \
40+ "Tensor? x_active_mask=None, Tensor? shared_expert_x=None, Tensor? ori_x=None, " \
41+ "Tensor? const_expert_alpha_1=None, Tensor? const_expert_alpha_2=None, Tensor? const_expert_v=None, " \
42+ "Tensor? performance_info=None, " \
43+ "int expert_shard_type=0, int shared_expert_num=0, int shared_expert_rank_num=0, " \
44+ "int global_bs=0, int comm_quant_mode=0, " \
45+ "str comm_alg=\"\", int zero_expert_num=0, int copy_expert_num=0, int const_expert_num=0) " \
46+ "-> Tensor");
47+}
48+ 
49+constexpr uint32_t DIM_ONE = 1UL;
50+constexpr uint32_t DIM_TWO = 2UL;
51+constexpr uint32_t TILINGKEY_XTYPE = 10;
52+constexpr uint32_t WORKSPACESIZE = 16 * 1024 *1024;
53+constexpr int64_t INT8_COMM_QUANT = 2;
54+constexpr uint32_t BUFFER_SINGLE = 1;
55+constexpr uint32_t BUFFER_DOUBLE = 2;
56+constexpr uint64_t UB_ALIGN = 32UL;
57+constexpr uint32_t DTYPE_SIZE_HALF = 2;
58+constexpr uint32_t ALIGNED_LEN = 256U;
59+constexpr uint32_t STATE_OFFSET = 32U;
60+ 
61+static void calculate_tilingkey(int32_t &tilingKey, at::ScalarType xType, const uint32_t quantMode)
62+{
63+ tilingKey += static_cast<uint64_t>(quantMode);
64+ if (xType == at::kBFloat16) {
65+ tilingKey += static_cast<uint64_t>(TILINGKEY_XTYPE);
66+ }
67+ return;
68+}
69+ 
70+void MoeDistributeCombineV2_api(aclrtStream stream,
71+ const at::Tensor &expand_x, const at::Tensor &expert_ids,
72+ const at::Tensor &assist_info_for_combine,
73+ const at::Tensor &ep_send_counts, const at::Tensor &expert_scales,
74+ const at::Tensor &new_workspace, const at::Tensor &mc2_context,
75+ const c10::optional<at::Tensor> &x_active_mask,
76+ const c10::optional<at::Tensor> &shared_expert_x,
77+ const c10::optional<at::Tensor> &ori_x,
78+ const c10::optional<at::Tensor> &const_expert_alpha_1,
79+ const c10::optional<at::Tensor> &const_expert_alpha_2,
80+ const c10::optional<at::Tensor> &const_expert_v,
81+ const c10::optional<at::Tensor> &performance_info,
82+ at::Tensor &x_out, int64_t comm_quant_mode,
83+ MoeDistributeCombineV2Info tilingData)
84+{
85+ auto expandX_ptr = get_first_tensor_address<at::Tensor>(expand_x.scalar_type(), expand_x, false);
86+ auto expertIds_ptr = get_first_tensor_address<at::Tensor>(expert_ids.scalar_type(), expert_ids, false);
87+ auto expandIdx_ptr = get_first_tensor_address<at::Tensor>(assist_info_for_combine.scalar_type(),
88+ assist_info_for_combine, false);
89+ auto epSendCount_ptr = get_first_tensor_address<at::Tensor>(ep_send_counts.scalar_type(), ep_send_counts, false);
90+ auto expertScales_ptr = get_first_tensor_address<at::Tensor>(expert_scales.scalar_type(), expert_scales, false);
91+ auto workspace_ptr = get_first_tensor_address<at::Tensor>(new_workspace.scalar_type(), new_workspace, false);
92+ auto mc2Context_ptr = get_first_tensor_address<at::Tensor>(mc2_context.scalar_type(), mc2_context, false);
93+ 
94+ 
95+ void* xActiveMask_ptr = nullptr;
96+ if(x_active_mask.has_value()) {
97+ xActiveMask_ptr = get_first_tensor_address<c10::optional<at::Tensor>>(x_active_mask->scalar_type(),
98+ x_active_mask, false);
99+ }
100+ 
101+ void* sharedExpertX_ptr = nullptr;
102+ if(shared_expert_x.has_value()) {
103+ sharedExpertX_ptr = get_first_tensor_address<c10::optional<at::Tensor>>(shared_expert_x->scalar_type(),
104+ shared_expert_x, false);
105+ }
106+ 
107+ void* oriX_ptr = nullptr;
108+ if(ori_x.has_value()) {
109+ oriX_ptr = get_first_tensor_address<c10::optional<at::Tensor>>(ori_x->scalar_type(), ori_x, false);
110+ }
111+ 
112+ void* constExpertAlpha1_ptr = nullptr;
113+ if(const_expert_alpha_1.has_value()) {
114+ constExpertAlpha1_ptr = get_first_tensor_address<c10::optional<at::Tensor>>(const_expert_alpha_1->scalar_type(),
115+ const_expert_alpha_1, false);
116+ }
117+ 
118+ void* constExpertAlpha2_ptr = nullptr;
119+ if(const_expert_alpha_2.has_value()) {
120+ constExpertAlpha2_ptr = get_first_tensor_address<c10::optional<at::Tensor>>(const_expert_alpha_2->scalar_type(),
121+ const_expert_alpha_2, false);
122+ }
123+ 
124+ void* constExpertV_ptr = nullptr;
125+ if(const_expert_v.has_value()) {
126+ constExpertV_ptr = get_first_tensor_address<c10::optional<at::Tensor>>(const_expert_v->scalar_type(),
127+ const_expert_v, false);
128+ }
129+ 
130+ void* performanceInfo_ptr = nullptr;
131+ if(performance_info.has_value()) {
132+ performanceInfo_ptr = get_first_tensor_address<c10::optional<at::Tensor>>(performance_info->scalar_type(),
133+ performance_info, false);
134+ }
135+ void* residualX_ptr = nullptr;
136+ void* gamma_ptr = nullptr;
137+ auto XOut_ptr = get_first_tensor_address<at::Tensor>(x_out.scalar_type(), x_out, false);
138+ 
139+
140+ int32_t tilingKey = 100;
141+ calculate_tilingkey(tilingKey, expand_x.scalar_type(), comm_quant_mode);
142+ 
143+ moe_distribute_combine_v2_entry(tilingKey, tilingData.aivNum, stream, (GM_ADDR)expandX_ptr, (GM_ADDR)expertIds_ptr,
144+ (GM_ADDR)expandIdx_ptr, (GM_ADDR)epSendCount_ptr, (GM_ADDR)residualX_ptr,
145+ (GM_ADDR)gamma_ptr, (GM_ADDR)expertScales_ptr, (GM_ADDR)xActiveMask_ptr, (GM_ADDR)sharedExpertX_ptr,
146+ (GM_ADDR)oriX_ptr, (GM_ADDR)constExpertAlpha1_ptr, (GM_ADDR)constExpertAlpha2_ptr,
147+ (GM_ADDR)constExpertV_ptr, (GM_ADDR)performanceInfo_ptr, (GM_ADDR)XOut_ptr,
148+ (GM_ADDR)workspace_ptr, (GM_ADDR)mc2Context_ptr, tilingData);
149+}
150+ 
151+void calculate_buffernum(MoeDistributeCombineV2Info &tilingData, const at::Tensor &expand_x, int64_t comm_quant_mode)
152+{
153+ uint32_t axisH = tilingData.h;
154+ uint32_t axisBS = tilingData.bs;
155+ uint32_t axisK = tilingData.k;
156+ uint32_t zeroExpertNum = tilingData.zeroExpertNum;
157+ uint32_t copyExpertNum = tilingData.copyExpertNum;
158+ uint32_t constExpertNum = tilingData.constExpertNum;
159+ bool isInputExpertMaskFlag = tilingData.isExpertMask;
160+ bool isInputTokenMaskFlag = tilingData.isTokenMask;
161+ bool enableSpecialExpert = (constExpertNum + zeroExpertNum + copyExpertNum > 0U);
162+ uint32_t maxSizeTokenBuf = (axisH * expand_x.element_size() + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
163+ uint32_t hExpandXTypeSize = axisH * expand_x.element_size();
164+ uint32_t activeMaskAlignSize = axisBS * ((axisK * sizeof(bool) + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN);
165+ uint32_t hExpandXAlign32Size = (hExpandXTypeSize + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
166+ uint32_t hFloatSize = axisH * static_cast<uint32_t>(sizeof(float));
167+ uint32_t hFloatAlign32Size = (hFloatSize + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
168+ uint32_t maxSizeRowTmpFloatBuf = hFloatAlign32Size;
169+ uint32_t flagRcvCount = axisK + tilingData.sharedExpertNum;
170+ uint32_t hFloatAlign256Size = (hFloatSize + ALIGNED_LEN - 1) / ALIGNED_LEN * ALIGNED_LEN;
171+ uint32_t bsKNum = axisBS * axisK;
172+ uint32_t bsKFloatAlign = (bsKNum * sizeof(float) + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
173+ uint32_t mulBufSize = hFloatAlign256Size > bsKFloatAlign ? hFloatAlign256Size : bsKFloatAlign;
174+ 
175+ if (isInputExpertMaskFlag || enableSpecialExpert) {
176+ uint32_t activeMaskAlignHalfSize = activeMaskAlignSize * sizeof(DTYPE_SIZE_HALF);
177+ maxSizeTokenBuf = (activeMaskAlignSize > hExpandXAlign32Size ? activeMaskAlignSize : hExpandXAlign32Size);
178+ maxSizeRowTmpFloatBuf = (activeMaskAlignHalfSize > hFloatAlign32Size ? activeMaskAlignHalfSize : hFloatAlign32Size);
179+ }
180+ 
181+ // LocalWindowCopy的ub使用总量
182+ uint32_t totalBufferSize = 0;
183+ totalBufferSize = maxSizeTokenBuf + maxSizeRowTmpFloatBuf + mulBufSize + hFloatAlign32Size + hExpandXAlign32Size
184+ * BUFFER_DOUBLE + flagRcvCount * STATE_OFFSET * BUFFER_DOUBLE + UB_ALIGN;
185+ if (comm_quant_mode == INT8_COMM_QUANT) {
186+ uint32_t scaleNum = (hExpandXAlign32Size / expand_x.element_size()) / static_cast<uint32_t>(UB_ALIGN / sizeof(float));
187+ uint32_t scaleNumAlignSize = (scaleNum * sizeof(float) + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
188+ totalBufferSize += scaleNumAlignSize;
189+ }
190+ if (isInputTokenMaskFlag) {
191+ uint32_t axisBsAlignSize = (axisBS * sizeof(bool) + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
192+ totalBufferSize += axisBsAlignSize + axisBsAlignSize * sizeof(DTYPE_SIZE_HALF) * BUFFER_DOUBLE;
193+ }
194+ if (isInputExpertMaskFlag) {
195+ totalBufferSize += (axisBS * sizeof(DTYPE_SIZE_HALF) + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN + (axisBS * sizeof(int32_t) +
196+ UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN + (axisBS * axisK * sizeof(bool) + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
197+ }
198+ if (enableSpecialExpert && !isInputExpertMaskFlag) {
199+ totalBufferSize += (axisBS * sizeof(DTYPE_SIZE_HALF) + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
200+ }
201+ tilingData.bufferNum = totalBufferSize > tilingData.totalUbSize ? BUFFER_SINGLE : BUFFER_DOUBLE;
202+}
203+ 
204+void calculate_tilingdata(MoeDistributeCombineV2Info &tilingData, int64_t ep_world_size, int64_t ep_rank_id,
205+ int64_t moe_expert_num, int64_t total_winsize_ep,
206+ int64_t expert_shard_type, int64_t shared_expert_num, int64_t shared_expert_rank_num,
207+ int64_t global_bs, int64_t bs, int64_t h, int64_t k, int64_t a,
208+ int64_t zero_expert_num, int64_t copy_expert_num, int64_t const_expert_num, int64_t comm_quant_mode,
209+ const at::Tensor &expand_x,
210+ const c10::optional<at::Tensor> &x_active_mask,
211+ const c10::optional<at::Tensor> &shared_expert_x,
212+ const c10::optional<at::Tensor> &performance_info)
213+{
214+ auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
215+ uint64_t ubSizePlatFrom;
216+ ascendcPlatform->GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatFrom);
217+ 
218+ tilingData.epWorldSize = ep_world_size;
219+ tilingData.epRankId = ep_rank_id;
220+ tilingData.expertShardType = expert_shard_type;
221+ tilingData.sharedExpertNum = shared_expert_num;
222+ tilingData.sharedExpertRankNum = shared_expert_rank_num;
223+ tilingData.moeExpertNum = moe_expert_num;
224+ tilingData.moeExpertPerRankNum = moe_expert_num / (ep_world_size - shared_expert_rank_num);
225+ tilingData.zeroExpertNum = zero_expert_num;
226+ tilingData.copyExpertNum = copy_expert_num;
227+ tilingData.constExpertNum = const_expert_num;
228+ tilingData.globalBs = global_bs;
229+ tilingData.bs = bs;
230+ tilingData.k = k;
231+ tilingData.h = h;
232+ tilingData.a = a;
233+ tilingData.aivNum = ascendcPlatform->GetCoreNumAiv();
234+ tilingData.isTokenMask = (x_active_mask.has_value() && x_active_mask->dim() == DIM_ONE);
235+ tilingData.isExpertMask = (x_active_mask.has_value() && x_active_mask->dim() == DIM_TWO);
236+ tilingData.hasSharedExpertX = shared_expert_x.has_value(); // 或根据实际变量名调整
237+ tilingData.isPerformance = performance_info.has_value();
238+ tilingData.reserved0 = false;
239+ tilingData.reserved1 = false;
240+ tilingData.reserved2 = false;
241+ tilingData.totalUbSize = ubSizePlatFrom;
242+ tilingData.totalWinSizeEp = total_winsize_ep;
243+ calculate_buffernum(tilingData, expand_x, comm_quant_mode);
244+}
245+ 
246+at::Tensor npu_moe_distribute_combine_v2(
247+ const at::Tensor &expand_x, const at::Tensor &expert_ids,
248+ const at::Tensor &assist_info_for_combine,
249+ const at::Tensor &ep_send_counts, const at::Tensor &expert_scales, const at::Tensor &mc2_context,
250+ c10::string_view group_ep, int64_t ep_world_size, int64_t ep_rank_id,
251+ int64_t moe_expert_num, int64_t total_winsize_ep,
Y
Yyangzeheng3月16日

模仿deepep计算buffersize

likedislike
252+ const c10::optional<at::Tensor> &x_active_mask,
253+ const c10::optional<at::Tensor> &shared_expert_x,
254+ const c10::optional<at::Tensor> &ori_x,
255+ const c10::optional<at::Tensor> &const_expert_alpha_1,
256+ const c10::optional<at::Tensor> &const_expert_alpha_2,
257+ const c10::optional<at::Tensor> &const_expert_v,
258+ const c10::optional<at::Tensor> &performance_info,
259+ int64_t expert_shard_type, int64_t shared_expert_num, int64_t shared_expert_rank_num,
260+ int64_t global_bs, int64_t comm_quant_mode,
261+ c10::string_view comm_alg, int64_t zero_expert_num, int64_t copy_expert_num, int64_t const_expert_num)
262+{
263+ auto expand_x_size = expand_x.sizes();
264+ auto expert_ids_size = expert_ids.sizes();
265+
266+ int64_t bs = expert_ids_size[0];
267+ int64_t h = expand_x_size[1];
268+ int64_t k = expert_ids_size[1];
269+ 
270+ bool is_shared_default = ((shared_expert_num == 1) && (shared_expert_rank_num == 0));
271+ bool is_no_shared = ((shared_expert_num == 0) && (shared_expert_rank_num == 0));
272+
273+ bool shared_front = (expert_shard_type == 0);
274+ int64_t local_moe_expert_num = 1;
275+ int64_t global_bs_real = (global_bs == 0) ? (bs * ep_world_size) : global_bs;
276+ int64_t a = 0;
277+ if (shared_front) {
278+ if (ep_rank_id < shared_expert_rank_num) {
279+ local_moe_expert_num = 1;
280+ int64_t max_bs = global_bs_real / ep_world_size; // 前面已有拦截，保证ep_world_size > 0
281+ int64_t rank_num_per_shared_expert = shared_expert_rank_num / shared_expert_num; // 前面已有拦截, 保证进入该分支时shared_expert_num > 0
282+ int64_t max_shared_group_num = (ep_world_size + rank_num_per_shared_expert - 1) / rank_num_per_shared_expert;
283+ a = max_bs * max_shared_group_num;
284+ } else {
285+ local_moe_expert_num = moe_expert_num / (ep_world_size - shared_expert_rank_num);
286+ a = global_bs_real * std::min(local_moe_expert_num, k);
287+ }
288+ }
289+ TORCH_CHECK((expand_x.dim() == DIM_TWO) && (expert_ids.dim() == DIM_TWO), "The x and expert_ids should be 2D");
290+ TORCH_CHECK((expand_x.scalar_type() == at::kBFloat16) || (expand_x.scalar_type() == at::kHalf)
291+ || (expand_x.scalar_type() == at::kInt),
292+ "dtype of expand_x should be BFloat16, Float16 or Int, but got " + std::string(c10::toString(expand_x.scalar_type())));
293+ TORCH_CHECK(expert_ids.scalar_type() == at::kInt,
294+ "dtype of expert_ids should be Int, but got " + std::string(c10::toString(expert_ids.scalar_type())));
295+ 
296+ char *group_ep_ptr = const_cast<char *>(group_ep.data());
297+ 
298+ at::Tensor output;
299+ at::Tensor new_workspace = at::empty({WORKSPACESIZE / 4}, expert_ids.options().dtype(at::kInt));
300+ if (expand_x.scalar_type() != at::kInt) {
301+ output = at::empty({bs, h}, expert_ids.options().dtype(expand_x.scalar_type()));
302+ } else {
303+ output = at::empty({bs, h}, expert_ids.options().dtype(at::kBFloat16));
304+ }
305+ 
306+ c10::optional<at::Tensor> nulltensor = c10::nullopt;
307+ int64_t out_dtype = 0;
308+ int64_t group_list_type = 0;
309+ 
310+ std::string comm_alg_str = std::string(comm_alg);
311+ char *comm_alg_ptr = const_cast<char *>(comm_alg_str.c_str());
312+ 
313+ MoeDistributeCombineV2Info tilingData;
314+ calculate_tilingdata(tilingData, ep_world_size, ep_rank_id, moe_expert_num, total_winsize_ep,
315+ expert_shard_type, shared_expert_num, shared_expert_rank_num, global_bs, bs, h, k, a, zero_expert_num,
316+ copy_expert_num, const_expert_num, comm_quant_mode, expand_x, x_active_mask, shared_expert_x, performance_info);
317+ 
318+ auto stream = c10_npu::getCurrentNPUStream().stream(false);
319+ auto acl_call = [=]() mutable -> int {
320+ MoeDistributeCombineV2_api(stream, expand_x, expert_ids, assist_info_for_combine, ep_send_counts,
321+ expert_scales, new_workspace, mc2_context, x_active_mask,
322+ shared_expert_x, ori_x, const_expert_alpha_1, const_expert_alpha_2,
323+ const_expert_v, performance_info, output, comm_quant_mode, tilingData);
324+ return 0;
325+ };
326+ at_npu::native::OpCommand::RunOpApiV2("moeDistributeCombineV2", acl_call);
327+
328+ return output;
329+}
330+ 
331+at::Tensor npu_moe_distribute_combine_v2_meta(
332+ const at::Tensor &expand_x, const at::Tensor &expert_ids,
333+ const at::Tensor &assist_info_for_combine,
334+ const at::Tensor &ep_send_counts, const at::Tensor &expert_scales, const at::Tensor &mc2_context,
335+ c10::string_view group_ep, int64_t ep_world_size, int64_t ep_rank_id,
336+ int64_t moe_expert_num, int64_t total_winsize_ep,
337+ const c10::optional<at::Tensor> &x_active_mask,
338+ const c10::optional<at::Tensor> &shared_expert_x,
339+ const c10::optional<at::Tensor> &ori_x,
340+ const c10::optional<at::Tensor> &const_expert_alpha_1,
341+ const c10::optional<at::Tensor> &const_expert_alpha_2,
342+ const c10::optional<at::Tensor> &const_expert_v,
343+ const c10::optional<at::Tensor> &performance_info,
344+ int64_t expert_shard_type, int64_t shared_expert_num, int64_t shared_expert_rank_num,
345+ int64_t global_bs, int64_t comm_quant_mode,
346+ c10::string_view comm_alg, int64_t zero_expert_num, int64_t copy_expert_num, int64_t const_expert_num)
347+{
348+ auto expert_ids_size = expert_ids.sizes();
349+ auto expand_x_size = expand_x.sizes();
350+
351+ int64_t bs = expert_ids_size[0];
352+ int64_t h = expand_x_size[1];
353+
354+ at::Tensor output;
355+ 
356+ if (expand_x.scalar_type() != at::kInt) {
357+ output = at::empty({bs, h}, expert_ids.options().dtype(expand_x.scalar_type()));
358+ } else {
359+ output = at::empty({bs, h}, expert_ids.options().dtype(at::kBFloat16));
360+ }
361+ return output;
362+}
363+ 
364+TORCH_LIBRARY_IMPL(ascend_ops, PrivateUse1, m)
365+{
366+ m.impl("MoeDistributeCombineV2", TORCH_FN(npu_moe_distribute_combine_v2));
367+}
368+ 
369+TORCH_LIBRARY_IMPL(ascend_ops, Meta, m)
370+{
371+ m.impl("MoeDistributeCombineV2", &npu_moe_distribute_combine_v2_meta);
372+}
373+ 
374+} // namespace MoeDistributeCombineV2
375+} // namespace ascend_ops
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_combine_v2/ascend910_93/op_kernel/moe_distribute_combine_v2.h+1200-0
@@ -0,0 +1,1200 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file moe_distribute_combine_v2.h
13+ * \brief
14+ */
15+#ifndef MOE_DISTRIBUTE_COMBINE_V2_H
16+#define MOE_DISTRIBUTE_COMBINE_V2_H
17+ 
18+#if ASC_DEVKIT_MAJOR >= 9
19+#include "basic_api/kernel_basic_intf.h"
20+#else
21+#include "kernel_operator.h"
22+#endif
23+#include "adv_api/reduce/sum.h"
24+#include "kernel_tiling/kernel_tiling.h"
25+#include "moe_distribute_combine_v2_tiling.h"
26+#include "moe_distribute_combine_v2_quant.h"
27+#include "../../../common/inc/kernel/mc2_moe_context.h"
28+#include "../../../moe_distribute_dispatch_v2/ascend910_93/op_kernel/moe_distribute_v2_constant.h"
29+#include "../../../moe_distribute_dispatch_v2/ascend910_93/op_kernel/moe_distribute_v2_base.h"
30+ 
31+namespace MoeDistributeCombineV2Impl {
32+using namespace MoeDistributeV2Base;
33+using namespace Mc2Kernel;
34+ 
35+#define CombineMC2TypeClass typename ExpandXType, bool IsInt8Quant
36+#define CombineMC2TypeFunc ExpandXType, IsInt8Quant
37+ 
38+using namespace AscendC;
39+template <CombineMC2TypeClass>
40+class MoeDistributeCombineV2 {
41+public:
42+ __aicore__ inline MoeDistributeCombineV2() {};
43+ __aicore__ inline void Init(GM_ADDR mc2Context, GM_ADDR expandX, GM_ADDR expertIds, GM_ADDR expandIdx, GM_ADDR epSendCount, GM_ADDR residualX,
44+ GM_ADDR gamma, GM_ADDR expertScales, GM_ADDR xActiveMask, GM_ADDR sharedExpertX, GM_ADDR oriX,
45+ GM_ADDR constExpertAlpha1, GM_ADDR constExpertAlpha2, GM_ADDR constExpertV, GM_ADDR performanceInfo, GM_ADDR XOut,
46+ GM_ADDR workspaceGM, MoeDistributeCombineV2Info tilingData, TPipe *pipe);
47+ __aicore__ inline void Process();
48+private:
49+ __aicore__ inline void InitInputAndOutput(GM_ADDR residualX, GM_ADDR gamma, GM_ADDR expandX, GM_ADDR expertIds, GM_ADDR expandIdx,
50+ GM_ADDR epSendCount, GM_ADDR expertScales, GM_ADDR xActiveMask,
51+ GM_ADDR sharedExpertX, GM_ADDR oriX,
52+ GM_ADDR constExpertAlpha1, GM_ADDR constExpertAlpha2,
53+ GM_ADDR constExpertV, GM_ADDR performanceInfo, GM_ADDR XOut);
54+ __aicore__ inline void InitAttrs(GM_ADDR mc2Context, const MoeDistributeCombineV2Info tilingData);
55+ __aicore__ inline void InitTilingAttrs(const MoeDistributeCombineV2Info tilingData);
56+ __aicore__ inline void AlltoAllBuffInitAndMaskCal();
57+ __aicore__ inline void TokenMaskCalCnt();
58+ __aicore__ inline void ExpertMaskCalCnt();
59+ __aicore__ inline void GenerateActiveMask(half val);
60+ __aicore__ inline void MaskSpecialExpert();
61+ __aicore__ inline void MaskAlign();
62+ __aicore__ inline void SetWaitTpStatusAndDisPatch();
63+ __aicore__ inline void CustomAdd(LocalTensor<ExpandXType> &dst, LocalTensor<ExpandXType> &src0, LocalTensor<ExpandXType> &src1);
64+ __aicore__ inline void ExpertAlltoAllDispatchInnerCopyAdd(uint32_t toRankId, uint32_t tokenId, uint32_t topkId, uint32_t tkIndex);
65+ __aicore__ inline void ExpertAlltoAllDispatchCopyAdd();
66+ __aicore__ inline void ProcessConstantExpert(uint32_t tokenIndex, uint32_t const_expert_idx, float scaleVal);
67+ __aicore__ inline void ProcessCopyExpert(uint32_t tokenIndex, float scaleVal);
68+ __aicore__ inline void ProcessMoeExpert(uint32_t tokenIndexOffset, uint32_t topkId, float scaleVal);
69+ __aicore__ inline void ProcessExpert(uint32_t tokenIndex, uint32_t processLen);
70+ __aicore__ inline void ExpertScaleCopy(const uint32_t beginIndex, const uint32_t endIndex, const uint32_t tokenPerAivNum);
71+ __aicore__ inline void CalConstExpertAlpha(GlobalTensor<ExpandXType> constExpertAlphaGM, uint32_t const_expert_idx, float &alphaFloat);
72+ __aicore__ inline void LocalWindowCopy();
73+ __aicore__ inline void BuffInit();
74+ __aicore__ inline void SplitCoreCal();
75+ __aicore__ inline bool WaitDispatch(uint32_t tokenIndex, uint64_t performanceTimeStart, uint32_t copyCount, uint32_t beginIndex);
76+ __aicore__ inline void PerformanceInfoPerToken(uint32_t tokenIndex, uint64_t performanceTimeStart, uint32_t beginIndex, LocalTensor<float> stateTensor);
77+ __aicore__ inline void AddRmsNormAddCompute(uint32_t tokenIndex, uint32_t tokenOffset, uint32_t numCol,
78+ LocalTensor<float>& x1TmpFloatLocal,
79+ LocalTensor<float>& x2TmpFloatLocal,
80+ LocalTensor<float>& addOutTmpFloatLocal,
81+ const DataCopyExtParams& copyExtParams,
82+ const DataCopyPadExtParams<ExpandXType>& copyPadExtParams);
83+ __aicore__ GM_ADDR GetWinAddrByRankId(const int32_t rankId, const uint8_t domain)
84+ {
85+ return (GM_ADDR)mc2Context_->epHcclBuffer[rankId] + STATE_SIZE + winDataSizeOffsetEp_;
86+ }
87+ 
88+ __aicore__ GM_ADDR GetWinStateAddrByRankId(const int32_t rankId, const uint8_t domain)
89+ {
90+ return (GM_ADDR)mc2Context_->epHcclBuffer[rankId] + winStatusOffset_;
91+ }
92+ 
93+ __aicore__ inline uint32_t MIN(uint32_t x, uint32_t y)
94+ {
95+ return (x < y) ? x : y;
96+ }
97+ 
98+ TPipe *tpipe_{nullptr};
99+ GlobalTensor<ExpandXType> expandXGM_;
100+ GlobalTensor<bool> xActiveMaskGM_;
101+ GlobalTensor<int32_t> expertIdsGM_;
102+ GlobalTensor<int32_t> expandIdxGM_;
103+ GlobalTensor<int32_t> epSendCountGM_;
104+ GlobalTensor<int32_t> performanceInfoGM_;
105+ GlobalTensor<float> expertScalesGM_;
106+ GlobalTensor<ExpandXType> sharedExpertXGM_;
107+ GlobalTensor<ExpandXType> residualXGM_;
108+ GlobalTensor<ExpandXType> gammaGM_;
109+ GlobalTensor<ExpandXType> yOutGlobal_;
110+ GlobalTensor<float> rstdOutGlobal_;
111+ GlobalTensor<ExpandXType> expandOutGlobal_;
112+ GlobalTensor<ExpandXType> rankWindow_; // 用于存对端window的变量
113+ GlobalTensor<ExpandXType> rowTmpGlobal_;
114+ GlobalTensor<ExpandXType> oriXGM_;
115+ GlobalTensor<ExpandXType> constExpertAlpha1GM_;
116+ GlobalTensor<ExpandXType> constExpertAlpha2GM_;
117+ GlobalTensor<ExpandXType> constExpertVGM_;
118+ GlobalTensor<uint32_t> selfDataStatusGMTensor_;
119+ 
120+ GM_ADDR epWindowGM_;
121+ GM_ADDR stateGM_;
122+ GM_ADDR maskCalcWorkspaceGM_;
123+ GM_ADDR statusDataSpaceGm_;
124+ 
125+ __gm__ Mc2MoeContext* mc2Context_{nullptr};
126+ 
127+ LocalTensor<ExpandXType> gmTpSendCountTensor_;
128+ LocalTensor<ExpandXType> outTensor_;
129+ LocalTensor<float> winTpSendCountFloatTensor_;
130+ LocalTensor<float> gmTpSendCountFloatTensor_;
131+ LocalTensor<int32_t> performanceInfoTensor_;
132+ LocalTensor<int32_t> performanceInfoTmpTensor_;
133+ LocalTensor<int32_t> firstRecordTensor_;;
134+ LocalTensor<bool> maskStrideTensor_;
135+ LocalTensor<bool> maskGenerateTensor_;
136+ LocalTensor<uint32_t> dataStateLocalTensor_;
137+ LocalTensor<float> stateResetTensor_;
138+ 
139+ // tiling侧已确保数据上限， 相乘不会越界，因此统一采用uin32_t进行处理
140+ uint32_t axisBS_{0};
141+ uint32_t axisH_{0};
142+ uint32_t axisK_{0};
143+ uint32_t aivNum_{0};
144+ uint32_t epWorldSize_{0};
145+ uint32_t epWorldSizeOriginal_{0};
146+ uint32_t epRankId_{0};
147+ uint32_t epRankIdOriginal_{0};
148+ uint32_t coreIdx_{0}; // aiv id
149+ uint32_t sharedExpertNum_{0};
150+ uint32_t sharedExpertRankNum_{0};
151+ uint32_t rankNumPerShareExpert_{0};
152+ uint32_t moeExpertPerRankNum_{0}; // 每张卡部署的moe专家数
153+ uint32_t moeSendNum_{0}; // moeExpertPerRankNum_ * epWorldSize_
154+ uint32_t bufferNum_{0};
155+ uint32_t zeroExpertNum_{0};
156+ uint32_t copyExpertNum_{0};
157+ uint32_t constExpertNum_{0};
158+ uint32_t moeExpertNum_{0};
159+ uint32_t moeExpertOriginalNum_{0};
160+ uint32_t globalBS_{0};
161+ uint32_t bsKNum_{0};
162+ uint32_t startTokenId_{0};
163+ uint32_t endTokenId_{0};
164+ uint32_t sendCntNum_{0};
165+ uint32_t ubSize_{0};
166+ uint32_t dataState_{0};
167+ uint32_t stateOffset_{0};
168+ uint64_t activeMaskBsCnt_{0};
169+ uint64_t winStatusOffset_{0};
170+ uint64_t totalWinSizeEp_{0};
171+ uint64_t winDataSizeOffsetEp_{0};
172+ uint32_t selfSendCnt_{0};
173+ uint32_t tpRemoteSendCnt_{0};
174+ uint32_t activeMaskAlignSize_{0};
175+ uint32_t hExpandXTypeSize_{0};
176+ uint32_t hFloatAlign32Size_{0};
177+ uint32_t hFloatAlign256Size_{0};
178+ uint32_t hExpandXAlign32Size_{0};
179+ uint32_t hAlignWinSize_{0};
180+ uint32_t hAlignWinCnt_{0};
181+ uint32_t tokenScaleCnt_{0};
182+ uint32_t scaleNumAlignSize_{0};
183+ uint32_t flagRcvCount_{0};
184+ uint32_t axisBsAlignSize_{0};
185+ uint32_t expertScaleBeginIdx_{0};
186+ uint32_t performanceInfoSizeAlign_{0};
187+ uint32_t tokenNumPerCoreAlign_{0};
188+ float armAvgFactor_{0.0};
189+ float epsilon_{0.0};
190+ 
191+ TQueBind<QuePosition::VECIN, QuePosition::VECOUT, 1> moeQueue_;
192+ TQue<QuePosition::VECIN, 1> moeSumQueue_;
193+ TQueBind<QuePosition::VECIN, QuePosition::VECOUT, 1> gmTpSendCountQueue_;
194+ TQue<QuePosition::VECOUT, 1> xOutQueue_;
195+ TBuf<> readStateBuf_;
196+ TBuf<> expertScalesBuf_;
197+ TBuf<> rowTmpFloatBuf_;
198+ TBuf<> sumFloatBuf_;
199+ TBuf<> mulBuf_;
200+ TBuf<> indexCountsBuf_;
201+ TBuf<> winTpSendCountFloatBuf_;
202+ TBuf<> tokenBuf_;
203+ TBuf<> gammaBuf_;
204+ TBuf<TPosition::VECCALC> reduceFp32Buf_;
205+ TBuf<> xActMaskTBuf_;
206+ TBuf<> xActMaskCastTBuf_;
207+ TBuf<> tokenTargetTBuf_;
208+ TBuf<> validBsIndexTBuf_;
209+ TBuf<> xActMaskSumTBuf_;
210+ TBuf<> stateBuf_;
211+ TBuf<> stateResetBuf_;
212+ TBuf<> expertMaskBuf_;
213+ TBuf<> performanceInfoBuf_;
214+ TBuf<> performanceInfoTmpBuf_;
215+ TBuf<> firstRecordBuf_;
216+ bool isInputTokenMaskFlag_ = false;
217+ bool isInputExpertMaskFlag_ = false;
218+ bool hasSharedExpertX_ = false;
219+ bool isPerformanceFlag_ = false;
220+ bool isScalingDownFlag_ = false;
221+ bool isShareExpertRankFlag_ = false;
222+ bool enableSpecialExpert_ = false;
223+ 
224+ // int8量化
225+ TBuf<> xAbsBuf_;
226+ TBuf<> xMaxBuf_;
227+ TBuf<> xScaleMulBuf_;
228+ 
229+ LocalTensor<half> fp16CastTensor_;
230+ LocalTensor<float> absFloatTensor_;
231+ LocalTensor<float> reduceMaxFloatTensor_;
232+ LocalTensor<float> scaleDivFloatTensor_;
233+ LocalTensor<float> scaleDupLocalTensor_;
234+ LocalTensor<ExpandXType> sendLocalTensor_;
235+ LocalTensor<half> tokenTargetTensor_;
236+ LocalTensor<int32_t> validBsIndexTensor_;
237+ LocalTensor<bool> expertMaskTensor_;
238+ LocalTensor<float> expertScalesLocal_;
239+ LocalTensor<float> rowTmpFloatLocal_;
240+ LocalTensor<float> mulBufLocal_;
241+ LocalTensor<float> sumFloatBufLocal_;
242+ 
243+ uint32_t scaleNum_{0};
244+ MoeDistributeCombineQuant<CombineMC2TypeFunc> quantInst_;
245+};
246+ 
247+template <CombineMC2TypeClass>
248+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::TokenMaskCalCnt()
249+{
250+ // 一维mask, 计算得到有效bs数量
251+ LocalTensor<bool> xActiveMaskTensor = xActMaskTBuf_.Get<bool>();
252+ LocalTensor<half> tempTensor = xActMaskCastTBuf_.Get<half>();
253+ LocalTensor<half> sumOutTensor = xActMaskSumTBuf_.Get<half>();
254+ DataCopyExtParams xActiveMaskParams{1U, static_cast<uint32_t>(axisBS_ * sizeof(bool)), 0U, 0U, 0U};
255+ DataCopyPadExtParams<bool> xActiveMaskCopyPadParams{false, 0U, 0U, 0U};
256+ DataCopyPad(xActiveMaskTensor, xActiveMaskGM_, xActiveMaskParams, xActiveMaskCopyPadParams);
257+ SyncFunc<AscendC::HardEvent::MTE2_V>();
258+ LocalTensor<int8_t> xActiveMaskInt8Tensor = xActiveMaskTensor.ReinterpretCast<int8_t>();
259+ Cast(tempTensor, xActiveMaskInt8Tensor, RoundMode::CAST_NONE, axisBS_);
260+ PipeBarrier<PIPE_V>();
261+ SumParams params{1, axisBsAlignSize_, axisBS_};
262+ Sum(sumOutTensor, tempTensor, params);
263+ SyncFunc<AscendC::HardEvent::V_S>();
264+ activeMaskBsCnt_ = static_cast<int32_t>(sumOutTensor.GetValue(0));
265+}
266+ 
267+template <CombineMC2TypeClass>
268+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::ExpertMaskCalCnt()
269+{
270+ // 二维mask, 挑选有效token
271+ uint64_t rsvdCnt = 0;
272+ uint32_t mask = axisBS_;
273+ LocalTensor<bool> maskStrideTensor = tokenBuf_.Get<bool>();
274+ LocalTensor<half> tempTensor = rowTmpFloatBuf_.Get<half>();
275+ LocalTensor<half> maskTempTensor = sumFloatBuf_.Get<half>();
276+ LocalTensor<uint8_t> maskTensor = tokenBuf_.Get<uint8_t>();
277+ LocalTensor<int32_t> bsIndexTensor = mulBuf_.Get<int32_t>();
278+ LocalTensor<uint32_t> maskTensorInt32 = tokenBuf_.Get<uint32_t>();
279+ DataCopyExtParams xActiveMaskParams{
280+ static_cast<uint16_t>(axisBS_), static_cast<uint32_t>(axisK_ * sizeof(bool)), 0U, 0U, 0U};
281+ DataCopyPadExtParams<bool> xActiveMaskCopyPadParams{false, 0U, 0U, 0U};
282+ SumParams axisBsSumParams{
283+ 1, static_cast<uint32_t>(Ceil(axisBS_ * sizeof(half), UB_ALIGN) * UB_ALIGN / sizeof(half)), axisBS_};
284+ uint32_t calCnt = Ceil(axisBS_ * sizeof(half), ALIGNED_LEN) * ALIGNED_LEN / sizeof(half);
285+ 
286+ Duplicate<half>(maskTempTensor, (half)0, calCnt);
287+ DataCopyPad(maskStrideTensor, xActiveMaskGM_, xActiveMaskParams, xActiveMaskCopyPadParams);
288+ SyncFunc<AscendC::HardEvent::MTE2_V>();
289+ LocalTensor<int8_t> maskStrideInt8Tensor = maskStrideTensor.ReinterpretCast<int8_t>();
290+ Cast(tempTensor, maskStrideInt8Tensor, RoundMode::CAST_NONE, activeMaskAlignSize_);
291+ PipeBarrier<PIPE_V>();
292+ uint32_t innerAlign = Ceil(axisK_ * sizeof(half), UB_ALIGN) * UB_ALIGN / sizeof(half) * BUFFER_NUM;
293+ SumParams axisKSumParams{axisBS_, innerAlign, axisK_};
294+ Sum(tokenTargetTensor_, tempTensor, axisKSumParams);
295+ PipeBarrier<PIPE_V>();
296+ Mins(maskTempTensor, tokenTargetTensor_, static_cast<half>(1), axisBS_);
297+ PipeBarrier<PIPE_V>();
298+ CompareScalar(maskTensor, maskTempTensor, static_cast<half>(1), AscendC::CMPMODE::EQ, calCnt);
299+ CreateVecIndex(bsIndexTensor, 0, axisBS_);
300+ PipeBarrier<PIPE_V>();
301+ GatherMask(validBsIndexTensor_, bsIndexTensor, maskTensorInt32, true, mask, {1, 1, 0, 0}, activeMaskBsCnt_);
302+}
303+ 
304+template <CombineMC2TypeClass>
305+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::InitInputAndOutput(
306+ GM_ADDR residualX, GM_ADDR gamma, GM_ADDR expandX, GM_ADDR expertIds, GM_ADDR expandIdx, GM_ADDR epSendCount, GM_ADDR expertScales,
307+ GM_ADDR xActiveMask, GM_ADDR sharedExpertX, GM_ADDR oriX, GM_ADDR constExpertAlpha1,
308+ GM_ADDR constExpertAlpha2, GM_ADDR constExpertV, GM_ADDR performanceInfo, GM_ADDR XOut)
309+{
310+ expandXGM_.SetGlobalBuffer((__gm__ ExpandXType*)expandX);
311+ expertIdsGM_.SetGlobalBuffer((__gm__ int32_t*)expertIds);
312+ expandIdxGM_.SetGlobalBuffer((__gm__ int32_t*)expandIdx);
313+ epSendCountGM_.SetGlobalBuffer((__gm__ int32_t*)epSendCount);
314+ expertScalesGM_.SetGlobalBuffer((__gm__ float*)expertScales);
315+ xActiveMaskGM_.SetGlobalBuffer((__gm__ bool*)xActiveMask);
316+ sharedExpertXGM_.SetGlobalBuffer((__gm__ ExpandXType*)sharedExpertX);
317+ oriXGM_.SetGlobalBuffer((__gm__ ExpandXType*)oriX);
318+ constExpertAlpha1GM_.SetGlobalBuffer((__gm__ ExpandXType*)constExpertAlpha1);
319+ constExpertAlpha2GM_.SetGlobalBuffer((__gm__ ExpandXType*)constExpertAlpha2);
320+ constExpertVGM_.SetGlobalBuffer((__gm__ ExpandXType*)constExpertV);
321+ performanceInfoGM_.SetGlobalBuffer((__gm__ int32_t*)performanceInfo);
322+ 
323+ expandOutGlobal_.SetGlobalBuffer((__gm__ ExpandXType*)XOut);
324+}
325+ 
326+template <CombineMC2TypeClass>
327+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::InitTilingAttrs(const MoeDistributeCombineV2Info tilingData)
328+{
329+ axisBS_ = tilingData.bs;
330+ axisH_ = tilingData.h;
331+ axisK_ = tilingData.k;
332+ aivNum_ = tilingData.aivNum;
333+ ubSize_ = tilingData.totalUbSize;
334+ globalBS_ = tilingData.globalBs;
335+ isPerformanceFlag_ = tilingData.isPerformance;
336+ epWorldSizeOriginal_ = tilingData.epWorldSize;
337+ epRankId_ = tilingData.epRankId;
338+ epRankIdOriginal_ = tilingData.epRankId;
339+ epWorldSize_ = tilingData.epWorldSize;
340+ moeExpertPerRankNum_ = tilingData.moeExpertPerRankNum;
341+ totalWinSizeEp_ = tilingData.totalWinSizeEp;
342+ isInputTokenMaskFlag_ = tilingData.isTokenMask;
343+ isInputExpertMaskFlag_ = tilingData.isExpertMask;
344+ hasSharedExpertX_ = tilingData.hasSharedExpertX;
345+ bufferNum_ = tilingData.bufferNum;
346+ zeroExpertNum_ = tilingData.zeroExpertNum;
347+ copyExpertNum_ = tilingData.copyExpertNum;
348+ constExpertNum_ = tilingData.constExpertNum;
349+ moeExpertNum_ = tilingData.moeExpertNum;
350+ moeExpertOriginalNum_ = tilingData.moeExpertNum;
351+ sharedExpertRankNum_ = tilingData.sharedExpertRankNum;
352+ enableSpecialExpert_ = (constExpertNum_ + zeroExpertNum_ + copyExpertNum_ > 0U);
353+}
354+ 
355+template <CombineMC2TypeClass>
356+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::InitAttrs(GM_ADDR mc2Context, const MoeDistributeCombineV2Info tilingData)
357+{
358+ InitTilingAttrs(tilingData);
359+ uint32_t epRankIdHccl{0};
360+ uint32_t epWorldSizeHccl{0};
361+ //Using Mc2Context instead of hccl context
362+ mc2Context_ = (__gm__ Mc2MoeContext*)mc2Context;
363+ epRankIdHccl = mc2Context_->epRankId;
364+ epWorldSizeHccl = tilingData.epWorldSize;
365+ statusDataSpaceGm_ = (GM_ADDR)(mc2Context_->epHcclBuffer[epRankIdHccl]);
366+ selfDataStatusGMTensor_.SetGlobalBuffer((__gm__ uint32_t*)(statusDataSpaceGm_ + COMBINE_STATE_WIN_OFFSET + coreIdx_ * WIN_ADDR_ALIGN));
367+ TBuf<> dataStateBuf;
368+ tpipe_->InitBuffer(dataStateBuf, UB_ALIGN);
369+ dataState_ = InitWinState(selfDataStatusGMTensor_, epRankIdHccl, epWorldSizeHccl, epRankIdOriginal_, moeExpertNum_, epWorldSizeOriginal_, globalBS_, dataStateBuf);
370+ 
371+ sharedExpertNum_ = tilingData.sharedExpertNum;
372+ moeSendNum_ = epWorldSize_ * moeExpertPerRankNum_;
373+ if (epRankId_ < sharedExpertRankNum_) {
374+ isShareExpertRankFlag_ = true;
375+ }
376+ 
377+ rankNumPerShareExpert_ = sharedExpertRankNum_ / sharedExpertNum_;
378+ 
379+ stateOffset_ = STATE_OFFSET;
380+ uint32_t hFloatSize = axisH_ * static_cast<uint32_t>(sizeof(float));
381+ hFloatAlign32Size_ = Ceil(hFloatSize, UB_ALIGN) * UB_ALIGN;
382+ hFloatAlign256Size_ = Ceil(hFloatSize, ALIGNED_LEN) * ALIGNED_LEN;
383+ hExpandXTypeSize_ = axisH_ * sizeof(ExpandXType);
384+ hExpandXAlign32Size_ = Ceil(hExpandXTypeSize_, UB_ALIGN) * UB_ALIGN;
385+ hAlignWinSize_ = Ceil(hExpandXTypeSize_, WIN_ADDR_ALIGN) * WIN_ADDR_ALIGN;
386+ hAlignWinCnt_ = hAlignWinSize_ / sizeof(ExpandXType);
387+ bsKNum_ = axisBS_ * axisK_;
388+}
389+ 
390+template <CombineMC2TypeClass>
391+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::Init(
392+ GM_ADDR mc2Context, GM_ADDR expandX, GM_ADDR expertIds, GM_ADDR expandIdx,
393+ GM_ADDR epSendCount, GM_ADDR residualX, GM_ADDR gamma, GM_ADDR expertScales,
394+ GM_ADDR xActiveMask, GM_ADDR sharedExpertX, GM_ADDR oriX,
395+ GM_ADDR constExpertAlpha1, GM_ADDR constExpertAlpha2, GM_ADDR constExpertV, GM_ADDR performanceInfo, GM_ADDR XOut,
396+ GM_ADDR workspaceGM, MoeDistributeCombineV2Info tilingData, TPipe *pipe)
397+{
398+ tpipe_ = pipe;
399+ 
400+ coreIdx_ = GetBlockIdx();
401+ 
402+ maskCalcWorkspaceGM_ = workspaceGM + coreIdx_ * MASK_CALC_NEED_WORKSPACE;
403+ 
404+ InitInputAndOutput(
405+ residualX, gamma, expandX, expertIds, expandIdx, epSendCount, expertScales, xActiveMask, sharedExpertX, oriX, constExpertAlpha1,
406+ constExpertAlpha2, constExpertV, performanceInfo, XOut);
407+ 
408+ // 检查hcclwinsize是否越界
409+ mc2Context_ = (__gm__ Mc2MoeContext*)mc2Context;
410+ 
411+ InitAttrs(mc2Context, tilingData);
412+ 
413+ if constexpr (IsInt8Quant) {
414+ quantInst_.SetQuantInitParams(axisH_);
415+ quantInst_.InitInt8Quant(scaleNum_, hExpandXAlign32Size_, hFloatAlign256Size_, tokenScaleCnt_);
416+ }
417+ 
418+ PipeBarrier<PIPE_ALL>();
419+ 
420+ // 当前win区划分为前后两半区，连续两次dispatch，切换半区
421+ winDataSizeOffsetEp_ = static_cast<uint64_t>(dataState_) * (tilingData.totalWinSizeEp / 2UL);
422+ winStatusOffset_ = COMBINE_STATE_OFFSET + dataState_ * WIN_STATE_OFFSET; // 前面的预留给dispatch使用
423+ epWindowGM_ = GetWinAddrByRankId(epRankIdOriginal_, EP_DOMAIN);
424+#if defined(ASCENDC_OOM) && ASCENDC_OOM == 1
425+ for (int tempepRankId = 0; tempepRankId < epWorldSize_; tempepRankId++) {
426+ OOMCheckAddrRange<ExpandXType>((__gm__ ExpandXType*)(GetWinAddrByRankId(tempepRankId, EP_DOMAIN)), totalWinSizeEp_);
427+ OOMCheckAddrRange<float>((__gm__ float*)(GetWinStateAddrByRankId(tempepRankId, EP_DOMAIN)), STATE_SIZE);
428+ }
429+#endif
430+ if (isShareExpertRankFlag_) {
431+ DataCacheCleanAndInvalid<int32_t, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(epSendCountGM_[epWorldSize_ - 1]);
432+ selfSendCnt_ = epSendCountGM_(epWorldSize_ - 1);
433+ } else {
434+ DataCacheCleanAndInvalid<int32_t, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(epSendCountGM_[moeSendNum_ - 1]);
435+ selfSendCnt_ = epSendCountGM_(moeSendNum_ - 1);
436+ }
437+ SplitCoreCal();
438+ tpipe_->InitBuffer(moeQueue_, BUFFER_NUM, hExpandXAlign32Size_);
439+ flagRcvCount_ = axisK_ + sharedExpertNum_;
440+}
441+ 
442+template <CombineMC2TypeClass>
443+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::BuffInit()
444+{
445+ tpipe_->Reset();
446+ tpipe_->InitBuffer(readStateBuf_, UB_ALIGN); // 32
447+ tpipe_->InitBuffer(gmTpSendCountQueue_, BUFFER_NUM, hExpandXAlign32Size_); // 28K 存储搬入token
448+ if constexpr (IsInt8Quant) {
449+ uint32_t tokenScaleAlign32Size = Ceil(tokenScaleCnt_ * sizeof(ExpandXType), UB_ALIGN) * UB_ALIGN;
450+ tpipe_->InitBuffer(xOutQueue_, BUFFER_NUM, tokenScaleAlign32Size); // 28K 输出token搬运
451+ tpipe_->InitBuffer(xAbsBuf_, hFloatAlign256Size_); // 28K blockReduceMax计算及后续Cast计算，256对齐
452+ uint32_t hFloatAlign256Cnt = hFloatAlign256Size_ / sizeof(float);
453+ tpipe_->InitBuffer(xMaxBuf_, (hFloatAlign256Cnt / REDUCE_NUM) * sizeof(float)); // 3.5K 存储ReduceMax结果
454+ tpipe_->InitBuffer(xScaleMulBuf_, hFloatAlign256Size_); // 28K 参与Brcb计算，256对齐
455+ tpipe_->InitBuffer(winTpSendCountFloatBuf_, hFloatAlign32Size_); // 28K 参与Div等token v核运算
456+ 
457+ winTpSendCountFloatTensor_ = winTpSendCountFloatBuf_.Get<float>();
458+ absFloatTensor_ = xAbsBuf_.Get<float>();
459+ reduceMaxFloatTensor_ = xMaxBuf_.Get<float>();
460+ scaleDupLocalTensor_ = xScaleMulBuf_.Get<float>();
461+ fp16CastTensor_ = xAbsBuf_.Get<half>();
462+ Duplicate(absFloatTensor_, float(0), hFloatAlign256Cnt); // 统一写0
463+ }
464+ 
465+ tpipe_->InitBuffer(indexCountsBuf_, sendCntNum_ * EXPAND_IDX_INFO * sizeof(int32_t));
466+}
467+ 
468+template <CombineMC2TypeClass>
469+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::MaskAlign()
470+{
471+ // 扩展后的二维mask通过GM对齐内轴元素个数
472+ uint32_t calcCnt = Ceil(axisBS_ * axisK_ * sizeof(half), ALIGNED_LEN_256) * ALIGNED_LEN_256 / sizeof(half);
473+ GlobalTensor<bool> MaskGMTensor;
474+ MaskGMTensor.SetGlobalBuffer((__gm__ bool*)maskCalcWorkspaceGM_);
475+ DataCopyExtParams maskCalcParams = {1U, static_cast<uint32_t>(calcCnt * sizeof(bool)), 0U, 0U, 0U};
476+ SyncFunc<AscendC::HardEvent::V_MTE3>();
477+ DataCopyPad(MaskGMTensor, maskGenerateTensor_, maskCalcParams);
478+ SyncFunc<AscendC::HardEvent::MTE3_MTE2>();
479+ DataCopyExtParams xActiveMaskParams{
480+ static_cast<uint16_t>(axisBS_), static_cast<uint32_t>(axisK_ * sizeof(bool)), 0U, 0U, 0U};
481+ DataCopyPadExtParams<bool> xActiveMaskCopyPadParams{true, 0U, static_cast<uint8_t>(UB_ALIGN - axisK_), 0U};
482+ DataCopyPad(maskStrideTensor_, MaskGMTensor, xActiveMaskParams, xActiveMaskCopyPadParams);
483+ SyncFunc<AscendC::HardEvent::MTE2_V>();
484+ SyncFunc<AscendC::HardEvent::MTE2_S>();
485+}
486+ 
487+template <CombineMC2TypeClass>
488+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::GenerateActiveMask(half val)
489+{
490+ maskStrideTensor_ = tokenBuf_.Get<bool>();
491+ LocalTensor<half> maskCalcTensor = tokenBuf_.Get<half>();
492+
493+ if (isInputTokenMaskFlag_) {
494+ // 根据一维场景下的activeMaskBsCnt_，构造出二维mask
495+ uint32_t calcCnt = Ceil(axisBS_ * axisK_ * sizeof(half), ALIGNED_LEN_256) * ALIGNED_LEN_256 / sizeof(half);
496+ Duplicate<half>(maskCalcTensor, static_cast<half>(0), calcCnt);
497+ PipeBarrier<PIPE_V>();
498+ uint32_t activeCalcCnt = Ceil(activeMaskBsCnt_ * axisK_ * sizeof(half), ALIGNED_LEN_256) * ALIGNED_LEN_256 / sizeof(half);
499+ Duplicate<half>(maskCalcTensor, static_cast<half>(1), activeCalcCnt);
500+ PipeBarrier<PIPE_V>();
501+ Cast(maskGenerateTensor_.ReinterpretCast<uint8_t>(), maskCalcTensor, RoundMode::CAST_NONE, calcCnt);
502+ } else {
503+ // 构造二维全true的mask
504+ uint32_t calcCnt = Ceil(axisBS_ * axisK_ * sizeof(half), ALIGNED_LEN_256) * ALIGNED_LEN_256 / sizeof(half);
505+ Duplicate<half>(maskCalcTensor, val, calcCnt);
506+ PipeBarrier<PIPE_V>();
507+ Cast(maskGenerateTensor_.ReinterpretCast<uint8_t>(), maskCalcTensor, RoundMode::CAST_NONE, calcCnt);
508+ }
509+}
510+ 
511+template <CombineMC2TypeClass>
512+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::MaskSpecialExpert()
513+{
514+ LocalTensor<int32_t> expertIdsTensor_ = mulBuf_.Get<int32_t>();
515+ LocalTensor<float> expertIdsFloat = rowTmpFloatBuf_.Get<float>();
516+ LocalTensor<uint8_t> maskTensor = mulBuf_.Get<uint8_t>();
517+ LocalTensor<half> maskCalcTensor = tokenBuf_.Get<half>();
518+ LocalTensor<half> maskCalcSelectedTensor = rowTmpFloatBuf_.Get<half>();
519+ maskStrideTensor_ = tokenBuf_.Get<bool>();
520+ LocalTensor<half> tempTensor = rowTmpFloatBuf_.Get<half>();
521+ 
522+ // 拷入expertIds
523+ uint32_t mask = axisBS_ * axisK_;
524+ DataCopyExtParams expertIdsCntParams = {1U, static_cast<uint32_t>(mask * sizeof(int32_t)), 0U, 0U, 0U};
525+ DataCopyPadExtParams<int32_t> expertIdsCntCopyPadParams{false, 0U, 0U, 0U};
526+ DataCopyPad(expertIdsTensor_, expertIdsGM_, expertIdsCntParams, expertIdsCntCopyPadParams);
527+ SyncFunc<AscendC::HardEvent::MTE2_V>();
528+ SyncFunc<AscendC::HardEvent::MTE2_S>();
529+ 
530+ // 根据expertId小于moeExpertNum，得到考虑特殊专家后的mask
531+ uint32_t calcCnt = Ceil(mask * sizeof(int32_t), ALIGNED_LEN_256) * ALIGNED_LEN_256 / sizeof(int32_t);
532+ Cast(expertIdsFloat, expertIdsTensor_, RoundMode::CAST_NONE, calcCnt);
533+ PipeBarrier<PIPE_V>();
534+ int32_t moeExpertNumInt32 = static_cast<int32_t>(moeExpertOriginalNum_);
535+ CompareScalar(maskTensor, expertIdsFloat, static_cast<float>(moeExpertNumInt32), AscendC::CMPMODE::LT, calcCnt);
536+ PipeBarrier<PIPE_V>();
537+ if (isInputExpertMaskFlag_) {
538+ Cast(maskCalcTensor, expertMaskTensor_.ReinterpretCast<uint8_t>(), RoundMode::CAST_NONE, calcCnt);
539+ } else {
540+ Cast(maskCalcTensor, maskGenerateTensor_.ReinterpretCast<uint8_t>(), RoundMode::CAST_NONE, calcCnt);
541+ }
542+ PipeBarrier<PIPE_V>();
543+ Select(
544+ maskCalcSelectedTensor, maskTensor, maskCalcTensor, static_cast<half>(0), SELMODE::VSEL_TENSOR_SCALAR_MODE,
545+ calcCnt);
546+ PipeBarrier<PIPE_V>();
547+ Cast(maskGenerateTensor_.ReinterpretCast<uint8_t>(), maskCalcSelectedTensor, RoundMode::CAST_NONE, calcCnt);
548+
549+ // 通过GM对齐内轴元素个数
550+ MaskAlign();
551+ 
552+ // 更新考虑特殊专家后的
553+ uint32_t calCnt = Ceil(axisBS_ * sizeof(half), ALIGNED_LEN) * ALIGNED_LEN / sizeof(half);
554+ LocalTensor<int8_t> maskStrideInt8Tensor = maskStrideTensor_.ReinterpretCast<int8_t>();
555+ activeMaskAlignSize_ = axisBS_ * (Ceil(axisK_ * sizeof(bool), UB_ALIGN) * UB_ALIGN);
556+ Cast(tempTensor, maskStrideInt8Tensor, RoundMode::CAST_NONE, activeMaskAlignSize_);
557+ PipeBarrier<PIPE_V>();
558+ uint32_t innerAlign = Ceil(axisK_ * sizeof(half), UB_ALIGN) * UB_ALIGN / sizeof(half) * BUFFER_NUM;
559+ SumParams axisKSumParams{axisBS_, innerAlign, axisK_};
560+ Sum(tokenTargetTensor_, tempTensor, axisKSumParams);
561+ PipeBarrier<PIPE_V>();
562+ SyncFunc<AscendC::HardEvent::V_S>();
563+}
564+ 
565+template <CombineMC2TypeClass>
566+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::AlltoAllBuffInitAndMaskCal()
567+{
568+ tpipe_->Reset();
569+ activeMaskBsCnt_ = axisBS_;
570+ activeMaskAlignSize_ = axisBS_ * (Ceil(axisK_ * sizeof(bool), UB_ALIGN) * UB_ALIGN);
571+ uint32_t maxSizeTokenBuf = hExpandXAlign32Size_;
572+ uint32_t maxSizeRowTmpFloatBuf = hFloatAlign32Size_;
573+ uint32_t bsKFloatAlign = Ceil(bsKNum_ * sizeof(float), UB_ALIGN) * UB_ALIGN;
574+ uint32_t mulBufSize = hFloatAlign256Size_ > bsKFloatAlign ? hFloatAlign256Size_ : bsKFloatAlign;
575+ if (isInputExpertMaskFlag_ || enableSpecialExpert_) {
576+ uint32_t activeMaskAlignHalfSize = activeMaskAlignSize_ * sizeof(half);
577+ maxSizeTokenBuf = (activeMaskAlignSize_ > hExpandXAlign32Size_ ? activeMaskAlignSize_ : hExpandXAlign32Size_);
578+ maxSizeRowTmpFloatBuf = (activeMaskAlignHalfSize > hFloatAlign32Size_ ? activeMaskAlignHalfSize : hFloatAlign32Size_);
579+ }
580+ // InitBuffer需要在tiling中计算ub总量
581+ tpipe_->InitBuffer(tokenBuf_, maxSizeTokenBuf); // 16K 用于搬入输入token
582+ tpipe_->InitBuffer(rowTmpFloatBuf_, maxSizeRowTmpFloatBuf); // 32K 用于存储cast之后的fp32 token数据
583+ tpipe_->InitBuffer(mulBuf_, mulBufSize); // 32K buffer复用， 最大用于存储Brcb之后的token，需要256对齐
584+ tpipe_->InitBuffer(sumFloatBuf_, hFloatAlign32Size_); // 32K add
585+ tpipe_->InitBuffer(moeSumQueue_, bufferNum_, hExpandXAlign32Size_); // 32K 搬入
586+ tpipe_->InitBuffer(stateBuf_, (flagRcvCount_) * STATE_OFFSET);
587+ tpipe_->InitBuffer(stateResetBuf_, (flagRcvCount_) * STATE_OFFSET); // 清理状态区
588+ stateResetTensor_ = stateResetBuf_.Get<float>();
589+ Duplicate<float>(stateResetTensor_, (float)0.0, static_cast<uint32_t>(flagRcvCount_ * FLOAT_PER_UB_ALIGN));
590+ SyncFunc<AscendC::HardEvent::V_MTE3>();
591+ if constexpr (IsInt8Quant) {
592+ scaleNumAlignSize_ = Ceil(scaleNum_ * sizeof(float), UB_ALIGN) * UB_ALIGN;
593+ tpipe_->InitBuffer(xAbsBuf_, scaleNumAlignSize_);
594+ fp16CastTensor_ = mulBuf_.Get<half>();
595+ absFloatTensor_ = rowTmpFloatBuf_.Get<float>();
596+ scaleDupLocalTensor_ = mulBuf_.Get<float>();
597+ scaleDivFloatTensor_ = xAbsBuf_.Get<float>();
598+ }
599+ if (isInputTokenMaskFlag_) {
600+ axisBsAlignSize_ = Ceil(axisBS_ * sizeof(bool), UB_ALIGN) * UB_ALIGN;
601+ tpipe_->InitBuffer(xActMaskTBuf_, axisBsAlignSize_);
602+ tpipe_->InitBuffer(xActMaskCastTBuf_, axisBsAlignSize_ * sizeof(half));
603+ tpipe_->InitBuffer(xActMaskSumTBuf_, axisBsAlignSize_ * sizeof(half));
604+ TokenMaskCalCnt(); // 计算一维mask
605+ }
606+ if (isInputExpertMaskFlag_) {
607+ tpipe_->InitBuffer(tokenTargetTBuf_, Ceil(axisBS_ * sizeof(half), UB_ALIGN) * UB_ALIGN);
608+ tpipe_->InitBuffer(validBsIndexTBuf_, Ceil(axisBS_ * sizeof(int32_t), UB_ALIGN) * UB_ALIGN);
609+ tpipe_->InitBuffer(expertMaskBuf_, Ceil(axisBS_ * axisK_ * sizeof(bool), UB_ALIGN) * UB_ALIGN);
610+ tokenTargetTensor_ = tokenTargetTBuf_.Get<half>();
611+ validBsIndexTensor_ = validBsIndexTBuf_.Get<int32_t>();
612+ ExpertMaskCalCnt(); // 计算二维mask
613+ expertMaskTensor_ = expertMaskBuf_.Get<bool>();
614+ DataCopyPadExtParams<bool> maskCopyPadParams{false, 0U, 0U, 0U};
615+ DataCopyExtParams maskParams{1U, static_cast<uint32_t>(axisBS_ * axisK_ * sizeof(bool)), 0U, 0U, 0U};
616+ DataCopyPad(expertMaskTensor_, xActiveMaskGM_, maskParams, maskCopyPadParams);
617+ SyncFunc<AscendC::HardEvent::V_S>();
618+ }
619+ if (enableSpecialExpert_) {
620+ maskGenerateTensor_ = sumFloatBuf_.Get<bool>();
621+ if (!isInputExpertMaskFlag_) {
622+ tpipe_->InitBuffer(tokenTargetTBuf_, Ceil(axisBS_ * sizeof(half), UB_ALIGN) * UB_ALIGN);
623+ tokenTargetTensor_ = tokenTargetTBuf_.Get<half>();
624+ GenerateActiveMask(static_cast<half>(1));
625+ }
626+ MaskSpecialExpert();
627+ }
628+ if (isPerformanceFlag_) {
629+ uint32_t performanceInfoSize = JUMP_WRITE * epWorldSizeOriginal_ * sizeof(int32_t);
630+ performanceInfoSizeAlign_ = Ceil(performanceInfoSize, UB_ALIGN) * UB_ALIGN;
631+ tpipe_->InitBuffer(performanceInfoBuf_, performanceInfoSizeAlign_);
632+ performanceInfoTensor_ = performanceInfoBuf_.Get<int32_t>();
633+ Duplicate<int32_t>(performanceInfoTensor_, 0, JUMP_WRITE * epWorldSizeOriginal_);
634+ }
635+}
636+ 
637+template <CombineMC2TypeClass>
638+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::SplitCoreCal()
639+{
640+ // 对需要发送的token数平均分核，得到每个核上处理的卡的数量
641+ sendCntNum_ = selfSendCnt_ / aivNum_;
642+ uint32_t remainderRankNum = selfSendCnt_ % aivNum_;
643+
644+ startTokenId_ = sendCntNum_ * coreIdx_;
645+
646+ if (coreIdx_ < remainderRankNum) {
647+ sendCntNum_++;
648+ startTokenId_ += coreIdx_;
649+ } else {
650+ startTokenId_ += remainderRankNum;
651+ }
652+ endTokenId_ = startTokenId_ + sendCntNum_;
653+}
654+ 
655+// 流水流程
656+// 46 -> gm -> ub syncall win->gm add -> alltoall
657+// 2 -> win wait syncall gm -> ub win ->gm add -> alltoall
658+template <CombineMC2TypeClass>
659+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::SetWaitTpStatusAndDisPatch()
660+{
661+ PipeBarrier<PIPE_ALL>();
662+ if ((coreIdx_ >= tpRemoteSendCnt_) && (coreIdx_ >= selfSendCnt_)) {
663+ return;
664+ }
665+ 
666+ // Copy win gm->ub add ->alltoall send
667+ ExpertAlltoAllDispatchCopyAdd();
668+ SyncFunc<AscendC::HardEvent::MTE3_S>();
669+}
670+ 
671+template <CombineMC2TypeClass>
672+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::ExpertAlltoAllDispatchCopyAdd()
673+{
674+ if (sendCntNum_ == 0U) { // 空闲核，直接返回
675+ return;
676+ }
677+ 
678+ LocalTensor<int32_t> expandIdxLocal = indexCountsBuf_.Get<int32_t>();
679+ const DataCopyExtParams bskParams{1U, static_cast<uint32_t>(sendCntNum_ * EXPAND_IDX_INFO * sizeof(uint32_t)), 0U,
680+ 0U, 0U};
681+ const DataCopyPadExtParams<int32_t> copyPadParams{false, 0U, 0U, 0U};
682+ DataCopyPad(expandIdxLocal, expandIdxGM_[startTokenId_ * EXPAND_IDX_INFO], bskParams, copyPadParams);
683+ LocalTensor<float> statusTensor = readStateBuf_.AllocTensor<float>();
684+ Duplicate<float>(statusTensor, (float)1, FLOAT_PER_UB_ALIGN);
685+ SyncFunc<AscendC::HardEvent::V_MTE3>();
686+ SyncFunc<AscendC::HardEvent::MTE2_S>();
687+ for (uint32_t loop = 0; loop < sendCntNum_; loop++) {
688+ uint32_t tkIndex = startTokenId_ + ((loop + epRankId_) % sendCntNum_); // 错位发送
689+ uint32_t baseOffset = (tkIndex - startTokenId_) * EXPAND_IDX_INFO;
690+ uint32_t rankIdExpandIdx = static_cast<uint32_t>(expandIdxLocal(baseOffset)); // 位置0是rank_id
691+ uint32_t toRankId = rankIdExpandIdx; // 位置0是rank_id
692+ uint32_t tokenId = static_cast<uint32_t>(expandIdxLocal(baseOffset + 1)); // 位置1是token_id
693+ uint32_t topkId = static_cast<uint32_t>(expandIdxLocal(baseOffset + 2)); // 位置2是topk_id
694+ 
695+ ExpertAlltoAllDispatchInnerCopyAdd(toRankId, tokenId, topkId, tkIndex);
696+ PipeBarrier<PIPE_MTE3>();
697+ GM_ADDR stateGM = GetWinStateAddrByRankId(toRankId, EP_DOMAIN) + tokenId * flagRcvCount_ * stateOffset_ +
698+ topkId * stateOffset_; // 计算地址偏移
699+ GlobalTensor<float> stateGMTensor;
700+ stateGMTensor.SetGlobalBuffer((__gm__ float*)stateGM);
701+ DataCopy<float>(stateGMTensor, statusTensor, FLOAT_PER_UB_ALIGN); // 8是数据大小，按32对齐拷贝
702+ }
703+}
704+ 
705+template <CombineMC2TypeClass>
706+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::ExpertAlltoAllDispatchInnerCopyAdd(
707+ uint32_t toRankId, uint32_t tokenId, uint32_t topkId, uint32_t tkIndex)
708+{
709+ uint32_t dataCnt = axisH_;
710+ uint32_t epOffset = tokenId * (axisK_ + sharedExpertNum_) + topkId;
711+ uint32_t tokenGMOffset = tkIndex * axisH_;
712+ uint32_t tokenWinOffset = tkIndex * hAlignWinCnt_;
713+ GM_ADDR rankGM = GetWinAddrByRankId(toRankId, EP_DOMAIN) + epOffset * hAlignWinSize_;
714+ rankWindow_.SetGlobalBuffer((__gm__ ExpandXType*)rankGM);
715+ DataCopyPadExtParams<ExpandXType> copyPadExtParams{false, 0U, 0U, 0U};
716+ DataCopyExtParams expandXCopyParams{1U, static_cast<uint32_t>(hExpandXTypeSize_), 0U, 0U, 0U};
717+ DataCopyExtParams xScaleCopyParams{1U, static_cast<uint32_t>(tokenScaleCnt_ * sizeof(ExpandXType)), 0U, 0U, 0U};
718+ if constexpr (IsInt8Quant) {
719+ gmTpSendCountTensor_ = gmTpSendCountQueue_.AllocTensor<ExpandXType>();
720+ DataCopyPad(gmTpSendCountTensor_, expandXGM_[tokenGMOffset], expandXCopyParams, copyPadExtParams);
721+ gmTpSendCountQueue_.EnQue(gmTpSendCountTensor_);
722+ gmTpSendCountTensor_ = gmTpSendCountQueue_.DeQue<ExpandXType>();
723+ sendLocalTensor_ = xOutQueue_.AllocTensor<ExpandXType>();
724+ quantInst_.Int8QuantProcess(sendLocalTensor_, winTpSendCountFloatTensor_, gmTpSendCountTensor_,
725+ fp16CastTensor_, absFloatTensor_, reduceMaxFloatTensor_, scaleDupLocalTensor_);
726+ xOutQueue_.EnQue(sendLocalTensor_);
727+ sendLocalTensor_ = xOutQueue_.DeQue<ExpandXType>();
728+ DataCopyPad(rankWindow_, sendLocalTensor_, xScaleCopyParams);
729+ gmTpSendCountQueue_.FreeTensor<ExpandXType>(gmTpSendCountTensor_);
730+ xOutQueue_.FreeTensor<ExpandXType>(sendLocalTensor_);
731+ } else {
732+ gmTpSendCountTensor_ = gmTpSendCountQueue_.AllocTensor<ExpandXType>();
733+ DataCopyPad(gmTpSendCountTensor_, expandXGM_[tokenGMOffset], expandXCopyParams, copyPadExtParams);
734+ gmTpSendCountQueue_.EnQue(gmTpSendCountTensor_);
735+ gmTpSendCountTensor_ = gmTpSendCountQueue_.DeQue<ExpandXType>();
736+ DataCopyPad(rankWindow_, gmTpSendCountTensor_, expandXCopyParams);
737+ gmTpSendCountQueue_.FreeTensor<ExpandXType>(gmTpSendCountTensor_);
738+ }
739+}
740+ 
741+template <CombineMC2TypeClass>
742+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::CustomAdd(LocalTensor<ExpandXType> &dst,
743+ LocalTensor<ExpandXType> &src0, LocalTensor<ExpandXType> &src1)
744+{
745+ if constexpr (AscendC::IsSameType<ExpandXType, bfloat16_t>::value) {
746+ Cast(winTpSendCountFloatTensor_, src0, RoundMode::CAST_NONE, axisH_);
747+ Cast(gmTpSendCountFloatTensor_, src1, RoundMode::CAST_NONE, axisH_);
748+ PipeBarrier<PIPE_V>();
749+ Add(winTpSendCountFloatTensor_, winTpSendCountFloatTensor_, gmTpSendCountFloatTensor_, axisH_);
750+ PipeBarrier<PIPE_V>();
751+ Cast(dst, winTpSendCountFloatTensor_, RoundMode::CAST_RINT, axisH_);
752+ } else {
753+ Add(dst, src0, src1, axisH_);
754+ }
755+}
756+ 
757+template <CombineMC2TypeClass>
758+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::PerformanceInfoPerToken(uint32_t tokenIndex,
759+ uint64_t performanceTimeStart, uint32_t beginIndex, LocalTensor<float> stateTensor)
760+{
761+ SyncFunc<AscendC::HardEvent::MTE2_S>();
762+ for (uint32_t i = 0; i < flagRcvCount_; i ++) {
763+ uint32_t fromRankId;
764+ if (i < axisK_) {
765+ uint32_t moeExpertId = expertIdsGM_.GetValue(tokenIndex * axisK_ + i);
766+ if (moeExpertId >= moeExpertNum_) {
767+ continue;
768+ }
769+ if (isInputExpertMaskFlag_) {
770+ bool maskExpertFlag = expertMaskTensor_.GetValue(tokenIndex * axisK_ + i);
771+ if (!maskExpertFlag) {
772+ continue;
773+ }
774+ }
775+ fromRankId = moeExpertId / moeExpertPerRankNum_ + sharedExpertRankNum_;
776+ } else {
777+ fromRankId = (i - axisK_) * rankNumPerShareExpert_ + epRankId_ % rankNumPerShareExpert_;
778+ }
779+ 
780+ if (stateTensor.GetValue(i * FLAG_OFFSET) > float(0.5) && firstRecordTensor_.GetValue((tokenIndex - beginIndex) * flagRcvCount_ + i) == 0) {
781+ uint64_t performanceTimeCheck = static_cast<uint64_t>(GetSystemCycle());
782+ int32_t performanceTimeWait = static_cast<int32_t>((performanceTimeCheck - performanceTimeStart) / CYCLES_PER_US);
783+ uint32_t fromRankIdTime = performanceInfoTensor_.GetValue(JUMP_WRITE * fromRankId);
784+ uint32_t maxTimeValue = (fromRankIdTime < performanceTimeWait) ? performanceTimeWait : fromRankIdTime;
785+ performanceInfoTensor_.SetValue(JUMP_WRITE * fromRankId, maxTimeValue);
786+ firstRecordTensor_.SetValue((tokenIndex - beginIndex) * flagRcvCount_ + i, 1);
787+ }
788+ }
789+}
790+ 
791+template <CombineMC2TypeClass>
792+__aicore__ inline bool MoeDistributeCombineV2<CombineMC2TypeFunc>::WaitDispatch(uint32_t tokenIndex,
793+ uint64_t performanceTimeStart, uint32_t copyCount, uint32_t beginIndex)
794+{
795+ uint32_t targetCount = copyCount;
796+ if (isInputExpertMaskFlag_ || ((zeroExpertNum_ + copyExpertNum_ + constExpertNum_) > 0U)) {
797+ int32_t tokenTarget = static_cast<int32_t>(tokenTargetTensor_.GetValue(tokenIndex)) + sharedExpertNum_;
798+ targetCount = tokenTarget * FLOAT_PER_UB_ALIGN;
799+ }
800+ float target = (float)1.0 * targetCount;
801+ float minTarget = target - (float)0.5;
802+ float maxTarget = target + (float)0.5;
803+ // 计算地址偏移
804+ GM_ADDR stateGM = GetWinStateAddrByRankId(epRankIdOriginal_, EP_DOMAIN) + tokenIndex * flagRcvCount_ * stateOffset_;
805+ GlobalTensor<float> stateGMTensor;
806+ stateGMTensor.SetGlobalBuffer((__gm__ float*)stateGM);
807+ float localState = 0;
808+ SumParams sumParams{1, copyCount, copyCount};
809+ LocalTensor<float> stateTensor = stateBuf_.Get<float>();
810+ SyncFunc<AscendC::HardEvent::S_MTE2>();
811+ DataCopy<float>(stateTensor, stateGMTensor, copyCount);
812+ SyncFunc<AscendC::HardEvent::MTE2_V>();
813+ if (isPerformanceFlag_) {
814+ PerformanceInfoPerToken(tokenIndex, performanceTimeStart, beginIndex, stateTensor);
815+ }
816+ Sum(stateTensor, stateTensor, sumParams);
817+ SyncFunc<AscendC::HardEvent::V_S>();
818+ localState = stateTensor(0);
819+ if (((minTarget < localState) && (localState < maxTarget))) {
820+ // 计算地址偏移，清状态
821+ GM_ADDR stateGM = GetWinStateAddrByRankId(epRankIdOriginal_, EP_DOMAIN) + tokenIndex * flagRcvCount_ * stateOffset_;
822+ GlobalTensor<float> stateGMTensor;
823+ stateGMTensor.SetGlobalBuffer((__gm__ float*)stateGM);
824+ DataCopy<float>(stateGMTensor, stateResetTensor_, copyCount);
825+ return true;
826+ }
827+ return false;
828+}
829+ 
830+template <CombineMC2TypeClass>
831+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::AddRmsNormAddCompute(
832+ uint32_t tokenIndex, uint32_t tokenOffset, uint32_t numCol, LocalTensor<float>& x1TmpFloatLocal,
833+ LocalTensor<float>& x2TmpFloatLocal, LocalTensor<float>& addOutTmpFloatLocal,
834+ const DataCopyExtParams& copyExtParams, const DataCopyPadExtParams<ExpandXType>& copyPadExtParams)
835+{
836+ // 计算x + residual_x
837+ LocalTensor<ExpandXType> x2 = tokenBuf_.Get<ExpandXType>();
838+ SyncFunc<AscendC::HardEvent::V_MTE2>();
839+ DataCopyPad(x2, residualXGM_[tokenIndex * axisH_ + tokenOffset], copyExtParams, copyPadExtParams);
840+ SyncFunc<AscendC::HardEvent::MTE2_V>();
841+ Cast(x2TmpFloatLocal, x2, AscendC::RoundMode::CAST_NONE, numCol);
842+ PipeBarrier<PIPE_V>();
843+ AscendC::Add(addOutTmpFloatLocal, x1TmpFloatLocal, x2TmpFloatLocal, numCol);
844+}
845+ 
846+template <CombineMC2TypeClass>
847+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::CalConstExpertAlpha(
848+ GlobalTensor<ExpandXType> constExpertAlphaGM, uint32_t const_expert_idx, float &alphaFloat)
849+{
850+ LocalTensor<ExpandXType> weightLocal = moeSumQueue_.AllocTensor<ExpandXType>();
851+ LocalTensor<float> weightFloatLocal = mulBuf_.Get<float>();
852+ DataCopyPadExtParams<ExpandXType> copyPadExtParams{false, 0U, 0U, 0U};
853+ DataCopyExtParams expandXCopyParams{1U, static_cast<uint32_t>(hExpandXTypeSize_), 0U, 0U, 0U};
854+ 
855+ // 使用moeSumQueue_分配缓冲区来存储alpha1对应的权重矩阵Wc
856+ DataCopyPad(weightLocal, constExpertAlphaGM[const_expert_idx * axisH_], expandXCopyParams, copyPadExtParams);
857+ moeSumQueue_.EnQue(weightLocal);
858+ weightLocal = moeSumQueue_.DeQue<ExpandXType>();
859+ Cast(weightFloatLocal, weightLocal, AscendC::RoundMode::CAST_NONE, axisH_);
860+ PipeBarrier<PIPE_V>();
861+ 
862+ // 计算Wc * x
863+ Mul(weightFloatLocal, weightFloatLocal, rowTmpFloatLocal_, axisH_);
864+ PipeBarrier<PIPE_V>();
865+ uint32_t innerAlign = Ceil(axisH_ * sizeof(float), UB_ALIGN) * UB_ALIGN / sizeof(float);
866+ SumParams params{1, innerAlign, axisH_};
867+ Sum(weightFloatLocal, weightFloatLocal, params);
868+ SyncFunc<AscendC::HardEvent::V_S>();
869+ alphaFloat = weightFloatLocal.GetValue(0);
870+ moeSumQueue_.FreeTensor<ExpandXType>(weightLocal);
871+}
872+ 
873+// 处理常量专家
874+template <CombineMC2TypeClass>
875+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::ProcessConstantExpert(
876+ uint32_t tokenIndex, uint32_t const_expert_idx, float scaleVal)
877+{
878+ PipeBarrier<PIPE_ALL>();
879+ LocalTensor<ExpandXType> rowTmpLocal = tokenBuf_.Get<ExpandXType>();
880+ LocalTensor<float> alphaFloatLocal = tokenBuf_.Get<float>();
881+ DataCopyPadExtParams<ExpandXType> copyPadExtParams{false, 0U, 0U, 0U};
882+ DataCopyExtParams expandXCopyParams{1U, static_cast<uint32_t>(hExpandXTypeSize_), 0U, 0U, 0U};
883+ float alpha1Float = static_cast<float>(0.0);
884+ float alpha2Float = static_cast<float>(0.0);
885+ 
886+ // 读取输入token
887+ DataCopyPad(rowTmpLocal, oriXGM_[tokenIndex * axisH_], expandXCopyParams, copyPadExtParams);
888+ SyncFunc<AscendC::HardEvent::MTE2_V>();
889+ Cast(rowTmpFloatLocal_, rowTmpLocal, AscendC::RoundMode::CAST_NONE, axisH_);
890+ PipeBarrier<PIPE_V>();
891+ 
892+ // 计算Wc * x
893+ CalConstExpertAlpha(constExpertAlpha1GM_, const_expert_idx, alpha1Float);
894+ CalConstExpertAlpha(constExpertAlpha2GM_, const_expert_idx, alpha2Float);
895+ 
896+ // 计算softmax(Wc * x)
897+ float maxAlphaFloat = (alpha1Float > alpha2Float) ? alpha1Float : alpha2Float;
898+ alphaFloatLocal.SetValue(0, alpha1Float - maxAlphaFloat);
899+ alphaFloatLocal.SetValue(1, alpha2Float - maxAlphaFloat);
900+ SyncFunc<AscendC::HardEvent::S_V>();
901+ Exp(alphaFloatLocal, alphaFloatLocal, 2);
902+ SyncFunc<AscendC::HardEvent::V_S>();
903+ float alphaSumFloat = alphaFloatLocal.GetValue(0) + alphaFloatLocal.GetValue(1);
904+ alpha1Float = alphaFloatLocal.GetValue(0) / alphaSumFloat;
905+ alpha2Float = alphaFloatLocal.GetValue(1) / alphaSumFloat;
906+ 
907+ // 使用moeSumQueue_分配缓冲区来存储常量专家向量v
908+ LocalTensor<float> constVFloatLocal = mulBuf_.Get<float>();
909+ LocalTensor<ExpandXType> const_v_ub = moeSumQueue_.AllocTensor<ExpandXType>();
910+ DataCopyPad(const_v_ub, constExpertVGM_[const_expert_idx * axisH_], expandXCopyParams, copyPadExtParams);
911+ moeSumQueue_.EnQue(const_v_ub);
912+ const_v_ub = moeSumQueue_.DeQue<ExpandXType>();
913+ 
914+ Cast(constVFloatLocal, const_v_ub, AscendC::RoundMode::CAST_NONE, axisH_);
915+ PipeBarrier<PIPE_V>();
916+ moeSumQueue_.FreeTensor<ExpandXType>(const_v_ub);
917+ 
918+ // 计算 alpha1 * x + alpha2 * v
919+ SyncFunc<AscendC::HardEvent::S_V>();
920+ Muls(rowTmpFloatLocal_, rowTmpFloatLocal_, alpha1Float, axisH_);
921+ Muls(constVFloatLocal, constVFloatLocal, alpha2Float, axisH_);
922+ PipeBarrier<PIPE_V>();
923+ Add(rowTmpFloatLocal_, rowTmpFloatLocal_, constVFloatLocal, axisH_);
924+ PipeBarrier<PIPE_V>();
925+ 
926+ // 乘以专家权重
927+ Muls(mulBufLocal_, rowTmpFloatLocal_, scaleVal, axisH_);
928+ PipeBarrier<PIPE_V>();
929+ Add(sumFloatBufLocal_, sumFloatBufLocal_, mulBufLocal_, axisH_);
930+ PipeBarrier<PIPE_V>();
931+}
932+ 
933+// 处理拷贝专家
934+template <CombineMC2TypeClass>
935+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::ProcessCopyExpert(
936+ uint32_t tokenIndex, float scaleVal)
937+{
938+ DataCopyPadExtParams<ExpandXType> copyPadExtParams{false, 0U, 0U, 0U};
939+ DataCopyExtParams expandXCopyParams{1U, static_cast<uint32_t>(hExpandXTypeSize_), 0U, 0U, 0U};
940+ LocalTensor<ExpandXType> tmpUb = moeSumQueue_.AllocTensor<ExpandXType>();
941+ DataCopyPad(tmpUb, oriXGM_[tokenIndex * axisH_], expandXCopyParams, copyPadExtParams);
942+ moeSumQueue_.EnQue(tmpUb);
943+ tmpUb = moeSumQueue_.DeQue<ExpandXType>();
944+ 
945+ Cast(rowTmpFloatLocal_, tmpUb, AscendC::RoundMode::CAST_NONE, axisH_);
946+ PipeBarrier<PIPE_V>();
947+ moeSumQueue_.FreeTensor<ExpandXType>(tmpUb);
948+ Muls(mulBufLocal_, rowTmpFloatLocal_, scaleVal, axisH_);
949+ PipeBarrier<PIPE_V>();
950+ Add(sumFloatBufLocal_, sumFloatBufLocal_, mulBufLocal_, axisH_);
951+ PipeBarrier<PIPE_V>();
952+}
953+ 
954+// 处理Moe专家
955+template <CombineMC2TypeClass>
956+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::ProcessMoeExpert(
957+ uint32_t tokenIndexOffset, uint32_t topkId, float scaleVal)
958+{
959+ uint32_t processLen = axisH_;
960+ const DataCopyExtParams xScaleCopyParams{
961+ 1U, static_cast<uint32_t>(tokenScaleCnt_ * sizeof(ExpandXType)), 0U, 0U, 0U};
962+ const DataCopyExtParams expandXCopyParams{1U, static_cast<uint32_t>(hExpandXTypeSize_), 0U, 0U, 0U};
963+ const DataCopyPadExtParams<ExpandXType> copyPadExtParams{false, 0U, 0U, 0U};
964+ 
965+ GM_ADDR wAddr = (__gm__ uint8_t*)(epWindowGM_) + (tokenIndexOffset + topkId) * hAlignWinSize_;
966+ rowTmpGlobal_.SetGlobalBuffer((__gm__ ExpandXType*)wAddr);
967+ LocalTensor<ExpandXType> tmpUb = moeSumQueue_.AllocTensor<ExpandXType>();
968+ if constexpr (IsInt8Quant) {
969+ DataCopyPad(tmpUb, rowTmpGlobal_, xScaleCopyParams, copyPadExtParams);
970+ } else {
971+ DataCopyPad(tmpUb, rowTmpGlobal_, expandXCopyParams, copyPadExtParams);
972+ }
973+ moeSumQueue_.EnQue(tmpUb);
974+ tmpUb = moeSumQueue_.DeQue<ExpandXType>();
975+ if constexpr (IsInt8Quant) {
976+ quantInst_.Int8DequantProcess(tmpUb, scaleDivFloatTensor_, fp16CastTensor_, absFloatTensor_, scaleDupLocalTensor_);
977+ }
978+ Cast(rowTmpFloatLocal_, tmpUb, AscendC::RoundMode::CAST_NONE, processLen);
979+ PipeBarrier<PIPE_V>();
980+ AscendC::Muls(mulBufLocal_, rowTmpFloatLocal_, scaleVal, processLen);
981+ PipeBarrier<PIPE_V>();
982+ AscendC::Add(sumFloatBufLocal_, sumFloatBufLocal_, mulBufLocal_, processLen);
983+ moeSumQueue_.FreeTensor<ExpandXType>(tmpUb);
984+}
985+ 
986+template <CombineMC2TypeClass>
987+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::ExpertScaleCopy(
988+ const uint32_t beginIndex, const uint32_t endIndex, const uint32_t tokenPerAivNum)
989+{
990+ expertScaleBeginIdx_ = beginIndex;
991+ uint32_t expertScaleEndIdx = endIndex;
992+ uint32_t expertScaleCntPerCore = tokenPerAivNum * axisK_;
993+ if (isInputExpertMaskFlag_) {
994+ expertScaleBeginIdx_ = validBsIndexTensor_.GetValue(beginIndex);
995+ expertScaleEndIdx = validBsIndexTensor_.GetValue(endIndex - 1);
996+ expertScaleCntPerCore = (expertScaleEndIdx - expertScaleBeginIdx_ + 1) * axisK_;
997+ }
998+ tpipe_->InitBuffer(expertScalesBuf_, Ceil(expertScaleCntPerCore * sizeof(float), UB_ALIGN) * UB_ALIGN);
999+ expertScalesLocal_ = expertScalesBuf_.Get<float>();
1000+ const DataCopyExtParams tokenScaleParams{1U, static_cast<uint32_t>(expertScaleCntPerCore * sizeof(float)), 0U, 0U, 0U};
1001+ const DataCopyPadExtParams<float> copyPadFloatParams{false, 0U, 0U, 0U};
1002+ DataCopyPad(expertScalesLocal_, expertScalesGM_[expertScaleBeginIdx_ * axisK_], tokenScaleParams, copyPadFloatParams);
1003+ SyncFunc<AscendC::HardEvent::MTE2_S>();
1004+}
1005+ 
1006+template <CombineMC2TypeClass>
1007+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::ProcessExpert(uint32_t tokenIndex, uint32_t processLen)
1008+{
1009+ uint32_t index = (tokenIndex - expertScaleBeginIdx_) * axisK_;
1010+ float scaleVal = 0.0;
1011+ GM_ADDR wAddr;
1012+ const DataCopyExtParams xScaleCopyParams{
1013+ 1U, static_cast<uint32_t>(tokenScaleCnt_ * sizeof(ExpandXType)), 0U, 0U, 0U};
1014+ const DataCopyPadExtParams<ExpandXType> copyPadExtParams{false, 0U, 0U, 0U};
1015+ const DataCopyExtParams expandXCopyParams{1U, static_cast<uint32_t>(hExpandXTypeSize_), 0U, 0U, 0U};
1016+ SyncFunc<AscendC::HardEvent::MTE3_V>(); // 与结果搬出datacopy同tensor
1017+ Duplicate(sumFloatBufLocal_, static_cast<float>(0), axisH_);
1018+ LocalTensor<ExpandXType> tmpUb;
1019+ uint32_t tokenIndexOffset = tokenIndex * (axisK_ + sharedExpertNum_);
1020+ 
1021+ if ((zeroExpertNum_ + copyExpertNum_ + constExpertNum_) == 0U) {
1022+ for (uint32_t topkId = 0U; topkId < axisK_; topkId++) {
1023+ if (isInputExpertMaskFlag_) {
1024+ bool maskExpertFlag = expertMaskTensor_.GetValue(tokenIndex * axisK_ + topkId);
1025+ if (!maskExpertFlag) {
1026+ index++;
1027+ continue;
1028+ }
1029+ }
1030+ scaleVal = expertScalesLocal_.GetValue(index);
1031+ ProcessMoeExpert(tokenIndexOffset, topkId, scaleVal);
1032+ index++;
1033+ }
1034+ } else {
1035+ for (uint32_t topkId = 0U; topkId < axisK_; topkId++) {
1036+ // 读取expert_id
1037+ DataCacheCleanAndInvalid<int32_t, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(expertIdsGM_[tokenIndex * axisK_ + topkId]);
1038+ uint32_t expert_id = expertIdsGM_.GetValue(tokenIndex * axisK_ + topkId);
1039+ if (isInputExpertMaskFlag_) {
1040+ bool maskExpertFlag = expertMaskTensor_.GetValue(tokenIndex * axisK_ + topkId);
1041+ if (!maskExpertFlag) {
1042+ index++;
1043+ continue;
1044+ }
1045+ }
1046+ scaleVal = expertScalesLocal_.GetValue(index);
1047+ 
1048+ if (expert_id < moeExpertOriginalNum_) {
1049+ ProcessMoeExpert(tokenIndexOffset, topkId, scaleVal);
1050+ index++;
1051+ } else if (expert_id < moeExpertOriginalNum_ + zeroExpertNum_) {
1052+ // 零专家不需要任何操作
1053+ index++;
1054+ } else if (expert_id < moeExpertOriginalNum_ + zeroExpertNum_ + copyExpertNum_) {
1055+ ProcessCopyExpert(tokenIndex, scaleVal);
1056+ index++;
1057+ } else if (expert_id < moeExpertOriginalNum_ + zeroExpertNum_ + copyExpertNum_ + constExpertNum_) {
1058+ uint32_t const_expert_idx = expert_id - (moeExpertOriginalNum_ + zeroExpertNum_ + copyExpertNum_);
1059+ ProcessConstantExpert(tokenIndex, const_expert_idx, scaleVal);
1060+ index++;
1061+ }
1062+ }
1063+ }
1064+ 
1065+ for (uint32_t topkId = axisK_; topkId < (axisK_ + sharedExpertNum_); topkId++) {
1066+ wAddr = (__gm__ uint8_t*)(epWindowGM_) + (tokenIndexOffset + topkId) * hAlignWinSize_;
1067+ rowTmpGlobal_.SetGlobalBuffer((__gm__ ExpandXType*)wAddr);
1068+ tmpUb = moeSumQueue_.AllocTensor<ExpandXType>();
1069+ if constexpr (IsInt8Quant) {
1070+ DataCopyPad(tmpUb, rowTmpGlobal_, xScaleCopyParams, copyPadExtParams);
1071+ } else {
1072+ DataCopyPad(tmpUb, rowTmpGlobal_, expandXCopyParams, copyPadExtParams);
1073+ }
1074+ moeSumQueue_.EnQue(tmpUb);
1075+ tmpUb = moeSumQueue_.DeQue<ExpandXType>();
1076+ if constexpr (IsInt8Quant) {
1077+ quantInst_.Int8DequantProcess(tmpUb, scaleDivFloatTensor_, fp16CastTensor_, absFloatTensor_, scaleDupLocalTensor_);
1078+ }
1079+ Cast(rowTmpFloatLocal_, tmpUb, AscendC::RoundMode::CAST_NONE, processLen);
1080+ PipeBarrier<PIPE_V>();
1081+ AscendC::Add(sumFloatBufLocal_, sumFloatBufLocal_, rowTmpFloatLocal_, processLen);
1082+ PipeBarrier<PIPE_V>();
1083+ moeSumQueue_.FreeTensor<ExpandXType>(tmpUb);
1084+ }
1085+ 
1086+ if (hasSharedExpertX_) {
1087+ LocalTensor<ExpandXType> rowTmpLocal = tokenBuf_.Get<ExpandXType>();
1088+ SyncFunc<AscendC::HardEvent::V_MTE2>(); // 与结果搬出Cast同地址
1089+ DataCopyPad(rowTmpLocal, sharedExpertXGM_[tokenIndex * axisH_], expandXCopyParams, copyPadExtParams);
1090+ SyncFunc<AscendC::HardEvent::MTE2_V>();
1091+ Cast(rowTmpFloatLocal_, rowTmpLocal, AscendC::RoundMode::CAST_NONE, processLen);
1092+ PipeBarrier<PIPE_V>();
1093+ AscendC::Add(sumFloatBufLocal_, sumFloatBufLocal_, rowTmpFloatLocal_, processLen);
1094+ }
1095+}
1096+ 
1097+template <CombineMC2TypeClass>
1098+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::LocalWindowCopy()
1099+{
1100+ if (activeMaskBsCnt_ == 0U) {
1101+ return;
1102+ }
1103+ uint32_t beginIndex = 0U;
1104+ uint32_t endIndex = 0U;
1105+ uint32_t processLen = 0U;
1106+ uint32_t tokenOffset = 0U;
1107+ uint32_t statePos = 1U;
1108+ uint32_t tokenPerAivNum = activeMaskBsCnt_ / aivNum_;
1109+ uint32_t remainderToken = activeMaskBsCnt_ % aivNum_;
1110+ 
1111+ beginIndex = tokenPerAivNum * coreIdx_;
1112+ if (coreIdx_ < remainderToken) {
1113+ tokenPerAivNum++;
1114+ beginIndex += coreIdx_;
1115+ } else {
1116+ beginIndex += remainderToken;
1117+ }
1118+ endIndex = beginIndex + tokenPerAivNum;
1119+ if (tokenPerAivNum == 0U) {
1120+ return;
1121+ }
1122+ processLen = axisH_;
1123+ TBuf<> opPosDfxBuf;
1124+ tpipe_->InitBuffer(opPosDfxBuf, UB_ALIGN);
1125+ dataStateLocalTensor_ = opPosDfxBuf.Get<uint32_t>();
1126+ rowTmpFloatLocal_ = rowTmpFloatBuf_.Get<float>();
1127+ mulBufLocal_ = mulBuf_.Get<float>();
1128+ sumFloatBufLocal_ = sumFloatBuf_.Get<float>();
1129+ const DataCopyPadExtParams<ExpandXType> copyPadXTypeParams{false, 0U, 0U, 0U};
1130+ DataCopyParams dataStateParams{1U, sizeof(uint32_t), 0U, 0U};
1131+ const DataCopyExtParams expandXCopyParams{1U, static_cast<uint32_t>(hExpandXTypeSize_), 0U, 0U, 0U};
1132+ ExpertScaleCopy(beginIndex, endIndex, tokenPerAivNum);
1133+ TBuf<> tokenStatusBuf;
1134+ tpipe_->InitBuffer(tokenStatusBuf, Ceil(tokenPerAivNum * sizeof(int32_t), UB_ALIGN) * UB_ALIGN);
1135+ LocalTensor tokenStatusTensor = tokenStatusBuf.Get<int32_t>();
1136+ Duplicate<int32_t>(tokenStatusTensor, static_cast<int32_t>(0), tokenPerAivNum);
1137+ uint32_t tokenNumCompleted = static_cast<uint32_t>(0);
1138+ 
1139+ if (isPerformanceFlag_) {
1140+ uint32_t tokenNumPerCore = tokenPerAivNum * flagRcvCount_ * sizeof(int32_t);
1141+ tokenNumPerCoreAlign_ = Ceil(tokenNumPerCore, UB_ALIGN) * UB_ALIGN;
1142+ tpipe_->InitBuffer(firstRecordBuf_, tokenNumPerCoreAlign_);
1143+ firstRecordTensor_ = firstRecordBuf_.Get<int32_t>();
1144+ Duplicate<int32_t>(firstRecordTensor_, static_cast<int32_t>(0), tokenPerAivNum * flagRcvCount_);
1145+ }
1146+ 
1147+ SyncFunc<AscendC::HardEvent::V_S>();
1148+ uint64_t performanceTimeStart = static_cast<uint64_t>(GetSystemCycle());
1149+ while (tokenNumCompleted != tokenPerAivNum) {
1150+ for (uint32_t curIdx = beginIndex; curIdx < endIndex; curIdx++) {
1151+ if (tokenStatusTensor(curIdx - beginIndex) == 1) {
1152+ continue;
1153+ }
1154+ uint32_t tokenIndex = curIdx;
1155+ if (isInputExpertMaskFlag_) {
1156+ tokenIndex = validBsIndexTensor_.GetValue(curIdx);
1157+ }
1158+ uint32_t copyCount = flagRcvCount_ * FLOAT_PER_UB_ALIGN;
1159+ if (!WaitDispatch(tokenIndex, performanceTimeStart, copyCount, beginIndex)) {
1160+ continue;
1161+ }
1162+ tokenNumCompleted++;
1163+ tokenStatusTensor.SetValue(curIdx - beginIndex, 1);
1164+ 
1165+ statePos++;
1166+ dataStateLocalTensor_.SetValue(0, statePos);
1167+ SyncFunc<AscendC::HardEvent::S_MTE3>();
1168+ DataCopyPad(selfDataStatusGMTensor_[1], dataStateLocalTensor_, dataStateParams);
1169+ ProcessExpert(tokenIndex, processLen);
1170+ 
1171+ // 结果搬出
1172+ PipeBarrier<PIPE_V>();
1173+ LocalTensor<ExpandXType> sumBufLocal = tokenBuf_.Get<ExpandXType>();
1174+ Cast(sumBufLocal, sumFloatBufLocal_, AscendC::RoundMode::CAST_RINT, processLen);
1175+ SyncFunc<AscendC::HardEvent::V_MTE3>();
1176+ DataCopyPad(expandOutGlobal_[tokenIndex * axisH_ + tokenOffset], sumBufLocal, expandXCopyParams);
1177+ }
1178+ }
1179+ if (isPerformanceFlag_) {
1180+ SyncFunc<AscendC::HardEvent::V_MTE3>();
1181+ SetAtomicMax<int32_t>();
1182+ DataCopyExtParams performanceInfoCopyParams{1U, static_cast<uint32_t>(JUMP_WRITE * epWorldSizeOriginal_ * sizeof(int32_t)), 0U, 0U, 0U};
1183+ DataCopyPad(performanceInfoGM_, performanceInfoTensor_, performanceInfoCopyParams);
1184+ SetAtomicNone();
1185+ }
1186+}
1187+ 
1188+template <CombineMC2TypeClass>
1189+__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::Process()
1190+{
1191+ if ASCEND_IS_AIV { // 全aiv处理
1192+ BuffInit();
1193+ SetWaitTpStatusAndDisPatch();
1194+ AlltoAllBuffInitAndMaskCal();
1195+ LocalWindowCopy();
1196+ }
1197+}
1198+ 
1199+} // MoeDistributeCombineV2Impl
1200+#endif // MOE_DISTRIBUTE_COMBINE_IMPL_H
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_combine_v2/ascend910_93/op_kernel/moe_distribute_combine_v2_quant.h+106-0
@@ -0,0 +1,106 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file moe_distribute_v2_quant.h
13+ * \brief
14+ */
15+ 
16+#ifndef MOE_DISTRIBUTE_COMBINE_V2_QUANT_H
17+#define MOE_DISTRIBUTE_COMBINE_V2_QUANT_H
18+ 
19+#include "../../../moe_distribute_dispatch_v2/ascend910_93/op_kernel/moe_distribute_v2_constant.h"
20+#include "../../../moe_distribute_dispatch_v2/ascend910_93/op_kernel/moe_distribute_v2_base.h"
21+ 
22+namespace Mc2Kernel {
23+using namespace AscendC;
24+using namespace MoeDistributeV2Base;
25+ 
26+template <typename ExpandXType, bool IsInt8Quant>
27+class MoeDistributeCombineQuant{
28+public:
29+ float scaleValFloat_;
30+ uint32_t axisH_{0};
31+ uint32_t mask_{0};
32+ uint32_t repeatNum_{0};
33+ uint32_t hAlign32Size_{0};
34+ uint32_t quantScaleNum_{0};
35+ LocalTensor<int8_t> castLocalTensor_;
36+ LocalTensor<ExpandXType> scaleDivTensor_;
37+ 
38+ __aicore__ inline MoeDistributeCombineQuant() = default;
39+ 
40+ __aicore__ inline void SetQuantInitParams(uint32_t axisH)
41+ {
42+ axisH_ = axisH;
43+ }
44+ 
45+ __aicore__ inline void InitInt8Quant(uint32_t &scaleNum_, uint32_t &hExpandXAlign32Size_,
46+ uint32_t &hFloatAlign256Size_, uint32_t &tokenScaleCnt_)
47+ {
48+ hAlign32Size_ = Ceil(axisH_, UB_ALIGN) * UB_ALIGN;
49+ scaleValFloat_ = static_cast<float>(1.0f / SCALE_PARAM);
50+ uint32_t scaleGranu = static_cast<uint32_t>(UB_ALIGN / sizeof(float)); // 计算每个block得到的reducemax结果数量
51+ quantScaleNum_ = (hExpandXAlign32Size_ / sizeof(ExpandXType)) / scaleGranu; // 得到有效scale的个数
52+ scaleNum_ = quantScaleNum_;
53+ repeatNum_ = static_cast<uint32_t>(hFloatAlign256Size_ / ALIGNED_LEN); // BlockReduceMax 与 Brcb的重复迭代次数，每次256b参与计算
54+ mask_ = static_cast<uint32_t>(ALIGNED_LEN / sizeof(float));
55+ tokenScaleCnt_ = hAlign32Size_ / sizeof(ExpandXType) + quantScaleNum_; // int8_align + scale有效个数
56+ }
57+ 
58+ __aicore__ inline void Int8QuantProcess(LocalTensor<ExpandXType> &sendLocalTensor_, LocalTensor<float> &winTpSendCountFloatTensor_, LocalTensor<ExpandXType> &gmTpSendCountTensor_,
59+ LocalTensor<half> &fp16CastTensor_, LocalTensor<float> &absFloatTensor_, LocalTensor<float> reduceMaxFloatTensor_, LocalTensor<float> scaleDupLocalTensor_)
60+ {
61+ SyncFunc<AscendC::HardEvent::MTE2_V>();
62+ castLocalTensor_ = sendLocalTensor_.template ReinterpretCast<int8_t>(); // 长度为int8H_Align + scaleNum
63+ scaleDivTensor_ = castLocalTensor_[hAlign32Size_].template ReinterpretCast<ExpandXType>(); // 偏移前面的int8
64+ 
65+ Cast(winTpSendCountFloatTensor_, gmTpSendCountTensor_, RoundMode::CAST_NONE, axisH_);
66+ PipeBarrier<PIPE_V>();
67+ Abs(absFloatTensor_, winTpSendCountFloatTensor_, axisH_); // absFloatTensor_ align到256并写0，支持ReduceMax与Brcb
68+ PipeBarrier<PIPE_V>();
69+ BlockReduceMax(reduceMaxFloatTensor_, absFloatTensor_, repeatNum_, mask_, 1, 1, BLOCK_NUM); // 32->1 256->8
70+ PipeBarrier<PIPE_V>();
71+ Muls(reduceMaxFloatTensor_, reduceMaxFloatTensor_, scaleValFloat_, quantScaleNum_); // 有效个数
72+ PipeBarrier<PIPE_V>();
73+ Cast(scaleDivTensor_, reduceMaxFloatTensor_, RoundMode::CAST_RINT, quantScaleNum_); // 有效个数
74+ PipeBarrier<PIPE_V>();
75+ Brcb(scaleDupLocalTensor_, reduceMaxFloatTensor_, repeatNum_, {1, BLOCK_NUM}); // 一次256
76+ PipeBarrier<PIPE_V>();
77+ Div(winTpSendCountFloatTensor_, winTpSendCountFloatTensor_, scaleDupLocalTensor_, axisH_); // 有效个数
78+ PipeBarrier<PIPE_V>();
79+ Cast(fp16CastTensor_, winTpSendCountFloatTensor_, RoundMode::CAST_RINT, axisH_);
80+ PipeBarrier<PIPE_V>();
81+ Cast(castLocalTensor_, fp16CastTensor_, RoundMode::CAST_RINT, axisH_);
82+ SyncFunc<AscendC::HardEvent::V_MTE3>();
83+ }
84+ 
85+ __aicore__ inline void Int8DequantProcess(LocalTensor<ExpandXType>& src, LocalTensor<float> scaleDivFloatTensor_, LocalTensor<half> &fp16CastTensor_,
86+ LocalTensor<float> &absFloatTensor_, LocalTensor<float> scaleDupLocalTensor_)
87+ {
88+ SyncFunc<AscendC::HardEvent::MTE2_V>();
89+ castLocalTensor_ = src.template ReinterpretCast<int8_t>();
90+ scaleDivTensor_ = src[hAlign32Size_ / INT8_DIVIVE];
91+ 
92+ SyncFunc<AscendC::HardEvent::S_V>();
93+ Cast(scaleDivFloatTensor_, scaleDivTensor_, RoundMode::CAST_NONE, quantScaleNum_);
94+ Cast(fp16CastTensor_, castLocalTensor_, RoundMode::CAST_NONE, axisH_);
95+ PipeBarrier<PIPE_V>();
96+ Cast(absFloatTensor_, fp16CastTensor_, RoundMode::CAST_NONE, axisH_);
97+ Brcb(scaleDupLocalTensor_, scaleDivFloatTensor_, repeatNum_, {1, BLOCK_NUM});
98+ PipeBarrier<PIPE_V>();
99+ Mul(absFloatTensor_, absFloatTensor_, scaleDupLocalTensor_, axisH_);
100+ PipeBarrier<PIPE_V>();
101+ Cast(src, absFloatTensor_, RoundMode::CAST_RINT, axisH_);
102+ PipeBarrier<PIPE_V>();
103+ }
104+};
105+}
106+#endif // MOE_DISTRIBUTE_V2_QUANT_H
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_combine_v2/ascend910_93/op_kernel/moe_distribute_combine_v2_tiling.h+51-0
@@ -0,0 +1,51 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file moe_distribute_combine_v2_tiling.h
13+ * \brief
14+ */
15+#ifndef MOE_DISTRIBUTE_COMBINE_V2_TILING_H
16+#define MOE_DISTRIBUTE_COMBINE_V2_TILING_H
17+ 
18+#include <cstdint>
19+#include "kernel_tiling/kernel_tiling.h"
20+ 
21+// a3
22+struct MoeDistributeCombineV2Info {
23+ uint32_t epWorldSize;
24+ uint32_t epRankId;
25+ uint32_t expertShardType;
26+ uint32_t sharedExpertNum;
27+ uint32_t sharedExpertRankNum;
28+ uint32_t moeExpertNum;
29+ uint32_t moeExpertPerRankNum;
30+ uint32_t zeroExpertNum;
31+ uint32_t copyExpertNum;
32+ uint32_t constExpertNum;
33+ uint32_t globalBs;
34+ uint32_t bs;
35+ uint32_t k;
36+ uint32_t h;
37+ uint32_t a;
38+ uint32_t aivNum;
39+ bool isTokenMask; // input active mask 1dims or not
40+ bool isExpertMask; // input active mask 2dims or not
41+ bool hasSharedExpertX; // input shared expert x or not
42+ bool isPerformance; // whether performance or not
43+ bool reserved0;
44+ bool reserved1;
45+ bool reserved2;
46+ uint64_t totalUbSize;
47+ uint64_t totalWinSizeEp;
48+ uint32_t bufferNum;
49+};
50+ 
51+#endif //__MOE_DISTRIBUTE_COMBINE_V2_TILING_H__
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_dispatch_v2/ascend910_93/CMakeLists.txt+11-0
@@ -0,0 +1,11 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# ----------------------------------------------------------------------------
10+ 
11+add_sources("--npu-arch=dav-2201")
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_dispatch_v2/ascend910_93/moe_distribute_dispatch_v2_entry.cpp+222-0
@@ -0,0 +1,222 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file moe_distribute_dispatch_v2.cpp
13+ * \brief
14+ */
15+#if ASC_DEVKIT_MAJOR >= 9
16+#include "basic_api/kernel_basic_intf.h"
17+#else
18+#include "kernel_operator.h"
19+#endif
20+#include "op_kernel/moe_distribute_dispatch_v2.h"
21+#include "op_kernel/moe_distribute_dispatch_v2_full_mesh.h"
22+#include "op_kernel/moe_distribute_dispatch_v2_tiling.h"
23+ 
24+using namespace MoeDistributeDispatchV2Impl;
25+using namespace MoeDistributeDispatchV2FullMeshImpl;
26+using namespace Mc2Kernel;
27+using namespace AscendC;
28+ 
29+/*
30+* A3 tilingkey说明
31+* 5位的十进制数
32+* 第1位（个位）：quantMode:
33+* 0: 不量化, 1: 静态量化, 2: 动态量化
34+* 第2位（十位）：x输入类型:
35+* 0: float16, 1: bfloat16
36+* 第3位（百位）：是否有smoothScale:
37+* 0: 无, 1: 有
38+* 第4位（千位）：是否走fullmesh_v2模板:
39+* 0: 不做, 1: 做
40+* 第5位（万位）：无实际含义
41+*/
42+ 
43+template<typename XType, typename ExpandxType, int32_t QuantMode, bool IsSmoothScaleExist>
44+__attribute__((always_inline)) __aicore__ __inline__ void moe_distribute_dispatch_v2(
45+ GM_ADDR x, GM_ADDR expertIds, GM_ADDR scales, GM_ADDR xActiveMask, GM_ADDR expertScales,
46+ GM_ADDR performanceInfo, GM_ADDR expandXOut, GM_ADDR dynamicScalesOut,
47+ GM_ADDR assistInfoOut, GM_ADDR expertTokenNumsOut, GM_ADDR epSendCountsOut,
48+ GM_ADDR expandScalesOut, GM_ADDR workspaceGM, GM_ADDR mc2Context, MoeDistributeDispatchV2Info tilingData)
49+{
50+ 
51+ TPipe pipe;
52+ MoeDistributeDispatchV2<XType, ExpandxType, QuantMode, IsSmoothScaleExist> op;
53+ op.Init(mc2Context, x, expertIds, scales, xActiveMask, performanceInfo, expandXOut, dynamicScalesOut,
54+ assistInfoOut, expertTokenNumsOut, epSendCountsOut, workspaceGM, tilingData, &pipe);
55+
56+ op.Process();
57+ return;
58+}
59+ 
60+template<typename XType, typename ExpandxType, int32_t QuantMode, bool IsSmoothScaleExist>
61+__attribute__((always_inline)) __aicore__ __inline__ void moe_distribute_dispatch_v2_full_mesh(
62+ GM_ADDR x, GM_ADDR expertIds, GM_ADDR scales, GM_ADDR xActiveMask, GM_ADDR expertScales,
63+ GM_ADDR performanceInfo, GM_ADDR expandXOut, GM_ADDR dynamicScalesOut,
64+ GM_ADDR assistInfoOut, GM_ADDR expertTokenNumsOut, GM_ADDR epSendCountsOut,
65+ GM_ADDR expandScalesOut, GM_ADDR workspaceGM, GM_ADDR mc2Context, MoeDistributeDispatchV2Info tilingData)
66+{
67+ 
68+ TPipe pipe;
69+ MoeDistributeDispatchV2FullMesh<XType, ExpandxType, QuantMode, IsSmoothScaleExist> op;
70+ op.Init(x, expertIds, scales, xActiveMask, performanceInfo, expandXOut, dynamicScalesOut,
71+ assistInfoOut, expertTokenNumsOut, epSendCountsOut, workspaceGM, mc2Context, tilingData, &pipe);
72+
73+ op.Process();
74+ return;
75+}
76+ 
77+ 
78+extern "C" __global__ __aicore__ void moe_distribute_dispatch_v2_generic(
79+ int32_t tilingKey,
80+ GM_ADDR x, GM_ADDR expertIds, GM_ADDR scales, GM_ADDR xActiveMask, GM_ADDR expertScales,
81+ GM_ADDR performanceInfo, GM_ADDR expandXOut, GM_ADDR dynamicScalesOut,
82+ GM_ADDR assistInfoOut, GM_ADDR expertTokenNumsOut, GM_ADDR epSendCountsOut,
83+ GM_ADDR expandScalesOut, GM_ADDR workspaceGM, GM_ADDR mc2Context, MoeDistributeDispatchV2Info tilingData)
84+{
85+ // 根据不同的数据类型调用不同的模板
86+ switch (tilingKey) {
87+ 
88+ case 10000:
89+ moe_distribute_dispatch_v2<float16_t, float16_t, MoeDistributeDispatchV2Impl::UNQUANT, false>(
90+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
91+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
92+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
93+ break;
94+ 
95+ case 10002:
96+ moe_distribute_dispatch_v2<float16_t, int8_t, MoeDistributeDispatchV2Impl::PERTOKEN_DYNAMIC_QUANT, false>(
97+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
98+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
99+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
100+ break;
101+ 
102+ case 10010:
103+ moe_distribute_dispatch_v2<bfloat16_t, bfloat16_t, MoeDistributeDispatchV2Impl::UNQUANT, false>(
104+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
105+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
106+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
107+ break;
108+ 
109+ case 10012:
110+ moe_distribute_dispatch_v2<bfloat16_t, int8_t, MoeDistributeDispatchV2Impl::PERTOKEN_DYNAMIC_QUANT, false>(
111+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
112+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
113+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
114+ break;
115+ 
116+ case 10100:
117+ moe_distribute_dispatch_v2<float16_t, float16_t, MoeDistributeDispatchV2Impl::UNQUANT, true>(
118+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
119+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
120+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
121+ break;
122+ 
123+ case 10102:
124+ moe_distribute_dispatch_v2<float16_t, int8_t, MoeDistributeDispatchV2Impl::PERTOKEN_DYNAMIC_QUANT, true>(
125+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
126+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
127+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
128+ break;
129+ 
130+ case 10110:
131+ moe_distribute_dispatch_v2<bfloat16_t, bfloat16_t, MoeDistributeDispatchV2Impl::UNQUANT, true>(
132+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
133+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
134+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
135+ break;
136+ 
137+ case 10112:
138+ moe_distribute_dispatch_v2<bfloat16_t, int8_t, MoeDistributeDispatchV2Impl::PERTOKEN_DYNAMIC_QUANT, true>(
139+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
140+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
141+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
142+ break;
143+ 
144+ /* ---------- FullMesh ---------- */
145+ 
146+ case 11000:
147+ moe_distribute_dispatch_v2_full_mesh<float16_t, float16_t, MoeDistributeDispatchV2Impl::UNQUANT, false>(
148+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
149+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
150+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
151+ break;
152+ 
153+ case 11002:
154+ moe_distribute_dispatch_v2_full_mesh<float16_t, int8_t, MoeDistributeDispatchV2Impl::PERTOKEN_DYNAMIC_QUANT, false>(
155+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
156+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
157+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
158+ break;
159+ 
160+ case 11010:
161+ moe_distribute_dispatch_v2_full_mesh<bfloat16_t, bfloat16_t, MoeDistributeDispatchV2Impl::UNQUANT, false>(
162+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
163+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
164+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
165+ break;
166+ 
167+ case 11012:
168+ moe_distribute_dispatch_v2_full_mesh<bfloat16_t, int8_t, MoeDistributeDispatchV2Impl::PERTOKEN_DYNAMIC_QUANT, false>(
169+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
170+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
171+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
172+ break;
173+ 
174+ case 11100:
175+ moe_distribute_dispatch_v2_full_mesh<float16_t, float16_t, MoeDistributeDispatchV2Impl::UNQUANT, true>(
176+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
177+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
178+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
179+ break;
180+ 
181+ case 11102:
182+ moe_distribute_dispatch_v2_full_mesh<float16_t, int8_t, MoeDistributeDispatchV2Impl::PERTOKEN_DYNAMIC_QUANT, true>(
183+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
184+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
185+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
186+ break;
187+ 
188+ case 11110:
189+ moe_distribute_dispatch_v2_full_mesh<bfloat16_t, bfloat16_t, MoeDistributeDispatchV2Impl::UNQUANT, true>(
190+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
191+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
192+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
193+ break;
194+ 
195+ case 11112:
196+ moe_distribute_dispatch_v2_full_mesh<bfloat16_t, int8_t, MoeDistributeDispatchV2Impl::PERTOKEN_DYNAMIC_QUANT, true>(
197+ x, expertIds, scales, xActiveMask, expertScales, performanceInfo,
198+ expandXOut, dynamicScalesOut, assistInfoOut, expertTokenNumsOut,
199+ epSendCountsOut, expandScalesOut, workspaceGM, mc2Context, tilingData);
200+ break;
201+ 
202+ default:
203+ AscendC::PRINTF("moe_distribute_dispatch_v2 Error: invalid tilingKey = %d\n", tilingKey);
204+ return;
205+ }
206+ return;
207+}
208+ 
209+ 
210+// <<<>>>调用函数
211+void moe_distribute_dispatch_v2_entry(int32_t tilingKey, uint32_t blockDim, void* stream, GM_ADDR x, GM_ADDR expertIds,
212+ GM_ADDR scales, GM_ADDR xActiveMask, GM_ADDR expertScales, GM_ADDR performanceInfo, GM_ADDR expandXOut,
213+ GM_ADDR dynamicScalesOut, GM_ADDR assistInfoOut, GM_ADDR expertTokenNumsOut, GM_ADDR epSendCountsOut,
214+ GM_ADDR expandScalesOut, GM_ADDR workspaceGM, GM_ADDR mc2Context,
215+ MoeDistributeDispatchV2Info tilingData)
216+{
217+ moe_distribute_dispatch_v2_generic<<<blockDim, nullptr, stream>>>(
218+ tilingKey, x, expertIds, scales, xActiveMask, expertScales, performanceInfo, expandXOut, dynamicScalesOut,
219+ assistInfoOut, expertTokenNumsOut, epSendCountsOut, expandScalesOut, workspaceGM,
220+ mc2Context, tilingData
221+ );
222+}
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_dispatch_v2/ascend910_93/moe_distribute_dispatch_v2_entry.h+21-0
@@ -0,0 +1,21 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+#ifndef MOE_DISTRIBUTE_DISPATCH_V2_ENTRY_H
11+#define MOE_DISTRIBUTE_DISPATCH_V2_ENTRY_H
12+#include "op_kernel/moe_distribute_dispatch_v2_tiling.h"
13+ 
14+// <<<>>>调用函数声明
15+void moe_distribute_dispatch_v2_entry(int32_t tilingKey, uint32_t blockDim, void* stream, GM_ADDR x, GM_ADDR expertIds,
16+ GM_ADDR scales, GM_ADDR xActiveMask, GM_ADDR expertScales, GM_ADDR performanceInfo, GM_ADDR expandXOut,
17+ GM_ADDR dynamicScalesOut, GM_ADDR assistInfoOut, GM_ADDR expertTokenNumsOut, GM_ADDR epSendCountsOut,
18+ GM_ADDR expandScalesOut, GM_ADDR workspaceGM, GM_ADDR mc2Context,
19+ MoeDistributeDispatchV2Info tilingData);
20+ 
21+#endif //MOE_DISTRIBUTE_DISPATCH_V2_ENTRY_H
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_dispatch_v2/ascend910_93/moe_distribute_dispatch_v2_torch.cpp+353-0
@@ -0,0 +1,353 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify it.
3+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
8+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+ * See LICENSE in the root of the software repository for the full text of the License.
10+ */
11+ 
12+/*!
13+ * \file moe_distribute_dispatch_v2_torch.cpp
14+ * \brief
15+ */
16+ 
17+#include <ATen/ATen.h>
18+#include <vector>
19+#include <torch/all.h>
20+ 
21+#include "acl/acl.h"
22+#include "tiling/platform/platform_ascendc.h"
23+#include "kernel_operator.h"
24+#include "moe_distribute_dispatch_v2_entry.h"
25+#include "op_kernel/moe_distribute_dispatch_v2_tiling.h"
26+#include "moe_distribute_dispatch_v2_torch.h"
27+ 
28+#include "torch_npu/csrc/core/npu/NPUStream.h"
29+#include "torch_npu/csrc/framework/OpCommand.h"
30+ 
31+namespace ascend_ops {
32+ 
33+namespace MoeDistributeDispatchV2 {
34+ 
35+TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m)
36+{
37+ m.def("MoeDistributeDispatchV2(Tensor x, Tensor expert_ids, Tensor mc2_context, str group_ep, int ep_world_size, " \
38+ "int ep_rank_id, int moe_expert_num, int total_winsize_ep, *, Tensor? scales=None, Tensor? x_active_mask=None, " \
39+ "Tensor? expert_scales=None, Tensor? performance_info=None, "\
40+ "int expert_shard_type=0, int shared_expert_num=0, " \
41+ "int shared_expert_rank_num=0, int quant_mode=0, int global_bs=0, int expert_token_nums_type=0, " \
42+ "str comm_alg=\"\", int zero_expert_num=0, int copy_expert_num=0, int const_expert_num=0 " \
43+ ") " \
44+ "-> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)");
45+}
46+ 
47+constexpr uint32_t DIM_ONE = 1UL;
48+constexpr uint32_t DIM_TWO = 2UL;
49+constexpr uint32_t TILINGKEY_XTYPE = 10;
50+constexpr uint32_t TILINGKEY_SCALES = 100;
51+constexpr uint32_t TILINGKEY_COMM_ALG = 1000;
52+constexpr uint32_t WORKSPACESIZE = 16 * 1024 *1024;
53+ 
54+static void calculate_tilingkey(int32_t &tilingKey, at::ScalarType xType, const bool isScales, const uint32_t quantMode,
55+ const bool isSetCommAlg)
56+{
57+ tilingKey += static_cast<uint64_t>(quantMode);
58+ // 检查是否为 bfloat16 (kBFloat16)
59+ if (xType == at::kBFloat16) {
60+ tilingKey += static_cast<uint64_t>(TILINGKEY_XTYPE);
61+ }
62+ if (isScales) {
63+ tilingKey += static_cast<uint64_t>(TILINGKEY_SCALES);
64+ }
65+ if (isSetCommAlg) {
66+ tilingKey += static_cast<uint64_t>(TILINGKEY_COMM_ALG);
67+ }
68+ 
69+ return;
70+}
71+ 
72+void moe_distribute_dispatch_v2_api(
73+ aclrtStream stream, bool is_fullmesh_v2, const at::Tensor &x, const at::Tensor &expert_ids,
74+ const at::Tensor &new_workspace, const at::Tensor &mc2_context,
75+ c10::string_view group_ep, int64_t ep_world_size, int64_t ep_rank_id, int64_t moe_expert_num,
76+ const c10::optional<at::Tensor> &scales, const c10::optional<at::Tensor> &x_active_mask,
77+ const c10::optional<at::Tensor> &expert_scales, const c10::optional<at::Tensor> &performance_info,
78+ at::Tensor &expand_x, at::Tensor &dynamic_scales, at::Tensor &assist_info_forcombine,
79+ at::Tensor &expert_token_nums, at::Tensor &ep_recv_counts,
80+ at::Tensor &expand_scales, MoeDistributeDispatchV2Info tilingData)
81+{
82+ auto x_ptr = get_first_tensor_address<at::Tensor>(x.scalar_type(), x, false);
83+ auto expertIds_ptr = get_first_tensor_address<at::Tensor>(expert_ids.scalar_type(), expert_ids, false);
84+ auto workspace_ptr = get_first_tensor_address<at::Tensor>(new_workspace.scalar_type(), new_workspace, false);
85+ auto mc2Context_ptr = get_first_tensor_address<at::Tensor>(mc2_context.scalar_type(), mc2_context, false);
86+ void* scales_ptr = nullptr;
87+ if(scales.has_value()) {
88+ scales_ptr = get_first_tensor_address<c10::optional<at::Tensor>>(scales->scalar_type(), scales, false);
89+ }
90+ void* xActiveMask_ptr = nullptr;
91+ if(x_active_mask.has_value()) {
92+ xActiveMask_ptr = get_first_tensor_address<c10::optional<at::Tensor>>(x_active_mask->scalar_type(), x_active_mask,
93+ false);
94+ }
95+ void* expertScales_ptr = nullptr;
96+ if(expert_scales.has_value()) {
97+ expertScales_ptr = get_first_tensor_address<c10::optional<at::Tensor>>(expert_scales->scalar_type(),
98+ expert_scales, false);
99+ }
100+ void* performanceInfo_ptr = nullptr;
101+ if(performance_info.has_value()) {
102+ performanceInfo_ptr = get_first_tensor_address<c10::optional<at::Tensor>>(performance_info->scalar_type(),
103+ performance_info, false);
104+ }
105+ auto expandXOut_ptr = get_first_tensor_address<at::Tensor>(expand_x.scalar_type(), expand_x, false);
106+ auto dynamicScalesOut_ptr =
107+ get_first_tensor_address<at::Tensor>(dynamic_scales.scalar_type(), dynamic_scales, false);
108+ auto assistInfoOut_ptr =
109+ get_first_tensor_address<at::Tensor>(assist_info_forcombine.scalar_type(), assist_info_forcombine, false);
110+ auto expertTokenNumsOut_ptr =
111+ get_first_tensor_address<at::Tensor>(expert_token_nums.scalar_type(), expert_token_nums, false);
112+ auto epSendCountsOut_ptr =
113+ get_first_tensor_address<at::Tensor>(ep_recv_counts.scalar_type(), ep_recv_counts, false);
114+ auto expandScalesOut_ptr =
115+ get_first_tensor_address<at::Tensor>(assist_info_forcombine.scalar_type(), assist_info_forcombine, false);
116+ int32_t tilingKey = 10000;
117+ calculate_tilingkey(tilingKey, x.scalar_type(), scales.has_value(), tilingData.quantMode, is_fullmesh_v2);
118+ moe_distribute_dispatch_v2_entry(tilingKey, tilingData.aivNum, (void*)stream, (GM_ADDR)x_ptr,
119+ (GM_ADDR)expertIds_ptr, (GM_ADDR)scales_ptr, (GM_ADDR)xActiveMask_ptr,
120+ (GM_ADDR)expertScales_ptr, (GM_ADDR)performanceInfo_ptr,
121+ (GM_ADDR)expandXOut_ptr, (GM_ADDR)dynamicScalesOut_ptr, (GM_ADDR)assistInfoOut_ptr,
122+ (GM_ADDR)expertTokenNumsOut_ptr, (GM_ADDR)epSendCountsOut_ptr,
123+ (GM_ADDR)expandScalesOut_ptr, (GM_ADDR)workspace_ptr, (GM_ADDR)mc2Context_ptr, tilingData);
124+}
125+ 
126+void calculate_tilingdata(MoeDistributeDispatchV2Info &tilingData, int64_t ep_world_size, int64_t ep_rank_id,
127+ int64_t moe_expert_num, int64_t total_winsize_ep,
128+ int64_t expert_shard_type, int64_t shared_expert_num, int64_t shared_expert_rank_num,
129+ int64_t global_bs, int64_t bs, int64_t h, int64_t k, int64_t a, int64_t quant_mode,
130+ int64_t zero_expert_num, int64_t copy_expert_num, int64_t const_expert_num, int64_t expert_token_nums_type,
131+ const c10::optional<at::Tensor> &scales,
132+ const c10::optional<at::Tensor> &x_active_mask,
133+ const c10::optional<at::Tensor> &performance_info)
134+{
135+ auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
136+ uint64_t ubSizePlatFrom;
137+ ascendcPlatform->GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatFrom);
138+ // 初始化TilingData结构体的成员变量
139+ tilingData.epWorldSize = ep_world_size; // epWorldSize
140+ tilingData.epRankId = ep_rank_id; // epRankId
141+ tilingData.expertShardType = expert_shard_type; // expert type
142+ tilingData.sharedExpertNum = shared_expert_num; // shared expert number
143+ tilingData.sharedExpertRankNum = shared_expert_rank_num; // shared expert rank number
144+ tilingData.moeExpertNum = moe_expert_num; // moe expert number
145+ tilingData.quantMode = quant_mode; // quant mode
146+ tilingData.globalBs = global_bs; // globalBs = gBS * worldSize
147+ tilingData.bs = bs; // bs
148+ tilingData.k = k; // k
149+ tilingData.h = h; // h
150+ tilingData.a = a; // a
151+ tilingData.aivNum = ascendcPlatform->GetCoreNumAiv(); // aivNum
152+ tilingData.isTokenMask = (x_active_mask.has_value() && x_active_mask->dim() == DIM_ONE); // input active mask 1dims or not
153+ tilingData.isExpertMask = (x_active_mask.has_value() && x_active_mask->dim() == DIM_TWO); // input active mask 2dims or not
154+ tilingData.isPerformance = performance_info.has_value(); // whether performance or not
155+ tilingData.isQuant = (quant_mode != 0); // whether quant or not
156+ tilingData.reserved0 = false;
157+ tilingData.reserved1 = false;
158+ tilingData.reserved2 = false;
159+ tilingData.totalUbSize = ubSizePlatFrom; // epWorldSize
160+ tilingData.totalWinSizeEp = total_winsize_ep;
161+ tilingData.expertTokenNumsType = expert_token_nums_type; // expert token nums type, support 0: cumsum mode, 1: count mode
162+ tilingData.zeroComputeExpertNum = zero_expert_num + copy_expert_num + const_expert_num; // sum of zero, copy and const expert nums
163+ tilingData.scalesRow = scales.has_value() ? scales.value()[0].item().toLong() : 0;
164+ tilingData.scalesCol = scales.has_value() ? scales.value()[1].item().toLong() : 0;
165+ tilingData.scalesTypeSize = scales.has_value() ? at::elementSize(scales.value().scalar_type()) : 0;
166+ tilingData.scalesCount = scales.has_value() ? (scales.value()[0].item().toLong() * scales.value()[1].item().toLong()) : 0;
167+}
168+ 
169+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_moe_distribute_dispatch_v2(
170+ const at::Tensor &x, const at::Tensor &expert_ids, const at::Tensor &mc2_context, c10::string_view group_ep,
171+ int64_t ep_world_size, int64_t ep_rank_id, int64_t moe_expert_num, int64_t total_winsize_ep,
172+ const c10::optional<at::Tensor> &scales, const c10::optional<at::Tensor> &x_active_mask,
173+ const c10::optional<at::Tensor> &expert_scales, const c10::optional<at::Tensor> &performance_info,
174+ int64_t expert_shard_type, int64_t shared_expert_num,
175+ int64_t shared_expert_rank_num, int64_t quant_mode, int64_t global_bs,
176+ int64_t expert_token_nums_type, c10::string_view comm_alg,
177+ int64_t zero_expert_num, int64_t copy_expert_num, int64_t const_expert_num)
178+{
179+ TORCH_CHECK((x.dim() == DIM_TWO) && (expert_ids.dim() == DIM_TWO), "The x and expert_ids should be 2D");
180+ TORCH_CHECK(x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kHalf,
181+ "dtype of x should be BFloat16 or Half, but got " + std::string(c10::toString(x.scalar_type())));
182+ TORCH_CHECK(expert_ids.scalar_type() == at::kInt,
183+ "dtype of expert_ids should be Int, but got " + std::string(c10::toString(expert_ids.scalar_type())));
184+ TORCH_CHECK((ep_rank_id >= 0) && (ep_rank_id < ep_world_size),
185+ "ep_rank_id should be in [0, ep_world_size), but got",
186+ " ep_world_size: ", ep_world_size,
187+ ", ep_rank_id: ", ep_rank_id);
188+ TORCH_CHECK((shared_expert_rank_num >= 0) && (shared_expert_rank_num < ep_world_size),
189+ "shared_expert_rank_num should be in [0, ep_world_size), but got",
190+ " ep_world_size: ", ep_world_size,
191+ ", shared_expert_rank_num: ", shared_expert_rank_num);
192+ bool is_shared_default = ((shared_expert_num == 1) && (shared_expert_rank_num == 0));
193+ bool is_no_shared = ((shared_expert_num == 0) && (shared_expert_rank_num == 0));
194+ bool is_valid_shared = ((shared_expert_num > 0)
195+ && ((shared_expert_rank_num / shared_expert_num) > 0)
196+ && ((shared_expert_rank_num % shared_expert_num) == 0));
197+ TORCH_CHECK(is_shared_default || is_no_shared || is_valid_shared,
198+ "shared_expert_num and shared_expertrank_num have obvious value situations: "
199+ "1. shared_expert_num is 1, shared_expert_rank_num is 0; 2. shared_expert num is 0, "
200+ "shared_expert_rank_num is 0; 3. shared_expert_num in (0, shared_expert_rank_num] and "
201+ "shared_expert_rank_num % shared_expert_num = 0. but the current input value is ",
202+ " shared_expert_num: ", shared_expert_num,
203+ ", shared_expert_rank_num: ", shared_expert_rank_num);
204+ TORCH_CHECK((expert_token_nums_type == 0) || (expert_token_nums_type == 1),
205+ "The expert_token_nums_type should be 0 or 1.");
206+ auto x_size = x.sizes();
207+ auto expert_ids_size = expert_ids.sizes();
208+ 
209+ int64_t bs = x_size[0];
210+ int64_t h = x_size[1];
211+ int64_t k = expert_ids_size[1];
212+ 
213+ bool shared_front = (expert_shard_type == 0);
214+ int64_t local_moe_expert_num = 1;
215+ int64_t global_bs_real = (global_bs == 0) ? (bs * ep_world_size) : global_bs;
216+ int64_t a = 0;
217+ int64_t ep_recv_cnt_num = 0;
218+ if (shared_front) {
219+ if (ep_rank_id < shared_expert_rank_num) {
220+ local_moe_expert_num = 1;
221+ int64_t max_bs = global_bs_real / ep_world_size; // 前面已有拦截，保证ep_world_size > 0
222+ int64_t rank_num_per_shared_expert = shared_expert_rank_num / shared_expert_num; // 前面已有拦截, 保证进入该分支时shared_expert_num > 0
223+ int64_t max_shared_group_num = (ep_world_size + rank_num_per_shared_expert - 1) / rank_num_per_shared_expert;
224+ a = max_bs * max_shared_group_num;
225+ } else {
226+ local_moe_expert_num = moe_expert_num / (ep_world_size - shared_expert_rank_num);
227+ a = global_bs_real * std::min(local_moe_expert_num, k);
228+ }
229+ }
230+ ep_recv_cnt_num = ep_world_size * local_moe_expert_num;
231+ 
232+ 
233+ auto output_dtype = (!scales.has_value() && quant_mode == 0) ? x.scalar_type() : at::kChar;
234+ char *group_ep_ptr = const_cast<char *>(group_ep.data());
235+ at::Tensor expand_x;
236+ at::Tensor dynamic_scales;
237+ 
238+ expand_x = at::empty({a, h}, x.options().dtype(output_dtype));
239+ dynamic_scales = at::empty({a}, x.options().dtype(at::kFloat));
240+
241+ at::Tensor expert_token_nums = at::empty({local_moe_expert_num}, x.options().dtype(at::kLong));
242+ at::Tensor ep_recv_counts = at::empty({ep_recv_cnt_num}, x.options().dtype(at::kInt));
243+ at::Tensor assist_info_forcombine{nullptr};
244+ at::Tensor new_workspace = at::empty({WORKSPACESIZE / 4}, x.options().dtype(at::kInt));
245+ 
246+ at::Tensor expand_scales = at::empty({a}, x.options().dtype(at::kFloat));
247+ if (expert_scales.has_value() && expert_scales.value().defined()) {
248+ ep_recv_cnt_num = ep_world_size * local_moe_expert_num + 2 * global_bs_real * k * (ep_world_size / 8); // 2: 2 buffer, 8 ranknum per server
249+ ep_recv_counts = at::empty({ep_recv_cnt_num}, x.options().dtype(at::kInt));
250+ }
251+ 
252+ bool is_fullmesh_v2 = (comm_alg == "fullmesh_v2");
253+ 
254+ assist_info_forcombine = at::empty({std::max(bs * k, a * 128)}, x.options().dtype(at::kInt));
255+ 
256+ MoeDistributeDispatchV2Info tilingData;
257+ calculate_tilingdata(tilingData, ep_world_size, ep_rank_id, moe_expert_num, total_winsize_ep,
258+ expert_shard_type, shared_expert_num, shared_expert_rank_num, global_bs, bs, h, k, a, quant_mode,
259+ zero_expert_num, copy_expert_num, const_expert_num, expert_token_nums_type, scales, x_active_mask, performance_info);
260+ 
261+ auto stream = c10_npu::getCurrentNPUStream().stream(false);
262+ auto acl_call = [=]() mutable -> int {
263+ moe_distribute_dispatch_v2_api(stream, is_fullmesh_v2, x, expert_ids, new_workspace, mc2_context, group_ep,
264+ ep_world_size, ep_rank_id, moe_expert_num, scales, x_active_mask, expert_scales, performance_info, expand_x,
265+ dynamic_scales, assist_info_forcombine, expert_token_nums, ep_recv_counts,
266+ expand_scales, tilingData);
267+ return 0;
268+ };
269+ at_npu::native::OpCommand::RunOpApiV2("moeDistributeDispatchV2", acl_call);
270+
271+ return std::tie(expand_x, dynamic_scales, assist_info_forcombine, expert_token_nums, ep_recv_counts,
272+ expand_scales);
273+}
274+ 
275+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_moe_distribute_dispatch_v2_meta(
276+ const at::Tensor &x, const at::Tensor &expert_ids, const at::Tensor &mc2_context, c10::string_view group_ep,
277+ int64_t ep_world_size, int64_t ep_rank_id, int64_t moe_expert_num, int64_t total_winsize_ep,
278+ const c10::optional<at::Tensor> &scales, const c10::optional<at::Tensor> &x_active_mask,
279+ const c10::optional<at::Tensor> &expert_scales, const c10::optional<at::Tensor> &performance_info,
280+ int64_t expert_shard_type, int64_t shared_expert_num, int64_t shared_expert_rank_num,
281+ int64_t quant_mode, int64_t global_bs, int64_t expert_token_nums_type,
282+ c10::string_view comm_alg, int64_t zero_expert_num, int64_t copy_expert_num, int64_t const_expert_num)
283+{
284+ bool is_shared_default = ((shared_expert_num == 1) && (shared_expert_rank_num == 0));
285+ bool is_no_shared = ((shared_expert_num == 0) && (shared_expert_rank_num == 0));
286+ bool is_valid_shared = ((shared_expert_num > 0)
287+ && ((shared_expert_rank_num / shared_expert_num) > 0)
288+ && ((shared_expert_rank_num % shared_expert_num) == 0));
289+ auto x_size = x.sizes();
290+ auto expert_ids_size = expert_ids.sizes();
291+ 
292+ int64_t bs = x_size[0];
293+ int64_t h = x_size[1];
294+ int64_t k = expert_ids_size[1];
295+ 
296+ bool shared_front = (expert_shard_type == 0);
297+ int64_t local_moe_expert_num = 1;
298+ int64_t global_bs_real = (global_bs == 0) ? (bs * ep_world_size) : global_bs;
299+ int64_t a = 0;
300+ int64_t ep_recv_cnt_num = 0;
301+ if (shared_front) {
302+ if (ep_rank_id < shared_expert_rank_num) {
303+ local_moe_expert_num = 1;
304+ int64_t max_bs = global_bs_real / ep_world_size; // 前面已有拦截，保证ep_world_size > 0
305+ int64_t rank_num_per_shared_expert = shared_expert_rank_num / shared_expert_num; // 前面已有拦截, 保证进入该分支时shared_expert_num > 0
306+ int64_t max_shared_group_num = (ep_world_size + rank_num_per_shared_expert - 1) / rank_num_per_shared_expert;
307+ a = max_bs * max_shared_group_num;
308+ } else {
309+ local_moe_expert_num = moe_expert_num / (ep_world_size - shared_expert_rank_num);
310+ a = global_bs_real * std::min(local_moe_expert_num, k);
311+ }
312+ }
313+ ep_recv_cnt_num = ep_world_size * local_moe_expert_num;
314+ 
315+ 
316+ auto output_dtype = (!scales.has_value() && quant_mode == 0) ? x.scalar_type() : at::kChar;
317+ char *group_ep_ptr = const_cast<char *>(group_ep.data());
318+ at::Tensor expand_x;
319+ at::Tensor dynamic_scales;
320+ 
321+ expand_x = at::empty({a, h}, x.options().dtype(output_dtype));
322+ dynamic_scales = at::empty({a}, x.options().dtype(at::kFloat));
323+
324+ at::Tensor expert_token_nums = at::empty({local_moe_expert_num}, x.options().dtype(at::kLong));
325+ at::Tensor ep_recv_counts = at::empty({ep_recv_cnt_num}, x.options().dtype(at::kInt));
326+ at::Tensor assist_info_forcombine{nullptr};
327+ 
328+ at::Tensor expand_scales = at::empty({a}, x.options().dtype(at::kFloat));
329+ if (expert_scales.has_value() && expert_scales.value().defined()) {
330+ ep_recv_cnt_num = ep_world_size * local_moe_expert_num + 2 * global_bs_real * k * (ep_world_size / 8); // 2: 2 buffer, 8 ranknum per server
331+ ep_recv_counts = at::empty({ep_recv_cnt_num}, x.options().dtype(at::kInt));
332+ }
333+ 
334+ bool is_fullmesh_v2 = (comm_alg == "fullmesh_v2");
335+ 
336+ assist_info_forcombine = at::empty({std::max(bs * k, a * 128)}, x.options().dtype(at::kInt));
337+ return std::tie(expand_x, dynamic_scales, assist_info_forcombine, expert_token_nums, ep_recv_counts,
338+ expand_scales);
339+}
340+ 
341+// Register Ascend implementations for MoeDistributeDispatchV2
342+TORCH_LIBRARY_IMPL(ascend_ops, PrivateUse1, m)
343+{
344+ m.impl("MoeDistributeDispatchV2", TORCH_FN(npu_moe_distribute_dispatch_v2));
345+}
346+ 
347+TORCH_LIBRARY_IMPL(ascend_ops, Meta, m)
348+{
349+ m.impl("MoeDistributeDispatchV2", &npu_moe_distribute_dispatch_v2_meta);
350+}
351+ 
352+} // namespace MoeDistributeDispatchV2
353+} // namespace ascend_ops
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_dispatch_v2/ascend910_93/moe_distribute_dispatch_v2_torch.h+109-0
@@ -0,0 +1,109 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify it.
3+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
8+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+ * See LICENSE in the root of the software repository for the full text of the License.
10+ */
11+ 
12+/*!
13+ * \file moe_distribute_dispatch_v2_torch.h
14+ * \brief
15+ */
16+#ifndef ASCEND_OPS_MOE_DISTRIBUTE_DISPATCH_TORCH_H
17+#define ASCEND_OPS_MOE_DISTRIBUTE_DISPATCH_TORCH_H
18+ 
19+#include <ATen/ATen.h>
20+#include <vector>
21+#include <torch/all.h>
22+ 
23+template <typename TensorType, typename ElementType>
24+ElementType *get_first_tensor_address_by_type(const TensorType &input, bool allow_empty)
25+{
26+ const auto &get_tensor = [&]() -> const torch::Tensor * {
27+ // 处理 optional<torch::Tensor>
28+ if constexpr (std::is_same_v<TensorType, c10::optional<torch::Tensor>>) {
29+ if (!input.has_value()) {
30+ if (!allow_empty)
31+ TORCH_CHECK(false, "optional<Tensor> has no value");
32+ return nullptr;
33+ }
34+ if (!input->defined()) {
35+ if (!allow_empty)
36+ TORCH_CHECK(false, "optional<Tensor> is undefined");
37+ return nullptr;
38+ }
39+ return &input.value();
40+ }
41+ // 处理 optional<TensorList>
42+ else if constexpr (std::is_same_v<TensorType, c10::optional<at::TensorList>>) {
43+ if (!input.has_value()) {
44+ if (!allow_empty)
45+ TORCH_CHECK(false, "optional<TensorList> has no value");
46+ return nullptr;
47+ }
48+ if (input->empty()) {
49+ if (!allow_empty)
50+ TORCH_CHECK(false, "optional<TensorList> is empty");
51+ return nullptr;
52+ }
53+ const auto &tensor = input.value()[0];
54+ if (!tensor.defined()) {
55+ if (!allow_empty)
56+ TORCH_CHECK(false, "First tensor in optional<TensorList> is undefined");
57+ return nullptr;
58+ }
59+ return &tensor;
60+ }
61+ // 处理 TensorList
62+ else if constexpr (std::is_same_v<TensorType, at::TensorList>) {
63+ if (input.empty()) {
64+ if (!allow_empty)
65+ TORCH_CHECK(false, "TensorList is empty");
66+ return nullptr;
67+ }
68+ const auto &tensor = input[0];
69+ if (!tensor.defined()) {
70+ if (!allow_empty)
71+ TORCH_CHECK(false, "First tensor in TensorList is undefined");
72+ return nullptr;
73+ }
74+ return &tensor;
75+ }
76+ // 处理 torch::Tensor
77+ else if constexpr (std::is_same_v<TensorType, torch::Tensor>) {
78+ if (!input.defined()) {
79+ if (!allow_empty)
80+ TORCH_CHECK(false, "Tensor is undefined");
81+ return nullptr;
82+ }
83+ return &input;
84+ }
85+ // 不支持的类型
86+ else {
87+ static_assert(std::is_same_v<TensorType, void>, "Unsupported tensor type");
88+ return nullptr;
89+ }
90+ };
91+ 
92+ const torch::Tensor *tensor_ptr = get_tensor();
93+ if (!tensor_ptr) {
94+ return nullptr;
95+ }
96+ 
97+ // 修复：使用非模板版本的data_ptr()，然后进行类型转换
98+ void *raw_ptr = tensor_ptr->data_ptr();
99+ return reinterpret_cast<ElementType *>(raw_ptr);
100+}
101+ 
102+template <typename TensorType>
103+void *get_first_tensor_address(c10::ScalarType dataType, const TensorType &input, bool allow_empty = false)
104+{
105+ return get_first_tensor_address_by_type<TensorType, void>(input, allow_empty);
106+}
107+ 
108+ 
109+#endif
Aexamples/fast_kernel_launch_example/csrc/moe_distribute_dispatch_v2/ascend910_93/op_kernel/common.h+183-0
@@ -0,0 +1,183 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file common.h
13+ * \brief
14+ */
15+ 
16+#ifndef MC2_MOE_DISPATCH_COMM_H
17+#define MC2_MOE_DISPATCH_COMM_H
18+ 
19+constexpr uint32_t NEED_ONE_HUNDRED_AND_TWENTY_SEVEN = 127;
20+constexpr uint32_t RIGHT_SHIFT_BIT_SEVEN = 7;
21+constexpr uint32_t NEED_THIRTY_FIRST = 31;
22+constexpr uint32_t ALIGN_UP_TO_2_MASK = 1;
23+constexpr uint32_t ALIGN_UP_TO_32_MASK = 31;
24+constexpr uint32_t ALIGN_UP_TO_64_MASK = 64;
25+constexpr uint32_t ALIGN_UP_TO_128_MASK = 127;
26+constexpr uint32_t ALIGN_UP_TO_256_MASK = 255;
27+constexpr uint32_t ALIGN_UP_TO_512_MASK = 511;
28+constexpr uint32_t RIGHT_SHIFT_BIT_FIVE = 5;
29+constexpr uint32_t FIVE_HUNDRED_AND_ELEVEN = 511;
30+constexpr uint32_t RIGHT_SHIFT_BIT_NINE = 9;
31+ 
32+namespace AscendC {
33+template <typename T1, typename T2>
34+__aicore__ inline T2 Ceil(T1 x, T1 y)
35+{
36+ return (x + y - 1) / y;
37+}
38+ 
39+template <typename T>
40+__aicore__ inline T Ceil32(T x)
41+{
42+ return (x + NEED_THIRTY_FIRST) >> RIGHT_SHIFT_BIT_FIVE;
43+}
44+ 
45+template <typename T>
46+__aicore__ inline T Ceil128(T x)
47+{
48+ return (x + NEED_ONE_HUNDRED_AND_TWENTY_SEVEN) >> RIGHT_SHIFT_BIT_SEVEN;
49+}
50+ 
51+template <typename T>
52+__aicore__ inline T Ceil512(T x)
53+{
54+ return (x + FIVE_HUNDRED_AND_ELEVEN) >> RIGHT_SHIFT_BIT_NINE;
55+}
56+ 
57+template <typename T1, typename T2>
58+__aicore__ inline T2 Align(T1 x, T1 y)
59+{
60+ return Ceil<T1, T2>(x, y) * y;
61+}
62+ 
63+template <typename T>
64+__aicore__ inline T Align2(T x)
65+{
66+ return (x + ALIGN_UP_TO_2_MASK) & (~ALIGN_UP_TO_2_MASK);
67+}
68+ 
69+template <typename T>
70+__aicore__ inline T Align32(T x)
71+{
72+ return (x + ALIGN_UP_TO_32_MASK) & (~ALIGN_UP_TO_32_MASK);
73+}
74+ 
75+template <typename T>
76+__aicore__ inline T Align64(T x)
77+{
78+ return (x + ALIGN_UP_TO_64_MASK) & (~ALIGN_UP_TO_64_MASK);
79+}
80+ 
81+template <typename T>
82+__aicore__ inline T Align128(T x)
83+{
84+ return (x + ALIGN_UP_TO_128_MASK) & (~ALIGN_UP_TO_128_MASK);
85+}
86+ 
87+template <typename T>
88+__aicore__ inline T Align256(T x)
89+{
90+ return (x + ALIGN_UP_TO_256_MASK) & (~ALIGN_UP_TO_256_MASK);
91+}
92+ 
93+template <typename T>
94+__aicore__ inline T Align512(T x)
95+{
96+ return (x + ALIGN_UP_TO_512_MASK) & (~ALIGN_UP_TO_512_MASK);
97+}
98+ 
99+template <MicroAPI::HistogramsType htype, typename T, typename U>
100+static __aicore__ inline void HistogramsVf(__local_mem__ U* dst, __local_mem__ T* src, uint16_t repeatElm,
101+ uint16_t halfRepeat, uint32_t totalElm, uint16_t repeatTimes)
102+{
103+ AscendC::MicroAPI::RegTensor<T> srcReg;
104+ AscendC::MicroAPI::RegTensor<U> dst0Reg;
105+ AscendC::MicroAPI::RegTensor<U> dst1Reg;
106+ AscendC::MicroAPI::MaskReg pregOut = AscendC::MicroAPI::CreateMask<T>();
107+ MicroAPI::Duplicate(dst0Reg, 0);
108+ MicroAPI::Duplicate(dst1Reg, 0);
109+ for (uint16_t i = 0; i < repeatTimes; ++i) {
110+ MicroAPI::MaskReg preg = MicroAPI::UpdateMask<T>(totalElm);
111+ MicroAPI::DataCopy(srcReg, src + repeatElm * i);
112+ MicroAPI::Histograms<T, U, MicroAPI::HistogramsBinType::BIN0, htype>(dst0Reg, srcReg, preg);
113+ MicroAPI::Histograms<T, U, MicroAPI::HistogramsBinType::BIN1, htype>(dst1Reg, srcReg, preg);
114+ }
115+ MicroAPI::DataCopy(dst, dst0Reg, pregOut);
116+ MicroAPI::DataCopy(dst + halfRepeat, dst1Reg, pregOut);
117+}
118+ 
119+__aicore__ inline void GetExpertFreq(LocalTensor<uint16_t>& dstLocal, LocalTensor<uint8_t>& srcLocal, uint32_t totalElm)
120+{
121+ uint32_t repeatElm = GetVecLen();
122+ uint16_t repeatTimes = Ceil<uint32_t, uint16_t>(totalElm, repeatElm);
123+ __local_mem__ uint8_t* src = (__local_mem__ uint8_t*)srcLocal.GetPhyAddr();
124+ __local_mem__ uint16_t* dst = (__local_mem__ uint16_t*)dstLocal.GetPhyAddr();
125+ VF_CALL<HistogramsVf<MicroAPI::HistogramsType::FREQUENCY, uint8_t, uint16_t>>(dst, src, repeatElm, repeatElm >> 1,
126+ totalElm, repeatTimes);
127+ PipeBarrier<PIPE_V>();
128+}
129+ 
130+__aicore__ inline void GetExpertCumSum(LocalTensor<uint16_t>& dstLocal, LocalTensor<uint8_t>& srcLocal,
131+ uint32_t totalElm)
132+{
133+ uint32_t repeatElm = GetVecLen();
134+ uint16_t repeatTimes = Ceil<uint32_t, uint16_t>(totalElm, repeatElm);
135+ __local_mem__ uint8_t *src = (__local_mem__ uint8_t *)srcLocal.GetPhyAddr();
136+ __local_mem__ uint16_t *dst = (__local_mem__ uint16_t *)dstLocal.GetPhyAddr();
137+ VF_CALL<HistogramsVf<MicroAPI::HistogramsType::ACCUMULATE, uint8_t, uint16_t>>(dst, src, repeatElm, repeatElm >> 1,
138+ totalElm, repeatTimes);
139+ PipeBarrier<PIPE_V>();
140+}
141+ 
142+static __aicore__ inline void ReduceLoop(__local_mem__ int32_t* dst, __local_mem__ int32_t* src, uint16_t repeat0Times,
143+ uint32_t repeat0SrcStride, uint16_t repeat1Times, uint32_t repeat1Stride,
144+ uint32_t repeat1Element, uint32_t repeat0DstStride)
145+{
146+ MicroAPI::MaskReg maskFirst = MicroAPI::CreateMask<int32_t, MicroAPI::MaskPattern::VL1>();
147+ for (uint16_t i0 = 0; i0 < repeat0Times; ++i0) {
148+ MicroAPI::RegTensor<int32_t> sumReg;
149+ uint32_t elements = repeat1Element;
150+ MicroAPI::Duplicate(sumReg, 0);
151+ for (uint16_t i1 = 0; i1 < repeat1Times; ++i1) {
152+ MicroAPI::RegTensor<int32_t> srcReg;
153+ MicroAPI::RegTensor<int32_t> dstReg;
154+ MicroAPI::MaskReg mask = MicroAPI::UpdateMask<int32_t>(elements);
155+ MicroAPI::DataCopy(srcReg, src + i0 * repeat0SrcStride + i1 * repeat1Stride);
156+ MicroAPI::ReduceSum(dstReg, srcReg, mask);
157+ MicroAPI::Add(sumReg, sumReg, dstReg, maskFirst);
158+ }
159+ MicroAPI::DataCopy<int32_t, MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dst + i0 * repeat0DstStride, sumReg,
160+ maskFirst);
161+ }
162+}
163+ 
164+__aicore__ inline void GetReduceSum(LocalTensor<int32_t>& dstLocal, LocalTensor<int32_t>& srcLocal,
165+ uint16_t repeat0Times, uint32_t repeat0SrcStride, uint32_t repeat1Element,
166+ uint32_t repeat0DstStride)
167+{
168+ if (repeat0Times == 0 || repeat1Element == 0) {
169+ return;
170+ }
171+ 
172+ uint32_t repeat1Stride = GetVecLen() / sizeof(int32_t);
173+ uint16_t repeat1Times = Ceil<uint32_t, uint16_t>(repeat1Element, repeat1Stride);
174+ __local_mem__ int32_t* src = (__local_mem__ int32_t*)srcLocal.GetPhyAddr();
175+ __local_mem__ int32_t* dst = (__local_mem__ int32_t*)dstLocal.GetPhyAddr();
176+ VF_CALL<ReduceLoop>(dst, src, repeat0Times, repeat0SrcStride, repeat1Times, repeat1Stride, repeat1Element,
177+ repeat0DstStride);
178+ PipeBarrier<PIPE_V>();
179+}
180+ 
181+} // namespace AscendC
182+ 
183+#endif