init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

csrc/mc2/CMakeLists.txt

@@ -0,0 +1,17 @@
+# -----------------------------------------------------------------------------------------------------------
+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+# CANN Open Software License Agreement Version 2.0 (the "License").
+# Please refer to the License for details. You may not use this file except in compliance with the License.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+# See LICENSE in the root of the software repository for the full text of the License.
+# -----------------------------------------------------------------------------------------------------------
+
+file(GLOB SUBDIRECTORIES LIST_DIRECTORIES true RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
+foreach(SUBDIR ${SUBDIRECTORIES})
+  # 检查子目录中是否存在 CMakeLists.txt
+  if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${SUBDIR}/CMakeLists.txt)
+    add_subdirectory(${SUBDIR})
+  endif()
+endforeach()

csrc/mc2/dispatch_ffn_combine/CMakeLists.txt

@@ -0,0 +1,19 @@
+# -----------------------------------------------------------------------------------------------------------
+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+# CANN Open Software License Agreement Version 2.0 (the "License").
+# Please refer to the License for details. You may not use this file except in compliance with the License.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+# See LICENSE in the root of the software repository for the full text of the License.
+# -----------------------------------------------------------------------------------------------------------
+
+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
+if(NOT ENABLE_TEST AND NOT BENCHMARK)
+    list(REMOVE_ITEM CURRENT_DIRS tests)
+endif()
+foreach(SUB_DIR ${CURRENT_DIRS})
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
+        add_subdirectory(${SUB_DIR})
+    endif()
+endforeach()

csrc/mc2/dispatch_ffn_combine/dispatch_ffn_combine_torch_adpt.h

@@ -0,0 +1,89 @@
+/*
+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OF CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef DISPATCH_FFN_COMBINE_TORCH_ADPT_H
+#define DISPATCH_FFN_COMBINE_TORCH_ADPT_H
+
+namespace vllm_ascend {
+std::tuple<at::Tensor&, at::Tensor&> dispatch_ffn_combine(
+    const at::Tensor& x,
+    const at::TensorList& weight1,
+    const at::TensorList& weight2,
+    const at::Tensor& expert_idx,
+    const at::TensorList& scale1,
+    const at::TensorList& scale2,
+    const at::TensorList& bias1,
+    const at::TensorList& bias2,
+    const at::Tensor& probs,
+    c10::string_view group,
+    int64_t max_output_size,
+    at::Tensor& out,
+    at::Tensor& expert_token_nums,
+    const c10::optional<at::Tensor>& x_active_mask,
+    double swiglu_limit
+) {
+    char *group_ep_ptr = const_cast<char *>(group.data());
+    bool is_int8 = weight1[0].dtype() == at::kChar;
+    bool is_int4 = weight1[0].dtype() == at::kInt;
+    if (is_int8) {
+        EXEC_NPU_CMD(aclnnDispatchFFNCombine,
+                 x,
+                 weight1,
+                 weight2,
+                 expert_idx,
+                 scale1,
+                 scale2,
+                 probs,
+                 x_active_mask.has_value() ? x_active_mask.value() : at::Tensor(),
+                 group_ep_ptr,
+                 max_output_size,
+                 swiglu_limit,
+                 out,
+                 expert_token_nums);
+    } else if (is_int4){
+        EXEC_NPU_CMD(aclnnDispatchFFNCombineW4A8,
+                 x,
+                 weight1,
+                 weight2,
+                 expert_idx,
+                 scale1,
+                 scale2,
+                 bias1,
+                 bias2,
+                 probs,
+                 x_active_mask.has_value() ? x_active_mask.value() : at::Tensor(),
+                 group_ep_ptr,
+                 max_output_size,
+                 swiglu_limit,
+                 out,
+                 expert_token_nums);
+    } else {
+        EXEC_NPU_CMD(aclnnDispatchFFNCombineBF16,
+                 x,
+                 weight1,
+                 weight2,
+                 expert_idx,
+                 scale1,
+                 scale2,
+                 probs,
+                 group_ep_ptr,
+                 max_output_size,
+                 out,
+                 expert_token_nums);
+    }    
+    return {out, expert_token_nums};
+}
+}
+#endif

csrc/mc2/dispatch_ffn_combine/op_host/CMakeLists.txt

@@ -0,0 +1,38 @@
+set(_DISPATCH_FFN_INC_OPTS)
+if (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/aarch64-linux/ascendc/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/aarch64-linux/ascendc/include)
+elseif (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/arm64-linux/ascendc/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/arm64-linux/ascendc/include)
+elseif (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/ascendc/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/ascendc/include)
+endif()
+if (EXISTS ${CMAKE_SOURCE_DIR}/third_party/catlass/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${CMAKE_SOURCE_DIR}/third_party/catlass/include)
+endif()
+
+add_op_to_compiled_list()
+
+if (BUILD_OPEN_PROJECT)
+    target_sources(op_host_aclnnInner PRIVATE
+        dispatch_ffn_combine_def.cpp
+    )
+endif()
+
+add_ops_compile_options(
+    OP_NAME DispatchFFNCombine
+    OPTIONS
+        --cce-auto-sync=on
+        -Wno-deprecated-declarations
+        -DHCCL_COMM
+        -DCATLASS_ARCH=2201
+        ${_DISPATCH_FFN_INC_OPTS}
+)
+
+if (NOT BUILD_OPS_RTY_KERNEL)
+    add_modules_sources(OPTYPE dispatch_ffn_combine ACLNNTYPE aclnn_inner)
+    target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
+        ${CMAKE_CURRENT_SOURCE_DIR}
+        ${CMAKE_CURRENT_SOURCE_DIR}/../op_kernel
+    )
+endif()
+

csrc/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_def.cpp

@@ -0,0 +1,99 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file dispatch_ffn_combine_def.cpp
+ * \brief
+ */
+#include "register/op_def_registry.h"
+
+namespace ops {
+class DispatchFFNCombine : public OpDef {
+ public:
+  explicit DispatchFFNCombine(const char *name) : OpDef(name) {
+    this->Input("a")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_BF16})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("w1")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ})
+        .IgnoreContiguous();
+    this->Input("w2")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ})
+        .IgnoreContiguous();
+    this->Input("expertIdx")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("scale1")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("scale2")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("probs")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("xActiveMaskOptional")
+        .ParamType(OPTIONAL)
+        .DataType({ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+
+    // Output
+    this->Output("out")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_BF16})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND,ge::FORMAT_ND});
+    this->Output("expert_token_nums")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+
+    this->Attr("group").AttrType(REQUIRED).String();
+    this->Attr("M").AttrType(OPTIONAL).Int();
+    this->Attr("transB").AttrType(OPTIONAL).Bool(false);
+    this->Attr("weightNz").AttrType(OPTIONAL).Bool(false);
+    this->Attr("swigluLimit").AttrType(OPTIONAL).Float(0.0f);
+
+    OpAICoreConfig aicore_config;
+    aicore_config.DynamicCompileStaticFlag(true)
+        .DynamicFormatFlag(true)
+        .DynamicRankSupportFlag(true)
+        .DynamicShapeSupportFlag(true)
+        .NeedCheckSupportFlag(false)
+        .PrecisionReduceFlag(true)
+        .ExtendCfgInfo("aclnnSupport.value", "support_aclnn")
+        .ExtendCfgInfo("jitCompile.flag", "static_false")
+        .ExtendCfgInfo("multiKernelSupportDynamicGraph.value", "multi_kernel");
+    this->AICore().AddConfig("ascend910_93", aicore_config);
+    this->AICore().AddConfig("ascend910b", aicore_config);
+    this->MC2().HcclGroup("group");
+  }
+};
+
+OP_ADD(DispatchFFNCombine);
+}  // namespace ops

csrc/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_proto.cpp

@@ -0,0 +1,42 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file dispatch_ffn_proto.cpp
+ * \brief
+ */
+#include <graph/utils/type_utils.h>
+#include <register/op_impl_registry.h>
+// #include "../../common/ophost/op_util.h"
+// #include "../../common/ophost/hcom_topo_info.h"
+// #include "log/ops_log.h"
+
+using namespace ge;
+namespace ops {
+const size_t ATTR_GROUP = 0;
+const size_t ATTR_RANK_SIZE = 1;
+const size_t SUPPORT_DIM_SIZE = 2;
+
+static ge::graphStatus InferShapeDispatchFFNCombine(gert::InferShapeContext* context) {
+  (void) context;
+  return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus InferDataTypeDispatchFFNCombine(gert::InferDataTypeContext* context) {
+  (void) context;
+  // auto d_type = context->GetInputDataType(0);
+  // context->SetOutputDataType(0, d_type);
+  return ge::GRAPH_SUCCESS;
+}
+
+IMPL_OP_INFERSHAPE(DispatchFFNCombine)
+  .InferShape(InferShapeDispatchFFNCombine)
+  .InferDataType(InferDataTypeDispatchFFNCombine);
+}  // namespace ops

csrc/mc2/dispatch_ffn_combine/op_host/dispatch_ffn_combine_tiling.cpp

@@ -0,0 +1,339 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+/*!
+ * \file dispatch_ffn_tiling.cpp
+ * \brief
+ */
+#include "vector"
+#include "register/tilingdata_base.h"
+#include "tiling/tiling_api.h"
+#include "tiling_base/error_log.h"
+#include "hcom_topo_info.h"
+#include "register/op_def_registry.h"
+#include "../op_kernel/dispatch_ffn_combine_tiling.h"
+#include <vector>
+#include <map>
+#include <algorithm>
+#include "../op_kernel/moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h"
+
+using namespace AscendC;
+using namespace ge;
+
+#define HCCL_BUFFSIZE  "HCCL_BUFFSIZE"
+
+namespace {
+    // 1. Constant definitions
+    const char *K_INNER_DEBUG = "DispatchFFNCombine Tiling Debug";
+    constexpr uint32_t ATTR_GROUP_INDEX = 0;
+    constexpr uint32_t ATTR_MAX_OUTPUT_SIZE_INDEX = 1;
+    constexpr uint32_t ATTR_IS_TRANS_B = 2;
+    constexpr uint32_t ATTR_WEIGHT_NZ = 3;
+    constexpr uint32_t ATTR_SWIGLU_LIMIT = 4;
+    constexpr uint64_t INIT_TILINGKEY = 1000000;
+    constexpr uint64_t TILINGKEY_TRANS_B = 1U;
+    constexpr uint64_t TILINGKEY_WEIGHT_NZ = 10;
+    constexpr uint32_t X_INDEX = 0;
+    constexpr uint32_t WEIGHT_INDEX = 1;
+    constexpr uint32_t WEIGHT2_INDEX = 2;
+    constexpr uint32_t EXPERTID_INDEX = 3;
+    constexpr uint32_t X_ACTIVE_MASK_INDEX = 7;
+    constexpr uint32_t BLOCK_NUM = 20;
+    constexpr uint32_t SYSTEM_NEED_WORKSPACE = 16 * 1024 * 1024;
+    constexpr uint64_t MB_SIZE = 1024 * 1024UL;
+}
+
+namespace optiling {
+
+static int32_t CeilDev(int32_t num, int32_t div)
+{
+    if (div == 0) {
+        return 0;
+    }
+    return (num + div - 1) / div;
+}
+
+static uint64_t GetMaxWindowSize()
+{
+    uint16_t defaultWindowSize = 200;
+    const char* hccl_buffsize_env = getenv(HCCL_BUFFSIZE);
+    if (hccl_buffsize_env != nullptr) {
+        try {
+            std::string envStr(hccl_buffsize_env);
+            unsigned long val = std::stoul(envStr);
+            if (val <= std::numeric_limits<uint16_t>::max()) {
+                defaultWindowSize = static_cast<uint16_t>(val);
+            } else {
+                OP_LOGW(K_INNER_DEBUG, "HCCL_BUFFSIZE value %lu is out of range, using default.", val);
+            }
+        } catch (const std::exception& e) {
+            OP_LOGE(K_INNER_DEBUG, "Exception encountered when parsing env HCCL_BUFFSIZE: %s", e.what());
+        }
+    } else {
+        OP_LOGD(K_INNER_DEBUG, "Env HCCL_BUFFSIZE not set");
+    }
+    const uint64_t maxWindowSize = static_cast<uint64_t>(defaultWindowSize) * MB_SIZE;
+    OP_LOGD(K_INNER_DEBUG, "Get maxWindowSize is %lu", maxWindowSize);
+    return maxWindowSize;
+}
+
+// Parse and validate rankId, group, worldSize, and isTransB attributes
+static ge::graphStatus DispatchFFNCombineCheckAttrAndSetTiling(gert::TilingContext *context, DispatchFFNCombineInfo& info)
+{
+    auto attrs = context->GetAttrs();
+    OP_TILING_CHECK(attrs == nullptr, OP_LOGE(K_INNER_DEBUG, "attrs is null."), return ge::GRAPH_FAILED);
+
+    // TODO: set, validate, and print tiling data related to attributes
+    auto groupPtr = attrs->GetAttrPointer<char>(static_cast<int>(ATTR_GROUP_INDEX));
+    auto maxOutputSizePtr = attrs->GetAttrPointer<int>(ATTR_MAX_OUTPUT_SIZE_INDEX);
+    auto is_trans_b = attrs->GetAttrPointer<bool>(ATTR_IS_TRANS_B);
+    auto weight_nz = attrs->GetAttrPointer<bool>(ATTR_WEIGHT_NZ);
+    auto swiglu_limit = attrs->GetAttrPointer<float>(ATTR_SWIGLU_LIMIT);
+    OP_TILING_CHECK(groupPtr == nullptr || strlen(groupPtr) == 0,
+    OP_LOGE(K_INNER_DEBUG, "group is invalid."), return GRAPH_FAILED);
+
+    OP_TILING_CHECK(is_trans_b == nullptr,
+        OP_LOGE(K_INNER_DEBUG, "is_trans_b is invalid."), return GRAPH_FAILED);
+    OP_TILING_CHECK(weight_nz == nullptr,
+        OP_LOGE(K_INNER_DEBUG, "weight_nz is invalid."), return GRAPH_FAILED);
+
+    info.maxOutputSize = *maxOutputSizePtr;
+    info.isTransposeB = *is_trans_b;
+    info.isWeightNz = *weight_nz;
+    info.swigluLimit = swiglu_limit != nullptr ? *swiglu_limit : 0.0f;
+
+    int64_t rankSize;
+    (void)ge::HcomTopoInfo::Instance().GetGroupRankSize(groupPtr, rankSize);
+    info.worldSize = rankSize;
+
+    OP_LOGD(K_INNER_DEBUG, "maxOutputSize=%d ", info.maxOutputSize);
+    OP_LOGD(K_INNER_DEBUG, "rankSize=%d ", info.worldSize);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+// Extract shapes of input tensors A and B to compute M, K, N
+static ge::graphStatus DispatchFFNCombineCheckShapeAndSetTiling(gert::TilingContext *context, DispatchFFNCombineInfo &info)
+{
+    const char *nodeName = context->GetNodeName();
+
+    const gert::StorageShape *aStorageShape = context->GetInputShape(X_INDEX);
+    auto expertIdxTensor = context->GetDynamicInputTensor(EXPERTID_INDEX, 0);
+    uint32_t M = aStorageShape->GetStorageShape().GetDim(0);
+    uint32_t K = aStorageShape->GetStorageShape().GetDim(1);
+
+    auto wTensor = context->GetDynamicInputTensor(WEIGHT_INDEX, 0);
+    uint32_t wTensorDims = wTensor->GetOriginShape().GetDimNum();
+    uint32_t N = wTensor->GetStorageShape().GetDim(wTensorDims - 1);
+
+    uint32_t topK = expertIdxTensor->GetStorageShape().GetDim(1);
+    uint32_t listLen = 0;
+    while (true) {
+        auto wTensorT = context->GetDynamicInputTensor(WEIGHT_INDEX, ++listLen);
+        if (wTensorT == nullptr) {break;}
+    }
+
+    uint32_t expertPerRank;
+    if (listLen == 1) {
+        expertPerRank = wTensor->GetStorageShape().GetDim(0);
+    } else {
+        expertPerRank = listLen;
+    }
+
+    info.M = M;
+    info.N = N;
+    info.K = K;
+    info.expertPerRank = expertPerRank;
+    info.topK = topK;
+    info.listLen = listLen;
+    OP_LOGD(K_INNER_DEBUG, "M=%d ", info.M);
+    OP_LOGD(K_INNER_DEBUG, "K=%d ", info.K);
+    OP_LOGD(K_INNER_DEBUG, "N=%d ", info.N);
+    OP_LOGD(K_INNER_DEBUG, "expertPerRank=%d ", info.expertPerRank);
+    OP_LOGD(K_INNER_DEBUG, "topK=%d ", info.topK);
+    OP_LOGD(K_INNER_DEBUG, "listLen=%d ", info.listLen);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+// Get hardware info such as AI Core count and UB capacity for the current chip platform.
+static ge::graphStatus DispatchFFNCombineGetPlatformInfoAndSetTiling(gert::TilingContext *context, DispatchFFNCombineInfo& info)
+{
+    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
+    uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
+    uint64_t ubSize = 0U;
+    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
+    info.aivNum = aivNum;
+    info.totalUbSize = ubSize;
+
+    OP_LOGD(K_INNER_DEBUG, "aivNum=%d", info.aivNum);
+    OP_LOGD(K_INNER_DEBUG, "ubSize=%lu", info.totalUbSize);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus CheckXActiveMaskShape(gert::TilingContext *context, const char *nodeName, DispatchFFNCombineInfo &info)
+{
+    const gert::StorageShape* xActiveMaskStorageShape = context->GetOptionalInputShape(X_ACTIVE_MASK_INDEX);
+    if (xActiveMaskStorageShape != nullptr) {
+        OP_TILING_CHECK(xActiveMaskStorageShape->GetStorageShape().GetDimNum() != 1,
+            OP_LOGE(nodeName, "xActiveMask shape dims must be 1, but current dim num is %lu.",
+                xActiveMaskStorageShape->GetStorageShape().GetDimNum()), return ge::GRAPH_FAILED);
+        const int64_t xActiveMaskDim0 = xActiveMaskStorageShape->GetStorageShape().GetDim(0);
+        OP_TILING_CHECK(xActiveMaskDim0 != static_cast<int64_t>(info.M),
+            OP_LOGE(nodeName, "xActiveMask Dim0 must be M(%u), but current dim is %lu.", info.M, xActiveMaskDim0),
+            return ge::GRAPH_FAILED);
+    }
+    return ge::GRAPH_SUCCESS;
+}
+
+void SetTilingData(CoCTiling &cocTilingData, DispatchFFNCombineInfo &info)
+{
+    cocTilingData.m0 = 128;
+    cocTilingData.k0 = 256;
+    cocTilingData.n0 = 256;
+    cocTilingData.swizzleDirect = 1;
+    cocTilingData.swizzleOffset = 7;
+    cocTilingData.ubMoveNum = 16 * 1024;
+    cocTilingData.pValue = 1;
+    cocTilingData.commNpuSplit = info.worldSize;
+    cocTilingData.commDataSplit = 1;
+    cocTilingData.lenPerLoop = cocTilingData.m0 * cocTilingData.n0 / 2;
+}
+
+// Main scheduling function:
+// Get tilingData ➝ check Attr ➝ check Shape ➝ get platform info 
+// ➝ call SetTilingData (based on rank count) ➝ set blockDim ➝ set tilingKey ➝ set workspace ➝ configure communication parameters
+
+static ge::graphStatus DispatchFFNCombineTilingFuncImpl(gert::TilingContext *context)
+{
+    const char *nodeName = context->GetNodeName();
+    OP_LOGI(nodeName, "Enter DispatchFFNCombine tiling func.");
+
+    // 1. tilingData
+    DispatchFFNCombineTilingData *tilingData = context->GetTilingData<DispatchFFNCombineTilingData>();
+    OP_TILING_CHECK(tilingData == nullptr, OP_LOGE(nodeName, "tilingData is nullptr."),
+        return ge::GRAPH_FAILED);
+    OP_LOGI(nodeName, "DispatchFFNCombine get tilingData.");
+    DispatchFFNCombineInfo& info = tilingData->dispatchFFNCombineInfo;
+    OP_LOGI(nodeName, "DispatchFFNCombine get tilingData info.");
+
+    OP_TILING_CHECK(DispatchFFNCombineCheckAttrAndSetTiling(context, info) != ge::GRAPH_SUCCESS,
+        OP_LOGE(context->GetNodeName(), "DispatchFFNCombine CheckAttrAndSetTiling Failed"),
+        return ge::GRAPH_FAILED);
+    OP_TILING_CHECK(DispatchFFNCombineCheckShapeAndSetTiling(context, info) != ge::GRAPH_SUCCESS,
+        OP_LOGE(context->GetNodeName(), "DispatchFFNCombine CheckShapeAndSetTiling Failed"),
+        return ge::GRAPH_FAILED);
+    OP_TILING_CHECK(DispatchFFNCombineGetPlatformInfoAndSetTiling(context, info) != ge::GRAPH_SUCCESS,
+        OP_LOGE(context->GetNodeName(), "DispatchFFNCombine GetPlatformInfoAndSetTiling Failed"),
+        return ge::GRAPH_FAILED);
+    OP_TILING_CHECK(CheckXActiveMaskShape(context, nodeName, info) != ge::GRAPH_SUCCESS,
+        OP_LOGE(context->GetNodeName(), "DispatchFFNCombine CheckXActiveMaskShape Failed"),
+        return ge::GRAPH_FAILED);
+
+    SetTilingData(tilingData->cocTiling, info);
+
+    // 2. set blockDim
+    uint32_t blockDim = 1U;
+    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
+    auto aicNum = ascendcPlatform.GetCoreNumAic();
+    auto aivNum = ascendcPlatform.GetCoreNumAiv();
+    blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum);
+    context->SetBlockDim(blockDim);
+
+    // 3. set tilingKey
+    uint64_t tilingKey = INIT_TILINGKEY;
+    tilingKey += info.isTransposeB ? TILINGKEY_TRANS_B : 0;
+    tilingKey += info.isWeightNz ? TILINGKEY_WEIGHT_NZ : 0;
+    context->SetTilingKey(tilingKey);
+
+    OP_LOGD(K_INNER_DEBUG, "tilingKey=%d", tilingKey);
+
+    optiling::MoeInitRoutingQuantV2TilingBase moeInitRoutingQuantV2TilingBase;
+    int64_t inuptXDtypeSize = sizeof(int16_t);
+    int64_t scaleDim0 = 0;
+    int64_t ubSize = 196352;
+    int64_t expertCapacity = 0;
+    int64_t expertNum = info.expertPerRank * info.worldSize + 1;    // enable expertId == expertNum
+    int64_t activeNum = 0;
+    int64_t dropPadMode = 0;
+     int64_t expertTokensCountOrCumsumFlag = 2;
+     bool expertTokensBeforeCapacityFlag = false;
+     int64_t quantMode = 1;
+     uint32_t aivNumInitRouting = 2 * BLOCK_NUM;
+    moeInitRoutingQuantV2TilingBase.DoTiling(info.M, info.K, info.topK, expertCapacity, expertNum, activeNum, dropPadMode, 
+        expertTokensCountOrCumsumFlag, expertTokensBeforeCapacityFlag, inuptXDtypeSize, quantMode, scaleDim0, aivNumInitRouting, ubSize);
+    uint64_t initRoutingQuantTilingKey = moeInitRoutingQuantV2TilingBase.tilingKey_;
+    size_t initRoutingWorkspace = moeInitRoutingQuantV2TilingBase.workspaceSize_;
+
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData = moeInitRoutingQuantV2TilingBase.quantTilingData;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.vbsComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.vbsComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.vmsMiddleComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.vmsMiddleComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.sortOutComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.sortOutComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.srcToDstComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.srcToDstComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.srcToDstCapacityComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.srcToDstCapacityComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.gatherOutComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.gatherOutComputeParamsOp;
+    tilingData->cocTiling.initRoutingQuantTilingKey = initRoutingQuantTilingKey;
+
+    uint64_t maxWindowSize = GetMaxWindowSize();
+    uint64_t actualSize = static_cast<uint64_t>(info.M) * info.topK * info.K * sizeof(int8_t) * 3 + 10 * MB_SIZE ;
+    OP_TILING_CHECK((actualSize > maxWindowSize),
+        OP_LOGE(nodeName, "HCCL_BUFFSIZE is too SMALL, m = %lu, k = %lu, topK = %lu"
+            " expected HCCL_BUFFSIZE is ((m * k * topK * sizeof(int8_t)) * 3 + 3MB)= %luMB, HCCL_BUFFSIZE=%luMB.",
+            info.M, info.K, info.topK, (actualSize + MB_SIZE - 1) / MB_SIZE, maxWindowSize / MB_SIZE),
+        return ge::GRAPH_FAILED);
+
+    // 4. workspace
+    size_t *workSpaces = context->GetWorkspaceSizes(1);
+    OP_TILING_CHECK(workSpaces == nullptr, OP_LOGE(nodeName, "workSpaces is nullptr."),
+        return ge::GRAPH_FAILED);
+
+    uint32_t n2 = info.K;
+    uint32_t k2 = info.N / 2;
+
+    uint64_t cocWorkspace = (info.M + 256 - 1) / 256 * 256 * info.topK *sizeof(int32_t) +
+                            info.worldSize * info.worldSize * info.expertPerRank * sizeof(int32_t) * 2 +
+                            info.maxOutputSize * sizeof(float) * 2 +
+                            info.maxOutputSize * n2 * sizeof(int16_t) +
+                            info.maxOutputSize * info.K * sizeof(int8_t) +
+                            info.worldSize  * sizeof(int32_t) * 16 +
+                            (info.expertPerRank + info.worldSize) * sizeof(int32_t) * 16;
+
+    workSpaces[0] = SYSTEM_NEED_WORKSPACE + std::max(cocWorkspace, initRoutingWorkspace);
+
+
+    // 5. communication
+    auto attrs = context->GetAttrs();
+    auto group = attrs->GetAttrPointer<char>(static_cast<int>(ATTR_GROUP_INDEX));
+    uint32_t opType = 8U;
+    std::string algConfig = "AlltoAll=level0:fullmesh;level1:pairwise";
+    AscendC::Mc2CcTilingConfig mc2CcTilingConfig(group, opType, algConfig);
+    mc2CcTilingConfig.GetTiling(tilingData->mc2InitTiling);
+    mc2CcTilingConfig.GetTiling(tilingData->mc2CcTiling);
+
+    OP_LOGI(nodeName, "Leave DispatchFFNCombine tiling func.");
+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus DispatchFFNCombineTilingFunc(gert::TilingContext* context)
+{
+    return DispatchFFNCombineTilingFuncImpl(context);
+}
+
+struct DispatchFFNCombineCompileInfo {};
+ge::graphStatus TilingParseForDispatchFFNCombine(gert::TilingParseContext *context)
+{
+    (void)context;
+    return ge::GRAPH_SUCCESS;
+}
+
+IMPL_OP_OPTILING(DispatchFFNCombine)
+    .Tiling(DispatchFFNCombineTilingFunc)
+    .TilingParse<DispatchFFNCombineCompileInfo>(TilingParseForDispatchFFNCombine);
+} // namespace optiling

csrc/mc2/dispatch_ffn_combine/op_host/hcom_topo_info.h

@@ -0,0 +1,72 @@
+/* Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ * ===================================================================================================================*/
+
+#ifndef METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_
+#define METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_
+
+#include <unordered_map>
+#include <mutex>
+
+using Status = int32_t;
+
+namespace ge {
+static constexpr uint32_t COMM_MESH = 0b1U;
+static constexpr uint32_t COMM_SWITCH = (COMM_MESH << 1U);
+static constexpr uint32_t COMM_RING = (COMM_MESH << 2U);
+static constexpr uint32_t COMM_PAIRWISE = (COMM_MESH << 3U);
+class HcomTopoInfo {
+ public:
+  enum class TopoLevel {
+    L0 = 0,
+    L1,
+    MAX,
+  };
+  struct TopoLevelDesc {
+    uint32_t comm_sets;
+    uint32_t rank_size;
+  };
+  using TopoDescs = TopoLevelDesc[static_cast<int32_t>(TopoLevel::MAX)];
+  struct TopoInfo {
+    int64_t rank_size;
+    void *notify_handle;
+    TopoDescs topo_level_descs;
+  };
+  static HcomTopoInfo &Instance();
+  bool TopoInfoHasBeenSet(const char_t *group);
+  bool TryGetGroupTopoInfo(const char_t *group, TopoInfo &info);
+  Status SetGroupTopoInfo(const char_t *group, const TopoInfo &info);
+  Status GetGroupRankSize(const char_t *group, int64_t &rank_size);
+  TopoDescs *GetGroupTopoDesc(const char_t *group);
+  Status GetGroupNotifyHandle(const char_t *group, void *&notify_handle);
+  void UnsetGroupTopoInfo(const char_t *group) {
+    const std::lock_guard<std::mutex> lock(mutex_);
+    (void) rank_info_.erase(group);
+  }
+
+  Status SetGroupOrderedStream(const char_t *group, void *stream);
+  Status GetGroupOrderedStream(const char_t *group, void *&stream);
+  void UnsetGroupOrderedStream(const char_t *group) {
+    const std::lock_guard<std::mutex> lock(mutex_);
+    (void) group_to_ordered_stream_.erase(group);
+  };
+
+  Status SetGroupOrderedStream(const int32_t device_id, const char_t *group, void *stream);
+  Status GetGroupOrderedStream(const int32_t device_id, const char_t *group, void *&stream);
+  void UnsetGroupOrderedStream(const int32_t device_id, const char_t *group);
+ private:
+  HcomTopoInfo() = default;
+  ~HcomTopoInfo() = default;
+  std::unordered_map<std::string, TopoInfo> rank_info_;
+  std::mutex mutex_;
+  std::unordered_map<std::string, void*> group_to_ordered_stream_; // Ordered stream for the communication domain
+  std::unordered_map<int32_t, std::unordered_map<std::string, void*>> device_id_to_group_to_ordered_stream_; // Ordered stream for the communication domain
+};
+}
+
+#endif // METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_

csrc/mc2/dispatch_ffn_combine/op_host/op_api/aclnn_dispatch_ffn_combine.cpp

@@ -0,0 +1,83 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+#include "aclnn_dispatch_ffn_combine.h"
+#include <algorithm>
+// #include "aclnn_kernels/common/op_error_check.h"
+// #include "opdev/op_log.h"
+// #include "opdev/common_types.h"
+// #include "opdev/platform.h"
+// #include "ophost/matmul_util.h"
+#include <unistd.h>
+#include <vector>
+#include <string>
+#include <iostream>
+#include <fcntl.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <sys/file.h>
+#include <climits>
+// using namespace op;
+
+// using namespace op;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+static constexpr size_t TWO_DIMS = 2;
+static constexpr int64_t KVALUE_MIN = 256;
+static constexpr int64_t KVALUE_MAX = 65535;
+static constexpr size_t HCCL_GROUP_NAME_MAX = 128U;
+enum NnopbaseHcclServerType {
+    NNOPBASE_HCCL_SERVER_TYPE_AICPU = 0,
+    NNOPBASE_HCCL_SERVER_TYPE_MTE,
+    NNOPBASE_HCCL_SERVER_TYPE_END
+};
+
+extern aclnnStatus aclnnInnerDispatchFFNCombineGetWorkspaceSize(const aclTensor* x, const aclTensorList* weight1, const aclTensorList* weight2,
+                                                         const aclTensor* expertId, const aclTensorList* scale1, const aclTensorList* scale2,
+                                                         const aclTensor* probs, const aclTensor* xActiveMask,
+                                                         const char* group, int64_t maxOutputSize,
+                                                         bool transB, bool weightNz, double swigluLimit,
+                                                         const aclTensor* out, const aclTensor* expertTokenNums,
+                                                         uint64_t* workspaceSize, aclOpExecutor** executor);
+extern aclnnStatus aclnnInnerDispatchFFNCombine(void *workspace, uint64_t workspaceSize,
+                                            aclOpExecutor *executor, aclrtStream stream);
+extern "C" void __attribute__((weak)) NnopbaseSetHcclServerType(void *executor, NnopbaseHcclServerType sType);
+
+
+
+aclnnStatus aclnnDispatchFFNCombineGetWorkspaceSize(const aclTensor* x, const aclTensorList* weight1, const aclTensorList* weight2,
+                                                    const aclTensor* expertId, const aclTensorList* scale1, const aclTensorList* scale2,
+                                                    const aclTensor* probs, const aclTensor* xActiveMask,
+                                                    const char* group, int64_t maxOutputSize, double swigluLimit,
+                                                    const aclTensor* out, const aclTensor* expertTokenNums,
+                                                    uint64_t* workspaceSize, aclOpExecutor** executor)
+{
+    bool transB = false;
+    bool weightNz = true;
+
+    aclnnStatus ret = aclnnInnerDispatchFFNCombineGetWorkspaceSize(x, weight1, weight2, expertId, scale1, scale2, probs, xActiveMask, group,
+                                                                    maxOutputSize, transB, weightNz, swigluLimit,
+                                                                    out, expertTokenNums, workspaceSize, executor);
+    return ret;
+}
+
+aclnnStatus aclnnDispatchFFNCombine(void* workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)
+{
+    if (NnopbaseSetHcclServerType) {
+        NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);
+    }
+    aclnnStatus ret = aclnnInnerDispatchFFNCombine(workspace, workspaceSize, executor, stream);
+    return ret;
+}
+#ifdef __cplusplus
+}
+#endif

csrc/mc2/dispatch_ffn_combine/op_host/op_api/aclnn_dispatch_ffn_combine.h

@@ -0,0 +1,64 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+#ifndef OP_API_INC_DISPATCH_FFN_COMBINE_
+#define OP_API_INC_DISPATCH_FFN_COMBINE_
+
+#include <string>
+
+#include "aclnn/aclnn_base.h"
+#include "hccl/hccl.h"
+#include "hccl/hccl_types.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * Operator function: fuse all distributed MoE ops from InitRouting through Unpermute.
+ * @brief First-stage interface of aclnnDispatchFFNCombine that calculates workspace size based on the specific compute flow.
+ * @domain aclnn_ops_infer
+ * @param [in] x: The input tensor.
+ * @param [in] weight1: The first weight tensor.
+ * @param [in] weight2: The second weight tensor.
+ * @param [in] expertId: The expert ID tensor.
+ * @param [in] scale1: The first scale tensor.
+ * @param [in] scale2: The second scale tensor.
+ * @param [in] probs: The probabilities tensor.
+ * @param [in] group: string identifying the communication domain name.
+ * @param [in] maxOutputSize: The maximum output size.
+ * @param [out] out: result of computation + communication; same dtype as input.
+ * @param [out] workspaceSize: workspace size to allocate on the NPU device side.
+ * @param [out] executor: op executor containing the operator compute flow.
+ * @return aclnnStatus: status code.
+ */
+__attribute__((visibility("default"))) aclnnStatus aclnnDispatchFFNCombineGetWorkspaceSize(const aclTensor* x, const aclTensorList* weight1, const aclTensorList* weight2,
+                                                                                        const aclTensor* expertId, const aclTensorList* scale1, const aclTensorList* scale2,
+                                                                                        const aclTensor* probs, const aclTensor* xActiveMask,
+                                                                                        const char* group, int64_t maxOutputSize, double swigluLimit,
+                                                                                        const aclTensor* out, const aclTensor* expertTokenNums,
+                                                                                        uint64_t* workspaceSize, aclOpExecutor** executor);
+
+/**
+ * @brief Second-stage interface of aclnnDispatchFFNCombine to execute computation.
+ * @param [in] workspace: workspace memory address allocated on the NPU device side.
+ * @param [in] workspace_size: workspace size allocated on the NPU device side, obtained from aclnnDispatchFFNCombineGetWorkspaceSize.
+ * @param [in] executor: op executor containing the operator compute flow.
+ * @param [in] stream: acl stream.
+ * @return aclnnStatus: status code.
+ */
+__attribute__((visibility("default"))) aclnnStatus aclnnDispatchFFNCombine(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor,
+                                           aclrtStream stream);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // OP_API_INC_DISPATCH_FFN_COMBINE_

csrc/mc2/dispatch_ffn_combine/op_host/tiling_args.h

@@ -0,0 +1,9 @@
+#ifndef TILING_ARGS_H
+#define TILING_ARGS_H
+#include <cstdint>
+
+namespace Moe {
+constexpr uint64_t COMBINE_STATE_WIN_OFFSET = 3U * 1024UL * 1024UL;
+constexpr uint64_t NOTIFY_DISPATCH_WIN_OFFSET = 204U * 1024UL * 1024UL;
+}  // namespace Moe
+#endif  // TILING_ARGS_H

csrc/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine.cpp

@@ -0,0 +1,33 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/* !
+ * \file dispatch_ffn_combine.cpp
+ * \brief
+ */
+#include "kernel_operator.h"
+#include "lib/matmul_intf.h"
+#include "dispatch_ffn_combine_tiling.h"
+#include "dispatch_ffn_combine.h"
+
+using namespace AscendC;
+using namespace DispatchFFNCombineImpl;
+extern "C" __global__ __aicore__ void dispatch_ffn_combine(GM_ADDR x, GM_ADDR w1, GM_ADDR w2,  GM_ADDR expertId, GM_ADDR scale1, GM_ADDR scale2, GM_ADDR probs,
+    GM_ADDR xActiveMask, GM_ADDR c, GM_ADDR expertTokenNums, GM_ADDR workspaceGM,  GM_ADDR tilingGM)
+{
+    REGISTER_TILING_DEFAULT(DispatchFFNCombineTilingData);
+    if (TILING_KEY_IS(1000010)) {
+        KERNEL_TASK_TYPE(1000010, KERNEL_TYPE_MIX_AIC_1_2);
+        GET_TILING_DATA_WITH_STRUCT(DispatchFFNCombineTilingData, tilingData, tilingGM);
+        DispatchFFNCombine<int8_t, DTYPE_W1, DTYPE_OUT, false, true> op;
+        op.Init(x, w1, w2, expertId, scale1, scale2, probs, xActiveMask, c, expertTokenNums, workspaceGM, tilingGM);
+        op.Process();
+    }
+}

csrc/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine.h

@@ -0,0 +1,291 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file dispatch_ffn_combine.h
+ * \brief
+ */
+
+#ifndef DISPATCH_FFN_COMBINE_H
+#define DISPATCH_FFN_COMBINE_H
+
+using namespace AscendC;
+
+#include "kernel_operator.h"
+
+#include "utils/moe_distribute_base.h"
+
+#include "dispatch_ffn_combine_tiling.h"
+
+#include "catlass/catlass.hpp"
+#include "catlass/arch/arch.hpp"
+#include "catlass/epilogue/dispatch_policy.hpp"
+#include "catlass/epilogue/block/block_epilogue.hpp"
+#include "catlass/epilogue/tile/tile_copy.hpp"
+#include "catlass/epilogue/tile/tile_elemwise_add.hpp"
+#include "catlass/epilogue/tile/tile_elemwise_muls.hpp"
+#include "catlass/gemm/block/block_mmad.hpp"
+#include "catlass/gemm/block/block_swizzle.hpp"
+#include "catlass/gemm/dispatch_policy.hpp"
+#include "catlass/gemm/kernel/matmul_epilogue.hpp"
+#include "catlass/gemm/gemm_type.hpp"
+#include "catlass/layout/layout.hpp"
+
+#include "utils/select_helper.hpp"
+#include "utils/const_args.hpp"
+#include "dispatch_ffn_combine_kernel.hpp"
+#include "moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h"
+
+using namespace Catlass;
+
+namespace DispatchFFNCombineImpl {
+#define TemplateMMA2AClass typename AType_, typename BType_, typename CType_, bool TB_, bool Nz_
+#define TemplateMMA2ACFunc AType_, BType_, CType_, TB_, Nz_
+
+using namespace AscendC;
+template <TemplateMMA2AClass>
+class DispatchFFNCombine {
+public:
+    __aicore__ inline DispatchFFNCombine() {};
+    __aicore__ inline void Init(GM_ADDR xGM, GM_ADDR weight1GM, GM_ADDR weight2GM, GM_ADDR expertIdGM, GM_ADDR scale1GM, GM_ADDR scale2GM,
+                                GM_ADDR probs, GM_ADDR xActiveMaskGM, GM_ADDR outGM, GM_ADDR expertTokenNums, GM_ADDR workspaceGM, GM_ADDR tilingGM);
+    __aicore__ inline void Process();
+
+
+private:
+    GM_ADDR xGM_;
+    GM_ADDR weight1GM_;
+    GM_ADDR weight2GM_;
+    GM_ADDR expertIdGM_;
+    GM_ADDR scale1GM_;
+    GM_ADDR scale2GM_;
+    GM_ADDR probs_;
+    GM_ADDR xActiveMaskGM_;
+    GM_ADDR outGM_;
+    GM_ADDR gmExpertTokenNums_;
+    GM_ADDR workspaceGM_;
+
+    GM_ADDR moeInitRoutingQuantV2Scale = nullptr;
+    GM_ADDR moeInitRoutingQuantV2Offset = nullptr;
+    GM_ADDR expertTokensBeforeCapacity = nullptr;
+
+
+    TBuf<AscendC::TPosition::VECCALC> uBuf_;
+
+    int32_t rank;
+    int32_t rankSize;
+    int32_t aivNum;
+    
+    int32_t m0;
+    int32_t k0;
+    int32_t n0;
+    int32_t swizzlOffset;
+    int32_t swizzlDirect;
+    int32_t ubMoveNum;
+    int32_t pValue;
+
+    int32_t commNpuSplit;
+    int32_t commDataSplit;
+    int32_t lenPerLoop;
+
+    int32_t m;
+    int32_t k;
+    int32_t n;
+    int32_t topK;
+    int32_t expertPerRank;
+    int32_t maxOutputSize;
+    int32_t EP;
+    int32_t listLen;
+    float swigluLimit;
+
+    optiling::MoeInitRoutingQuantV2TilingData moeInitRoutingQuantV2TilingData;
+    uint64_t initRoutingQuantTilingKey;
+
+    // Hccl<HCCL_SERVER_TYPE_AICPU> hccl_;
+
+};
+
+
+template <TemplateMMA2AClass>
+__aicore__ inline void DispatchFFNCombine<TemplateMMA2ACFunc>::Init(GM_ADDR xGM, GM_ADDR weight1GM, GM_ADDR weight2GM, GM_ADDR expertIdGM, GM_ADDR scale1GM, GM_ADDR scale2GM,
+                                                                    GM_ADDR probs, GM_ADDR xActiveMaskGM, GM_ADDR outGM, GM_ADDR expertTokenNums, GM_ADDR workspaceGM, GM_ADDR tilingGM)
+{
+    REGISTER_TILING_DEFAULT(DispatchFFNCombineTilingData);
+    auto tiling = (__gm__ DispatchFFNCombineTilingData*)tilingGM;
+    GET_TILING_DATA(tilingData, tilingGM);
+
+    xGM_ = xGM;
+    weight1GM_ = weight1GM;
+    weight2GM_ = weight2GM;
+    expertIdGM_ = expertIdGM;
+    scale1GM_ = scale1GM;
+    scale2GM_ = scale2GM;
+    probs_ = probs;
+    xActiveMaskGM_ = xActiveMaskGM;
+
+    outGM_ = outGM;
+    gmExpertTokenNums_ = expertTokenNums;
+
+    workspaceGM_ = workspaceGM;
+
+    aivNum = tilingData.dispatchFFNCombineInfo.aivNum;
+
+    m = tilingData.dispatchFFNCombineInfo.M;
+    k = tilingData.dispatchFFNCombineInfo.K;
+    n = tilingData.dispatchFFNCombineInfo.N;
+    EP =  tilingData.dispatchFFNCombineInfo.worldSize;
+    topK = tilingData.dispatchFFNCombineInfo.topK;
+    expertPerRank = tilingData.dispatchFFNCombineInfo.expertPerRank;
+    maxOutputSize = tilingData.dispatchFFNCombineInfo.maxOutputSize;
+    listLen = tilingData.dispatchFFNCombineInfo.listLen;
+    swigluLimit = tilingData.dispatchFFNCombineInfo.swigluLimit;
+
+    m0 = tilingData.cocTiling.m0;
+    k0 = tilingData.cocTiling.k0;
+    n0 = tilingData.cocTiling.n0;
+    swizzlDirect = tilingData.cocTiling.swizzleDirect;
+    swizzlOffset = tilingData.cocTiling.swizzleOffset;
+    ubMoveNum = tilingData.cocTiling.ubMoveNum;
+    pValue = tilingData.cocTiling.pValue;
+    commNpuSplit = tilingData.cocTiling.commNpuSplit;
+    commDataSplit = tilingData.cocTiling.commDataSplit;
+    lenPerLoop = tilingData.cocTiling.lenPerLoop;
+    moeInitRoutingQuantV2TilingData = tilingData.cocTiling.moeInitRoutingQuantV2TilingData;
+    initRoutingQuantTilingKey = tilingData.cocTiling.initRoutingQuantTilingKey;
+
+    auto contextGM0 = AscendC::GetHcclContext<HCCL_GROUP_ID_0>();
+    __gm__ HcclOpResParamCustom *WinContext_{nullptr};
+    WinContext_ = (__gm__ HcclOpResParamCustom *)contextGM0;
+
+    rank = WinContext_->localUsrRankId;
+    rankSize = WinContext_->rankSize;
+}
+
+template <TemplateMMA2AClass>
+__aicore__ inline void DispatchFFNCombine<TemplateMMA2ACFunc>::Process()
+{
+    // Define ArchTag
+    using ArchTag = Arch::AtlasA2;
+    constexpr bool enableUnitFlag = false;
+    constexpr bool enableShuffleK = true;
+
+    uint32_t k2 = n/2;
+    uint32_t n2 = k;
+
+    int64_t activeNum = 0;
+    int64_t expertCapacity = 0;
+    int64_t expertNum = expertPerRank * EP;
+    int64_t dropPadMode = 0;
+    int64_t expertTokensCountOrCumsumFlag = 2;
+    bool expertTokensBeforeCapacityFlag = false;
+    int64_t quantMode = 1;
+
+    using LayoutA = layout::RowMajor;
+    using LayoutB = typename std::conditional<
+        Nz_,
+        layout::zN,
+        typename std::conditional<TB_, layout::ColumnMajor, layout::RowMajor>::type
+    >::type;
+
+    LayoutB layoutB1 = LayoutBInitializer<LayoutB, BType_>::create(k, n);
+    LayoutB layoutB2 = LayoutBInitializer<LayoutB, BType_>::create(k2, n2);
+    using LayoutC = layout::RowMajor;
+    using L1TileShape = GemmShape<128, 256, 512>;   // M, N, K
+
+    constexpr uint32_t workspaceStages = 2;
+    constexpr uint32_t preloadStages = 1;
+    constexpr uint32_t l1Stages = 2;
+    constexpr uint32_t l0AStages = 2;
+    constexpr uint32_t l0BStages = 2;
+    constexpr uint32_t l0CStages = 1;
+
+    using DispatchPolicy = Gemm::MmadAtlasA2PreloadAsyncFixpipe<
+        preloadStages,
+        l1Stages, l0AStages, l0BStages, l0CStages,
+        enableUnitFlag, enableShuffleK
+    >;
+
+    using L0TileShape = GemmShape<128, 256, 128>;
+    using AType = Gemm::GemmType<int8_t, layout::RowMajor>;
+    using BType = Gemm::GemmType<int8_t, LayoutB>;
+    using CType = Gemm::GemmType<float16_t, layout::RowMajor>;
+    using D1Type = Gemm::GemmType<int8_t, layout::RowMajor>;
+
+    using D2Type = typename std::conditional<
+        std::is_same_v<CType_, bfloat16_t>, 
+        Gemm::GemmType<bfloat16_t, layout::RowMajor>,
+        Gemm::GemmType<CType_, layout::RowMajor>
+        >::type;
+
+    using BlockMmad = Gemm::Block::BlockMmad<DispatchPolicy, L1TileShape, L0TileShape, AType, BType, CType>;
+    constexpr uint32_t ubStages = 2;
+
+    using EpilogueDispatchPolicy1 = Epilogue::EpilogueAtlasA2PerTokenDequantSwigluQuant<ubStages>;
+    
+    using ScaleType = Gemm::GemmType<uint64_t, layout::VectorLayout>;
+    using PerTokenScaleType = Gemm::GemmType<float, layout::VectorLayout>;
+    using ElementMulType = Gemm::GemmType<float, layout::RowMajor>;
+    using TileElemWiseMuls = Epilogue::Tile::TileElemWiseMuls<ArchTag, ElementMulType, 0>;
+
+    using TileCopy1 = Epilogue::Tile::TileCopy<ArchTag, CType, ScaleType, PerTokenScaleType, D1Type>;
+    using BlockEpilogue1 = Epilogue::Block::BlockEpilogue<EpilogueDispatchPolicy1, CType, PerTokenScaleType,
+        D1Type, TileElemWiseMuls, TileCopy1>;
+
+    using EpilogueDispatchPolicy2 = Epilogue::EpilogueAtlasA2PerTokenDequant<ubStages>;
+    using EpilogueDispatchPolicy3 =  Epilogue::EpilogueAtlasA2PerTokenDequantV2<ubStages>;
+    
+    using TileCopy2 = Epilogue::Tile::TileCopy<ArchTag, CType, ScaleType, PerTokenScaleType, D2Type>;
+    using BlockEpilogue2 = Epilogue::Block::BlockEpilogue<EpilogueDispatchPolicy2, CType,PerTokenScaleType,
+        D2Type, TileCopy2>;
+    using BlockEpilogue3 = Epilogue::Block::BlockEpilogue<EpilogueDispatchPolicy3, CType,PerTokenScaleType,
+        D2Type, TileCopy2>;
+
+
+    using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<9, 1>;
+    using ElementGroupList = int64_t;
+    using MatmulKernel = Gemm::Kernel::DispatchFFNCombineKernel<BlockMmad,
+        BlockScheduler, ElementGroupList, BlockEpilogue1, BlockEpilogue2, BlockEpilogue3>;
+
+    LayoutA layoutA1{static_cast<uint32_t>(m), static_cast<uint32_t>(k)};
+    LayoutA layoutA2{static_cast<uint32_t>(m), static_cast<uint32_t>(k2)};
+    layout::VectorLayout layoutScale1{static_cast<uint32_t>(n)};
+    layout::VectorLayout layoutScale2{static_cast<uint32_t>(n2)};
+    layout::RowMajor layoutD1{static_cast<uint32_t>(maxOutputSize), static_cast<uint32_t>(k2)};
+    layout::RowMajor layoutD2{static_cast<uint32_t>(m*topK), static_cast<uint32_t>(n2)};
+    // Prepare params
+
+    GemmCoord problemShape{static_cast<uint32_t>(m), static_cast<uint32_t>(n), static_cast<uint32_t>(k)};
+
+    uint32_t epilogueCoreNum = aivNum;
+    uint32_t epilogueGranularity = expertPerRank - 3;
+    if (expertPerRank <= 4) {
+        epilogueGranularity = expertPerRank - 1;
+    }
+    typename MatmulKernel::Params params{
+        problemShape, static_cast<uint32_t>(EP), static_cast<uint32_t>(listLen), static_cast<uint32_t>(expertPerRank), static_cast<uint32_t>(maxOutputSize),
+        static_cast<uint32_t>(rank), static_cast<uint32_t>(rankSize),
+        static_cast<uint32_t>(topK), initRoutingQuantTilingKey,
+        epilogueCoreNum, epilogueGranularity,
+        xGM_, layoutA1, layoutA2,
+        weight1GM_, layoutB1,
+        weight2GM_, layoutB2,
+        scale1GM_, layoutScale1,
+        scale2GM_, layoutScale2,
+        outGM_, layoutD1, layoutD2,
+        expertIdGM_, moeInitRoutingQuantV2Scale, moeInitRoutingQuantV2Offset,
+        expertTokensBeforeCapacity, probs_,
+        workspaceGM_, gmExpertTokenNums_, ubMoveNum, xActiveMaskGM_, moeInitRoutingQuantV2TilingData, swigluLimit};
+    //Call kernel
+    MatmulKernel kernel(params);
+    kernel(params);
+}
+
+} // DispatchFFNCombineImpl
+#endif // DISPATCH_FFN_COMBINE_H

csrc/mc2/dispatch_ffn_combine/op_kernel/dispatch_ffn_combine_tiling.h

@@ -0,0 +1,58 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file dispatch_ffn_combine_tiling.h
+ * \brief
+ */
+
+#include "moe_init_routing_quant_v2/moe_init_routing_v2_tiling.h"
+#include "moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h"
+
+#ifndef ASCENDC_DISPATCH_FFN_COMBINE_TILING_H
+#define ASCENDC_DISPATCH_FFN_COMBINE_TILING_H
+struct DispatchFFNCombineInfo {
+    uint32_t M;
+    uint32_t K;
+    uint32_t N;
+    uint32_t expertPerRank;
+    uint32_t maxOutputSize;
+    uint32_t isTransposeB;
+    uint32_t isWeightNz;
+    uint32_t aivNum;
+    uint32_t totalUbSize;
+    uint32_t topK;
+    uint32_t worldSize;
+    uint32_t listLen;
+    float swigluLimit;
+};
+
+struct CoCTiling {
+    int32_t m0 = -1;
+    int32_t k0 = -1;
+    int32_t n0 = -1;
+    int32_t swizzleDirect = -1;
+    int32_t swizzleOffset = -1;
+    int32_t ubMoveNum = -1;
+    int32_t pValue = -1;
+    int32_t commNpuSplit = -1;
+    int32_t commDataSplit = -1;
+    int32_t lenPerLoop = -1;
+    uint64_t initRoutingQuantTilingKey;
+    optiling::MoeInitRoutingQuantV2TilingData moeInitRoutingQuantV2TilingData;
+};
+
+struct DispatchFFNCombineTilingData {
+    Mc2InitTiling mc2InitTiling;
+    Mc2CcTiling mc2CcTiling;
+    DispatchFFNCombineInfo dispatchFFNCombineInfo;
+    CoCTiling cocTiling;
+};
+#endif

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_quant_v2.cpp

@@ -0,0 +1,133 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_init_routing_quant_v2.cpp
+ * \brief
+ */
+#include "moe_v2_sort_one_core.h"
+#include "moe_v2_sort_multi_core.h"
+#include "moe_v2_mrgsort_out.h"
+#include "moe_v2_mrgsort.h"
+#include "moe_v2_expert_token_out.h"
+#include "moe_v2_src_to_dst_op.h"
+#include "moe_v2_src_to_dst_with_capacity.h"
+#include "moe_v2_fullload_quant.h"
+#include "moe_v2_fullload_dynamic_quant.h"
+#include "moe_v2_gather_quant.h"
+#include "moe_v2_gather_dynamic_quant.h"
+#include "moe_v2_src_to_dst_and_gather.h"
+
+using namespace AscendC;
+using namespace MoeInitRoutingQuantV2;
+using namespace optiling;
+
+template <class DTYPE_X = bfloat16_t>
+__aicore__ inline  void moe_init_routing_quant_v2(
+    GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedX, GM_ADDR expandedRowIdx,
+    GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity, GM_ADDR dynamicQuantScale, GM_ADDR workspace,
+    const MoeInitRoutingQuantV2TilingData* tilingData, uint64_t tilingKey) {
+
+  if (g_coreType == AIC) {
+    return;
+  }
+
+  if (workspace == nullptr) {
+    return;
+  }
+
+  if (tilingKey == 20000) {  // quant full load
+    TPipe sortPipe;
+    MoeV2FullLoadQuant<DTYPE_X> op;
+    op.Init(x, expertIdx, scale, offset, expandedX, expandedRowIdx, expertTokensCountOrCumsum, workspace, tilingData, &sortPipe);
+    op.Process();
+    sortPipe.Destroy();
+    return;
+  } 
+  
+  
+  else if (tilingKey == 21000) {  // dynamic quant full load
+    TPipe sortPipe;
+    MoeV2FullLoadDynamicQuant<DTYPE_X> op;
+    op.Init(x, expertIdx, expandedX, expandedRowIdx, expertTokensCountOrCumsum, scale, dynamicQuantScale, workspace, tilingData,
+            &sortPipe);
+    op.Process();
+    sortPipe.Destroy();
+    return;
+  }
+
+  // sort
+  if (tilingKey == 10000 || tilingKey == 10100 || tilingKey == 11000 || tilingKey == 11100) {
+    TPipe sortPipe;
+    MoeV2SortOneCore op;
+    op.Init<MoeInitRoutingQuantV2TilingData>(expertIdx, expertTokensCountOrCumsum, expertTokensBeforeCapacity, workspace,
+                                             tilingData, &sortPipe);
+    op.Process();
+    sortPipe.Destroy();
+  } else if (tilingKey == 10010 || tilingKey == 10110 || tilingKey == 11010  || tilingKey== 11110) {
+    TPipe sortPipe;
+    MoeV2SortMultiCore op;
+    op.Init<MoeInitRoutingQuantV2TilingData>(expertIdx, expertTokensCountOrCumsum, expertTokensBeforeCapacity, workspace,
+                                             tilingData, &sortPipe);
+    op.Process();
+    sortPipe.Destroy();
+  }
+
+  if (tilingKey == 10000 || tilingKey == 10010 || tilingKey ==11000 || tilingKey ==11010) { // No drop scenario
+    if (tilingData->expertTokensCountOrCumsumFlag != EXERPT_TOKENS_NONE) {
+      TPipe expertTokenOutPipe;
+      MoeV2ExpertTokenOut expertTokenOutOp;
+      expertTokenOutOp.Init<MoeInitRoutingQuantV2TilingData>(expertTokensCountOrCumsum, expertTokensBeforeCapacity,
+                                                             expandedRowIdx, workspace, tilingData, &expertTokenOutPipe);
+      expertTokenOutOp.Process();
+      expertTokenOutPipe.Destroy();
+    }
+    TPipe srcToDstPipe;
+    MoeV2SrcToDstOp srcToDstOp;
+    srcToDstOp.Init<MoeInitRoutingQuantV2TilingData>(expandedRowIdx, workspace, tilingData, &srcToDstPipe);
+    srcToDstOp.Process();
+    srcToDstPipe.Destroy();
+  } else if (tilingKey ==10100 || tilingKey ==10110 || tilingKey ==11100 || tilingKey ==11110) { // Drop scenario
+    TPipe expertTokenOutPipe;
+    MoeV2ExpertTokenOut expertTokenOutOp;
+    expertTokenOutOp.Init<MoeInitRoutingQuantV2TilingData>(expertTokensCountOrCumsum, expertTokensBeforeCapacity,
+                                                           expandedRowIdx, workspace, tilingData, &expertTokenOutPipe);
+    expertTokenOutOp.Process();
+    expertTokenOutPipe.Destroy();
+
+    if (tilingKey == 10100 || tilingKey == 10110) {
+      TPipe srcToDstPipe;
+      MoeV2SrcToDstWithCapacity<int8_t, MoeInitRoutingQuantV2TilingData> srcToDstWithCapacityOp;
+      srcToDstWithCapacityOp.Init(expandedRowIdx, expandedX, workspace, tilingData, &srcToDstPipe);
+      srcToDstWithCapacityOp.Process();
+      srcToDstPipe.Destroy();
+    } else {
+      TPipe srcToDstGatherPipe;
+      MoeV2SrcToDstAndGather<DTYPE_X, MoeInitRoutingQuantV2TilingData> srcToDstAndGatherOp;
+      srcToDstAndGatherOp.Init(x, scale, expandedRowIdx, expandedX, dynamicQuantScale, workspace, tilingData, &srcToDstGatherPipe);
+      srcToDstAndGatherOp.Process();
+      srcToDstGatherPipe.Destroy();
+    }
+  }
+
+  if (tilingKey == 10000 || tilingKey == 10010 || tilingKey == 10100 || tilingKey == 10110) {
+    TPipe gatherPipe;
+    MoeV2GatherQuant<DTYPE_X> gatherQuantOp;
+    gatherQuantOp.Init(x, scale, offset, expandedRowIdx, expandedX, workspace, tilingData, &gatherPipe);
+    gatherQuantOp.Process();
+    gatherPipe.Destroy();
+  } else if (tilingKey  == 11000 || tilingKey == 11010) {
+    TPipe gatherPipe;
+    MoeV2GatherDynamicQuant<DTYPE_X> gatherDynamicQuantOp;
+    gatherDynamicQuantOp.Init(x, scale, expandedRowIdx, expandedX, dynamicQuantScale, workspace, tilingData, &gatherPipe);
+    gatherDynamicQuantOp.Process();
+    gatherPipe.Destroy();
+  }
+}

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h

@@ -0,0 +1,434 @@
+#pragma once
+#ifndef ASCENDC_DISPATCH_FFN_COMBINE_INIT_ROUTING_QUANT_TILING_H
+#define ASCENDC_DISPATCH_FFN_COMBINE_INIT_ROUTING_QUANT_TILING_H
+
+#include "moe_init_routing_v2_tiling.h"
+
+namespace optiling {
+
+const static int64_t ATTR_QUANT_MODE = 6;
+const static int64_t TILING_KEY_BASE = 10000;
+const static int64_t TILING_KEY_PERF_BASE = 20000;
+const static int64_t TILING_KEY_QUANT_BASE = 1000;
+const static int64_t TILING_KEY_DROP_MODE_BASE = 100;
+const static int64_t TILING_KEY_SORT_BASE = 10;
+const static int64_t FOUR_BLOCK_BYTE = 128;
+const static int64_t MAX_COLS_ONE_LOOP_QUANT = 8192;
+const static int64_t INDEX_SCALE = 2;
+const static int64_t INDEX_OFFSET = 3;
+const static int64_t SMOOTH_NONE = 0;
+const static int64_t SMOOTH_1H = 1;
+const static int64_t SMOOTH_EH = 2;
+const static int64_t MAX_COLS_DYNAMIC_QUANT = 6144;
+const static int64_t DYNAMIC_QUANT_SRC_TO_DST_BUFFER = 15;
+const static int64_t DYNAMIC_QUANT_COLS_BUFFER = 21;
+const static int64_t DYNAMIC_QUANT_FULLLOAD_COLS_BUFFER = 13;
+const static int64_t DYNAMIC_QUANT_SCALE_SIZE_64 = 64;
+const static int64_t DYNAMIC_QUANT_SCALE_SIZE_128 = 128;
+const static int64_t OUTOUT_DYNAMIC_QUANT_SCALE = 4;
+const static int64_t FULLLOAD_H_LIMIT = 7168;
+
+
+inline static int64_t AlignOneBlockByte(int64_t x) {
+  return (x + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+}
+
+inline static int64_t AlignOneBlockByteCeil(int64_t x) {
+  return x / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+}
+
+struct MoeInitRoutingQuantV2TilingData {
+    int64_t coreNum;
+    int64_t n;
+    int64_t cols;
+    int64_t k;
+    int64_t expertCapacity;
+    int64_t expertNum;
+    int64_t dropPadMode;
+    int64_t expertTokensCountOrCumsumFlag;
+    int64_t expertTokensBeforeCapacityFlag;
+    int64_t smoothType;
+    InnerMoeV2VBSComputeTilingData vbsComputeParamsOp;
+    InnerMoeV2VMSMiddleComputeTilingData vmsMiddleComputeParamsOp;
+    InnerMoeV2SortOutComputeTilingData sortOutComputeParamsOp;
+    InnerMoeV2GatherOutComputeTilingData srcToDstComputeParamsOp;
+    InnerMoeV2GatherOutComputeTilingData srcToDstCapacityComputeParamsOp;
+    InnerMoeV2GatherOutComputeTilingData gatherOutComputeParamsOp;
+};
+
+
+
+class MoeInitRoutingQuantV2TilingBase : public InnerMoeInitRoutingV2TilingBase {
+public:
+protected:
+
+  bool GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, 
+  int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag,
+  bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) override;
+  uint64_t GetTilingKey() const override;
+  bool GetWorkspaceSize() override;
+  bool PostTiling() override;
+public:
+  //bool CheckOutShape() override;
+  bool IsFullLoadQuant(int64_t space);
+  bool IsFullLoadDynamicQuant(int64_t space);
+  bool IsFullLoad() override;
+  void SetGatherTilingData(InnerMoeV2GatherOutComputeTilingData* tilingData, int64_t perCoreRows, int64_t lastCoreRows,
+                           int64_t cols);
+  void SetGatherTilingDataCols(InnerMoeV2GatherOutComputeTilingData* tilingData, int64_t baseMaxCols, int64_t cols);
+  void SetGatherTilingDataRows(InnerMoeV2GatherOutComputeTilingData* tilingData, int64_t perCoreRows,
+                               int64_t lastCoreRows, int64_t basePerLoopMaxRows);
+  void Tiling4GatherQuant();
+  void Tiling4GatherDynamicQuant();
+  void Tiling4SrcToDstCapacityCompute() override;
+  void Tiling4GatherOutCompute() override;
+  void CopyGatherOutTiling(InnerMoeV2GatherOutComputeTilingData& dst, InnerMoeV2GatherOutComputeTilingData& src);
+  void CopyTilingData();
+
+
+  int64_t quantMode;
+  MoeInitRoutingQuantV2TilingData quantTilingData;
+};
+
+
+bool MoeInitRoutingQuantV2TilingBase::IsFullLoadQuant(int64_t space) {
+  int64_t perCoreXRows = moeInitRoutingTilingData.n / aivNum;
+  int64_t remainder = moeInitRoutingTilingData.n % aivNum;
+  // NUM_TWO is Max xRows need add 2 becauseof the left and right row may be another row.
+  perCoreXRows = remainder <= 1 ? perCoreXRows + 1 : perCoreXRows + NUM_TWO;
+  int64_t quantBaseSpace = AlignOneBlockByte(moeInitRoutingTilingData.cols);
+  int64_t quantSpace =
+      quantBaseSpace * (inuptXDtypeSize_ + sizeof(int8_t) + sizeof(float) + sizeof(int16_t)) * perCoreXRows;
+  int64_t remainUbAfterSort = aicoreParams_.ubSize - space - quantSpace;
+  return remainUbAfterSort > 0;
+}
+
+bool MoeInitRoutingQuantV2TilingBase::IsFullLoadDynamicQuant(int64_t space) {
+  int64_t quantSpace = AlignOneBlockByte(moeInitRoutingTilingData.cols) * DYNAMIC_QUANT_FULLLOAD_COLS_BUFFER;
+  int64_t scaleOutSpace = 64;
+  int64_t remainUbAfterSort = aicoreParams_.ubSize - space - scaleOutSpace - quantSpace;
+  return remainUbAfterSort > 0;
+}
+
+bool MoeInitRoutingQuantV2TilingBase::IsFullLoad() {
+  if (totalLength > sortLoopMaxElement || moeInitRoutingTilingData.cols > MAX_COLS_ONE_LOOP_QUANT ||
+      this->dropPadMode == 1) {
+    return false;
+  }
+  int64_t sortSpace = AlignOneBlockByte(this->totalLength) * sizeof(int32_t) * ONE_CORE_SORT_BUFFER;
+  int64_t otherSpace = AlignOneBlockByte(this->totalLength) * sizeof(int32_t) * NUM_THREE;
+  int64_t expertSpace = AlignOneBlockByte(this->expertNum * sizeof(int32_t));
+  if (quantMode == 0) {
+    return IsFullLoadQuant(sortSpace + otherSpace + expertSpace);
+  } else {
+    return IsFullLoadDynamicQuant(sortSpace + otherSpace + expertSpace);
+  }
+}
+
+bool MoeInitRoutingQuantV2TilingBase::GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, 
+  int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag,
+  bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) {
+  
+  InnerMoeInitRoutingV2TilingBase::GetShapeAttrsInfo(m, cols, topK, expertCapacity, expertNum, activeNum, dropPadMode, 
+            expertTokensCountOrCumsumFlag, expertTokensBeforeCapacityFlag, inuptXDtypeSize, quantMode, scaleDim0);
+  this -> quantMode = quantMode;
+  if (quantMode == 0) {
+  } else {
+    if (scaleDim0 > 0) {
+      quantTilingData.smoothType = ((scaleDim0 == 1) ? SMOOTH_1H : SMOOTH_EH);
+    } else {
+      quantTilingData.smoothType = SMOOTH_NONE;
+    }
+  }
+  return true;
+}
+
+
+uint64_t MoeInitRoutingQuantV2TilingBase::GetTilingKey() const {
+  if (isFullLoad) {
+    return TILING_KEY_PERF_BASE + quantMode * TILING_KEY_QUANT_BASE;
+  }
+  return TILING_KEY_BASE + quantMode * TILING_KEY_QUANT_BASE + dropPadMode * TILING_KEY_DROP_MODE_BASE +
+         (totalLength > sortLoopMaxElement) * TILING_KEY_SORT_BASE;
+}
+
+
+bool MoeInitRoutingQuantV2TilingBase::PostTiling() {
+  CopyTilingData();
+  return true;
+}
+void MoeInitRoutingQuantV2TilingBase::CopyGatherOutTiling(InnerMoeV2GatherOutComputeTilingData& dst,
+                                                          InnerMoeV2GatherOutComputeTilingData& src) {
+  dst.needCoreNum = (src.needCoreNum);
+  dst.activateRows = (src.activateRows);
+  dst.perCoreRows = (src.perCoreRows);
+  dst.perCorePerLoopRows = (src.perCorePerLoopRows);
+  dst.perCoreLastLoopRows = (src.perCoreLastLoopRows);
+  dst.lastCoreRows = (src.lastCoreRows);
+  dst.lastCorePerLoopRows = (src.lastCorePerLoopRows);
+  dst.lastCoreLastLoopRows = (src.lastCoreLastLoopRows);
+  dst.perCoreLoops = (src.perCoreLoops);
+  dst.lastCoreLoops = (src.lastCoreLoops);
+  dst.perLoopCols = (src.perLoopCols);
+  dst.lastLoopCols = (src.lastLoopCols);
+  dst.colLoops = (src.colLoops);
+}
+
+void MoeInitRoutingQuantV2TilingBase::CopyTilingData() {
+  quantTilingData.coreNum = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.coreNum);
+  quantTilingData.n = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.n);
+  quantTilingData.cols = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.cols);
+  quantTilingData.k = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.k);
+  quantTilingData.expertCapacity = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.expertCapacity);
+  quantTilingData.expertNum = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.expertNum);
+  quantTilingData.dropPadMode = (InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.dropPadMode);
+  quantTilingData.expertTokensCountOrCumsumFlag = (
+      InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.expertTokensCountOrCumsumFlag);
+  quantTilingData.expertTokensBeforeCapacityFlag = (
+      InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.expertTokensBeforeCapacityFlag);
+
+  auto vbsTilingData = &InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.vbsComputeParamsOp;
+  quantTilingData.vbsComputeParamsOp.needCoreNum = (vbsTilingData->needCoreNum);
+  quantTilingData.vbsComputeParamsOp.perCoreElements = (vbsTilingData->perCoreElements);
+  quantTilingData.vbsComputeParamsOp.perCoreLoops = (vbsTilingData->perCoreLoops);
+  quantTilingData.vbsComputeParamsOp.perCorePerLoopElements = (vbsTilingData->perCorePerLoopElements);
+  quantTilingData.vbsComputeParamsOp.perCoreLastLoopElements = (vbsTilingData->perCoreLastLoopElements);
+  quantTilingData.vbsComputeParamsOp.lastCoreElements = (vbsTilingData->lastCoreElements);
+  quantTilingData.vbsComputeParamsOp.lastCoreLoops = (vbsTilingData->lastCoreLoops);
+  quantTilingData.vbsComputeParamsOp.lastCorePerLoopElements = (vbsTilingData->lastCorePerLoopElements);
+  quantTilingData.vbsComputeParamsOp.lastCoreLastLoopElements = (vbsTilingData->lastCoreLastLoopElements);
+  quantTilingData.vbsComputeParamsOp.oneLoopMaxElements = (vbsTilingData->oneLoopMaxElements);
+
+  quantTilingData.vmsMiddleComputeParamsOp.needCoreNum = (
+      InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.vmsMiddleComputeParamsOp.needCoreNum);
+  quantTilingData.sortOutComputeParamsOp.oneLoopMaxElements = (
+      InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.sortOutComputeParamsOp.oneLoopMaxElements);
+
+  CopyGatherOutTiling(quantTilingData.srcToDstComputeParamsOp,
+                      InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.srcToDstComputeParamsOp);
+  CopyGatherOutTiling(quantTilingData.srcToDstCapacityComputeParamsOp,
+                      InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp);
+}
+
+
+bool MoeInitRoutingQuantV2TilingBase::GetWorkspaceSize() {
+  InnerMoeInitRoutingV2TilingBase::GetWorkspaceSize();
+  bool useCols =
+      (dropPadMode == 0 && quantTilingData.gatherOutComputeParamsOp.colLoops > 1) ||
+      (dropPadMode == 1 &&
+       InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp.colLoops > 1);
+  if (quantMode == 1 && useCols) {
+    workspaceSize_ += aivNum * InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.cols * sizeof(float);
+  }
+  return true;
+}
+
+void MoeInitRoutingQuantV2TilingBase::SetGatherTilingData(InnerMoeV2GatherOutComputeTilingData* tilingData,
+                                                          int64_t perCoreRows, int64_t lastCoreRows, int64_t cols) {
+  tilingData->perCorePerLoopRows = perCoreRows;
+  tilingData->perCoreLastLoopRows = perCoreRows;
+  tilingData->lastCorePerLoopRows = lastCoreRows;
+  tilingData->lastCoreLastLoopRows = lastCoreRows;
+  tilingData->perCoreLoops = 1;
+  tilingData->lastCoreLoops = 1;
+  tilingData->perLoopCols = cols;
+  tilingData->lastLoopCols = cols;
+  tilingData->colLoops = 1;
+}
+
+void MoeInitRoutingQuantV2TilingBase::SetGatherTilingDataCols(InnerMoeV2GatherOutComputeTilingData* tilingData,
+                                                              int64_t baseMaxCols, int64_t cols) {
+  tilingData->perLoopCols = (std::min(baseMaxCols, cols));
+  tilingData->lastLoopCols = (GetPerOrLastValue(cols, baseMaxCols));
+  tilingData->colLoops = (baseMaxCols == 0 ? 0 : (cols + baseMaxCols - 1) / baseMaxCols);
+}
+
+void MoeInitRoutingQuantV2TilingBase::SetGatherTilingDataRows(InnerMoeV2GatherOutComputeTilingData* tilingData,
+                                                              int64_t perCoreRows, int64_t lastCoreRows,
+                                                              int64_t basePerLoopMaxRows) {
+  tilingData->perCorePerLoopRows = (std::min(perCoreRows, basePerLoopMaxRows));
+  tilingData->perCoreLastLoopRows = (GetPerOrLastValue(perCoreRows, basePerLoopMaxRows));
+  tilingData->perCoreLoops = (basePerLoopMaxRows == 0 ? 0
+                                                      : (perCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows);
+  tilingData->lastCorePerLoopRows = (std::min(lastCoreRows, basePerLoopMaxRows));
+  tilingData->lastCoreLastLoopRows = (GetPerOrLastValue(lastCoreRows, basePerLoopMaxRows));
+  tilingData->lastCoreLoops = (basePerLoopMaxRows == 0 ? 0
+                                                          : (lastCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows);
+}
+
+void MoeInitRoutingQuantV2TilingBase::Tiling4SrcToDstCapacityCompute() {
+  if (quantMode == 0 || dropPadMode == 0) {
+    InnerMoeInitRoutingV2TilingBase::Tiling4SrcToDstCapacityCompute();
+    return;
+  }
+  
+  auto tilingData = &moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp;
+  int64_t perCoreRows = CeilDiv(totalLength, aivNum);
+  if (perCoreRows <= 0) {
+    tilingData->needCoreNum = 0;
+    return;
+  }
+
+  tilingData->needCoreNum = CeilDiv(totalLength, perCoreRows);
+  int64_t cols = moeInitRoutingTilingData.cols;
+  tilingData->perCoreRows = perCoreRows;
+  int64_t lastCoreRows = totalLength - perCoreRows * (tilingData->needCoreNum - 1);
+  tilingData->lastCoreRows = lastCoreRows;
+
+  int64_t rowSize = AlignOneBlockByte(perCoreRows * sizeof(int32_t)) * NUM_FOUR;
+  int64_t colSize = AlignOneBlockByte(cols * sizeof(int8_t)) * DYNAMIC_QUANT_SRC_TO_DST_BUFFER;
+  int64_t scaleSize = DYNAMIC_QUANT_SCALE_SIZE_64;
+  if (rowSize + colSize + scaleSize < static_cast<int64_t>(aicoreParams_.ubSize)) {
+    
+    SetGatherTilingData(tilingData, perCoreRows, lastCoreRows, cols);
+  } else {
+    
+    int64_t baseMaxCols = MAX_COLS_DYNAMIC_QUANT;
+    int64_t totalColSize = AlignOneBlockByte(baseMaxCols * sizeof(int8_t)) * DYNAMIC_QUANT_SRC_TO_DST_BUFFER;
+    int64_t ubSize = static_cast<int64_t>(aicoreParams_.ubSize);
+    int64_t basePerLoopMaxRows =
+        AlignOneBlockByteCeil((ubSize - totalColSize - scaleSize) / sizeof(int32_t)) / NUM_FOUR;
+    if (cols < MAX_COLS_DYNAMIC_QUANT) {
+      basePerLoopMaxRows = AlignOneBlockByteCeil((ubSize - colSize - scaleSize) / sizeof(int32_t)) / NUM_FOUR;
+    } else if (perCoreRows < basePerLoopMaxRows) {
+      baseMaxCols = AlignOneBlockByteCeil(ubSize - rowSize - scaleSize) / DYNAMIC_QUANT_SRC_TO_DST_BUFFER;
+    }
+    SetGatherTilingDataCols(tilingData, baseMaxCols, cols);
+    SetGatherTilingDataRows(tilingData, perCoreRows, lastCoreRows, basePerLoopMaxRows);
+  }
+}
+
+
+void MoeInitRoutingQuantV2TilingBase::Tiling4GatherQuant() {
+  auto tilingData = &quantTilingData.gatherOutComputeParamsOp;
+  tilingData->activateRows = totalLength;
+  if (dropPadMode == 0 && activateNum > 0) {
+      tilingData->activateRows = (std::min(activateNum, totalLength));
+  }
+  int64_t perCoreRows = CeilDiv(totalLength, aivNum);
+
+  if (perCoreRows <= 0) {
+    tilingData->needCoreNum = 0;
+    return;
+  }
+
+  tilingData->needCoreNum = (CeilDiv(totalLength, perCoreRows));
+  int64_t cols = moeInitRoutingTilingData.cols;
+  tilingData->perCoreRows = perCoreRows;
+  int64_t lastCoreRows = totalLength - perCoreRows * (tilingData->needCoreNum - 1);
+  tilingData->lastCoreRows = lastCoreRows;
+  int64_t sizeOfCol = sizeof(int8_t) * NUM_TWO + sizeof(float) + sizeof(int16_t) + inuptXDtypeSize_ * NUM_TWO;
+  int64_t rowSize = AlignOneBlockByte((perCoreRows * sizeof(int32_t) * NUM_TWO));
+  int64_t colSize = AlignOneBlockByte(cols * sizeOfCol);
+  if (rowSize + colSize < static_cast<int64_t>(aicoreParams_.ubSize) / NUM_TWO) {
+    SetGatherTilingData(tilingData, perCoreRows, lastCoreRows, cols);
+  } else {
+    int64_t baseMaxCols = MAX_COLS_ONE_LOOP_QUANT;
+    int64_t baseMaxColsSize = AlignOneBlockByte(baseMaxCols * sizeOfCol);
+    int64_t ubSize = static_cast<int64_t>(aicoreParams_.ubSize);
+    int64_t basePerLoopMaxRows = AlignOneBlockByteCeil((ubSize - baseMaxColsSize) / NUM_TWO / sizeof(int32_t));
+    if (cols < MAX_COLS_ONE_LOOP_QUANT) {
+      basePerLoopMaxRows = AlignOneBlockByteCeil((ubSize - colSize) / NUM_TWO / sizeof(int32_t));
+    } else if (perCoreRows < basePerLoopMaxRows) {
+      baseMaxCols = AlignOneBlockByteCeil((ubSize - rowSize) / sizeOfCol);
+    }
+    SetGatherTilingDataCols(tilingData, baseMaxCols, cols);
+    SetGatherTilingDataRows(tilingData, perCoreRows, lastCoreRows, basePerLoopMaxRows);
+  }
+}
+
+
+
+void SetGatherTilingDatawithloop(InnerMoeV2GatherOutComputeTilingData* tilingData,
+                                 int64_t perCorePerLoopRows, int64_t lastCorePerLoopRows, int64_t cols,
+                                 int64_t perCoreLastLoopRows = 1, int64_t lastCoreLastLoopRows = 1,
+                                 int64_t perCoreLoops = 1, int64_t lastCoreLoops = 1) {
+    tilingData-> perCorePerLoopRows = perCorePerLoopRows;
+    tilingData-> perCoreLastLoopRows = perCoreLastLoopRows;
+    tilingData-> lastCorePerLoopRows = lastCorePerLoopRows;
+    tilingData-> lastCoreLastLoopRows = lastCoreLastLoopRows;
+    tilingData-> perCoreLoops = perCoreLoops;
+    tilingData-> lastCoreLoops = lastCoreLoops;
+    tilingData-> perLoopCols = cols;
+    tilingData-> lastLoopCols = cols;
+    tilingData-> colLoops = 1;
+}
+
+void MoeInitRoutingQuantV2TilingBase::Tiling4GatherDynamicQuant() {
+
+  auto tilingData = &quantTilingData.gatherOutComputeParamsOp;
+  tilingData->activateRows = totalLength;
+  if (dropPadMode == 0 && activateNum > 0) {
+      tilingData->activateRows = (std::min(activateNum, totalLength));
+  }
+  int64_t perCoreRows = CeilDiv(totalLength, aivNum);
+
+  if (perCoreRows <= 0) {
+    tilingData->needCoreNum = 0;
+    return;
+  }
+
+  tilingData->needCoreNum = (CeilDiv(totalLength, perCoreRows));
+
+  int64_t cols = InnerMoeInitRoutingV2TilingBase::moeInitRoutingTilingData.cols;
+
+  tilingData->perCoreRows = perCoreRows;
+  int64_t lastCoreRows = totalLength - perCoreRows * (tilingData->needCoreNum - 1);
+  tilingData->lastCoreRows = lastCoreRows;
+
+
+  int64_t rowSize = AlignOneBlockByte(perCoreRows * sizeof(int32_t)) * NUM_FOUR;
+  int64_t colSize = AlignOneBlockByte(cols * sizeof(int8_t)) * DYNAMIC_QUANT_COLS_BUFFER;
+  int64_t scaleSize = DYNAMIC_QUANT_SCALE_SIZE_64;
+  int64_t onceRowSize = (static_cast<int64_t>(aicoreParams_.ubSize) - 
+                          colSize - scaleSize - 
+                          ONE_BLOCK_BYTE * NUM_FOUR * NUM_THREE) /
+                        (sizeof(int32_t) * NUM_FOUR);
+  int64_t oneBlockNumInt = static_cast<int64_t>(ONE_BLOCK_BYTE) / static_cast<int64_t>(sizeof(int32_t));
+  onceRowSize = onceRowSize / oneBlockNumInt * oneBlockNumInt;
+  bool ifOneLoop = ((static_cast<int64_t>(aicoreParams_.ubSize) > colSize +
+            scaleSize + ONE_BLOCK_BYTE * NUM_FOUR * NUM_FOUR) && 
+            quantTilingData.smoothType == SMOOTH_NONE &&
+            cols == FULLLOAD_H_LIMIT);
+
+  int64_t perCoreOnceRowSize = ifOneLoop ? std::min(onceRowSize, perCoreRows) : perCoreRows;
+  int64_t lastCoreOnceRowSize = ifOneLoop ? std::min(onceRowSize, lastCoreRows) : lastCoreRows;
+  int64_t perCoreLoops = ifOneLoop ? CeilDiv(perCoreRows, perCoreOnceRowSize) : 1;
+  int64_t lastCoreLoops = ifOneLoop ? CeilDiv(lastCoreRows, lastCoreOnceRowSize) : 1;
+  int64_t perCoreLastLoopRows = ifOneLoop ? GetPerOrLastValue(perCoreRows, perCoreOnceRowSize) : perCoreRows;
+  int64_t lastCoreLastLoopRows = ifOneLoop ?  GetPerOrLastValue(lastCoreRows, lastCoreOnceRowSize) : lastCoreRows;
+ 
+  if (rowSize + colSize + scaleSize < static_cast<int64_t>(aicoreParams_.ubSize) || ifOneLoop) {
+    
+    SetGatherTilingDatawithloop(tilingData, perCoreOnceRowSize, lastCoreOnceRowSize, cols,
+                                perCoreLastLoopRows, lastCoreLastLoopRows,
+                                perCoreLoops, lastCoreLoops);
+  } else {
+    int64_t baseMaxCols = MAX_COLS_DYNAMIC_QUANT;
+    int64_t totalColSize = AlignOneBlockByte(baseMaxCols * sizeof(int8_t)) * DYNAMIC_QUANT_COLS_BUFFER;
+    int64_t ubSize = static_cast<int64_t>(aicoreParams_.ubSize);
+    int64_t basePerLoopMaxRows =
+        AlignOneBlockByteCeil((ubSize - totalColSize - scaleSize) / sizeof(int32_t)) / NUM_FOUR;
+    if (cols < MAX_COLS_DYNAMIC_QUANT) {
+      basePerLoopMaxRows = AlignOneBlockByteCeil((ubSize - colSize - scaleSize) / sizeof(int32_t)) / NUM_FOUR;
+    } else if (perCoreRows < basePerLoopMaxRows) {
+      baseMaxCols = AlignOneBlockByteCeil(ubSize - rowSize - scaleSize) / DYNAMIC_QUANT_COLS_BUFFER;
+    }
+    SetGatherTilingDataCols(tilingData, baseMaxCols, cols);
+    SetGatherTilingDataRows(tilingData, perCoreRows, lastCoreRows, basePerLoopMaxRows);
+  }
+}
+
+
+void MoeInitRoutingQuantV2TilingBase::Tiling4GatherOutCompute() {
+  if (quantMode == 0) {
+    Tiling4GatherQuant();
+  } else {
+    Tiling4GatherDynamicQuant();
+  }
+}
+
+
+}
+
+#endif

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_init_routing_v2_tiling.h

@@ -0,0 +1,415 @@
+#pragma once
+
+#ifndef ASCENDC_DISPATCH_FFN_COMBINE_INIT_ROUTING_TILING_H
+#define ASCENDC_DISPATCH_FFN_COMBINE_INIT_ROUTING_TILING_H
+
+#include "tiling_base.h"
+
+
+namespace optiling {
+const static int64_t TILING_KEY_DROPLESS_SORT_ONE_CORE = 10001;
+const static int64_t TILING_KEY_DROPLESS_SORT_MULTI_CORE = 10002;
+const static int64_t TILING_KEY_DROP_PAD_MODE_SORT_ONE_CORE = 10011;
+const static int64_t TILING_KEY_DROP_PAD_MODE_SORT_MULTI_CORE = 10012;
+const static int64_t TILING_KEY_HIGH_PERFORMANCE = 20000;
+const static int64_t NUM_TWO = 2;
+const static int64_t NUM_THREE = 3;
+const static int64_t NUM_FOUR = 4;
+const static int64_t MRG_LIST_NUM = 4;
+const static int64_t SORT32_ALIGN_ELEMENT = 32;
+const static int64_t ONE_BLOCK_BYTE = 32;
+const static size_t DIM_ONE = 1;
+const static size_t DIM_TWO = 2;
+const static size_t DIM_THREE = 3;
+const static int32_t SIZE_16 = 16;
+const static int32_t LENGTH_1024 = 1024;
+const static int64_t MAX_COLS_ONE_LOOP = 16376;
+const static int64_t ASSIST_NUM = 256;
+const static int64_t INDEX_INPUT_X = 0;
+const static int64_t INDEX_INPUT_EXPERT_IDX = 1;
+const static int64_t ATTR_ACTIVE_ROWS = 0;
+const static int64_t ATTR_EXPERT_CAPACITY = 1;
+const static int64_t ATTR_EXPERT_NUM = 2;
+const static int64_t ATTR_DROP_PAD_MODE = 3;
+const static int64_t ATTR_EXPERT_TOKENS_COUNT_OR_CUMSUM_FLAG = 4;
+const static int64_t ATTR_EXPERT_TOKENS_BEFORE_CAPACITY_FLAG = 5;
+const static int64_t OUTOUT_EXPANDED_X = 0;
+const static int64_t OUTOUT_EXPANDED_ROW_IDX = 1;
+const static int64_t OUTOUT_EXPERT_TOKENS_COUNT_OR_CUMSUM = 2;
+const static int64_t OUTOUT_EXPERT_TOKENS_BEFORE_CAPACITY = 3;
+const static int64_t KV_FACTOR = 2;
+const static int64_t ONE_CORE_SORT_BUFFER = 6;
+const static int64_t EXPERT_TOKENS_COUNT = 2;
+
+
+inline static int64_t CeilLog4(int64_t x) {
+  return static_cast<int64_t>(std::ceil(std::log(x) / std::log(NUM_FOUR)));
+}
+
+inline static int64_t GetPerOrLastValue(int64_t x, int64_t y) {
+  if (y == 0) {
+    return 0;
+  }
+  return x <= y ? x : x % y;
+}
+
+template <class T>
+constexpr T CeilDiv(const T dividend, const T divisor)
+{
+    return (dividend + divisor - 1) / divisor;
+}
+
+
+struct InnerMoeV2VBSComputeTilingData {
+    int64_t needCoreNum = 0;
+    int64_t perCoreElements = 0;
+    int64_t perCoreLoops = 0;
+    int64_t perCorePerLoopElements = 0;
+    int64_t perCoreLastLoopElements = 0;
+    int64_t lastCoreElements = 0;
+    int64_t lastCoreLoops = 0;
+    int64_t lastCorePerLoopElements = 0;
+    int64_t lastCoreLastLoopElements = 0;
+    int64_t oneLoopMaxElements = 0;
+};
+
+struct InnerMoeV2VMSMiddleComputeTilingData {
+    int64_t needCoreNum = 0;
+};
+
+struct InnerMoeV2SortOutComputeTilingData {
+    int64_t oneLoopMaxElements = 0;
+};
+
+struct InnerMoeV2GatherOutComputeTilingData {
+    int64_t needCoreNum = 0;
+    int64_t activateRows = 0;
+    int64_t perCoreRows = 0;
+    int64_t perCorePerLoopRows = 0;
+    int64_t perCoreLastLoopRows = 0;
+    int64_t lastCoreRows = 0;
+    int64_t lastCorePerLoopRows = 0;
+    int64_t lastCoreLastLoopRows = 0;
+    int64_t perCoreLoops = 0;
+    int64_t lastCoreLoops = 0;
+    int64_t perLoopCols = 0;
+    int64_t lastLoopCols = 0;
+    int64_t colLoops = 0;
+};
+
+struct InnerMoeInitRoutingV2TilingData {
+    int64_t coreNum;
+    int64_t n;
+    int64_t cols;
+    int64_t k;
+    int64_t expertCapacity;
+    int64_t expertNum;
+    int64_t dropPadMode;
+    int64_t expertTokensCountOrCumsumFlag;
+    int64_t expertTokensBeforeCapacityFlag;
+    InnerMoeV2VBSComputeTilingData vbsComputeParamsOp;
+    InnerMoeV2VMSMiddleComputeTilingData vmsMiddleComputeParamsOp;
+    InnerMoeV2SortOutComputeTilingData sortOutComputeParamsOp;
+    InnerMoeV2GatherOutComputeTilingData srcToDstComputeParamsOp;
+    InnerMoeV2GatherOutComputeTilingData srcToDstCapacityComputeParamsOp;
+    InnerMoeV2GatherOutComputeTilingData gatherOutComputeParamsOp;
+};
+
+
+class InnerMoeInitRoutingV2TilingBase : public TilingBaseClass {
+
+protected:
+  bool GetPlatformInfo(int64_t aivCoreNum, int64_t ubSizePlatForm) override;
+  bool GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, 
+  int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag,
+  bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) override;
+
+  bool DoOpTiling() override;
+  uint64_t GetTilingKey() const override;
+  bool GetWorkspaceSize() override;
+
+
+protected:
+  bool CheckTokenCount(int64_t num, const char* tag);
+
+  virtual void Tiling4GatherOutCompute() = 0;
+  void Tiling4SrcToDstCompute();
+  virtual void Tiling4SrcToDstCapacityCompute();
+  void Tiling4SortOutCompute();
+  void Tiling4VMSMiddleCompute();
+  void Tiling4VBSCompute();
+  void ShowTilingData();
+  void Tiling4VBSMultiCoreCompute(InnerMoeV2VBSComputeTilingData* tilingData);
+  void Tiling4VBSOneCoreCompute(InnerMoeV2VBSComputeTilingData* tilingData);
+  virtual bool IsFullLoad() = 0;
+
+
+
+
+  int64_t aivNum = 0;
+  int64_t sortLoopMaxElement = 0;
+  int64_t mrgSortListMaxElement = 2040;
+  int64_t totalLength = 0;
+  int64_t activateNum = 0;
+  int64_t expertCapacity = 0;
+  int64_t expertNum = 0;
+  int64_t dropPadMode = 0;
+  int64_t expertTokensCountOrCumsumFlag = 0;
+  bool expertTokensBeforeCapacityFlag = false;
+  int64_t inuptXDtypeSize_ = 0;
+  bool isFullLoad = false;
+
+  InnerMoeInitRoutingV2TilingData moeInitRoutingTilingData;
+};
+
+
+bool InnerMoeInitRoutingV2TilingBase::DoOpTiling() {
+  sortLoopMaxElement =
+      (aicoreParams_.ubSize) / (sizeof(int32_t) * NUM_TWO * NUM_FOUR) / SORT32_ALIGN_ELEMENT * SORT32_ALIGN_ELEMENT;
+  isFullLoad = IsFullLoad();
+  Tiling4VBSCompute();
+  Tiling4VMSMiddleCompute();
+  Tiling4SortOutCompute();
+  Tiling4SrcToDstCompute();
+  Tiling4SrcToDstCapacityCompute();
+  Tiling4GatherOutCompute();
+  return true;
+};
+
+uint64_t InnerMoeInitRoutingV2TilingBase::GetTilingKey() const {
+  if (isFullLoad) {
+    return TILING_KEY_HIGH_PERFORMANCE;
+  }
+  if (dropPadMode == 0) {
+    if (totalLength <= sortLoopMaxElement) {  // Sorting uses only one core
+      return TILING_KEY_DROPLESS_SORT_ONE_CORE;
+    } else {
+      return TILING_KEY_DROPLESS_SORT_MULTI_CORE;
+    }
+  } else {
+    if (totalLength <= sortLoopMaxElement) {
+      return TILING_KEY_DROP_PAD_MODE_SORT_ONE_CORE;
+    } else {
+      return TILING_KEY_DROP_PAD_MODE_SORT_MULTI_CORE;
+    }
+  }
+  return tilingKey_;
+}
+
+
+
+bool InnerMoeInitRoutingV2TilingBase::GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, 
+  int64_t expertNum, int64_t activateNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag,
+  bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) {
+  
+  this->activateNum = activateNum;
+  this->expertCapacity = expertCapacity;
+  this->expertNum = expertNum;
+  this->dropPadMode = dropPadMode;
+  this->expertTokensCountOrCumsumFlag = expertTokensCountOrCumsumFlag;
+  this->expertTokensBeforeCapacityFlag = expertTokensBeforeCapacityFlag;
+  if (dropPadMode == 1) {
+    // Do not output expertTokensCountOrCumsum in drop-pad mode
+    expertTokensCountOrCumsumFlag = 0;
+  } else {
+    // Do not output expertTokensBeforeCapacity in dropless mode
+    expertTokensBeforeCapacityFlag = false;
+  }
+  moeInitRoutingTilingData.cols = cols;
+  moeInitRoutingTilingData.n = m;
+  moeInitRoutingTilingData.k = topK;
+  moeInitRoutingTilingData.expertCapacity = expertCapacity;
+  moeInitRoutingTilingData.expertNum = expertNum;
+  moeInitRoutingTilingData.dropPadMode = dropPadMode;
+  moeInitRoutingTilingData.expertTokensCountOrCumsumFlag = expertTokensCountOrCumsumFlag;
+  moeInitRoutingTilingData.expertTokensBeforeCapacityFlag = expertTokensBeforeCapacityFlag;
+  totalLength = moeInitRoutingTilingData.n * moeInitRoutingTilingData.k;
+  inuptXDtypeSize_ = inuptXDtypeSize;
+  return true;
+}
+
+bool InnerMoeInitRoutingV2TilingBase::GetPlatformInfo(int64_t aivCoreNum, int64_t ubSizePlatForm) {
+  aivNum = aivCoreNum;
+  aicoreParams_.blockDim = aivCoreNum;
+  aicoreParams_.ubSize = ubSizePlatForm;
+  moeInitRoutingTilingData.coreNum = aivCoreNum;
+  return true;
+}
+
+
+bool InnerMoeInitRoutingV2TilingBase::GetWorkspaceSize() {
+  // Calculate workspace size
+  size_t sortWorkspaceSize = totalLength * sizeof(float) * NUM_TWO * NUM_THREE;  // Space needed for sorting
+  size_t scatterWorkspaceSize = totalLength * sizeof(int32_t) * NUM_TWO;
+  size_t expertTokenFlagSize = aivNum * 2 * sizeof(int32_t);
+  workspaceSize_ = sortWorkspaceSize + scatterWorkspaceSize + expertTokenFlagSize + SIZE_16 * LENGTH_1024 * LENGTH_1024;
+  return true;
+}
+
+void InnerMoeInitRoutingV2TilingBase::Tiling4VBSOneCoreCompute(InnerMoeV2VBSComputeTilingData* tilingData) {
+    tilingData->needCoreNum = 1;
+    tilingData->perCoreElements = totalLength;
+    tilingData->perCoreLoops = 1;
+    tilingData->perCorePerLoopElements = tilingData->perCoreElements;
+    tilingData->perCoreLastLoopElements = tilingData->perCoreElements;
+    tilingData->lastCoreElements = tilingData->perCoreElements;
+    tilingData->lastCoreLoops = 1;
+    tilingData->lastCorePerLoopElements = tilingData->perCoreElements;
+    tilingData->lastCoreLastLoopElements = tilingData->perCoreElements;
+}
+
+void InnerMoeInitRoutingV2TilingBase::Tiling4VBSMultiCoreCompute(InnerMoeV2VBSComputeTilingData* tilingData) {
+      //Tiling4VBSMultiCoreCompute
+      int64_t needCoreNum = CeilDiv(totalLength, sortLoopMaxElement);  // Round up
+      needCoreNum = static_cast<int64_t>(std::pow(4, CeilLog4(needCoreNum)));
+      needCoreNum = std::min(needCoreNum, aivNum);                     // Cannot exceed physical core count
+      if (needCoreNum > 0) {
+          int64_t perCoreElements = totalLength / needCoreNum;  // Elements handled per core
+          int64_t alineFloorPerCoreElements = perCoreElements - perCoreElements % SORT32_ALIGN_ELEMENT;
+          int64_t lastCoreElement = totalLength - (needCoreNum - 1) * alineFloorPerCoreElements;
+          int64_t alineCeilPerCoreElements = perCoreElements + SORT32_ALIGN_ELEMENT - perCoreElements % SORT32_ALIGN_ELEMENT;
+          if (lastCoreElement > alineCeilPerCoreElements) {
+            perCoreElements = alineCeilPerCoreElements;
+            needCoreNum = CeilDiv(totalLength, perCoreElements);
+          } else {
+            perCoreElements = alineFloorPerCoreElements;
+          }
+          tilingData->needCoreNum = needCoreNum;
+          do {
+              tilingData->perCoreElements = perCoreElements;
+              tilingData->perCoreLoops = CeilDiv(tilingData->perCoreElements, sortLoopMaxElement);  // Loops handled per core
+              tilingData->perCorePerLoopElements = std::min(tilingData->perCoreElements, sortLoopMaxElement);
+              tilingData->perCoreLastLoopElements = tilingData->perCoreElements - (tilingData->perCoreLoops - 1) * tilingData->perCorePerLoopElements;
+              tilingData->lastCoreElements = totalLength - (tilingData->needCoreNum - 1) * tilingData->perCoreElements;
+              tilingData->lastCoreLoops = tilingData->perCoreLoops;
+              int64_t tmp = CeilDiv(tilingData->lastCoreElements, tilingData->lastCoreLoops);
+              int64_t lastCorePerLoopElements =
+                  CeilDiv(CeilDiv(tilingData->lastCoreElements, tilingData->lastCoreLoops), SORT32_ALIGN_ELEMENT) *
+                  SORT32_ALIGN_ELEMENT;
+              tilingData->lastCorePerLoopElements = lastCorePerLoopElements;
+              tilingData->lastCoreLastLoopElements = tilingData-> lastCoreElements - (tilingData->lastCoreLoops - 1) * tilingData->lastCorePerLoopElements;
+              perCoreElements -= SORT32_ALIGN_ELEMENT;
+          } while (tilingData->lastCoreLastLoopElements <= 0 && perCoreElements > 0);
+      }
+}
+
+
+void InnerMoeInitRoutingV2TilingBase::Tiling4VBSCompute() {
+  auto tilingData = &moeInitRoutingTilingData.vbsComputeParamsOp;
+  tilingData->oneLoopMaxElements = sortLoopMaxElement;
+  if (totalLength <= sortLoopMaxElement) {  // Only one core is used
+    Tiling4VBSOneCoreCompute(tilingData);
+    return;
+  }
+  Tiling4VBSMultiCoreCompute(tilingData);
+}
+
+void InnerMoeInitRoutingV2TilingBase::Tiling4VMSMiddleCompute() {
+  auto vbsComputeTilingData = &moeInitRoutingTilingData.vbsComputeParamsOp;
+  auto tilingData = &moeInitRoutingTilingData.vmsMiddleComputeParamsOp;
+  if (vbsComputeTilingData->needCoreNum <= MRG_LIST_NUM) {  // No intermediate merge if queue count fits one VMS
+      tilingData->needCoreNum = 0;                               // Required core count
+  } else {
+      int64_t needCoreNum = CeilDiv(vbsComputeTilingData->needCoreNum, MRG_LIST_NUM);
+      tilingData->needCoreNum = needCoreNum;  // Required core count
+  }
+}
+
+void InnerMoeInitRoutingV2TilingBase::Tiling4SortOutCompute() {
+  auto tilingData = &moeInitRoutingTilingData.sortOutComputeParamsOp;
+  tilingData->oneLoopMaxElements = mrgSortListMaxElement;
+}
+
+
+void InnerMoeInitRoutingV2TilingBase::Tiling4SrcToDstCompute() {
+  auto tilingData = &moeInitRoutingTilingData.srcToDstComputeParamsOp;
+
+  int64_t perLoopMaxRows = (aicoreParams_.ubSize - ASSIST_NUM * sizeof(float) - aivNum * SORT32_ALIGN_ELEMENT) /
+                           (SORT32_ALIGN_ELEMENT * NUM_TWO) / NUM_TWO;
+  int64_t perCoreRows = CeilDiv(totalLength, aivNum);
+  if (perCoreRows <= 0) {
+    tilingData->needCoreNum = 0;
+    return;
+  }
+
+  int64_t needCoreNum = CeilDiv(totalLength, perCoreRows);
+  tilingData->needCoreNum = needCoreNum;
+  int64_t lastCoreNum = totalLength - perCoreRows * (tilingData->needCoreNum - 1);
+  tilingData->perCoreRows = perCoreRows;
+  if (perLoopMaxRows >= tilingData->perCoreRows) {  // One loop completes
+      tilingData->perCorePerLoopRows = tilingData->perCoreRows;
+      tilingData->perCoreLastLoopRows = tilingData->perCoreRows;
+  } else {
+      tilingData->perCorePerLoopRows = perLoopMaxRows;
+      tilingData->perCoreLastLoopRows = tilingData->perCoreRows - (CeilDiv(tilingData->perCoreRows, perLoopMaxRows) - 1) * perLoopMaxRows;
+  }
+  tilingData->lastCoreRows = lastCoreNum;
+  if (perLoopMaxRows >= tilingData->lastCoreRows) {
+      tilingData->lastCorePerLoopRows = tilingData->lastCoreRows;
+      tilingData->lastCoreLastLoopRows = tilingData->lastCoreRows;
+  } else {
+      tilingData->lastCorePerLoopRows = perLoopMaxRows;
+      tilingData->lastCoreLastLoopRows = tilingData->lastCoreRows - (CeilDiv(tilingData->lastCoreRows, perLoopMaxRows) - 1) * perLoopMaxRows;
+  }
+}
+
+
+void InnerMoeInitRoutingV2TilingBase::Tiling4SrcToDstCapacityCompute() {
+  auto tilingData = &moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp;
+  int64_t perCoreRows = CeilDiv(totalLength, aivNum);
+
+  if (perCoreRows <= 0) {
+    tilingData->needCoreNum = 0;
+    return;
+  }
+
+  int64_t needCoreNum = CeilDiv(totalLength, perCoreRows);
+  tilingData->needCoreNum = needCoreNum;
+  int64_t cols = moeInitRoutingTilingData.cols;
+  tilingData->perCoreRows = perCoreRows;
+  int64_t lastCoreRows = totalLength - perCoreRows * (needCoreNum - 1);
+  tilingData->lastCoreRows = lastCoreRows;
+
+
+  int64_t rowSize =
+      (perCoreRows * sizeof(int32_t) * 2 + ONE_BLOCK_BYTE + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+  int64_t colSize = (cols * inuptXDtypeSize_ + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+
+  if (rowSize + colSize < static_cast<int64_t>(aicoreParams_.ubSize)) {
+    tilingData->perCorePerLoopRows = perCoreRows;
+    tilingData->perCoreLastLoopRows = perCoreRows;
+    tilingData->lastCorePerLoopRows = lastCoreRows;
+    tilingData->lastCoreLastLoopRows = lastCoreRows;
+    tilingData->perCoreLoops = 1;
+    tilingData->lastCoreLoops = 1;
+    tilingData->perLoopCols = cols;
+    tilingData->lastLoopCols = cols;
+    tilingData->colLoops = 1;
+
+  } else {
+    int64_t baseMaxCols = MAX_COLS_ONE_LOOP;
+    int64_t baseMaxColsSize = (baseMaxCols * inuptXDtypeSize_ + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+    int64_t basePerLoopMaxRows = (static_cast<int64_t>(aicoreParams_.ubSize) - baseMaxColsSize - ONE_BLOCK_BYTE) /
+                                 static_cast<int64_t>(sizeof(int32_t)) / NUM_TWO / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+    if (cols < MAX_COLS_ONE_LOOP) {
+      basePerLoopMaxRows = (static_cast<int64_t>(aicoreParams_.ubSize) - colSize - ONE_BLOCK_BYTE) /
+                           static_cast<int64_t>(sizeof(int32_t)) / NUM_TWO / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+    } else if (perCoreRows < basePerLoopMaxRows) {
+      baseMaxCols =
+          (static_cast<int64_t>(aicoreParams_.ubSize) - rowSize) / inuptXDtypeSize_ / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+    }
+    tilingData->perLoopCols = (std::min(baseMaxCols, cols));
+    tilingData->lastLoopCols = (GetPerOrLastValue(cols, baseMaxCols));
+    tilingData->colLoops = ((cols + baseMaxCols - 1) / baseMaxCols);
+    tilingData->perCorePerLoopRows = (std::min(perCoreRows, basePerLoopMaxRows));
+    tilingData->perCoreLastLoopRows = (GetPerOrLastValue(perCoreRows, basePerLoopMaxRows));
+    tilingData->perCoreLoops = ((perCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows);
+    tilingData->lastCorePerLoopRows = (std::min(lastCoreRows, basePerLoopMaxRows));
+    tilingData->lastCoreLastLoopRows = (GetPerOrLastValue(lastCoreRows, basePerLoopMaxRows));
+    tilingData->lastCoreLoops = ((lastCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows);
+  }
+}
+
+}
+
+#endif

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_common.h

@@ -0,0 +1,96 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_common.h
+ * \brief
+ */
+#ifndef INNER_MOE_V2_COMMON_H
+#define INNER_MOE_V2_COMMON_H
+
+#include "kernel_operator.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+constexpr int64_t SPLIT_N = 0;
+constexpr int64_t SPLIT_K = 1;
+constexpr float MIN_FP32 = -3.4e38;
+constexpr int64_t ONE_REPEAT_SORT_NUM = 32;
+constexpr int64_t BLOCK_BYTES = 32;
+constexpr int64_t INT32_ONE_BLOCK_NUM = 8;
+
+constexpr int64_t ASSIST_NUM = 256;
+constexpr int64_t ASSIST_INDEX_NUM = 32;
+
+constexpr int64_t MERGE_LIST_TWO = 2;
+constexpr int64_t MERGE_LIST_THREE = 3;
+constexpr int64_t MERGE_LIST_FOUR = 4;
+
+constexpr int64_t MERGE_LIST_IDX_TWO = 2;
+constexpr int64_t MERGE_LIST_IDX_THREE = 3;
+
+constexpr int64_t MAX_EXPERT_NUM = 5120;
+constexpr int64_t DROPLESS_MODE = 0;
+constexpr int64_t DROP_PAD_MODE = 1;
+constexpr int64_t EXERPT_TOKENS_COUNT = 2;
+constexpr int64_t EXERPT_TOKENS_CUMSUM = 1;
+constexpr int64_t EXERPT_TOKENS_NONE = 0;
+constexpr int64_t EXERPT_TOKENS_BEFORE_CAPACITY = 1;
+constexpr int64_t ALIGN_512 = 512;
+constexpr int64_t ALIGN_128 = 128;
+
+const __gm__ int32_t assist[256] = {
+    0,  0, 0, 0, 0, 0, 0, 0, 1,  0, 0, 0, 0, 0, 0, 0, 2,  0, 0, 0, 0, 0, 0, 0, 3,  0, 0, 0, 0, 0, 0, 0,
+    4,  0, 0, 0, 0, 0, 0, 0, 5,  0, 0, 0, 0, 0, 0, 0, 6,  0, 0, 0, 0, 0, 0, 0, 7,  0, 0, 0, 0, 0, 0, 0,
+    8,  0, 0, 0, 0, 0, 0, 0, 9,  0, 0, 0, 0, 0, 0, 0, 10, 0, 0, 0, 0, 0, 0, 0, 11, 0, 0, 0, 0, 0, 0, 0,
+    12, 0, 0, 0, 0, 0, 0, 0, 13, 0, 0, 0, 0, 0, 0, 0, 14, 0, 0, 0, 0, 0, 0, 0, 15, 0, 0, 0, 0, 0, 0, 0,
+    16, 0, 0, 0, 0, 0, 0, 0, 17, 0, 0, 0, 0, 0, 0, 0, 18, 0, 0, 0, 0, 0, 0, 0, 19, 0, 0, 0, 0, 0, 0, 0,
+    20, 0, 0, 0, 0, 0, 0, 0, 21, 0, 0, 0, 0, 0, 0, 0, 22, 0, 0, 0, 0, 0, 0, 0, 23, 0, 0, 0, 0, 0, 0, 0,
+    24, 0, 0, 0, 0, 0, 0, 0, 25, 0, 0, 0, 0, 0, 0, 0, 26, 0, 0, 0, 0, 0, 0, 0, 27, 0, 0, 0, 0, 0, 0, 0,
+    28, 0, 0, 0, 0, 0, 0, 0, 29, 0, 0, 0, 0, 0, 0, 0, 30, 0, 0, 0, 0, 0, 0, 0, 31, 0, 0, 0, 0, 0, 0, 0};
+
+__aicore__ inline int64_t Ceil(int64_t a, int64_t b) {
+  if (b == 0) {
+    return 0;
+  }
+  return (a + b - 1) / b;
+}
+
+__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes) {
+  if (bytes == 0) {
+    return 0;
+  }
+  return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes;
+}
+
+__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes) {
+  return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES;
+}
+
+template <typename T>
+__aicore__ inline T Min(T a, T b) {
+  return a > b ? b : a;
+}
+
+template <typename T>
+__aicore__ inline T Max(T a, T b) {
+  return a < b ? b : a;
+}
+
+template <HardEvent event>
+__aicore__ inline void SetWaitFlag(HardEvent evt) {
+  event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(evt));
+  SetFlag<event>(eventId);
+  WaitFlag<event>(eventId);
+}
+
+}  // namespace MoeInitRoutingQuantV2
+#endif  // INNER_MOE_V2_COMMON_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_expert_token_out.h

@@ -0,0 +1,310 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_expert_token_out.h
+ * \brief
+ */
+#ifndef INNER_MOE_V2_EXPERT_TOKEN_OUT_H
+#define INNER_MOE_V2_EXPERT_TOKEN_OUT_H
+
+#include "moe_v2_common.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+constexpr int64_t EXPERT_ID_VALUE_NUM = 2;
+
+class MoeV2ExpertTokenOut {
+ public:
+  __aicore__ inline MoeV2ExpertTokenOut(){};
+  template <typename TilingData>
+  __aicore__ inline void Init(GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity,
+                              GM_ADDR expandedRowIdx, GM_ADDR workspace, const TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void CopyIn(int64_t progress);
+  __aicore__ inline void Compute(int64_t progress);
+  __aicore__ inline void SyncAll();
+  __aicore__ inline void InitLocal();
+  __aicore__ inline void GetExpertTokenCount(int32_t curExpertId);
+  __aicore__ inline void CopyOutTokenGm();
+  __aicore__ inline void CopyOutExpertTokensCumsum(bool isTail);
+  __aicore__ inline void CopyOutExpertTokensCount(bool isTail);
+
+ private:
+  TPipe* pipe;
+  TQue<QuePosition::VECIN, 1> copyInQueue;
+  TQue<QuePosition::VECIN, 1> expertTokenIdxCopyInQueue;
+  TQue<QuePosition::VECOUT, 1> expertTokenIdxCopyOutQueue;
+
+  GlobalTensor<int32_t> expertTokensCountOrCumsumGm;
+  GlobalTensor<int32_t> expertTokensBeforeCapacityGm;
+  GlobalTensor<int32_t> expandedExpertIdxGm;
+  GlobalTensor<int32_t> expertIdxValueGm;
+  GlobalTensor<int32_t> expandedRowIdxGm;
+
+  LocalTensor<int32_t> expertTokenIdxOutLocal;
+
+  const InnerMoeV2GatherOutComputeTilingData* srcToDstTilingData;
+
+  int64_t coreNum;
+  int64_t blockIdx;
+  int64_t totalLength;
+  int64_t currentLoopRows;
+  int64_t coreRows;
+  int64_t perLoopRows;
+  int64_t lastLoopRows;
+  int64_t expertNum;
+  int64_t expertNumUbAlign;
+  int64_t dropPadMode = 0;
+  int64_t expertTokensCountOrCumsumFlag = 0;
+  int64_t expertTokensBeforeCapacityFlag = 0;
+
+  int64_t tokenCount = 0;
+  int64_t expertIdx = 0;
+  int32_t lastExpertId = -1;
+  int32_t firstExpertId = -1;
+
+  int32_t expertTokenValue = 0;
+};
+
+__aicore__ inline void MoeV2ExpertTokenOut::InitLocal() {
+  LocalTensor<int32_t> tokenIdxLocal = expertTokenIdxCopyOutQueue.AllocTensor<int32_t>();
+  Duplicate<int32_t>(tokenIdxLocal, 0, this->expertNumUbAlign);
+  expertTokenIdxCopyOutQueue.EnQue<int32_t>(tokenIdxLocal);
+
+  // expandedRowIdx initialized to -1, which is used in the src_to_dst_with_capacity step.
+  // use this step SyncAll to synchronize every core data
+  if (this->dropPadMode == 0) {
+    return;
+  }
+  LocalTensor<int32_t> outLocal = copyInQueue.AllocTensor<int32_t>();
+  int64_t loops = (coreRows + perLoopRows - 1) / perLoopRows;
+  Duplicate<int32_t>(outLocal, -1, perLoopRows);
+  SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+  for (int64_t loop = 0; loop < loops; loop++) {
+    int64_t copyLength = perLoopRows;
+    if (loop == loops - 1) {
+      copyLength = lastLoopRows;
+    }
+    DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(copyLength * sizeof(int32_t)), 0, 0,
+                                 0};
+    DataCopyPad(expandedRowIdxGm[this->blockIdx * this->srcToDstTilingData->perCoreRows + loop * perLoopRows], outLocal,
+                copyParams);
+  }
+  SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
+  copyInQueue.FreeTensor(outLocal);
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::CopyIn(int64_t progress) {
+  LocalTensor<int32_t> inLocal = copyInQueue.AllocTensor<int32_t>();
+  DataCopy(inLocal, expandedExpertIdxGm[progress * perLoopRows], Align(currentLoopRows, sizeof(int32_t)));
+  copyInQueue.EnQue<int32_t>(inLocal);
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::GetExpertTokenCount(int32_t curExpertId) {
+  this->tokenCount++;
+  if (this->lastExpertId < curExpertId) {
+    this->expertTokenIdxOutLocal.SetValue(this->expertIdx, this->tokenCount - 1);
+    this->tokenCount = 1;
+    this->expertIdx += (curExpertId - this->lastExpertId);
+    while (curExpertId - this->firstExpertId + 1 > this->expertNumUbAlign) {
+      SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+      CopyOutExpertTokensCumsum(false);
+      CopyOutExpertTokensCount(false);
+      SetWaitFlag<HardEvent::MTE3_V>(HardEvent::MTE3_V);
+      Duplicate<int32_t>(this->expertTokenIdxOutLocal, 0, this->expertNumUbAlign);
+      SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
+      this->firstExpertId += this->expertNumUbAlign;
+      this->expertIdx = curExpertId - this->firstExpertId;
+    }
+    this->lastExpertId = curExpertId;
+  }
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::Compute(int64_t progress) {
+  LocalTensor<int32_t> inLocal = copyInQueue.DeQue<int32_t>();
+  SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+  if (this->lastExpertId == -1) {
+    this->lastExpertId = inLocal.GetValue(0);
+    this->firstExpertId = this->lastExpertId;
+  }
+  for (int64_t i = 0; i < currentLoopRows; i++) {
+    int32_t expertId = inLocal.GetValue(i);
+    GetExpertTokenCount(expertId);
+  }
+  this->expertTokenIdxOutLocal.SetValue(this->expertIdx, this->tokenCount);
+  copyInQueue.FreeTensor(inLocal);
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::CopyOutExpertTokensCumsum(bool isTail) {
+  if (this->dropPadMode != DROPLESS_MODE || expertTokensCountOrCumsumFlag != EXERPT_TOKENS_CUMSUM) {
+    return;
+  }
+#ifdef __CCE_KT_TEST__
+    // CPU twin debugging cannot use multi-core sync, so index may contain uninitialized dirty data; handle specially
+    if (this->firstExpertId > expertTokensCountOrCumsumGm.GetSize()) {
+        return;
+    }
+#endif
+  int64_t copyLength = isTail ? this->lastExpertId - this->firstExpertId + 1 : this->expertNumUbAlign;
+  int64_t end = this->expertNum - this->firstExpertId;
+  for (int64_t i = 0; i < copyLength; i++) {
+    this->expertTokenValue += this->expertTokenIdxOutLocal.GetValue(i);
+    this->expertTokenIdxOutLocal.SetValue(i, this->expertTokenValue);
+  }
+  // if the remaining UB is sufficient, use the UB space to copy
+  // otherwise, copy the calculated data first, and then copy the last tokenValue to remaining expert position
+  if (isTail && end <= this->expertNumUbAlign) {
+    int64_t startAlign = Min(Align(copyLength, sizeof(int32_t)), end);
+    for (int64_t i = copyLength; i < startAlign; i++) {
+      this->expertTokenIdxOutLocal.SetValue(i, this->expertTokenValue);
+    }
+    if (startAlign < end) {
+      Duplicate<int32_t>(this->expertTokenIdxOutLocal[startAlign], this->expertTokenValue, end - startAlign);
+    }
+    copyLength = end;
+    SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+  }
+  DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(copyLength * sizeof(int32_t)), 0, 0, 0};
+  SetAtomicAdd<int32_t>();
+#ifndef __CCE_KT_TEST__
+  DataCopyPad(expertTokensCountOrCumsumGm[this->firstExpertId], this->expertTokenIdxOutLocal, copyParams);
+#endif
+  SetAtomicNone();
+  if (isTail && end > this->expertNumUbAlign) {
+    int64_t remainderLength = end - copyLength;
+    SetWaitFlag<HardEvent::MTE3_V>(HardEvent::MTE3_V);
+    Duplicate<int32_t>(this->expertTokenIdxOutLocal, this->expertTokenValue, this->expertNumUbAlign);
+    SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+    int64_t loopTimes = remainderLength / this->expertNumUbAlign + 1;
+    for (int64_t i = 0; i < loopTimes; i++) {
+      copyLength = i == loopTimes - 1 ? remainderLength - this->expertNumUbAlign * i : this->expertNumUbAlign;
+      DataCopyExtParams params{static_cast<uint16_t>(1), static_cast<uint32_t>(copyLength * sizeof(int32_t)), 0, 0, 0};
+      SetAtomicAdd<int32_t>();
+      DataCopyPad(expertTokensCountOrCumsumGm[this->lastExpertId + 1 + this->expertNumUbAlign * i],
+                  this->expertTokenIdxOutLocal, params);
+      SetAtomicNone();
+    }
+  }
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::CopyOutExpertTokensCount(bool isTail) {
+  int64_t copyLength = isTail ? this->lastExpertId - this->firstExpertId + 1 : this->expertNumUbAlign;
+  DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(copyLength * sizeof(int32_t)), 0, 0, 0};
+#ifdef __CCE_KT_TEST__
+    // CPU twin debugging skips output copies
+    return;
+#endif
+  SetAtomicAdd<int32_t>();
+  if (this->dropPadMode == DROP_PAD_MODE && expertTokensBeforeCapacityFlag > EXERPT_TOKENS_NONE) {
+    DataCopyPad(expertTokensBeforeCapacityGm[this->firstExpertId], this->expertTokenIdxOutLocal, copyParams);
+  }
+  if (this->dropPadMode == DROPLESS_MODE && expertTokensCountOrCumsumFlag == EXERPT_TOKENS_COUNT) {
+    DataCopyPad(expertTokensCountOrCumsumGm[this->firstExpertId], this->expertTokenIdxOutLocal, copyParams);
+  }
+  SetAtomicNone();
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::CopyOutTokenGm() {
+  if (this->dropPadMode == DROPLESS_MODE) {
+    SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+    CopyOutExpertTokensCumsum(true);
+    CopyOutExpertTokensCount(true);
+    return;
+  }
+  this->expertTokenIdxOutLocal.SetValue(this->expertNumUbAlign, this->lastExpertId);
+  this->expertTokenIdxOutLocal.SetValue(this->expertNumUbAlign + 1, this->tokenCount);
+  DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(EXPERT_ID_VALUE_NUM * sizeof(int32_t)),
+                               0, 0, 0};
+  SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+  DataCopyPad(expertIdxValueGm[this->blockIdx * EXPERT_ID_VALUE_NUM],
+              this->expertTokenIdxOutLocal[this->expertNumUbAlign], copyParams);
+  CopyOutExpertTokensCount(true);
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::SyncAll() {
+  if (coreNum == 1) {
+    return;
+  }
+#ifndef __CCE_KT_TEST__
+  AscendC::SyncAll();
+#endif
+}
+
+template <typename TilingData>
+__aicore__ inline void MoeV2ExpertTokenOut::Init(GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity,
+                                                 GM_ADDR expandedRowIdx, GM_ADDR workspace,
+                                                 const TilingData* tilingData, TPipe* tPipe) {
+  int64_t blockNum = GetBlockNum();
+  this->pipe = tPipe;
+  //this->blockIdx = GetBlockIdx();
+  this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+  this->coreNum = tilingData->coreNum;
+  this->totalLength = tilingData->n * tilingData->k;
+  this->srcToDstTilingData = &(tilingData->srcToDstComputeParamsOp);
+  this->expertNum = tilingData->expertNum;
+  this->dropPadMode = tilingData->dropPadMode;
+  this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag;
+  this->expertTokensBeforeCapacityFlag = tilingData->expertTokensBeforeCapacityFlag;
+
+  if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) {
+    this->coreRows = this->srcToDstTilingData->lastCoreRows;
+    this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows;
+    this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows;
+  } else {
+    this->coreRows = this->srcToDstTilingData->perCoreRows;
+    this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows;
+    this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows;
+  }
+
+  expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, Align(this->totalLength, sizeof(int32_t)));
+  if (this->dropPadMode == DROPLESS_MODE && this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) {
+    expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensCountOrCumsum, this->expertNum);
+  }
+  if (this->dropPadMode == DROP_PAD_MODE && this->expertTokensBeforeCapacityFlag == EXERPT_TOKENS_BEFORE_CAPACITY) {
+    expertTokensBeforeCapacityGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensBeforeCapacity, this->expertNum);
+  }
+
+  expandedExpertIdxGm.SetGlobalBuffer(
+      (__gm__ int32_t*)workspace + this->blockIdx * this->srcToDstTilingData->perCoreRows,
+      Align(this->coreRows, sizeof(int32_t)));
+  expertIdxValueGm.SetGlobalBuffer((__gm__ int32_t*)workspace + Align(this->totalLength, sizeof(int32_t)) * 2,
+                                   this->coreNum * 2);
+
+  this->expertNumUbAlign = Min(Align(this->expertNum, sizeof(int32_t)), MAX_EXPERT_NUM);
+  pipe->InitBuffer(copyInQueue, 1, this->perLoopRows * BLOCK_BYTES);
+  pipe->InitBuffer(expertTokenIdxCopyInQueue, 1, this->expertNumUbAlign * sizeof(int32_t));
+  pipe->InitBuffer(expertTokenIdxCopyOutQueue, 1, (this->expertNumUbAlign + EXPERT_ID_VALUE_NUM) * sizeof(int32_t));
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::Process() {
+  if (this->blockIdx < this->srcToDstTilingData->needCoreNum) {
+    int64_t loops = (coreRows + perLoopRows - 1) / perLoopRows;
+    currentLoopRows = perLoopRows;
+    InitLocal();
+    this->expertTokenIdxOutLocal = expertTokenIdxCopyOutQueue.DeQue<int32_t>();
+    for (int64_t loop = 0; loop < loops - 1; loop++) {
+      CopyIn(loop);
+      Compute(loop);
+    }
+    currentLoopRows = lastLoopRows;
+    CopyIn(loops - 1);
+    Compute(loops - 1);
+    CopyOutTokenGm();
+    expertTokenIdxCopyOutQueue.FreeTensor(this->expertTokenIdxOutLocal);
+  }
+  this->SyncAll();
+}
+
+}  // namespace MoeInitRoutingQuantV2
+#endif  // INNER_MOE_V2_EXPERT_TOKEN_OUT_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_dynamic_quant.h

@@ -0,0 +1,398 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/* !
+ * \file moe_v2_fullload_dynamic_quant.h
+ * \brief
+ */
+#ifndef MOE_V2_FULL_LOAD_DYNAMIC_QUANT_H
+#define MOE_V2_FULL_LOAD_DYNAMIC_QUANT_H
+
+#include "moe_v2_mrgsort.h"
+#include "moe_v2_sort_base.h"
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+template <typename T>
+class MoeV2FullLoadDynamicQuant : public MoeV2SortBase {
+ public:
+  __aicore__ inline MoeV2FullLoadDynamicQuant(){};
+  __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, GM_ADDR expandedRowIdx,
+                              GM_ADDR expertTokensCountOrCumsum, GM_ADDR quantSmooth, GM_ADDR dynamicQuantScale,
+                              GM_ADDR workspace, const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void CopyIn();
+  __aicore__ inline void SortCompute();
+  __aicore__ inline void CopyOutIdx();
+  __aicore__ inline void CopyOutEmpty();
+  __aicore__ inline void CopyOutXQuant1H();
+  __aicore__ inline void ComputeExpertTokenCountOrCumsum();
+  __aicore__ inline void Compute(LocalTensor<float>& smoothLocal);
+
+ private:
+  int64_t sortNum_;
+  const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData_;
+  int64_t blockIdx_;
+  int64_t needCoreNum_;
+  int64_t coreRows_;
+  int64_t perCoreRows_;
+  int64_t k_;
+  int64_t n_;
+  int64_t cols_;
+  int64_t cols_scale_;
+  int64_t activateRows_;
+  int64_t expertNum;
+  int64_t expertCapacity;
+  int64_t smoothType;
+  int64_t colsAlign;
+
+  TQue<QuePosition::VECIN, 1> xCopyInQueue_;
+  TQue<QuePosition::VECOUT, 1> expandedRowIdxCopyOutQueue_;
+  TQue<QuePosition::VECOUT, 1> expandedExpertIdxCopyOutQueue_;
+  TQue<QuePosition::VECOUT, 1> expandDstToSrcRowQueue_;
+  TQue<QuePosition::VECOUT, 1> expertTokensCopyOutQueue_;
+  TQue<QuePosition::VECIN, 1> smoothInQueue;
+  TQue<QuePosition::VECOUT, 1> calcQueue;
+  TQue<QuePosition::VECOUT, 1> inputXOutQueue;
+
+  GlobalTensor<T> xGm_;
+  GlobalTensor<int32_t> expertIdxGm_;
+  GlobalTensor<float> quantSmoothGm;
+
+  GlobalTensor<int8_t> expandedXGm_;
+  GlobalTensor<int32_t> expandedRowIdxGm_;
+  GlobalTensor<int32_t> expandedExpertIdxGm_;
+  GlobalTensor<int32_t> expertTokensCountOrCumsumGm;
+  GlobalTensor<int32_t> expertTokensBeforeCapacityGm;
+
+  int64_t expertTokensCountOrCumsumFlag = 0;
+  int64_t expertTokensBeforeCapacityFlag = 0;
+  int64_t dropPadMode = 0;
+
+  LocalTensor<uint32_t> expandDstToSrcRowLocal;
+  LocalTensor<int32_t> expandedExpertIdxLocal;
+};
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadDynamicQuant<T>::CopyIn() {
+  LocalTensor<int32_t> inLocal = sortDataCopyInQueue.AllocTensor<int32_t>();
+  DataCopyExtParams dataCopyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(this->totalLength * sizeof(int32_t)),
+                                   0, 0, 0};
+  DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+  DataCopyPad(inLocal[0], expertIdxGm_, dataCopyParams, dataCopyPadParams);
+  ArithProgression<int32_t>(inLocal[this->sortNum_], 0, 1, this->totalLength);
+  sortDataCopyInQueue.EnQue(inLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadDynamicQuant<T>::SortCompute() {
+  LocalTensor<int32_t> inLocal = sortDataCopyInQueue.DeQue<int32_t>();
+  LocalTensor<int32_t> expertIdxLocal = inLocal[0];
+  LocalTensor<float> expertIdxLocalFp32 = expertIdxLocal.ReinterpretCast<float>();
+  Cast(expertIdxLocalFp32, expertIdxLocal, RoundMode::CAST_ROUND, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  Muls(expertIdxLocalFp32, expertIdxLocalFp32, (float)-1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM;
+  if (duplicateNum > 0) {
+    int duplicateIndex = this->totalLength - duplicateNum;
+    uint64_t mask0 = UINT64_MAX;
+    mask0 = mask0 << duplicateNum;
+    mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
+    uint64_t mask[2] = {mask0, 0};
+    Duplicate(expertIdxLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+  LocalTensor<float> concatLocal;
+  LocalTensor<float> tempTensor = tempBuffer.Get<float>(GetSortLen<float>(this->sortNum_));
+  Concat(concatLocal, expertIdxLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  LocalTensor<uint32_t> rowIdxLocal = inLocal[this->sortNum_].template ReinterpretCast<uint32_t>();
+  LocalTensor<float> sortedLocal = sortedBuffer.Get<float>(GetSortLen<float>(this->sortNum_));
+  Sort<float, true>(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  LocalTensor<float> expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue_.AllocTensor<float>();
+  expandDstToSrcRowLocal = expandDstToSrcRowQueue_.AllocTensor<uint32_t>();
+  LocalTensor<float> expandDstToSrcRowLocalFp32 = expandDstToSrcRowLocal.ReinterpretCast<float>();
+  Extract(expandedExpertIdxLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  Cast(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocal.ReinterpretCast<int32_t>(), RoundMode::CAST_ROUND,
+       this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  Muls(expandedExpertIdxLocal, expandedExpertIdxLocal, (float)-1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  LocalTensor<int32_t> expandedExpertIdxLocalInt32;
+  expandedExpertIdxLocalInt32 = expandedExpertIdxLocal.ReinterpretCast<int32_t>();
+  Cast(expandedExpertIdxLocalInt32, expandedExpertIdxLocal, RoundMode::CAST_ROUND, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  expandedExpertIdxCopyOutQueue_.EnQue<int32_t>(expandedExpertIdxLocalInt32);
+
+  LocalTensor<uint32_t> expandedRowIdx = expandedRowIdxCopyOutQueue_.AllocTensor<uint32_t>();
+  LocalTensor<uint32_t> expandedRowIdxU32 = expandedRowIdx.ReinterpretCast<uint32_t>();
+  Muls(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocalFp32, (float)-1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  ArithProgression<int32_t>(inLocal[this->sortNum_], 0, 1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  if (duplicateNum > 0) {
+    int duplicateIndex = this->totalLength - duplicateNum;
+    uint64_t mask0 = UINT64_MAX;
+    mask0 = mask0 << duplicateNum;
+    mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
+    uint64_t mask[2] = {mask0, 0};
+    Duplicate(expandDstToSrcRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+  Concat(concatLocal, expandDstToSrcRowLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  Sort<float, true>(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  Extract(tempTensor, expandedRowIdxU32, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  expandedRowIdxCopyOutQueue_.EnQue<uint32_t>(expandedRowIdx);
+  sortDataCopyInQueue.FreeTensor(inLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadDynamicQuant<T>::CopyOutIdx() {
+  LocalTensor<int32_t> expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue<int32_t>();
+  DataCopyParams intriParams;
+  intriParams.blockCount = 1;
+  intriParams.blockLen = this->totalLength * sizeof(int32_t);
+  DataCopyPad(expandedRowIdxGm_, expandedRowIdx, intriParams);
+  expandedRowIdxCopyOutQueue_.EnQue(expandedRowIdx);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadDynamicQuant<T>::ComputeExpertTokenCountOrCumsum() {
+  expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue_.DeQue<int32_t>();
+  LocalTensor<int32_t> expertTokensCount = expertTokensCopyOutQueue_.AllocTensor<int32_t>();
+
+  int64_t expertNumAlign = Align(this->expertNum, sizeof(int32_t));
+  Duplicate(expertTokensCount, 0, expertNumAlign);
+  SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
+
+  int32_t lastExpertId = expandedExpertIdxLocal.GetValue(0);
+  int64_t tokenCount = 0;
+  int64_t lastExpertCount = 0;
+  for (int64_t i = 0; i < this->totalLength; i++) {
+    int32_t curExpertId = expandedExpertIdxLocal.GetValue(i);
+    tokenCount++;
+    while (lastExpertId < curExpertId) {
+      expertTokensCount.SetValue(lastExpertId, tokenCount - 1);
+      if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_COUNT) {
+        tokenCount = 1;
+      }
+      lastExpertId++;
+    }
+  }
+#ifndef __CCE_KT_TEST__
+  expertTokensCount.SetValue(lastExpertId, tokenCount);
+  if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_CUMSUM) {
+    lastExpertId++;
+    while (lastExpertId < this->expertNum) {
+      expertTokensCount.SetValue(lastExpertId, tokenCount);
+      lastExpertId++;
+    }
+  }
+  DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(this->expertNum * sizeof(int32_t)), 0, 0,
+                               0};
+  if (this->expertTokensCountOrCumsumFlag > 0) {
+    DataCopyPad(expertTokensCountOrCumsumGm, expertTokensCount, copyParams);
+  }
+  expertTokensCopyOutQueue_.FreeTensor(expertTokensCount);
+#endif
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadDynamicQuant<T>::CopyOutEmpty() {
+  expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue_.DeQue<int32_t>();
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadDynamicQuant<T>::Compute(LocalTensor<float>& smoothLocal) {
+  LocalTensor<float> inLocal = xCopyInQueue_.DeQue<float>();
+
+  LocalTensor<float> tempLocal = calcQueue.AllocTensor<float>();
+  LocalTensor<int8_t> outLocal = inputXOutQueue.AllocTensor<int8_t>();
+  LocalTensor<float> dynamicQuantLocal = outLocal[this->cols_].template ReinterpretCast<float>();
+
+  if constexpr (!IsSameType<T, float>::value) {
+    Cast(inLocal, inLocal.ReinterpretCast<T>()[colsAlign], RoundMode::CAST_NONE, this->cols_);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  if (smoothType != 0) {
+    Mul(inLocal, inLocal, smoothLocal, this->cols_);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  Abs(tempLocal, inLocal, this->cols_);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  ReduceMax(dynamicQuantLocal, tempLocal, tempLocal, this->cols_);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  float maxValue = dynamicQuantLocal.GetValue(0) / 127.0f;
+
+  Duplicate<float>(dynamicQuantLocal, maxValue, 8);
+  Duplicate<float>(tempLocal, maxValue, this->cols_);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Div(tempLocal, inLocal, tempLocal, this->cols_);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Cast(tempLocal.ReinterpretCast<half>(), tempLocal, RoundMode::CAST_TRUNC, this->cols_);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Cast(outLocal, tempLocal.ReinterpretCast<half>(), RoundMode::CAST_ROUND, this->cols_);
+
+  calcQueue.FreeTensor(tempLocal);
+  inputXOutQueue.EnQue(outLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadDynamicQuant<T>::CopyOutXQuant1H() {
+  expandDstToSrcRowQueue_.FreeTensor(expandDstToSrcRowLocal);
+  expandedExpertIdxCopyOutQueue_.FreeTensor(expandedExpertIdxLocal);
+
+  LocalTensor<int32_t> expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue<int32_t>();
+  int64_t curRowsStart = this->blockIdx_ * this->perCoreRows_;
+  int64_t curRowsEnd = curRowsStart + this->coreRows_ - 1;
+  int64_t startXRow = curRowsStart / this->k_;
+  int64_t endXRow = curRowsEnd / this->k_;
+
+  DataCopyExtParams dataXCopyParams{1, static_cast<uint32_t>(this->cols_ * sizeof(T)), 0, 0, 0};
+  DataCopyExtParams smoothCopyParams{1, static_cast<uint32_t>(this->cols_ * sizeof(float)), 0, 0, 0};
+  DataCopyExtParams intriParams{1, static_cast<uint32_t>((this->cols_ + BLOCK_BYTES) * sizeof(int8_t)), 0, 0, 0};
+
+  LocalTensor<float> smoothLocal;
+  if (smoothType == 1) {
+    smoothLocal = smoothInQueue.AllocTensor<float>();
+    DataCopyPad(smoothLocal, quantSmoothGm, smoothCopyParams, {false, 0, 0, 0});
+    smoothInQueue.EnQue(smoothLocal);
+    smoothLocal = smoothInQueue.DeQue<float>();
+  }
+  for (int64_t row = startXRow; row <= endXRow; row++) {
+    LocalTensor<T> xLocal = xCopyInQueue_.AllocTensor<T>();
+    if constexpr (IsSameType<T, float>::value) {
+      DataCopyPad(xLocal, xGm_[row * this->cols_], dataXCopyParams, {false, 0, 0, 0});
+    } else {
+      DataCopyPad(xLocal[colsAlign], xGm_[row * this->cols_], dataXCopyParams, {false, 0, 0, 0});
+    }
+
+    xCopyInQueue_.EnQue<T>(xLocal);
+    Compute(smoothLocal);
+
+    LocalTensor<int8_t> outLocal = inputXOutQueue.DeQue<int8_t>();
+    while (curRowsStart <= curRowsEnd && curRowsStart / this->k_ == row) {
+      int32_t outIndex = expandedRowIdx.GetValue(curRowsStart);
+      curRowsStart++;
+      if (outIndex == -1 || (this->dropPadMode == DROPLESS_MODE && outIndex >= this->activateRows_)) {
+        continue;
+      }
+      DataCopyPad(expandedXGm_[outIndex * this->cols_scale_], outLocal, intriParams);
+    }
+
+    xCopyInQueue_.FreeTensor(xLocal);
+    inputXOutQueue.FreeTensor(outLocal);
+  }
+  expandedRowIdxCopyOutQueue_.FreeTensor(expandedRowIdx);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadDynamicQuant<T>::Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX,
+                                                          GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum,
+                                                          GM_ADDR quantSmooth, GM_ADDR dynamicQuantScale,
+                                                          GM_ADDR workspace,
+                                                          const MoeInitRoutingQuantV2TilingData* tilingData,
+                                                          TPipe* tPipe) {
+  this->gatherOutTilingData_ = &(tilingData->gatherOutComputeParamsOp);
+  //this->blockIdx_ = GetBlockIdx();
+  this->blockIdx_ = get_block_idx() + get_subblockid() * get_block_num();
+  this->k_ = tilingData->k;
+  this->n_ = tilingData->n;
+  this->cols_ = tilingData->cols;
+  this->cols_scale_ = this->cols_ + UB_ALIGN;
+  this->needCoreNum_ = this->gatherOutTilingData_->needCoreNum;
+  this->perCoreRows_ = this->gatherOutTilingData_->perCoreRows;
+  this->activateRows_ = this->gatherOutTilingData_->activateRows;
+  if (this->blockIdx_ == this->gatherOutTilingData_->needCoreNum - 1) {
+    this->coreRows_ = this->gatherOutTilingData_->lastCoreRows;
+  } else {
+    this->coreRows_ = this->gatherOutTilingData_->perCoreRows;
+  }
+  this->expertNum = tilingData->expertNum;
+  this->dropPadMode = tilingData->dropPadMode;
+  this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag;
+
+  this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t));
+  this->sortNum_ = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
+  this->totalLength = tilingData->n * tilingData->k;
+  this->smoothType = tilingData->smoothType;
+  this->colsAlign = Align(this->cols_, sizeof(T));
+  this->pipe = tPipe;
+
+  xGm_.SetGlobalBuffer((__gm__ T*)x);
+  expertIdxGm_.SetGlobalBuffer((__gm__ int32_t*)expertIdx, this->tileLength);
+
+  expandedXGm_.SetGlobalBuffer((__gm__ int8_t*)expandedX);
+  expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, this->tileLength);
+  if (this->expertTokensCountOrCumsumFlag > 0) {
+    // dropless
+    expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensCountOrCumsum,
+                                                Align(this->expertNum, sizeof(int32_t)));
+  }
+  quantSmoothGm.SetGlobalBuffer((__gm__ float*)quantSmooth);
+
+  int64_t kvFactor = 2;
+  int64_t buffSize = this->sortNum_ * sizeof(int32_t);
+
+  int64_t curRowsStart = this->blockIdx_ * this->perCoreRows_;
+  int64_t startXRow = curRowsStart / this->k_;
+  int64_t endXRow = (curRowsStart + this->coreRows_ - 1) / this->k_;
+
+  pipe->InitBuffer(expandedRowIdxCopyOutQueue_, bufferNum, buffSize);
+  pipe->InitBuffer(expandedExpertIdxCopyOutQueue_, bufferNum, buffSize);
+  pipe->InitBuffer(expertTokensCopyOutQueue_, bufferNum, AlignBytes(this->expertNum, sizeof(int32_t)));
+  pipe->InitBuffer(expandDstToSrcRowQueue_, bufferNum, buffSize);
+  pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize * kvFactor);
+  pipe->InitBuffer(tempBuffer, buffSize * kvFactor);
+  pipe->InitBuffer(sortedBuffer, buffSize * kvFactor);
+
+  if constexpr (IsSameType<T, float>::value) {
+    pipe->InitBuffer(xCopyInQueue_, 1, AlignBytes(this->cols_, sizeof(float)));
+  } else {
+    pipe->InitBuffer(xCopyInQueue_, 1, 2 * AlignBytes(this->cols_, sizeof(T)));
+  }
+  pipe->InitBuffer(smoothInQueue, 1, AlignBytes(this->cols_, sizeof(float)));
+  pipe->InitBuffer(calcQueue, 1, AlignBytes(this->cols_, sizeof(float)));
+  pipe->InitBuffer(inputXOutQueue, 1, AlignBytes(this->cols_scale_, sizeof(int8_t)));
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadDynamicQuant<T>::Process() {
+  if (this->blockIdx_ < this->needCoreNum_) {
+    CopyIn();
+    SortCompute();
+    if (this->blockIdx_ == 0) {
+      CopyOutIdx();
+    }
+    if (this->blockIdx_ == this->needCoreNum_ - 1 && this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) {
+      ComputeExpertTokenCountOrCumsum();
+    } else {
+      CopyOutEmpty();
+    }
+    CopyOutXQuant1H();
+  }
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // MOE_V2_DYNAMIC_QUANT_FULL_LOAD_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_quant.h

@@ -0,0 +1,155 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/* !
+ * \file moe_v2_fullload_quant.h
+ * \brief
+ */
+#ifndef MOE_V2_FULL_LOAD_QUANT_H
+#define MOE_V2_FULL_LOAD_QUANT_H
+
+#include "moe_v2_fullload_quant_base.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+template <typename T>
+class MoeV2FullLoadQuant : public MoeV2FullLoadQuantBase {
+ public:
+  __aicore__ inline MoeV2FullLoadQuant(){};
+  __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedX,
+                              GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace,
+                              const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void Compute(int64_t xLocalLength);
+  __aicore__ inline void CopyOutX();
+
+ private:
+  TQue<QuePosition::VECOUT, 1> floatQueue;
+  TQue<QuePosition::VECOUT, 1> halfQueue;
+  TQue<QuePosition::VECOUT, 1> inputXCopyOutQueue;
+
+  GlobalTensor<T> xGm;
+  GlobalTensor<float> scaleGm;
+  GlobalTensor<float> offsetGm;
+
+  float scale;
+  float offset;
+};
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadQuant<T>::Compute(int64_t xLocalLength) {
+  LocalTensor<T> inLocal = xCopyInQueue.DeQue<T>();
+  LocalTensor<int8_t> outLocal = inputXCopyOutQueue.AllocTensor<int8_t>();
+  LocalTensor<float> floatLocal = floatQueue.AllocTensor<float>();
+  LocalTensor<half> halfLocal = halfQueue.AllocTensor<half>();
+
+  uint32_t elements = Align(this->cols, sizeof(int8_t)) * xLocalLength;
+  if constexpr (IsSameType<T, bfloat16_t>::value) {
+    Cast(floatLocal, inLocal, RoundMode::CAST_NONE, elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Cast(halfLocal, floatLocal, RoundMode::CAST_NONE, elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Muls(halfLocal, halfLocal, static_cast<half>(this->scale), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Adds(halfLocal, halfLocal, static_cast<half>(this->offset), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    LocalTensor<int32_t> intLocal = floatLocal.ReinterpretCast<int32_t>();
+    Cast(intLocal, halfLocal, RoundMode::CAST_RINT, elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    SetDeqScale((half)1.000000e+00f);
+    AscendC::PipeBarrier<PIPE_V>();
+    Cast(halfLocal, intLocal, RoundMode::CAST_RINT, elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Cast(outLocal, halfLocal, RoundMode::CAST_RINT, elements);
+  } else if constexpr (IsSameType<T, float>::value) {
+    Cast(halfLocal, inLocal, RoundMode::CAST_NONE, elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Muls(halfLocal, halfLocal, static_cast<half>(this->scale), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Adds(halfLocal, halfLocal, static_cast<half>(this->offset), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Cast(outLocal, halfLocal, RoundMode::CAST_RINT, elements);
+  } else {
+    Muls(inLocal, inLocal, static_cast<T>(this->scale), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Adds(inLocal, inLocal, static_cast<T>(this->offset), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Cast(outLocal, inLocal, RoundMode::CAST_RINT, elements);
+  }
+  inputXCopyOutQueue.EnQue(outLocal);
+  xCopyInQueue.FreeTensor(inLocal);
+  floatQueue.FreeTensor(floatLocal);
+  halfQueue.FreeTensor(halfLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadQuant<T>::CopyOutX() {
+  LocalTensor<T> xLocal = xCopyInQueue.AllocTensor<T>();
+  LocalTensor<int32_t> expandedRowIdx = expandedRowIdxCopyOutQueue.DeQue<int32_t>();
+  int64_t inFactor = Align(this->cols, sizeof(int8_t));
+  int64_t curRowsStart = this->blockIdx * this->perCoreRows;
+  int64_t startXRow = curRowsStart / this->k;
+  int64_t endXRow = (curRowsStart + this->coreRows - 1) / this->k;
+
+  uint32_t dstStride = (inFactor * sizeof(T) - AlignBytes(this->cols, sizeof(T))) / BLOCK_BYTES;
+  DataCopyExtParams dataXCopyParams{static_cast<uint16_t>(endXRow - startXRow + 1),
+                                    static_cast<uint32_t>(this->cols * sizeof(T)), 0, dstStride, 0};
+  DataCopyPadExtParams<T> dataXCopyPadParams{false, 0, 0, 0};
+  DataCopyPad(xLocal, xGm[startXRow * this->cols], dataXCopyParams, dataXCopyPadParams);
+  xCopyInQueue.EnQue(xLocal);
+  Compute(endXRow - startXRow + 1);
+  LocalTensor<int8_t> outLocal = inputXCopyOutQueue.DeQue<int8_t>();
+  int64_t k = 0;
+  DataCopyExtParams intriParams{1, static_cast<uint32_t>(this->cols * sizeof(int8_t)), 0, 0, 0};
+  for (int64_t i = startXRow; i <= endXRow; i++) {
+    for (; k < this->perCoreRows && curRowsStart / this->k == i; curRowsStart++, k++) {
+      int32_t outIndex = expandedRowIdx.GetValue(curRowsStart);
+      if (outIndex < this->activateRows) {
+        DataCopyPad(expandedXGm[outIndex * this->cols], outLocal[(i - startXRow) * inFactor], intriParams);
+      }
+    }
+  }
+  expandedRowIdxCopyOutQueue.FreeTensor(expandedRowIdx);
+  inputXCopyOutQueue.FreeTensor(outLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadQuant<T>::Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR scale, GM_ADDR offset,
+                                                   GM_ADDR expandedX, GM_ADDR expandedRowIdx,
+                                                   GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace,
+                                                   const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe) {
+  this->InitBase(x, expertIdx, expandedX, expandedRowIdx, expertTokensCountOrCumsum, workspace, tilingData, tPipe);
+  xGm.SetGlobalBuffer((__gm__ T*)x);
+  scaleGm.SetGlobalBuffer((__gm__ float*)scale, 1);
+  offsetGm.SetGlobalBuffer((__gm__ float*)offset, 1);
+  this->scale = scaleGm.GetValue(0);
+  this->offset = offsetGm.GetValue(0);
+
+  int64_t curRowsStart = this->blockIdx * this->perCoreRows;
+  int64_t rowLength = (curRowsStart + this->coreRows - 1) / this->k - curRowsStart / this->k + 1;
+  int64_t xAlignedCount = Align(this->cols, sizeof(int8_t));
+  pipe->InitBuffer(xCopyInQueue, bufferNum, xAlignedCount * sizeof(T) * rowLength);
+  pipe->InitBuffer(inputXCopyOutQueue, 1, xAlignedCount * sizeof(int8_t) * rowLength);
+  pipe->InitBuffer(floatQueue, 1, xAlignedCount * sizeof(float) * rowLength);
+  pipe->InitBuffer(halfQueue, 1, xAlignedCount * sizeof(half) * rowLength);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoadQuant<T>::Process() {
+  if (this->blockIdx < this->needCoreNum) {
+    this->ProcessBase();
+    CopyOutX();
+  }
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_fullload_quant_base.h

@@ -0,0 +1,279 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/* !
+ * \file moe_v2_fullload_quant_base.h
+ * \brief
+ */
+#ifndef MOE_V2_FULL_LOAD_QUANT_BASE_H
+#define MOE_V2_FULL_LOAD_QUANT_BASE_H
+
+#include "kernel_operator.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+class MoeV2FullLoadQuantBase {
+ public:
+  __aicore__ inline MoeV2FullLoadQuantBase(){};
+
+ protected:
+  __aicore__ inline void InitBase(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, GM_ADDR expandedRowIdx,
+                                  GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace,
+                                  const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void ProcessBase();
+  __aicore__ inline void CopyIn();
+  __aicore__ inline void SortCompute();
+  __aicore__ inline void CopyOutIdx();
+  __aicore__ inline void CopyOutEmpty();
+  __aicore__ inline void ComputeExpertTokenCountOrCumsum();
+
+ protected:
+  const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData;
+
+  TPipe* pipe;
+  int64_t tileLength;
+  int64_t bufferNum = 1;
+  int64_t totalLength;
+  int64_t coreNum;
+  int64_t sortNum;
+  int64_t blockIdx;
+  int64_t needCoreNum;
+  int64_t coreRows;
+  int64_t perCoreRows;
+  int64_t k;
+  int64_t n;
+  int64_t cols;
+  int64_t activateRows;
+  int64_t expertNum;
+  int64_t expertCapacity;
+
+  TQue<QuePosition::VECIN, 1> sortDataCopyInQueue;
+  TBuf<TPosition::VECCALC> tempBuffer;
+  TBuf<TPosition::VECCALC> sortedBuffer;
+  TQue<QuePosition::VECIN, 1> xCopyInQueue;
+  TQue<QuePosition::VECOUT, 1> expandedRowIdxCopyOutQueue;
+  TQue<QuePosition::VECOUT, 1> expandedExpertIdxCopyOutQueue;
+  TQue<QuePosition::VECOUT, 1> expandDstToSrcRowQueue;
+  TQue<QuePosition::VECOUT, 1> expertTokensCopyOutQueue;
+
+  GlobalTensor<int32_t> expertIdxGm;
+  GlobalTensor<int8_t> expandedXGm;
+  GlobalTensor<int32_t> expandedRowIdxGm;
+  GlobalTensor<int32_t> expandedExpertIdxGm;
+  GlobalTensor<int32_t> expertTokensCountOrCumsumGm;
+  GlobalTensor<int32_t> expertTokensBeforeCapacityGm;
+
+  int64_t expertTokensCountOrCumsumFlag = 0;
+  int64_t expertTokensBeforeCapacityFlag = 0;
+  int64_t dropPadMode = 0;
+  static constexpr int64_t DST_BLK_STRIDE = 1;
+  static constexpr int64_t DST_REP_STRIDE = 8;
+  static constexpr int64_t FOUR_BLOCK_BYTES = 128;
+};
+
+__aicore__ inline void MoeV2FullLoadQuantBase::CopyIn() {
+  LocalTensor<int32_t> inLocal = sortDataCopyInQueue.AllocTensor<int32_t>();
+  DataCopyExtParams dataCopyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(this->totalLength * sizeof(int32_t)),
+                                   0, 0, 0};
+  DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+  DataCopyPad(inLocal[0], expertIdxGm, dataCopyParams, dataCopyPadParams);
+  ArithProgression<int32_t>(inLocal[this->sortNum], 0, 1, this->totalLength);
+  sortDataCopyInQueue.EnQue(inLocal);
+}
+
+__aicore__ inline void MoeV2FullLoadQuantBase::SortCompute() {
+  LocalTensor<int32_t> inLocal = sortDataCopyInQueue.DeQue<int32_t>();
+  LocalTensor<int32_t> expertIdxLocal = inLocal[0];
+  LocalTensor<float> expertIdxLocalFp32 = expertIdxLocal.ReinterpretCast<float>();
+  Cast(expertIdxLocalFp32, expertIdxLocal, RoundMode::CAST_ROUND, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  Muls(expertIdxLocalFp32, expertIdxLocalFp32, (float)-1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM;
+  if (duplicateNum > 0) {
+    int duplicateIndex = this->totalLength - duplicateNum;
+    uint64_t mask0 = UINT64_MAX;
+    mask0 = mask0 << duplicateNum;
+    mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
+    uint64_t mask[2] = {mask0, 0};
+    Duplicate(expertIdxLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+  LocalTensor<float> concatLocal;
+  LocalTensor<float> tempTensor = tempBuffer.Get<float>(GetSortLen<float>(this->sortNum));
+  Concat(concatLocal, expertIdxLocalFp32, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  LocalTensor<uint32_t> rowIdxLocal = inLocal[this->sortNum].template ReinterpretCast<uint32_t>();
+  LocalTensor<float> sortedLocal = sortedBuffer.Get<float>(GetSortLen<float>(this->sortNum));
+  Sort<float, true>(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  LocalTensor<float> expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue.AllocTensor<float>();
+  LocalTensor<uint32_t> expandDstToSrcRowLocal = expandDstToSrcRowQueue.AllocTensor<uint32_t>();
+  LocalTensor<float> expandDstToSrcRowLocalFp32 = expandDstToSrcRowLocal.ReinterpretCast<float>();
+  Extract(expandedExpertIdxLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  Cast(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocal.ReinterpretCast<int32_t>(), RoundMode::CAST_ROUND,
+       this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  Muls(expandedExpertIdxLocal, expandedExpertIdxLocal, (float)-1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  LocalTensor<int32_t> expandedExpertIdxLocalInt32;
+  expandedExpertIdxLocalInt32 = expandedExpertIdxLocal.ReinterpretCast<int32_t>();
+  Cast(expandedExpertIdxLocalInt32, expandedExpertIdxLocal, RoundMode::CAST_ROUND, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  expandedExpertIdxCopyOutQueue.EnQue<int32_t>(expandedExpertIdxLocalInt32);
+
+  LocalTensor<uint32_t> expandedRowIdx = expandedRowIdxCopyOutQueue.AllocTensor<uint32_t>();
+  LocalTensor<uint32_t> expandedRowIdxU32 = expandedRowIdx.ReinterpretCast<uint32_t>();
+  Muls(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocalFp32, (float)-1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  ArithProgression<int32_t>(inLocal[this->sortNum], 0, 1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  if (duplicateNum > 0) {
+    int duplicateIndex = this->totalLength - duplicateNum;
+    uint64_t mask0 = UINT64_MAX;
+    mask0 = mask0 << duplicateNum;
+    mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
+    uint64_t mask[2] = {mask0, 0};
+    Duplicate(expandDstToSrcRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+  Concat(concatLocal, expandDstToSrcRowLocalFp32, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  Sort<float, true>(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  Extract(tempTensor, expandedRowIdxU32, sortedLocal, this->sortNum / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  expandedRowIdxCopyOutQueue.EnQue<uint32_t>(expandedRowIdx);
+  sortDataCopyInQueue.FreeTensor(inLocal);
+
+  expandDstToSrcRowQueue.FreeTensor(expandDstToSrcRowLocal);
+}
+
+__aicore__ inline void MoeV2FullLoadQuantBase::CopyOutIdx() {
+  LocalTensor<int32_t> expandedRowIdx = expandedRowIdxCopyOutQueue.DeQue<int32_t>();
+  DataCopyParams intriParams;
+  intriParams.blockCount = 1;
+  intriParams.blockLen = this->totalLength * sizeof(int32_t);
+  DataCopyPad(expandedRowIdxGm, expandedRowIdx, intriParams);
+  expandedRowIdxCopyOutQueue.EnQue(expandedRowIdx);
+}
+
+__aicore__ inline void MoeV2FullLoadQuantBase::ComputeExpertTokenCountOrCumsum() {
+  LocalTensor<int32_t> expandedExpertIdx = expandedExpertIdxCopyOutQueue.DeQue<int32_t>();
+  LocalTensor<int32_t> expertTokensCount = expertTokensCopyOutQueue.AllocTensor<int32_t>();
+
+  int64_t expertNumAlign = Align(this->expertNum, sizeof(int32_t));
+  Duplicate(expertTokensCount, 0, expertNumAlign);
+  SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
+
+  int32_t lastExpertId = expandedExpertIdx.GetValue(0);
+  int64_t tokenCount = 0;
+  int64_t lastExpertCount = 0;
+  for (int64_t i = 0; i < this->totalLength; i++) {
+    int32_t curExpertId = expandedExpertIdx.GetValue(i);
+    tokenCount++;
+    while (lastExpertId < curExpertId) {
+      expertTokensCount.SetValue(lastExpertId, tokenCount - 1);
+      if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_COUNT) {
+        tokenCount = 1;
+      }
+      lastExpertId++;
+    }
+  }
+  expertTokensCount.SetValue(lastExpertId, tokenCount);
+  if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_CUMSUM) {
+    lastExpertId++;
+    while (lastExpertId < this->expertNum) {
+      expertTokensCount.SetValue(lastExpertId, tokenCount);
+      lastExpertId++;
+    }
+  }
+  DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(this->expertNum * sizeof(int32_t)), 0, 0,
+                               0};
+  if (this->expertTokensCountOrCumsumFlag > 0) {
+    DataCopyPad(expertTokensCountOrCumsumGm, expertTokensCount, copyParams);
+  }
+  expertTokensCopyOutQueue.FreeTensor(expertTokensCount);
+  expandedExpertIdxCopyOutQueue.FreeTensor(expandedExpertIdx);
+}
+
+__aicore__ inline void MoeV2FullLoadQuantBase::CopyOutEmpty() {
+  LocalTensor<int32_t> outLocal = expandedExpertIdxCopyOutQueue.DeQue<int32_t>();
+  expandedExpertIdxCopyOutQueue.FreeTensor(outLocal);
+}
+
+__aicore__ inline void MoeV2FullLoadQuantBase::InitBase(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX,
+                                                        GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum,
+                                                        GM_ADDR workspace,
+                                                        const MoeInitRoutingQuantV2TilingData* tilingData,
+                                                        TPipe* tPipe) {
+  this->gatherOutTilingData = &(tilingData->gatherOutComputeParamsOp);
+  this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+  this->k = tilingData->k;
+  this->n = tilingData->n;
+  this->cols = tilingData->cols;
+  this->needCoreNum = this->gatherOutTilingData->needCoreNum;
+  this->perCoreRows = this->gatherOutTilingData->perCoreRows;
+  this->activateRows = this->gatherOutTilingData->activateRows;
+  if (this->blockIdx == this->gatherOutTilingData->needCoreNum - 1) {
+    this->coreRows = this->gatherOutTilingData->lastCoreRows;
+  } else {
+    this->coreRows = this->gatherOutTilingData->perCoreRows;
+  }
+  this->expertNum = tilingData->expertNum;
+  this->dropPadMode = tilingData->dropPadMode;
+  this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag;
+
+  this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t));
+  this->sortNum = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
+  this->totalLength = tilingData->n * tilingData->k;
+  this->pipe = tPipe;
+
+  expertIdxGm.SetGlobalBuffer((__gm__ int32_t*)expertIdx, this->tileLength);
+
+  expandedXGm.SetGlobalBuffer((__gm__ int8_t*)expandedX);
+  expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, this->tileLength);
+  if (this->expertTokensCountOrCumsumFlag > 0) {
+    // dropless
+    expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensCountOrCumsum,
+                                                Align(this->expertNum, sizeof(int32_t)));
+  }
+
+  int64_t kvFactor = 2;
+  int64_t buffSize = this->sortNum * sizeof(int32_t);
+
+  pipe->InitBuffer(expandedRowIdxCopyOutQueue, bufferNum, buffSize);
+  pipe->InitBuffer(expandedExpertIdxCopyOutQueue, bufferNum, buffSize);
+  pipe->InitBuffer(expertTokensCopyOutQueue, bufferNum, AlignBytes(this->expertNum, sizeof(int32_t)));
+  pipe->InitBuffer(expandDstToSrcRowQueue, bufferNum, buffSize);
+  pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize * kvFactor);
+  pipe->InitBuffer(tempBuffer, buffSize * kvFactor);
+  pipe->InitBuffer(sortedBuffer, buffSize * kvFactor);
+}
+
+__aicore__ inline void MoeV2FullLoadQuantBase::ProcessBase() {
+  if (this->blockIdx < this->needCoreNum) {
+    CopyIn();
+    SortCompute();
+    if (this->blockIdx == 0) {
+      CopyOutIdx();
+    }
+    if (this->blockIdx == this->needCoreNum - 1 && this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) {
+      ComputeExpertTokenCountOrCumsum();
+    } else {
+      CopyOutEmpty();
+    }
+  }
+}
+
+}  // namespace MoeInitRoutingQuantV2
+#endif  // MOE_V2_FULL_LOAD_QUANT_BASE_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_dynamic_quant.h

@@ -0,0 +1,547 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_gather_dynamic_quant.h
+ * \brief
+ */
+#ifndef MOE_V2_GATHER_DYNAMIC_QUANT_H
+#define MOE_V2_GATHER_DYNAMIC_QUANT_H
+
+#include "moe_v2_common.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+template <typename T>
+class MoeV2GatherDynamicQuant {
+ public:
+  __aicore__ inline MoeV2GatherDynamicQuant(){};
+  __aicore__ inline void Init(GM_ADDR inputX, GM_ADDR quantSmooth, GM_ADDR expandedRowIdx, GM_ADDR expandedX,
+                              GM_ADDR dynamicQuantScale, GM_ADDR workspace,
+                              const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void CopyInExpandedRowIdx(int64_t progress);
+  __aicore__ inline void CopyInExpandedExpertIdx(int64_t progress);
+  __aicore__ inline void CopyOutXQuant1H(int64_t progress);
+  __aicore__ inline void CopyOutXQuantEH(int64_t progress);
+  __aicore__ inline void Compute(LocalTensor<float>& smoothLocal);
+  __aicore__ inline void CopyOutPartialXQuantEH(int64_t progress);
+  __aicore__ inline void CopyOutPartialXQuant1H(int64_t progress);
+  __aicore__ inline float ComputeMax(LocalTensor<float>& inLocal, LocalTensor<float>& tempLocal,
+                                     LocalTensor<float>& dynamicQuantLocal, int32_t srcIdx, int32_t expertIdx,
+                                     int64_t j);
+  __aicore__ inline void ComputeScale(LocalTensor<float>& inLocal, LocalTensor<float>& tempLocal, float scaleTemp,
+                                      int64_t dstIndex, int64_t j);
+
+ private:
+  TPipe* pipe;
+  TQue<QuePosition::VECIN, BUFFER_NUM> inputXInQueue;
+  TQue<QuePosition::VECIN, BUFFER_NUM> smoothInQueue;
+  TQue<QuePosition::VECIN, BUFFER_NUM> expandRowIdxInQueue;
+  TQue<QuePosition::VECOUT, 1> calcQueue;
+  TQue<QuePosition::VECOUT, 1> inputXOutQueue;
+  TQue<QuePosition::VECOUT, 1> scaleOutQueue;
+
+  GlobalTensor<T> inputXGm;
+  GlobalTensor<int8_t> expandedXGm;
+  GlobalTensor<int32_t> expandedRowIdxGm;
+  GlobalTensor<float> quantSmoothGm;
+  GlobalTensor<float> dynamicQuantScaleGm;
+  GlobalTensor<float> quantSrcGm;
+  GlobalTensor<int32_t> expandedExpertIdxGm;
+  GlobalTensor<int32_t> sortedRowIdxGm;
+
+  const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData;
+
+  int64_t needCoreNum;
+  int64_t blockIdx;
+  int64_t cols;
+  int64_t cols_scale_;
+  int64_t n;
+  int64_t k;
+  int64_t totalLength;
+  int64_t activateRows;
+  int64_t currentLoopRows;
+  int64_t currentLoopRowsAlign;
+  int64_t coreRows;
+  int64_t perLoopRows;
+  int64_t lastLoopRows;
+  int64_t rowLoops;
+  int64_t colsTileLength;
+  int64_t perLoopCols;
+  int64_t perLoopColsAlign;
+  int64_t lastLoopCols;
+  int64_t colLoops;
+  int64_t dropPadMode;
+  int64_t smoothType;
+
+  int64_t indicesOffset;
+  int64_t inputOffset;
+  int64_t outOffset;
+};
+
+template <typename T>
+__aicore__ inline void MoeV2GatherDynamicQuant<T>::CopyInExpandedRowIdx(int64_t progress) {
+  this->indicesOffset = progress * this->perLoopRows;
+  LocalTensor<int32_t> indicesLocal = expandRowIdxInQueue.AllocTensor<int32_t>();
+  DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(this->currentLoopRows * sizeof(int32_t)), 0, 0, 0};
+  DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+  DataCopyPad(indicesLocal, expandedRowIdxGm[indicesOffset], dataCopyParams, dataCopyPadParams);
+  expandRowIdxInQueue.EnQue<int32_t>(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherDynamicQuant<T>::CopyInExpandedExpertIdx(int64_t progress) {
+  this->indicesOffset = progress * this->perLoopRows;
+  LocalTensor<int32_t> indicesLocal = expandRowIdxInQueue.AllocTensor<int32_t>();
+  DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(this->currentLoopRows * sizeof(int32_t)), 0, 0, 0};
+  DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+  DataCopyPad(indicesLocal, sortedRowIdxGm[indicesOffset], dataCopyParams, dataCopyPadParams);
+  DataCopyPad(indicesLocal[currentLoopRowsAlign], expandedExpertIdxGm[indicesOffset], dataCopyParams,
+              dataCopyPadParams);
+  expandRowIdxInQueue.EnQue<int32_t>(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherDynamicQuant<T>::Compute(LocalTensor<float>& smoothLocal) {
+  LocalTensor<float> inLocal = inputXInQueue.DeQue<float>();
+
+  LocalTensor<float> tempLocal = calcQueue.AllocTensor<float>();
+  LocalTensor<int8_t> outLocal = inputXOutQueue.AllocTensor<int8_t>();
+  LocalTensor<float> dynamicQuantLocal = outLocal[this->cols].template ReinterpretCast<float>();
+
+  if constexpr (!IsSameType<T, float>::value) {
+    Cast(inLocal, inLocal.ReinterpretCast<T>()[perLoopColsAlign], RoundMode::CAST_NONE, this->cols);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  if (smoothType != 0) {
+    Mul(inLocal, inLocal, smoothLocal, this->cols);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  Abs(tempLocal, inLocal, this->cols);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  ReduceMax(dynamicQuantLocal, tempLocal, tempLocal, this->cols);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  float maxValue = dynamicQuantLocal.GetValue(0) / 127.0f;
+
+  Duplicate<float>(dynamicQuantLocal, maxValue, 8);
+  Duplicate<float>(tempLocal, maxValue, this->cols);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Div(tempLocal, inLocal, tempLocal, this->cols);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Cast(tempLocal.ReinterpretCast<half>(), tempLocal, RoundMode::CAST_TRUNC, this->cols);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Cast(outLocal, tempLocal.ReinterpretCast<half>(), RoundMode::CAST_ROUND, this->cols);
+
+  calcQueue.FreeTensor(tempLocal);
+  inputXOutQueue.EnQue(outLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherDynamicQuant<T>::CopyOutXQuant1H(int64_t progress) {
+  LocalTensor<int32_t> indicesLocal = expandRowIdxInQueue.DeQue<int32_t>();
+
+  int64_t initialRow = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress;
+  int64_t curLoopRow = 0;
+  int64_t currentLoopStartRow = initialRow / this->k;
+  int64_t currentLoopLastRow = (initialRow + this->currentLoopRows - 1) / this->k;
+  DataCopyExtParams copyInParams{1, static_cast<uint32_t>(this->cols * sizeof(T)), 0, 0, 0};
+  DataCopyExtParams copyOutParams{1, static_cast<uint32_t>((this->cols + BLOCK_BYTES) * sizeof(int8_t)), 0, 0, 0};
+  DataCopyExtParams smoothParams{1, static_cast<uint32_t>(this->cols * sizeof(float)), 0, 0, 0};
+
+  LocalTensor<float> smoothLocal;
+  if (smoothType == 1) {
+    smoothLocal = smoothInQueue.AllocTensor<float>();
+    DataCopyPad(smoothLocal, quantSmoothGm, smoothParams, {false, 0, 0, 0});
+    smoothInQueue.EnQue(smoothLocal);
+    smoothLocal = smoothInQueue.DeQue<float>();
+  }
+
+  for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) {
+    LocalTensor<T> inLocal = inputXInQueue.AllocTensor<T>();
+    if constexpr (IsSameType<T, float>::value) {
+      DataCopyPad(inLocal, inputXGm[row * this->cols], copyInParams, {false, 0, 0, 0});
+    } else {
+      DataCopyPad(inLocal[perLoopColsAlign], inputXGm[row * this->cols], copyInParams, {false, 0, 0, 0});
+    }
+
+    inputXInQueue.EnQue<T>(inLocal);
+
+    // Compute quantization
+    Compute(smoothLocal);
+
+    LocalTensor<int8_t> outLocal = inputXOutQueue.DeQue<int8_t>();
+
+    while (curLoopRow < this->currentLoopRows && initialRow / this->k == row) {
+      int32_t outIndex = indicesLocal.GetValue(curLoopRow);
+      curLoopRow++;
+      initialRow++;
+      if (outIndex == -1 || (this->dropPadMode == DROPLESS_MODE && outIndex >= this->activateRows)) {
+        continue;
+      }
+      // Scale is placed after the data position
+      DataCopyPad(expandedXGm[outIndex * cols_scale_], outLocal, copyOutParams);
+    }
+    inputXInQueue.FreeTensor(inLocal);
+    inputXOutQueue.FreeTensor(outLocal);
+  }
+  expandRowIdxInQueue.FreeTensor(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherDynamicQuant<T>::CopyOutXQuantEH(int64_t progress) {
+  LocalTensor<int32_t> indicesLocal = expandRowIdxInQueue.DeQue<int32_t>();
+  SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+
+  DataCopyExtParams copyInParams{1, static_cast<uint32_t>(this->perLoopCols * sizeof(T)), 0, 0, 0};
+  DataCopyExtParams smoothParams{1, static_cast<uint32_t>(this->perLoopCols * sizeof(float)), 0, 0, 0};
+  DataCopyExtParams copyOutParams{1, static_cast<uint32_t>(this->perLoopCols * sizeof(int8_t)), 0, 0, 0};
+
+  int32_t lastExpertIdx = -1;
+  LocalTensor<T> inLocal = inputXInQueue.AllocTensor<T>();
+  LocalTensor<float> smoothLocal = smoothInQueue.AllocTensor<float>();
+  for (int64_t i = 0; i < this->currentLoopRows; i++) {
+    int64_t rowOffset = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress;
+    if (this->dropPadMode == DROPLESS_MODE && rowOffset + i >= this->activateRows) {
+      break;
+    }
+    int32_t srcIdx = indicesLocal.GetValue(i);
+    int32_t expertIdx = indicesLocal.GetValue(currentLoopRowsAlign + i);
+
+    if constexpr (IsSameType<T, float>::value) {
+      DataCopyPad(inLocal, inputXGm[srcIdx / this->k * this->cols], copyInParams, {false, 0, 0, 0});
+    } else {
+      DataCopyPad(inLocal[perLoopColsAlign], inputXGm[srcIdx / this->k * this->cols], copyInParams, {false, 0, 0, 0});
+    }
+    inputXInQueue.EnQue<T>(inLocal);
+    
+    if (expertIdx != lastExpertIdx) {
+      DataCopyPad(smoothLocal, quantSmoothGm[expertIdx * this->cols], smoothParams, {false, 0, 0, 0});
+      smoothInQueue.EnQue(smoothLocal);
+      smoothLocal = smoothInQueue.DeQue<float>();
+      lastExpertIdx = expertIdx;
+    }
+
+    Compute(smoothLocal);
+
+    LocalTensor<float> quantScaleLocal = scaleOutQueue.DeQue<float>();
+    DataCopyPad(dynamicQuantScaleGm[(rowOffset + i)], quantScaleLocal, {1, 4, 0, 0, 0});
+
+    LocalTensor<int8_t> outLocal = inputXOutQueue.DeQue<int8_t>();
+    DataCopyPad(expandedXGm[(rowOffset + i) * this->cols], outLocal, copyOutParams);
+
+    inputXOutQueue.FreeTensor(outLocal);
+    scaleOutQueue.FreeTensor(quantScaleLocal);
+  }
+
+  inputXInQueue.FreeTensor(inLocal);
+  smoothInQueue.FreeTensor(smoothLocal);
+  expandRowIdxInQueue.FreeTensor(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline float MoeV2GatherDynamicQuant<T>::ComputeMax(LocalTensor<float>& inLocal,
+                                                               LocalTensor<float>& tempLocal,
+                                                               LocalTensor<float>& dynamicQuantLocal, int32_t srcIdx,
+                                                               int32_t expertIdx, int64_t j) {
+  LocalTensor<float> smoothLocal = smoothInQueue.AllocTensor<float>();
+
+  DataCopyExtParams intriParamsT{1, static_cast<uint32_t>(colsTileLength * sizeof(T)), 0, 0, 0};
+  DataCopyExtParams intriParamsFp32{1, static_cast<uint32_t>(colsTileLength * sizeof(float)), 0, 0, 0};
+
+  if constexpr (!IsSameType<T, float>::value) {
+    DataCopyPad(inLocal.ReinterpretCast<T>()[perLoopColsAlign], inputXGm[srcIdx * this->cols + j * this->perLoopCols],
+                intriParamsT, {false, 0, 0, 0});
+  } else {
+    DataCopyPad(inLocal, inputXGm[srcIdx * this->cols + j * this->perLoopCols], intriParamsT, {false, 0, 0, 0});
+  }
+
+  inputXInQueue.EnQue<float>(inLocal);
+  inLocal = inputXInQueue.DeQue<float>();
+
+  if (smoothType != 0) {
+    DataCopyPad(smoothLocal, quantSmoothGm[expertIdx * this->cols + j * this->perLoopCols], intriParamsFp32,
+                {false, 0, 0, 0});
+    smoothInQueue.EnQue(smoothLocal);
+    smoothLocal = smoothInQueue.DeQue<float>();
+  }
+
+  if constexpr (!IsSameType<T, float>::value) {
+    Cast(inLocal, inLocal.ReinterpretCast<T>()[perLoopColsAlign], RoundMode::CAST_NONE, colsTileLength);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  if (smoothType != 0) {
+    Mul(inLocal, inLocal, smoothLocal, colsTileLength);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  Abs(tempLocal, inLocal, colsTileLength);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  ReduceMax(dynamicQuantLocal[8], tempLocal, tempLocal, colsTileLength);
+
+  DataCopyPad(quantSrcGm[j * this->perLoopCols], inLocal, intriParamsFp32);
+  smoothInQueue.FreeTensor(smoothLocal);
+  SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
+
+  return dynamicQuantLocal.GetValue(8);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherDynamicQuant<T>::ComputeScale(LocalTensor<float>& inLocal,
+                                                                LocalTensor<float>& tempLocal, float scaleTemp,
+                                                                int64_t dstIndex, int64_t j) {
+  DataCopyExtParams copyInParams{1, static_cast<uint32_t>(colsTileLength * sizeof(float)), 0, 0, 0};
+  DataCopyExtParams copyOutParams{1, static_cast<uint32_t>(colsTileLength * sizeof(int8_t)), 0, 0, 0};
+
+  LocalTensor<int8_t> outLocal = inputXOutQueue.AllocTensor<int8_t>();
+
+  DataCopyPad(inLocal, quantSrcGm[j * this->perLoopCols], copyInParams, {false, 0, 0, 0});
+  inputXInQueue.EnQue<float>(inLocal);
+  inLocal = inputXInQueue.DeQue<float>();
+
+  Duplicate<float>(tempLocal, scaleTemp, colsTileLength);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Div(tempLocal, inLocal, tempLocal, colsTileLength);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Cast(tempLocal.ReinterpretCast<half>(), tempLocal, RoundMode::CAST_TRUNC, colsTileLength);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Cast(outLocal, tempLocal.ReinterpretCast<half>(), RoundMode::CAST_ROUND, colsTileLength);
+
+  inputXOutQueue.EnQue(outLocal);
+  outLocal = inputXOutQueue.DeQue<int8_t>();
+  DataCopyPad(expandedXGm[dstIndex * this->cols + j * this->perLoopCols], outLocal, copyOutParams);
+
+  inputXOutQueue.FreeTensor(outLocal);
+  SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherDynamicQuant<T>::CopyOutPartialXQuantEH(int64_t progress) {
+  LocalTensor<int32_t> indicesLocal = expandRowIdxInQueue.DeQue<int32_t>();
+  SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+
+  for (int64_t i = 0; i < this->currentLoopRows; i++) {
+    int64_t rowOffset = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress;
+    if (this->dropPadMode == DROPLESS_MODE && rowOffset + i >= this->activateRows) {
+      break;
+    }
+    int32_t srcIdx = indicesLocal.GetValue(i);
+    int32_t expertIdx = indicesLocal.GetValue(currentLoopRowsAlign + i);
+
+    LocalTensor<float> inLocal = inputXInQueue.AllocTensor<float>();
+    LocalTensor<float> tempLocal = calcQueue.AllocTensor<float>();
+    LocalTensor<float> quantScaleLocal = scaleOutQueue.AllocTensor<float>();
+
+    uint32_t tmp = 0xFF7FFFFF;
+    float reduceMax = *((float*)&tmp);
+    for (int64_t j = 0; j < this->colLoops; j++) {
+      colsTileLength = this->perLoopCols;
+      if (j == this->colLoops - 1) {
+        colsTileLength = this->lastLoopCols;
+      }
+      float tileMax = ComputeMax(inLocal, tempLocal, quantScaleLocal, srcIdx / this->k, expertIdx, j);
+      reduceMax = (reduceMax > tileMax) ? reduceMax : tileMax;
+    }
+
+    float scaleTemp = reduceMax / 127.0f;
+    Duplicate<float>(quantScaleLocal, scaleTemp, 8);
+    scaleOutQueue.EnQue(quantScaleLocal);
+    quantScaleLocal = scaleOutQueue.DeQue<float>();
+
+    DataCopyPad(dynamicQuantScaleGm[(rowOffset + i)], quantScaleLocal, {1, 4, 0, 0, 0});
+
+    for (int64_t j = 0; j < this->colLoops; j++) {
+      colsTileLength = this->perLoopCols;
+      if (j == this->colLoops - 1) {
+        colsTileLength = this->lastLoopCols;
+      }
+
+      ComputeScale(inLocal, tempLocal, scaleTemp, rowOffset + i, j);
+    }
+
+    inputXInQueue.FreeTensor(inLocal);
+    calcQueue.FreeTensor(tempLocal);
+    scaleOutQueue.FreeTensor(quantScaleLocal);
+  }
+
+  expandRowIdxInQueue.FreeTensor(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherDynamicQuant<T>::CopyOutPartialXQuant1H(int64_t progress) {
+  LocalTensor<int32_t> indicesLocal = expandRowIdxInQueue.DeQue<int32_t>();
+
+  int64_t initialRow = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress;
+  int64_t curLoopRow = 0;
+
+  int64_t currentLoopStartRow = initialRow / this->k;
+  int64_t currentLoopLastRow = (initialRow + this->currentLoopRows - 1) / this->k;
+
+  for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) {
+    LocalTensor<float> inLocal = inputXInQueue.AllocTensor<float>();
+    LocalTensor<float> tempLocal = calcQueue.AllocTensor<float>();
+    LocalTensor<float> quantScaleLocal = scaleOutQueue.AllocTensor<float>();
+
+    uint32_t tmp = 0xFF7FFFFF;
+    float reduceMax = *((float*)&tmp);
+    for (int64_t j = 0; j < this->colLoops; j++) {
+      colsTileLength = this->perLoopCols;
+      if (j == this->colLoops - 1) {
+        colsTileLength = this->lastLoopCols;
+      }
+
+      float tileMax = ComputeMax(inLocal, tempLocal, quantScaleLocal, row, 0, j);
+      reduceMax = (reduceMax > tileMax) ? reduceMax : tileMax;
+    }
+
+    float scaleTemp = reduceMax / 127.0f;
+    Duplicate<float>(quantScaleLocal, scaleTemp, 8);
+    scaleOutQueue.EnQue(quantScaleLocal);
+    quantScaleLocal = scaleOutQueue.DeQue<float>();
+
+    while (curLoopRow < this->currentLoopRows && initialRow / this->k == row) {
+      int32_t outIndex = indicesLocal.GetValue(curLoopRow);
+      curLoopRow++;
+      initialRow++;
+      if (outIndex == -1 || (this->dropPadMode == DROPLESS_MODE && outIndex >= this->activateRows)) {
+        continue;
+      }
+      DataCopyPad(dynamicQuantScaleGm[outIndex], quantScaleLocal, {1, 4, 0, 0, 0});
+      for (int64_t j = 0; j < this->colLoops; j++) {
+        colsTileLength = this->perLoopCols;
+        if (j == this->colLoops - 1) {
+          colsTileLength = this->lastLoopCols;
+        }
+
+        ComputeScale(inLocal, tempLocal, scaleTemp, outIndex, j);
+      }
+    }
+    inputXInQueue.FreeTensor(inLocal);
+    calcQueue.FreeTensor(tempLocal);
+    scaleOutQueue.FreeTensor(quantScaleLocal);
+  }
+
+  expandRowIdxInQueue.FreeTensor(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherDynamicQuant<T>::Init(GM_ADDR inputX, GM_ADDR quantSmooth, GM_ADDR expandedRowIdx,
+                                                        GM_ADDR expandedX, GM_ADDR dynamicQuantScale, GM_ADDR workspace,
+                                                        const MoeInitRoutingQuantV2TilingData* tilingData,
+                                                        TPipe* tPipe) {
+  this->pipe = tPipe;
+  this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+  this->gatherOutTilingData = &(tilingData->gatherOutComputeParamsOp);
+
+  this->needCoreNum = this->gatherOutTilingData->needCoreNum;
+  this->activateRows = this->gatherOutTilingData->activateRows;
+  this->cols = tilingData->cols;
+  this->cols_scale_ = this->cols + UB_ALIGN;
+  this->n = tilingData->n;
+  this->k = tilingData->k;
+  this->totalLength = tilingData->n * tilingData->k;
+  this->dropPadMode = tilingData->dropPadMode;
+  this->smoothType = tilingData->smoothType;
+
+  if (this->blockIdx == this->gatherOutTilingData->needCoreNum - 1) {
+    this->coreRows = this->gatherOutTilingData->lastCoreRows;
+    this->perLoopRows = this->gatherOutTilingData->lastCorePerLoopRows;
+    this->lastLoopRows = this->gatherOutTilingData->lastCoreLastLoopRows;
+    this->rowLoops = this->gatherOutTilingData->lastCoreLoops;
+  } else {
+    this->coreRows = this->gatherOutTilingData->perCoreRows;
+    this->perLoopRows = this->gatherOutTilingData->perCorePerLoopRows;
+    this->lastLoopRows = this->gatherOutTilingData->perCoreLastLoopRows;
+    this->rowLoops = this->gatherOutTilingData->perCoreLoops;
+  }
+  this->perLoopCols = this->gatherOutTilingData->perLoopCols;
+  this->lastLoopCols = this->gatherOutTilingData->lastLoopCols;
+  this->colLoops = this->gatherOutTilingData->colLoops;
+  this->perLoopColsAlign = Align(this->perLoopCols, sizeof(T));
+
+  inputXGm.SetGlobalBuffer((__gm__ T*)inputX);
+  expandedXGm.SetGlobalBuffer((__gm__ int8_t*)expandedX);
+
+  expandedRowIdxGm.SetGlobalBuffer(
+      (__gm__ int32_t*)expandedRowIdx + this->blockIdx * this->gatherOutTilingData->perCoreRows,
+      Align(this->coreRows, sizeof(int32_t)));
+
+  quantSmoothGm.SetGlobalBuffer((__gm__ float*)quantSmooth);
+  dynamicQuantScaleGm.SetGlobalBuffer((__gm__ float*)dynamicQuantScale);
+
+  expandedExpertIdxGm.SetGlobalBuffer(
+      (__gm__ int32_t*)workspace + this->blockIdx * this->gatherOutTilingData->perCoreRows,
+      Align(this->coreRows, sizeof(int32_t)));
+  sortedRowIdxGm.SetGlobalBuffer((__gm__ int32_t*)workspace + Align(this->totalLength, sizeof(int32_t)) +
+                                     this->blockIdx * this->gatherOutTilingData->perCoreRows,
+                                 Align(this->coreRows, sizeof(int32_t)));
+  if (this->cols > 1) {
+    quantSrcGm.SetGlobalBuffer(
+        (__gm__ float*)workspace + Align(this->totalLength, sizeof(int32_t)) * 2 + this->blockIdx * this->cols,
+        this->cols * sizeof(float));
+  }
+
+  this->currentLoopRowsAlign = Align(this->perLoopRows, sizeof(int32_t));
+
+  int64_t perLoopColsAlignBytes = AlignBytes(this->perLoopCols, sizeof(T));
+  perLoopColsAlignBytes =
+      Max(int64_t(perLoopColsAlignBytes * sizeof(float) / sizeof(T)), int64_t(BLOCK_BYTES + BLOCK_BYTES));
+
+  pipe->InitBuffer(expandRowIdxInQueue, BUFFER_NUM, 2 * AlignBytes(this->perLoopRows, sizeof(int32_t)));
+  pipe->InitBuffer(inputXInQueue, BUFFER_NUM, perLoopColsAlignBytes);
+  pipe->InitBuffer(smoothInQueue, BUFFER_NUM, AlignBytes(this->perLoopCols, sizeof(float)));
+  pipe->InitBuffer(calcQueue, 1, AlignBytes(this->perLoopCols, sizeof(float)));
+  pipe->InitBuffer(inputXOutQueue, 1, AlignBytes(this->perLoopCols, sizeof(int8_t)));
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherDynamicQuant<T>::Process() {
+  if (this->blockIdx < this->needCoreNum) {
+    currentLoopRows = perLoopRows;
+      if (colLoops > 1) {  // Cannot fit all data in one row, workspace is required
+        trap();   // Not supported
+      } else {  // All data can fit in one row
+      if (smoothType == 2) {
+        for (int64_t loop = 0; loop < this->rowLoops - 1; loop++) {
+          CopyInExpandedExpertIdx(loop);
+          CopyOutXQuantEH(loop);
+        }
+        currentLoopRows = lastLoopRows;
+        CopyInExpandedExpertIdx(this->rowLoops - 1);
+        CopyOutXQuantEH(this->rowLoops - 1);
+      } else {
+        for (int64_t loop = 0; loop < this->rowLoops - 1; loop++) {
+          CopyInExpandedRowIdx(loop);
+          CopyOutXQuant1H(loop);
+        }
+        currentLoopRows = lastLoopRows;
+        CopyInExpandedRowIdx(this->rowLoops - 1);
+        CopyOutXQuant1H(this->rowLoops - 1);
+      }
+    }
+  }
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // MOE_V2_GATHER_DYNAMIC_QUANT_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_out.h

@@ -0,0 +1,181 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_gather_out.h
+ * \brief
+ */
+#ifndef INNER_MOE_V2_GATHER_OUT_H
+#define INNER_MOE_V2_GATHER_OUT_H
+
+#include "moe_v2_common.h"
+#include "kernel_operator.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+constexpr int64_t BUFFER_NUM = 2;
+
+template <typename T>
+class MoeV2GatherOut {
+ public:
+  __aicore__ inline MoeV2GatherOut(){};
+  __aicore__ inline void Init(GM_ADDR inputX, GM_ADDR expandedRowIdx, GM_ADDR expandedX, GM_ADDR workspace,
+                              const InnerMoeInitRoutingV2TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void CopyInIndices(int64_t progress);
+  __aicore__ inline void CopyOut(int64_t progress);
+
+ private:
+  TPipe* pipe;
+  TQueBind<QuePosition::VECIN, QuePosition::VECOUT, BUFFER_NUM> inputActivationsCopyInQueue;
+  TQue<QuePosition::VECIN, BUFFER_NUM> expandDstToSrcRowCopyInQueue;
+
+  GlobalTensor<T> inputXGm;
+  GlobalTensor<T> expandedXGm;
+  GlobalTensor<int32_t> expandedRowIdxGm;
+
+  const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData;
+
+  int64_t needCoreNum;
+  int64_t blockIdx;
+  int64_t cols;
+  int64_t n;
+  int64_t k;
+  int64_t activateRows;
+  int64_t currentLoopRows;
+  int64_t coreRows;
+  int64_t perLoopRows;
+  int64_t lastLoopRows;
+  int64_t rowLoops;
+  int64_t colsTileLength;
+  int64_t perLoopCols;
+  int64_t lastLoopCols;
+  int64_t colLoops;
+  int64_t dropPadMode;
+
+  int64_t indicesOffset;
+  int64_t inputOffset;
+  int64_t outOffset;
+};
+
+template <typename T>
+__aicore__ inline void MoeV2GatherOut<T>::CopyInIndices(int64_t progress) {
+  this->indicesOffset = progress * this->perLoopRows;
+  LocalTensor<int32_t> indicesLocal = expandDstToSrcRowCopyInQueue.AllocTensor<int32_t>();
+  DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(this->currentLoopRows * sizeof(int32_t)), 0, 0, 0};
+  DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+  DataCopyPad(indicesLocal, expandedRowIdxGm[indicesOffset], dataCopyParams, dataCopyPadParams);
+
+  expandDstToSrcRowCopyInQueue.EnQue<int32_t>(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherOut<T>::CopyOut(int64_t progress) {
+  LocalTensor<int32_t> indicesLocal = expandDstToSrcRowCopyInQueue.DeQue<int32_t>();
+  SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+  colsTileLength = this->perLoopCols;
+  for (int64_t colsLoop = 0; colsLoop < this->colLoops; colsLoop++) {
+    int64_t initialRow = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress;
+    int64_t curLoopRow = 0;
+    if (colsLoop == this->colLoops - 1) {
+      colsTileLength = this->lastLoopCols;
+    }
+    int64_t currentLoopStartRow = initialRow / this->k;
+    int64_t currentLoopLastRow = (initialRow + this->currentLoopRows - 1) / this->k;
+    for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) {
+      LocalTensor<T> inLocal = inputActivationsCopyInQueue.AllocTensor<T>();
+      // input row position
+      inputOffset = row * this->cols + colsLoop * this->perLoopCols;
+      DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(this->colsTileLength * sizeof(T)), 0, 0, 0};
+      DataCopyPadExtParams<T> dataCopyPadParams{false, 0, 0, 0};
+      DataCopyPad(inLocal, inputXGm[inputOffset], dataCopyParams, dataCopyPadParams);
+      SetWaitFlag<HardEvent::MTE2_MTE3>(HardEvent::MTE2_MTE3);
+
+      DataCopyExtParams intriParams{1, static_cast<uint32_t>(this->colsTileLength * sizeof(T)), 0, 0, 0};
+      while (curLoopRow < this->currentLoopRows && initialRow / this->k == row) {
+        int32_t outIndex = indicesLocal.GetValue(curLoopRow);
+        curLoopRow++;
+        initialRow++;
+        if (outIndex == -1 || (this->dropPadMode == DROPLESS_MODE && outIndex >= this->activateRows)) {
+          continue;
+        }
+        outOffset = outIndex * cols + colsLoop * this->perLoopCols;
+#ifdef __CCE_KT_TEST__
+        // CPU twin debugging cannot use multi-core sync, so index may contain uninitialized dirty data; handle specially
+        if (outOffset > expandedXGm.GetSize()) {
+            continue;
+        }
+#endif
+        DataCopyPad(expandedXGm[outOffset], inLocal, intriParams);
+      }
+      inputActivationsCopyInQueue.FreeTensor(inLocal);
+    }
+  }
+  expandDstToSrcRowCopyInQueue.FreeTensor(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherOut<T>::Init(GM_ADDR inputX, GM_ADDR expandedRowIdx, GM_ADDR expandedX,
+                                               GM_ADDR workspace, const InnerMoeInitRoutingV2TilingData* tilingData,
+                                               TPipe* tPipe) {
+  this->pipe = tPipe;
+  this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+  this->gatherOutTilingData = &(tilingData->gatherOutComputeParamsOp);
+
+  this->needCoreNum = this->gatherOutTilingData->needCoreNum;
+  this->activateRows = this->gatherOutTilingData->activateRows;
+  this->cols = tilingData->cols;
+  this->n = tilingData->n;
+  this->k = tilingData->k;
+  this->dropPadMode = tilingData->dropPadMode;
+
+  if (this->blockIdx == this->gatherOutTilingData->needCoreNum - 1) {
+    this->coreRows = this->gatherOutTilingData->lastCoreRows;
+    this->perLoopRows = this->gatherOutTilingData->lastCorePerLoopRows;
+    this->lastLoopRows = this->gatherOutTilingData->lastCoreLastLoopRows;
+    this->rowLoops = this->gatherOutTilingData->lastCoreLoops;
+  } else {
+    this->coreRows = this->gatherOutTilingData->perCoreRows;
+    this->perLoopRows = this->gatherOutTilingData->perCorePerLoopRows;
+    this->lastLoopRows = this->gatherOutTilingData->perCoreLastLoopRows;
+    this->rowLoops = this->gatherOutTilingData->perCoreLoops;
+  }
+  this->perLoopCols = this->gatherOutTilingData->perLoopCols;
+  this->lastLoopCols = this->gatherOutTilingData->lastLoopCols;
+  this->colLoops = this->gatherOutTilingData->colLoops;
+
+  inputXGm.SetGlobalBuffer((__gm__ T*)inputX, this->coreRows * this->cols);
+  expandedXGm.SetGlobalBuffer((__gm__ T*)expandedX, tilingData->n * tilingData->k * this->cols);
+  expandedRowIdxGm.SetGlobalBuffer(
+      (__gm__ int32_t*)expandedRowIdx + this->blockIdx * this->gatherOutTilingData->perCoreRows,
+      Align(this->coreRows, sizeof(int32_t)));
+
+  pipe->InitBuffer(inputActivationsCopyInQueue, BUFFER_NUM, AlignBytes(this->perLoopCols, sizeof(T)));
+  pipe->InitBuffer(expandDstToSrcRowCopyInQueue, BUFFER_NUM, AlignBytes(this->perLoopRows, sizeof(int32_t)));
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherOut<T>::Process() {
+  if (this->blockIdx < this->needCoreNum) {
+    currentLoopRows = perLoopRows;
+    for (int64_t loop = 0; loop < this->rowLoops; loop++) {
+      if (loop == this->rowLoops - 1) {
+        currentLoopRows = lastLoopRows;
+      }
+      CopyInIndices(loop);
+      CopyOut(loop);
+    }
+  }
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // INNER_MOE_V2_GATHER_OUT_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_gather_quant.h

@@ -0,0 +1,235 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_gather_quant.h
+ * \brief
+ */
+#ifndef MOE_V2_GATHER_QUANT_H
+#define MOE_V2_GATHER_QUANT_H
+
+#include "moe_v2_common.h"
+#include "kernel_operator.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+constexpr int64_t BUFFER_NUM = 2;
+
+template <typename T>
+class MoeV2GatherQuant {
+ public:
+  __aicore__ inline MoeV2GatherQuant(){};
+  __aicore__ inline void Init(GM_ADDR inputX, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedRowIdx, GM_ADDR expandedX,
+                              GM_ADDR workspace, const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void CopyInIndices(int64_t progress);
+  __aicore__ inline void Compute();
+  __aicore__ inline void CopyOut(int64_t progress);
+
+ private:
+  TPipe* pipe;
+  TQue<QuePosition::VECIN, BUFFER_NUM> inputXCopyInQueue;
+  TQue<QuePosition::VECIN, BUFFER_NUM> expandRowIdxCopyInQueue;
+  TQue<QuePosition::VECOUT, BUFFER_NUM> inputXCopyOutQueue;
+  TQue<QuePosition::VECOUT, 1> floatQueue;
+  TQue<QuePosition::VECOUT, 1> halfQueue;
+
+  GlobalTensor<T> inputXGm;
+  GlobalTensor<int8_t> expandedXGm;
+  GlobalTensor<int32_t> expandedRowIdxGm;
+  GlobalTensor<float> scaleGm;
+  GlobalTensor<float> offsetGm;
+
+  const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData;
+
+  int64_t needCoreNum;
+  int64_t blockIdx;
+  int64_t cols;
+  int64_t n;
+  int64_t k;
+  int64_t activateRows;
+  int64_t currentLoopRows;
+  int64_t coreRows;
+  int64_t perLoopRows;
+  int64_t lastLoopRows;
+  int64_t rowLoops;
+  int64_t colsTileLength;
+  int64_t perLoopCols;
+  int64_t lastLoopCols;
+  int64_t colLoops;
+  int64_t dropPadMode;
+  float scale;
+  float offset;
+
+  int64_t indicesOffset;
+  int64_t inputOffset;
+  int64_t outOffset;
+};
+
+template <typename T>
+__aicore__ inline void MoeV2GatherQuant<T>::CopyInIndices(int64_t progress) {
+  this->indicesOffset = progress * this->perLoopRows;
+  LocalTensor<int32_t> indicesLocal = expandRowIdxCopyInQueue.AllocTensor<int32_t>();
+  DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(this->currentLoopRows * sizeof(int32_t)), 0, 0, 0};
+  DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+  DataCopyPad(indicesLocal, expandedRowIdxGm[indicesOffset], dataCopyParams, dataCopyPadParams);
+  expandRowIdxCopyInQueue.EnQue<int32_t>(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherQuant<T>::Compute() {
+  LocalTensor<T> inLocal = inputXCopyInQueue.DeQue<T>();
+  LocalTensor<int8_t> outLocal = inputXCopyOutQueue.AllocTensor<int8_t>();
+  LocalTensor<float> floatLocal = floatQueue.AllocTensor<float>();
+  LocalTensor<half> halfLocal = halfQueue.AllocTensor<half>();
+  uint32_t elements = Align(this->colsTileLength, sizeof(T));
+  if constexpr (IsSameType<T, bfloat16_t>::value) {
+    Cast(floatLocal, inLocal, RoundMode::CAST_NONE, elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Cast(halfLocal, floatLocal, RoundMode::CAST_NONE, elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Muls(halfLocal, halfLocal, static_cast<half>(this->scale), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Adds(halfLocal, halfLocal, static_cast<half>(this->offset), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    LocalTensor<int32_t> intLocal = floatLocal.ReinterpretCast<int32_t>();
+    Cast(intLocal, halfLocal, RoundMode::CAST_RINT, elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    SetDeqScale((half)1.000000e+00f);
+    AscendC::PipeBarrier<PIPE_V>();
+    Cast(halfLocal, intLocal, RoundMode::CAST_RINT, elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Cast(outLocal, halfLocal, RoundMode::CAST_RINT, elements);
+  } else if constexpr (IsSameType<T, float>::value) {
+    Cast(halfLocal, inLocal, RoundMode::CAST_NONE, elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Muls(halfLocal, halfLocal, static_cast<half>(this->scale), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Adds(halfLocal, halfLocal, static_cast<half>(this->offset), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Cast(outLocal, halfLocal, RoundMode::CAST_RINT, elements);
+  } else {
+    Muls(inLocal, inLocal, static_cast<T>(this->scale), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Adds(inLocal, inLocal, static_cast<T>(this->offset), elements);
+    AscendC::PipeBarrier<PIPE_V>();
+    Cast(outLocal, inLocal, RoundMode::CAST_RINT, elements);
+  }
+  inputXCopyOutQueue.EnQue(outLocal);
+  floatQueue.FreeTensor(floatLocal);
+  halfQueue.FreeTensor(halfLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherQuant<T>::CopyOut(int64_t progress) {
+  LocalTensor<int32_t> indicesLocal = expandRowIdxCopyInQueue.DeQue<int32_t>();
+  SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+  colsTileLength = this->perLoopCols;
+  for (int64_t colsLoop = 0; colsLoop < this->colLoops; colsLoop++) {
+    int64_t initialRow = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress;
+    int64_t curLoopRow = 0;
+    if (colsLoop == this->colLoops - 1) {
+      colsTileLength = this->lastLoopCols;
+    }
+    int64_t currentLoopStartRow = initialRow / this->k;
+    int64_t currentLoopLastRow = (initialRow + this->currentLoopRows - 1) / this->k;
+    for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) {
+      LocalTensor<T> inLocal = inputXCopyInQueue.AllocTensor<T>();
+      // input row position
+      inputOffset = row * this->cols + colsLoop * this->perLoopCols;
+      DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(this->colsTileLength * sizeof(T)), 0, 0, 0};
+      DataCopyPadExtParams<T> dataCopyPadParams{false, 0, 0, 0};
+      DataCopyPad(inLocal, inputXGm[inputOffset], dataCopyParams, dataCopyPadParams);
+      inputXCopyInQueue.EnQue<T>(inLocal);
+      Compute();
+      LocalTensor<int8_t> outLocal = inputXCopyOutQueue.DeQue<int8_t>();
+      DataCopyExtParams intriParams{1, static_cast<uint32_t>(this->colsTileLength * sizeof(int8_t)), 0, 0, 0};
+      while (curLoopRow < this->currentLoopRows && initialRow / this->k == row) {
+        int32_t outIndex = indicesLocal.GetValue(curLoopRow);
+        curLoopRow++;
+        initialRow++;
+        if (outIndex == -1 || (this->dropPadMode == DROPLESS_MODE && outIndex >= this->activateRows)) {
+          continue;
+        }
+        outOffset = outIndex * cols + colsLoop * this->perLoopCols;
+        DataCopyPad(expandedXGm[outOffset], outLocal, intriParams);
+      }
+      inputXCopyInQueue.FreeTensor(inLocal);
+      inputXCopyOutQueue.FreeTensor(outLocal);
+    }
+  }
+  expandRowIdxCopyInQueue.FreeTensor(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherQuant<T>::Init(GM_ADDR inputX, GM_ADDR scale, GM_ADDR offset, GM_ADDR expandedRowIdx,
+                                                 GM_ADDR expandedX, GM_ADDR workspace,
+                                                 const MoeInitRoutingQuantV2TilingData* tilingData, TPipe* tPipe) {
+  this->pipe = tPipe;
+  this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+  this->gatherOutTilingData = &(tilingData->gatherOutComputeParamsOp);
+
+  this->needCoreNum = this->gatherOutTilingData->needCoreNum;
+  this->activateRows = this->gatherOutTilingData->activateRows;
+  this->cols = tilingData->cols;
+  this->n = tilingData->n;
+  this->k = tilingData->k;
+  this->dropPadMode = tilingData->dropPadMode;
+
+  if (this->blockIdx == this->gatherOutTilingData->needCoreNum - 1) {
+    this->coreRows = this->gatherOutTilingData->lastCoreRows;
+    this->perLoopRows = this->gatherOutTilingData->lastCorePerLoopRows;
+    this->lastLoopRows = this->gatherOutTilingData->lastCoreLastLoopRows;
+    this->rowLoops = this->gatherOutTilingData->lastCoreLoops;
+  } else {
+    this->coreRows = this->gatherOutTilingData->perCoreRows;
+    this->perLoopRows = this->gatherOutTilingData->perCorePerLoopRows;
+    this->lastLoopRows = this->gatherOutTilingData->perCoreLastLoopRows;
+    this->rowLoops = this->gatherOutTilingData->perCoreLoops;
+  }
+  this->perLoopCols = this->gatherOutTilingData->perLoopCols;
+  this->lastLoopCols = this->gatherOutTilingData->lastLoopCols;
+  this->colLoops = this->gatherOutTilingData->colLoops;
+
+  inputXGm.SetGlobalBuffer((__gm__ T*)inputX);
+  expandedXGm.SetGlobalBuffer((__gm__ int8_t*)expandedX);
+  expandedRowIdxGm.SetGlobalBuffer(
+      (__gm__ int32_t*)expandedRowIdx + this->blockIdx * this->gatherOutTilingData->perCoreRows,
+      Align(this->coreRows, sizeof(int32_t)));
+  scaleGm.SetGlobalBuffer((__gm__ float*)scale, 1);
+  offsetGm.SetGlobalBuffer((__gm__ float*)offset, 1);
+  this->scale = scaleGm.GetValue(0);
+  this->offset = offsetGm.GetValue(0);
+
+  pipe->InitBuffer(inputXCopyInQueue, BUFFER_NUM, AlignBytes(this->perLoopCols, sizeof(T)));
+  pipe->InitBuffer(inputXCopyOutQueue, BUFFER_NUM, AlignBytes(this->perLoopCols, sizeof(int8_t)));
+  pipe->InitBuffer(expandRowIdxCopyInQueue, BUFFER_NUM, AlignBytes(this->perLoopRows, sizeof(int32_t)));
+  pipe->InitBuffer(floatQueue, 1, AlignBytes(this->perLoopCols, sizeof(float)));
+  pipe->InitBuffer(halfQueue, 1, AlignBytes(this->perLoopCols, sizeof(half)));
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherQuant<T>::Process() {
+  if (this->blockIdx < this->needCoreNum) {
+    currentLoopRows = perLoopRows;
+    for (int64_t loop = 0; loop < this->rowLoops; loop++) {
+      if (loop == this->rowLoops - 1) {
+        currentLoopRows = lastLoopRows;
+      }
+      CopyInIndices(loop);
+      CopyOut(loop);
+    }
+  }
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // MOE_V2_GATHER_QUANT_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_init_routing_fullload.h

@@ -0,0 +1,312 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/* !
+ * \file moe_v2_init_routing_fullload.h
+ * \brief
+ */
+#ifndef INNER_MOE_V2_FULL_LOAD_H
+#define INNER_MOE_V2_FULL_LOAD_H
+
+#include "moe_v2_mrgsort.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+template <typename T>
+class MoeV2FullLoad : public MoeV2SortBase {
+ public:
+  __aicore__ inline MoeV2FullLoad(){};
+  __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, GM_ADDR expandedRowIdx,
+                              GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace,
+                              const InnerMoeInitRoutingV2TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void CopyIn();
+  __aicore__ inline void SortCompute();
+  __aicore__ inline void CopyOutIdx();
+  __aicore__ inline void CopyOutEmpty();
+  __aicore__ inline void CopyOutX();
+  __aicore__ inline void ComputeExpertTokenCountOrCumsum();
+
+ private:
+  int64_t sortNum_;
+  const InnerMoeV2GatherOutComputeTilingData* gatherOutTilingData_;
+  int64_t blockIdx_;
+  int64_t needCoreNum_;
+  int64_t coreRows_;
+  int64_t perCoreRows_;
+  int64_t k_;
+  int64_t n_;
+  int64_t cols_;
+  int64_t activateRows_;
+  int64_t expertNum;
+  int64_t expertCapacity;
+
+  TQue<QuePosition::VECIN, 1> xCopyInQueue_;
+  TQue<QuePosition::VECOUT, 1> expandedRowIdxCopyOutQueue_;
+  TQue<QuePosition::VECOUT, 1> expandedExpertIdxCopyOutQueue_;
+  TQue<QuePosition::VECOUT, 1> expandDstToSrcRowQueue_;
+  TQue<QuePosition::VECOUT, 1> expertTokensCopyOutQueue_;
+
+  GlobalTensor<T> xGm_;
+  GlobalTensor<int32_t> expertIdxGm_;
+
+  GlobalTensor<T> expandedXGm_;
+  GlobalTensor<int32_t> expandedRowIdxGm_;
+  GlobalTensor<int32_t> expandedExpertIdxGm_;
+  GlobalTensor<int32_t> expertTokensCountOrCumsumGm;
+  GlobalTensor<int32_t> expertTokensBeforeCapacityGm;
+
+  int64_t expertTokensCountOrCumsumFlag = 0;
+  int64_t expertTokensBeforeCapacityFlag = 0;
+  int64_t dropPadMode = 0;
+};
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::CopyIn() {
+  LocalTensor<int32_t> inLocal = sortDataCopyInQueue.AllocTensor<int32_t>();
+  DataCopyExtParams dataCopyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(this->totalLength * sizeof(int32_t)),
+                                   0, 0, 0};
+  DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+  DataCopyPad(inLocal[0], expertIdxGm_, dataCopyParams, dataCopyPadParams);
+  ArithProgression<int32_t>(inLocal[this->sortNum_], 0, 1, this->totalLength);
+  sortDataCopyInQueue.EnQue(inLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::SortCompute() {
+  LocalTensor<int32_t> inLocal = sortDataCopyInQueue.DeQue<int32_t>();
+  LocalTensor<int32_t> expertIdxLocal = inLocal[0];
+  LocalTensor<float> expertIdxLocalFp32 = expertIdxLocal.ReinterpretCast<float>();
+  Cast(expertIdxLocalFp32, expertIdxLocal, RoundMode::CAST_ROUND, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  Muls(expertIdxLocalFp32, expertIdxLocalFp32, (float)-1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM;
+  if (duplicateNum > 0) {
+    int duplicateIndex = this->totalLength - duplicateNum;
+    uint64_t mask0 = UINT64_MAX;
+    mask0 = mask0 << duplicateNum;
+    mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
+    uint64_t mask[2] = {mask0, 0};
+    Duplicate(expertIdxLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+  LocalTensor<float> concatLocal;
+  LocalTensor<float> tempTensor = tempBuffer.Get<float>(GetSortLen<float>(this->sortNum_));
+  Concat(concatLocal, expertIdxLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  LocalTensor<uint32_t> rowIdxLocal = inLocal[this->sortNum_].template ReinterpretCast<uint32_t>();
+  LocalTensor<float> sortedLocal = sortedBuffer.Get<float>(GetSortLen<float>(this->sortNum_));
+  Sort<float, true>(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  LocalTensor<float> expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue_.AllocTensor<float>();
+  LocalTensor<uint32_t> expandDstToSrcRowLocal = expandDstToSrcRowQueue_.AllocTensor<uint32_t>();
+  LocalTensor<float> expandDstToSrcRowLocalFp32 = expandDstToSrcRowLocal.ReinterpretCast<float>();
+  Extract(expandedExpertIdxLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  Cast(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocal.ReinterpretCast<int32_t>(), RoundMode::CAST_ROUND,
+       this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  Muls(expandedExpertIdxLocal, expandedExpertIdxLocal, (float)-1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  LocalTensor<int32_t> expandedExpertIdxLocalInt32;
+  expandedExpertIdxLocalInt32 = expandedExpertIdxLocal.ReinterpretCast<int32_t>();
+  Cast(expandedExpertIdxLocalInt32, expandedExpertIdxLocal, RoundMode::CAST_ROUND, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  expandedExpertIdxCopyOutQueue_.EnQue<int32_t>(expandedExpertIdxLocalInt32);
+
+  LocalTensor<uint32_t> expandedRowIdx = expandedRowIdxCopyOutQueue_.AllocTensor<uint32_t>();
+  LocalTensor<uint32_t> expandedRowIdxU32 = expandedRowIdx.ReinterpretCast<uint32_t>();
+  Muls(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocalFp32, (float)-1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  ArithProgression<int32_t>(inLocal[this->sortNum_], 0, 1, this->totalLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  if (duplicateNum > 0) {
+    int duplicateIndex = this->totalLength - duplicateNum;
+    uint64_t mask0 = UINT64_MAX;
+    mask0 = mask0 << duplicateNum;
+    mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
+    uint64_t mask[2] = {mask0, 0};
+    Duplicate(expandDstToSrcRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+  Concat(concatLocal, expandDstToSrcRowLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  Sort<float, true>(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  Extract(tempTensor, expandedRowIdxU32, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  expandedRowIdxCopyOutQueue_.EnQue<uint32_t>(expandedRowIdx);
+  sortDataCopyInQueue.FreeTensor(inLocal);
+
+  expandDstToSrcRowQueue_.FreeTensor(expandDstToSrcRowLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::CopyOutIdx() {
+  LocalTensor<int32_t> expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue<int32_t>();
+  DataCopyParams intriParams;
+  intriParams.blockCount = 1;
+  intriParams.blockLen = this->totalLength * sizeof(int32_t);
+  DataCopyPad(expandedRowIdxGm_, expandedRowIdx, intriParams);
+  expandedRowIdxCopyOutQueue_.EnQue(expandedRowIdx);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::ComputeExpertTokenCountOrCumsum() {
+  LocalTensor<int32_t> expandedExpertIdx = expandedExpertIdxCopyOutQueue_.DeQue<int32_t>();
+  LocalTensor<int32_t> expertTokensCount = expertTokensCopyOutQueue_.AllocTensor<int32_t>();
+
+  int64_t expertNumAlign = Align(this->expertNum, sizeof(int32_t));
+  Duplicate(expertTokensCount, 0, expertNumAlign);
+  SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
+
+  int32_t lastExpertId = expandedExpertIdx.GetValue(0);
+  int64_t tokenCount = 0;
+  int64_t lastExpertCount = 0;
+  for (int64_t i = 0; i < this->totalLength; i++) {
+    int32_t curExpertId = expandedExpertIdx.GetValue(i);
+    tokenCount++;
+    while (lastExpertId < curExpertId) {
+      expertTokensCount.SetValue(lastExpertId, tokenCount - 1);
+      if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_COUNT) {
+        tokenCount = 1;
+      }
+      lastExpertId++;
+    }
+  }
+  expertTokensCount.SetValue(lastExpertId, tokenCount);
+  if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_CUMSUM) {
+    lastExpertId++;
+    while (lastExpertId < this->expertNum) {
+      expertTokensCount.SetValue(lastExpertId, tokenCount);
+      lastExpertId++;
+    }
+  }
+  DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(this->expertNum * sizeof(int32_t)), 0, 0,
+                               0};
+  if (this->expertTokensCountOrCumsumFlag > 0) {
+    DataCopyPad(expertTokensCountOrCumsumGm, expertTokensCount, copyParams);
+  }
+  expertTokensCopyOutQueue_.FreeTensor(expertTokensCount);
+  expandedExpertIdxCopyOutQueue_.FreeTensor(expandedExpertIdx);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::CopyOutX() {
+  LocalTensor<T> xLocal = xCopyInQueue_.AllocTensor<T>();
+  LocalTensor<int32_t> expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue<int32_t>();
+  DataCopyParams intriParams;
+  intriParams.blockCount = 1;
+  intriParams.blockLen = this->cols_ * sizeof(T);
+  int64_t inFactor = Align(this->cols_, sizeof(T));
+  int64_t curRowsStart = this->blockIdx_ * this->perCoreRows_;
+  int64_t startXRow = curRowsStart / this->k_;
+  int64_t endXRow = (curRowsStart + this->coreRows_ - 1) / this->k_;
+
+  DataCopyExtParams dataXCopyParams{static_cast<uint16_t>(endXRow - startXRow + 1),
+                                    static_cast<uint32_t>(this->cols_ * sizeof(T)), 0, 0, 0};
+  DataCopyPadExtParams<T> dataXCopyPadParams{false, 0, 0, 0};
+  DataCopyPad(xLocal, xGm_[startXRow * this->cols_], dataXCopyParams, dataXCopyPadParams);
+  SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+
+  int64_t k = 0;
+  for (int64_t i = startXRow; i <= endXRow; i++) {
+    for (; k < this->perCoreRows_ && curRowsStart / this->k_ == i; curRowsStart++, k++) {
+      int32_t outIndex = expandedRowIdx.GetValue(curRowsStart);
+      if (outIndex < this->activateRows_) {
+        DataCopyPad(expandedXGm_[outIndex * this->cols_], xLocal[(i - startXRow) * inFactor], intriParams);
+      }
+    }
+  }
+  expandedRowIdxCopyOutQueue_.FreeTensor(expandedRowIdx);
+  xCopyInQueue_.FreeTensor(xLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::CopyOutEmpty() {
+  LocalTensor<int32_t> outLocal = expandedExpertIdxCopyOutQueue_.DeQue<int32_t>();
+  expandedExpertIdxCopyOutQueue_.FreeTensor(outLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, GM_ADDR expandedRowIdx,
+                                              GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace,
+                                              const InnerMoeInitRoutingV2TilingData* tilingData, TPipe* tPipe) {
+  this->gatherOutTilingData_ = &(tilingData->gatherOutComputeParamsOp);
+  this->blockIdx_ = get_block_idx() + get_subblockid() * get_block_num();
+  this->k_ = tilingData->k;
+  this->n_ = tilingData->n;
+  this->cols_ = tilingData->cols;
+  this->needCoreNum_ = this->gatherOutTilingData_->needCoreNum;
+  this->perCoreRows_ = this->gatherOutTilingData_->perCoreRows;
+  this->activateRows_ = this->gatherOutTilingData_->activateRows;
+  if (this->blockIdx_ == this->gatherOutTilingData_->needCoreNum - 1) {
+    this->coreRows_ = this->gatherOutTilingData_->lastCoreRows;
+  } else {
+    this->coreRows_ = this->gatherOutTilingData_->perCoreRows;
+  }
+  this->expertNum = tilingData->expertNum;
+  this->dropPadMode = tilingData->dropPadMode;
+  this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag;
+
+  this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t));
+  this->sortNum_ = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
+  this->totalLength = tilingData->n * tilingData->k;
+  this->pipe = tPipe;
+
+  xGm_.SetGlobalBuffer((__gm__ T*)x);
+  expertIdxGm_.SetGlobalBuffer((__gm__ int32_t*)expertIdx, this->tileLength);
+
+  expandedXGm_.SetGlobalBuffer((__gm__ T*)expandedX);
+  expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, this->tileLength);
+  if (this->expertTokensCountOrCumsumFlag > 0) {
+    // dropless
+    expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensCountOrCumsum,
+                                                Align(this->expertNum, sizeof(int32_t)));
+  }
+
+  int64_t kvFactor = 2;
+  int64_t buffSize = this->sortNum_ * sizeof(int32_t);
+
+  int64_t curRowsStart = this->blockIdx_ * this->perCoreRows_;
+  int64_t startXRow = curRowsStart / this->k_;
+  int64_t endXRow = (curRowsStart + this->coreRows_ - 1) / this->k_;
+  pipe->InitBuffer(xCopyInQueue_, bufferNum, AlignBytes(this->cols_, sizeof(T)) * (endXRow - startXRow + 1));
+  pipe->InitBuffer(expandedRowIdxCopyOutQueue_, bufferNum, buffSize);
+  pipe->InitBuffer(expandedExpertIdxCopyOutQueue_, bufferNum, buffSize);
+  pipe->InitBuffer(expertTokensCopyOutQueue_, bufferNum, AlignBytes(this->expertNum, sizeof(int32_t)));
+  pipe->InitBuffer(expandDstToSrcRowQueue_, bufferNum, buffSize);
+  pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize * kvFactor);
+  pipe->InitBuffer(tempBuffer, buffSize * kvFactor);
+  pipe->InitBuffer(sortedBuffer, buffSize * kvFactor);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::Process() {
+  if (this->blockIdx_ < this->needCoreNum_) {
+    CopyIn();
+    SortCompute();
+    if (this->blockIdx_ == 0) {
+      CopyOutIdx();
+    }
+    if (this->blockIdx_ == this->needCoreNum_ - 1 && this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) {
+      ComputeExpertTokenCountOrCumsum();
+    } else {
+      CopyOutEmpty();
+    }
+    CopyOutX();
+  }
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // INNER_MOE_V2_FULL_LOAD_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_mrgsort.h

@@ -0,0 +1,189 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_mrgsort.h
+ * \brief
+ */
+#ifndef INNER_MOE_V2_MRGSORT_H
+#define INNER_MOE_V2_MRGSORT_H
+
+#include "moe_v2_common.h"
+#include "kernel_operator.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+struct MoeV2MrgsortParam {
+  int64_t perListElements;
+  int64_t lastListElements;
+  int64_t oneLoopMaxElements;
+};
+
+class MoeV2Mrgsort {
+ public:
+  __aicore__ inline MoeV2Mrgsort(){};
+  __aicore__ inline void Init(MoeV2MrgsortParam* param);
+  __aicore__ inline void Process();
+  __aicore__ inline void SetInput(GlobalTensor<float>& gmInput, LocalTensor<float>& ubInput);
+  __aicore__ inline void SetOutput(GlobalTensor<float>& gmOutput, LocalTensor<float>& ubOutput);
+
+ private:
+  __aicore__ inline void CopyIn();
+  __aicore__ inline void UpdateMrgParam();
+  __aicore__ inline void MrgsortCompute();
+  __aicore__ inline void UpdateSortInfo();
+  __aicore__ inline void CopyOut();
+  __aicore__ inline void ClearCache();
+
+ private:
+  MoeV2MrgsortParam* param = nullptr;
+
+  GlobalTensor<float> gmInputs[4];
+  GlobalTensor<float> gmOutput;
+
+  LocalTensor<float> ubInputs[4];
+  LocalTensor<float> ubOutput;
+
+  int64_t listNum{0};
+  int64_t remainListNum{0};
+  int64_t outOffset{0};
+  int64_t offsets[4];
+  int64_t listRemainElements[4];
+  int64_t lengths[4];
+  int64_t allRemainElements{0};
+  int64_t curLoopSortedNum{0};
+
+  // for MrgSort
+  uint16_t validBitTail{0};
+  uint16_t elementCountListTail[4];
+  uint32_t listSortedNums[4];
+  LocalTensor<float> tmpUbInputs[4];
+};
+
+__aicore__ inline void MoeV2Mrgsort::ClearCache() {
+  this->listNum = 0;
+  this->allRemainElements = 0;
+  this->outOffset = 0;
+}
+
+__aicore__ inline void MoeV2Mrgsort::SetInput(GlobalTensor<float>& gmInput, LocalTensor<float>& ubInput) {
+  this->gmInputs[listNum] = gmInput;
+  this->ubInputs[listNum] = ubInput;
+  this->listNum += 1;
+}
+
+__aicore__ inline void MoeV2Mrgsort::SetOutput(GlobalTensor<float>& gmOutput, LocalTensor<float>& ubOutput) {
+  this->gmOutput = gmOutput;
+  this->ubOutput = ubOutput;
+}
+
+__aicore__ inline void MoeV2Mrgsort::UpdateMrgParam() {
+  if (this->remainListNum == MERGE_LIST_TWO) {
+    elementCountListTail[MERGE_LIST_IDX_TWO] = 0;
+    elementCountListTail[MERGE_LIST_IDX_THREE] = 0;
+    validBitTail = 0b0011;
+  } else if (this->remainListNum == MERGE_LIST_THREE) {
+    elementCountListTail[MERGE_LIST_IDX_THREE] = 0;
+    validBitTail = 0b0111;
+  } else if (this->remainListNum == MERGE_LIST_FOUR) {
+    validBitTail = 0b1111;
+  } else {
+    validBitTail = 0b0001;
+  }
+}
+
+__aicore__ inline void MoeV2Mrgsort::CopyIn() {
+  this->remainListNum = 0;
+  SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
+  for (int64_t i = 0, j = 0; i < listNum; i++) {
+    lengths[i] = Min(param->oneLoopMaxElements, listRemainElements[i]);
+    if (lengths[i] > 0) {
+      DataCopy(this->ubInputs[i], this->gmInputs[i][offsets[i]], Align(GetSortLen<float>(lengths[i]), sizeof(float)));
+      tmpUbInputs[j] = this->ubInputs[i];
+      elementCountListTail[j] = lengths[i];
+      this->remainListNum += 1;
+      j++;
+    }
+  }
+}
+
+__aicore__ inline void MoeV2Mrgsort::MrgsortCompute() {
+  SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
+  if (this->remainListNum == MERGE_LIST_TWO) {
+    MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]);
+    MrgSort<float, true>(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+  } else if (this->remainListNum == MERGE_LIST_THREE) {
+    MrgSortSrcList sortListTail =
+        MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], tmpUbInputs[0]);
+    MrgSort<float, true>(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+  } else if (this->remainListNum == MERGE_LIST_FOUR) {
+    MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO],
+                                                 tmpUbInputs[MERGE_LIST_IDX_THREE]);
+    MrgSort<float, true>(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+  } else {
+    DataCopy(this->ubOutput, this->tmpUbInputs[0], Align(GetSortLen<float>(elementCountListTail[0]), sizeof(float)));
+    listSortedNums[0] = elementCountListTail[0];
+  }
+}
+
+__aicore__ inline void MoeV2Mrgsort::UpdateSortInfo() {
+  curLoopSortedNum = 0;
+  for (int64_t i = 0, j = 0; i < listNum; i++) {
+    if (lengths[i] > 0) {
+      // update remain size
+      listRemainElements[i] -= listSortedNums[j];
+      allRemainElements -= listSortedNums[j];
+      // update offset
+      offsets[i] += GetSortOffset<float>(listSortedNums[j]);
+      // update current loop sorted nums
+      curLoopSortedNum += listSortedNums[j];
+      j += 1;
+    }
+  }
+}
+
+__aicore__ inline void MoeV2Mrgsort::CopyOut() {
+  DataCopyParams intriParams;
+  intriParams.blockCount = 1;
+  intriParams.blockLen = GetSortLen<float>(curLoopSortedNum) * sizeof(float);
+  SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+  DataCopyPad(this->gmOutput[outOffset], this->ubOutput, intriParams);
+  outOffset += GetSortLen<float>(curLoopSortedNum);
+}
+
+__aicore__ inline void MoeV2Mrgsort::Init(MoeV2MrgsortParam* param) {
+  this->param = param;
+  this->remainListNum = listNum;
+
+  for (int64_t i = 0; i < listNum; i++) {
+    offsets[i] = GetSortOffset<float>(param->perListElements * i);
+    if (i == listNum - 1) {
+      listRemainElements[i] = param->lastListElements;
+    } else {
+      listRemainElements[i] = param->perListElements;
+    }
+    allRemainElements += listRemainElements[i];
+  }
+}
+
+__aicore__ inline void MoeV2Mrgsort::Process() {
+  for (; allRemainElements > 0;) {
+    CopyIn();
+    UpdateMrgParam();
+    MrgsortCompute();
+    UpdateSortInfo();
+    CopyOut();
+  }
+
+  ClearCache();
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // INNER_MOE_V2_MRGSORT_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_mrgsort_out.h

@@ -0,0 +1,213 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_mrgsort_out.h
+ * \brief
+ */
+#ifndef INNER_MOE_V2_MRGSORT_OUT_H
+#define INNER_MOE_V2_MRGSORT_OUT_H
+
+#include "moe_v2_mrgsort.h"
+#include "kernel_operator.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+class MoeV2MrgsortOut {
+ public:
+  __aicore__ inline MoeV2MrgsortOut(){};
+  __aicore__ inline void Init(MoeV2MrgsortParam* param, TPipe* tPipe);
+  __aicore__ inline void Process();
+  __aicore__ inline void SetInput(GlobalTensor<float>& gmInput, LocalTensor<float>& ubInput);
+  __aicore__ inline void SetOutput(GlobalTensor<int32_t>& gmOutput1, GlobalTensor<int32_t>& gmOutput2,
+                                   LocalTensor<float>& ubOutput1, LocalTensor<float>& ubOutput2);
+  __aicore__ inline void SetBuffer(LocalTensor<float>& tempBuffer);
+
+ private:
+  __aicore__ inline void CopyIn();
+  __aicore__ inline void UpdateMrgParam();
+  __aicore__ inline void MrgsortCompute();
+  __aicore__ inline void UpdateSortInfo();
+  __aicore__ inline void Extract();
+  __aicore__ inline void CopyOut();
+  __aicore__ inline void ClearCache();
+
+ private:
+  MoeV2MrgsortParam* param = nullptr;
+
+  GlobalTensor<float> gmInputs[4];
+  GlobalTensor<int32_t> gmOutput1;
+  GlobalTensor<int32_t> gmOutput2;
+
+  LocalTensor<float> ubInputs[4];
+  LocalTensor<float> tempBuffer;
+
+  // for extract
+  LocalTensor<float> ubOutput1;
+  LocalTensor<uint32_t> ubOutput2;
+
+  // for copy out
+  LocalTensor<int32_t> ubOutputInt1;
+  LocalTensor<int32_t> ubOutputInt2;
+
+  int64_t listNum{0};
+  int64_t remainListNum{0};
+  int64_t outOffset{0};
+  int64_t offsets[4];
+  int64_t listRemainElements[4];
+  int64_t lengths[4];
+  int64_t allRemainElements{0};
+  int64_t curLoopSortedNum{0};
+
+  // for MrgSort
+  uint16_t validBitTail;
+  uint16_t elementCountListTail[4];
+  uint32_t listSortedNums[4];
+  LocalTensor<float> tmpUbInputs[4];
+};
+
+__aicore__ inline void MoeV2MrgsortOut::ClearCache() {
+  this->listNum = 0;
+  this->allRemainElements = 0;
+  this->outOffset = 0;
+}
+
+__aicore__ inline void MoeV2MrgsortOut::SetInput(GlobalTensor<float>& gmInput, LocalTensor<float>& ubInput) {
+  this->gmInputs[listNum] = gmInput;
+  this->ubInputs[listNum] = ubInput;
+  this->listNum += 1;
+}
+
+__aicore__ inline void MoeV2MrgsortOut::SetOutput(GlobalTensor<int32_t>& gmOutput1, GlobalTensor<int32_t>& gmOutput2,
+                                                  LocalTensor<float>& ubOutput1, LocalTensor<float>& ubOutput2) {
+  this->gmOutput1 = gmOutput1;
+  this->ubOutput1 = ubOutput1;
+  this->ubOutputInt1 = ubOutput1.ReinterpretCast<int32_t>();
+
+  this->gmOutput2 = gmOutput2;
+  this->ubOutput2 = ubOutput2.ReinterpretCast<uint32_t>();
+  this->ubOutputInt2 = ubOutput2.ReinterpretCast<int32_t>();
+}
+
+__aicore__ inline void MoeV2MrgsortOut::SetBuffer(LocalTensor<float>& tempBuffer) {
+  this->tempBuffer = tempBuffer;
+}
+
+__aicore__ inline void MoeV2MrgsortOut::UpdateMrgParam() {
+  if (this->remainListNum == MERGE_LIST_TWO) {
+    elementCountListTail[MERGE_LIST_IDX_TWO] = 0;
+    elementCountListTail[MERGE_LIST_IDX_THREE] = 0;
+    validBitTail = 0b0011;
+  } else if (this->remainListNum == MERGE_LIST_THREE) {
+    elementCountListTail[MERGE_LIST_IDX_THREE] = 0;
+    validBitTail = 0b0111;
+  } else if (this->remainListNum == MERGE_LIST_FOUR) {
+    validBitTail = 0b1111;
+  } else {
+    validBitTail = 0b0001;
+  }
+}
+
+__aicore__ inline void MoeV2MrgsortOut::CopyIn() {
+  this->remainListNum = 0;
+  SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
+  for (int64_t i = 0, j = 0; i < listNum; i++) {
+    lengths[i] = Min(param->oneLoopMaxElements, listRemainElements[i]);
+    if (lengths[i] > 0) {
+      DataCopy(this->ubInputs[i], this->gmInputs[i][offsets[i]], Align(GetSortLen<float>(lengths[i]), sizeof(float)));
+      tmpUbInputs[j] = this->ubInputs[i];
+      elementCountListTail[j] = lengths[i];
+      this->remainListNum += 1;
+      j++;
+    }
+  }
+}
+
+__aicore__ inline void MoeV2MrgsortOut::MrgsortCompute() {
+  SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
+  if (this->remainListNum == MERGE_LIST_TWO) {
+    MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]);
+    MrgSort<float, true>(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+  } else if (this->remainListNum == MERGE_LIST_THREE) {
+    MrgSortSrcList sortListTail =
+        MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], tmpUbInputs[0]);
+    MrgSort<float, true>(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+  } else if (this->remainListNum == MERGE_LIST_FOUR) {
+    MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO],
+                                                 tmpUbInputs[MERGE_LIST_IDX_THREE]);
+    MrgSort<float, true>(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+  } else {
+    DataCopy(this->tempBuffer, this->tmpUbInputs[0], Align(GetSortLen<float>(elementCountListTail[0]), sizeof(float)));
+    listSortedNums[0] = elementCountListTail[0];
+  }
+}
+
+__aicore__ inline void MoeV2MrgsortOut::UpdateSortInfo() {
+  curLoopSortedNum = 0;
+  for (int64_t i = 0, j = 0; i < listNum; i++) {
+    if (lengths[i] > 0) {
+      // update remain size
+      listRemainElements[i] -= listSortedNums[j];
+      allRemainElements -= listSortedNums[j];
+      // update offset
+      offsets[i] += GetSortOffset<float>(listSortedNums[j]);
+      // update current loop sorted nums
+      curLoopSortedNum += listSortedNums[j];
+      j += 1;
+    }
+  }
+}
+
+__aicore__ inline void MoeV2MrgsortOut::Extract() {
+  AscendC::Extract(this->ubOutput1, this->ubOutput2, this->tempBuffer, Ceil(curLoopSortedNum, ONE_REPEAT_SORT_NUM));
+  AscendC::PipeBarrier<PIPE_V>();
+  Muls(this->ubOutput1, this->ubOutput1, (float)-1, Align(curLoopSortedNum, sizeof(float)));
+  AscendC::PipeBarrier<PIPE_V>();
+  Cast(this->ubOutputInt1, this->ubOutput1, RoundMode::CAST_ROUND, Align(curLoopSortedNum, sizeof(float)));
+}
+
+__aicore__ inline void MoeV2MrgsortOut::CopyOut() {
+  DataCopyParams intriParams;
+  intriParams.blockCount = 1;
+  intriParams.blockLen = curLoopSortedNum * sizeof(int32_t);
+  SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+  DataCopyPad(this->gmOutput1[outOffset], this->ubOutputInt1, intriParams);
+  DataCopyPad(this->gmOutput2[outOffset], this->ubOutputInt2, intriParams);
+  outOffset += curLoopSortedNum;
+}
+
+__aicore__ inline void MoeV2MrgsortOut::Init(MoeV2MrgsortParam* param, TPipe* tPipe) {
+  this->param = param;
+  this->allRemainElements = 0;
+  for (int64_t i = 0; i < listNum; i++) {
+    offsets[i] = GetSortOffset<float>(param->perListElements * i);
+    if (i == listNum - 1) {
+      listRemainElements[i] = param->lastListElements;
+    } else {
+      listRemainElements[i] = param->perListElements;
+    }
+    allRemainElements += listRemainElements[i];
+  }
+}
+
+__aicore__ inline void MoeV2MrgsortOut::Process() {
+  for (; allRemainElements > 0;) {
+    CopyIn();
+    UpdateMrgParam();
+    MrgsortCompute();
+    UpdateSortInfo();
+    Extract();
+    CopyOut();
+  }
+  ClearCache();
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // INNER_MOE_V2_MRGSORT_OUT_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_base.h

@@ -0,0 +1,70 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_sort_base.h
+ * \brief
+ */
+#ifndef INNER_MOE_V2_SORT_BASE_H
+#define INNER_MOE_V2_SORT_BASE_H
+
+#include "kernel_operator.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+class MoeV2SortBase {
+ public:
+  __aicore__ inline MoeV2SortBase(){};
+
+ protected:
+  __aicore__ inline void SyncAll();
+
+ protected:
+  TPipe* pipe;
+  TQue<QuePosition::VECIN, 1> sortDataCopyInQueue;
+  TQue<QuePosition::VECOUT, 1> sortDataCopyOutQueue;
+  TBuf<TPosition::VECCALC> tempBuffer;
+  TBuf<TPosition::VECCALC> sortedBuffer;
+
+  GlobalTensor<int32_t> expertIdxGm;
+  GlobalTensor<int32_t> sortedexpertIdxGm;
+  GlobalTensor<int32_t> expandDstToSrcRowGm;
+  GlobalTensor<int32_t> expertTokensCountOrCumsumGm;
+  GlobalTensor<int32_t> expertTokensBeforeCapacityGm;
+
+  int64_t tileLength;
+  int64_t bufferNum = 1;
+  int64_t totalLength;
+  int64_t coreNum;
+  int64_t n;
+  int64_t k;
+  int64_t existRowIdx;
+  int64_t expertNum;
+  int64_t expertTokensCountOrCumsumFlag = 0;
+  int64_t expertTokensBeforeCapacityFlag = 0;
+
+  static constexpr int64_t SYNC_GM_NUM = 2;
+  static constexpr int64_t WORK_GM_NUM = 2;
+  static constexpr int64_t DST_BLK_STRIDE = 1;
+  static constexpr int64_t DST_REP_STRIDE = 8;
+};
+
+__aicore__ inline void MoeV2SortBase::SyncAll() {
+  if (coreNum == 1) {
+    return;
+  }
+#ifndef __CCE_KT_TEST__
+  AscendC::SyncAll();
+#endif
+}
+
+}  // namespace MoeInitRoutingQuantV2
+#endif  // INNER_MOE_V2_SORT_BASE_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_multi_core.h

@@ -0,0 +1,373 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_sort_multi_core.h
+ * \brief
+ */
+#ifndef INNER_MOE_V2_VBS_ONE_CORE_H
+#define INNER_MOE_V2_VBS_ONE_CORE_H
+
+#include "moe_v2_sort_base.h"
+#include "moe_v2_mrgsort.h"
+#include "moe_v2_mrgsort_out.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+class MoeV2SortMultiCore : public MoeV2SortBase {
+ public:
+  __aicore__ inline MoeV2SortMultiCore(){};
+  template <typename TilingData>
+  __aicore__ inline void Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity,
+                              GM_ADDR workspace, const TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void VBSProcess();
+  __aicore__ inline void UBSortProcess(int64_t progress, int64_t size, int64_t sortNum);
+  __aicore__ inline void OneCoreVMSProcess(int64_t listNum, int64_t perListElements, int64_t lastListElements);
+  __aicore__ inline void VMSProcess();
+  __aicore__ inline void SortOutProcess();
+  __aicore__ inline void VBSCopyIn(int64_t progress, int64_t size, int64_t sortNum);
+  __aicore__ inline void UBSortCompute(int64_t progress, int64_t size, int64_t sortNum);
+  __aicore__ inline void VBSCopyOut(int64_t progress, int64_t size, int64_t sortNum);
+  __aicore__ inline void InitMoeMrgSort(MoeV2Mrgsort* sorter, int64_t listNum, int64_t coreOffset, int64_t loopOffset);
+  __aicore__ inline void InitMoeMrgSortOut(MoeV2MrgsortOut* sorter, int64_t listNum, int64_t coreOffset);
+  __aicore__ inline void InitExpertTokensGlobalMemory();
+
+ private:
+  GlobalTensor<float> workspaceGms[2];
+
+  const InnerMoeV2VBSComputeTilingData* vbsTilingData;
+  const InnerMoeV2VMSMiddleComputeTilingData* vmsTilingData;
+  const InnerMoeV2SortOutComputeTilingData* sortOutTilingData;
+
+  // for MoeMrgsort
+  MoeV2Mrgsort mrgsorter;
+  MoeV2MrgsortParam mrgsortParam;
+
+  int64_t coreNum;
+  int64_t blockIdx;
+  int64_t srcWsIndex = 0;
+
+  int64_t listNum;
+  int64_t perListElements;
+  int64_t lastListElements;
+
+  int64_t sortTotalLength;
+  int64_t sortCoreLoops;
+  int64_t sortCoreLoopElements;
+  int64_t sortCoreLastLoopElements;
+
+  int64_t perCoreExpert;
+  int64_t needInitExpertCore;
+  int64_t currentCoreExpert;
+
+  static constexpr int64_t MAX_MRGSORT_LIST = 4;
+};
+
+__aicore__ inline void MoeV2SortMultiCore::InitExpertTokensGlobalMemory() {
+  if (this->blockIdx < this->needInitExpertCore) {
+    if (this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) {
+      InitGlobalMemory(expertTokensCountOrCumsumGm, currentCoreExpert, 0);
+    }
+    if (this->expertTokensBeforeCapacityFlag == EXERPT_TOKENS_BEFORE_CAPACITY) {
+      InitGlobalMemory(expertTokensBeforeCapacityGm, currentCoreExpert, 0);
+    }
+  }
+}
+
+__aicore__ inline void MoeV2SortMultiCore::VBSCopyIn(int64_t progress, int64_t size, int64_t sortNum) {
+  LocalTensor<int32_t> inLocal = sortDataCopyInQueue.AllocTensor<int32_t>();
+  int64_t inOffset = progress * sortCoreLoopElements;
+  DataCopyExtParams dataCopyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(size * sizeof(int32_t)), 0, 0, 0};
+  DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+  DataCopyPad(inLocal[0], expertIdxGm[inOffset], dataCopyParams, dataCopyPadParams);
+
+  LocalTensor<int32_t> rowIdxLocal = inLocal[sortNum];
+  int64_t startValue = this->blockIdx * this->vbsTilingData->perCoreElements + inOffset;
+  SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+  ArithProgression<int32_t>(rowIdxLocal, startValue, 1, size);
+  sortDataCopyInQueue.EnQue(inLocal);
+}
+
+__aicore__ inline void MoeV2SortMultiCore::UBSortCompute(int64_t progress, int64_t size, int64_t sortNum) {
+  LocalTensor<int32_t> inLocal = sortDataCopyInQueue.DeQue<int32_t>();
+  LocalTensor<int32_t> expertForSourceRowLocal = inLocal[0];
+  LocalTensor<float> expertForSourceRowLocalFp32;
+
+  expertForSourceRowLocalFp32 = expertForSourceRowLocal.ReinterpretCast<float>();
+  Cast(expertForSourceRowLocalFp32, expertForSourceRowLocal, RoundMode::CAST_ROUND, sortNum);
+  AscendC::PipeBarrier<PIPE_V>();
+  Muls(expertForSourceRowLocalFp32, expertForSourceRowLocalFp32, (float)-1, sortNum);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  int64_t duplicateNum = size % ONE_REPEAT_SORT_NUM;
+  if (duplicateNum > 0) {
+    int duplicateIndex = size - duplicateNum;
+    uint64_t mask0 = UINT64_MAX;
+    mask0 = mask0 << duplicateNum;
+    mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
+    uint64_t mask[2] = {mask0, 0};
+    Duplicate(expertForSourceRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  LocalTensor<float> concatLocal = expertForSourceRowLocalFp32;
+  LocalTensor<float> sortedLocal = sortedBuffer.Get<float>(GetSortLen<float>(sortNum));
+  LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
+  LocalTensor<uint32_t> sourceRowLocal;
+  sourceRowLocal = inLocal[sortNum].ReinterpretCast<uint32_t>();
+  Sort<float, true>(outLocal, concatLocal, sourceRowLocal, sortedLocal, sortNum / ONE_REPEAT_SORT_NUM);
+
+  sortDataCopyOutQueue.EnQue<float>(outLocal);
+  sortDataCopyInQueue.FreeTensor(inLocal);
+}
+
+__aicore__ inline void MoeV2SortMultiCore::VBSCopyOut(int64_t progress, int64_t size, int64_t sortNum) {
+  LocalTensor<float> outLocal = sortDataCopyOutQueue.DeQue<float>();
+  DataCopy(workspaceGms[0][this->blockIdx * GetSortLen<float>(this->vbsTilingData->perCoreElements) +
+                           GetSortLen<float>(progress * sortCoreLoopElements)],
+           outLocal, Align(GetSortLen<float>(size), sizeof(float)));
+  sortDataCopyOutQueue.FreeTensor(outLocal);
+}
+
+__aicore__ inline void MoeV2SortMultiCore::InitMoeMrgSort(MoeV2Mrgsort* sorter, int64_t listNum, int64_t coreOffset,
+                                                          int64_t loopOffset) {
+  GlobalTensor<float> srcWsGm = workspaceGms[srcWsIndex][blockIdx * coreOffset + loopOffset];
+  LocalTensor<float> inLocal = sortDataCopyInQueue.AllocTensor<float>();
+  LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
+  for (int64_t i = 0; i < listNum; i++) {
+    LocalTensor<float> inLocalT = inLocal[GetSortLen<float>(this->sortOutTilingData->oneLoopMaxElements) * i];
+    sorter->SetInput(srcWsGm, inLocalT);
+  }
+  GlobalTensor<float> dstWsGm = workspaceGms[1 - srcWsIndex][blockIdx * coreOffset + loopOffset];
+  sorter->SetOutput(dstWsGm, outLocal);
+  sortDataCopyInQueue.FreeTensor(inLocal);
+  sortDataCopyOutQueue.FreeTensor(outLocal);
+}
+
+__aicore__ inline void MoeV2SortMultiCore::InitMoeMrgSortOut(MoeV2MrgsortOut* sorter, int64_t listNum,
+                                                             int64_t coreOffset) {
+  GlobalTensor<float> srcWsGm = workspaceGms[srcWsIndex];
+  LocalTensor<float> inLocal = sortDataCopyInQueue.AllocTensor<float>();
+  LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
+
+  for (int64_t i = 0; i < listNum; i++) {
+    LocalTensor<float> inLocalT = inLocal[GetSortLen<float>(this->sortOutTilingData->oneLoopMaxElements) * i];
+    sorter->SetInput(srcWsGm, inLocalT);
+  }
+
+  LocalTensor<float> outLocalV = outLocal[this->sortOutTilingData->oneLoopMaxElements * MAX_MRGSORT_LIST];
+  sorter->SetOutput(this->sortedexpertIdxGm, this->expandDstToSrcRowGm, outLocal, outLocalV);
+
+  LocalTensor<float> tempBuffer =
+      sortedBuffer.Get<float>(GetSortLen<float>(this->sortOutTilingData->oneLoopMaxElements) * MAX_MRGSORT_LIST);
+  sorter->SetBuffer(tempBuffer);
+  sortDataCopyInQueue.FreeTensor(inLocal);
+  sortDataCopyOutQueue.FreeTensor(outLocal);
+}
+
+__aicore__ inline void MoeV2SortMultiCore::OneCoreVMSProcess(int64_t listNum, int64_t perListElements,
+                                                             int64_t lastListElements) {
+  int64_t coreOffset = GetSortLen<float>(this->vbsTilingData->perCoreElements);
+  mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements;
+
+  for (int64_t i = 0; listNum >= 1; i++) {
+    int64_t loops = (listNum + MAX_MRGSORT_LIST - 1) / MAX_MRGSORT_LIST;
+    int64_t remainListNum = listNum - (loops - 1) * MAX_MRGSORT_LIST;
+
+    mrgsortParam.perListElements = perListElements;
+    mrgsortParam.lastListElements = perListElements;
+
+    int64_t loopOffset = GetSortLen<float>(mrgsortParam.perListElements * MAX_MRGSORT_LIST);
+    for (int64_t loop = 0; loop < loops - 1; loop++) {
+      InitMoeMrgSort(&mrgsorter, MAX_MRGSORT_LIST, coreOffset, loop * loopOffset);
+      mrgsorter.Init(&mrgsortParam);
+      mrgsorter.Process();
+    }
+
+    mrgsortParam.perListElements = perListElements;
+    mrgsortParam.lastListElements = lastListElements;
+    InitMoeMrgSort(&mrgsorter, remainListNum, coreOffset, (loops - 1) * loopOffset);
+    mrgsorter.Init(&mrgsortParam);
+    mrgsorter.Process();
+
+    listNum = loops;
+    lastListElements = perListElements * (remainListNum - 1) + lastListElements;
+    perListElements = perListElements * MAX_MRGSORT_LIST;
+    srcWsIndex = (srcWsIndex + 1) % WORK_GM_NUM;
+
+    if (loops == 1) {
+      break;
+    }
+  }
+}
+
+__aicore__ inline void MoeV2SortMultiCore::UBSortProcess(int64_t progress, int64_t size, int64_t sortNum) {
+  VBSCopyIn(progress, size, sortNum);
+  UBSortCompute(progress, size, sortNum);
+  VBSCopyOut(progress, size, sortNum);
+}
+
+__aicore__ inline void MoeV2SortMultiCore::VBSProcess() {
+  if (this->blockIdx < this->vbsTilingData->needCoreNum) {
+    int64_t sortNum = Ceil(sortCoreLoopElements, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
+    for (int64_t loop = 0; loop < sortCoreLoops - 1; loop++) {
+      UBSortProcess(loop, sortCoreLoopElements, sortNum);
+    }
+
+    sortNum = Ceil(sortCoreLastLoopElements, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
+    UBSortProcess(sortCoreLoops - 1, sortCoreLastLoopElements, sortNum);
+    if (sortCoreLoops > 1) {
+      OneCoreVMSProcess(sortCoreLoops, sortCoreLoopElements, sortCoreLastLoopElements);
+    }
+  }
+#ifndef __CCE_KT_TEST__
+  AscendC::SyncAll();
+#endif
+}
+
+__aicore__ inline void MoeV2SortMultiCore::VMSProcess() {
+  int64_t currentStageNeedCoreNum = this->vmsTilingData->needCoreNum;
+  perListElements = this->vbsTilingData->perCoreElements;
+  lastListElements = this->vbsTilingData->lastCoreElements;
+  listNum = this->vbsTilingData->needCoreNum;
+
+  for (; listNum > MAX_MRGSORT_LIST;) {
+    currentStageNeedCoreNum = Ceil(listNum, MAX_MRGSORT_LIST);
+    int64_t coreOffset = GetSortLen<float>(perListElements * MAX_MRGSORT_LIST);
+    int64_t remainListNum = listNum - (currentStageNeedCoreNum - 1) * MAX_MRGSORT_LIST;
+
+    if (this->blockIdx < currentStageNeedCoreNum - 1) {
+      mrgsortParam.perListElements = perListElements;
+      mrgsortParam.lastListElements = perListElements;
+      mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements;
+      InitMoeMrgSort(&mrgsorter, MAX_MRGSORT_LIST, coreOffset, 0);
+      mrgsorter.Init(&mrgsortParam);
+      mrgsorter.Process();
+    } else if (this->blockIdx == currentStageNeedCoreNum - 1) {
+      mrgsortParam.perListElements = perListElements;
+      mrgsortParam.lastListElements = lastListElements;
+      mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements;
+      InitMoeMrgSort(&mrgsorter, remainListNum, coreOffset, 0);
+      mrgsorter.Init(&mrgsortParam);
+      mrgsorter.Process();
+    }
+    listNum = currentStageNeedCoreNum;
+    currentStageNeedCoreNum = Ceil(listNum, MAX_MRGSORT_LIST);
+    srcWsIndex = (srcWsIndex + 1) % WORK_GM_NUM;
+
+    lastListElements = perListElements * (remainListNum - 1) + lastListElements;
+    perListElements = perListElements * MAX_MRGSORT_LIST;
+#ifndef __CCE_KT_TEST__
+    AscendC::SyncAll();
+#endif
+  }
+}
+
+__aicore__ inline void MoeV2SortMultiCore::SortOutProcess() {
+  if (this->blockIdx < 1) {
+    mrgsortParam.perListElements = perListElements;
+    mrgsortParam.lastListElements = lastListElements;
+    mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements;
+
+    MoeV2MrgsortOut sorter;
+    InitMoeMrgSortOut(&sorter, listNum, GetSortLen<float>(perListElements));
+    sorter.Init(&mrgsortParam, pipe);
+    sorter.Process();
+  }
+#ifndef __CCE_KT_TEST__
+  AscendC::SyncAll();
+#endif
+}
+
+template <typename TilingData>
+__aicore__ inline void MoeV2SortMultiCore::Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum,
+                                                GM_ADDR expertTokensBeforeCapacity, GM_ADDR workspace,
+                                                const TilingData* tilingData, TPipe* tPipe) {
+  this->totalLength = tilingData->n * tilingData->k;
+  this->coreNum = tilingData->coreNum;
+  this->vbsTilingData = &(tilingData->vbsComputeParamsOp);
+  this->vmsTilingData = &(tilingData->vmsMiddleComputeParamsOp);
+  this->sortOutTilingData = &(tilingData->sortOutComputeParamsOp);
+
+  this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+  this->tileLength = this->vbsTilingData->perCorePerLoopElements;
+  this->sortTotalLength = this->vbsTilingData->perCoreElements;
+  if (this->blockIdx == tilingData->vbsComputeParamsOp.needCoreNum - 1) {
+    this->tileLength = this->vbsTilingData->lastCorePerLoopElements;
+    this->sortTotalLength = this->vbsTilingData->lastCoreElements;
+  }
+  this->n = tilingData->n;
+  this->k = tilingData->k;
+  this->expertNum = tilingData->expertNum;
+  this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag;
+  this->expertTokensBeforeCapacityFlag = tilingData->expertTokensBeforeCapacityFlag;
+
+  // VBS param init
+  if (this->blockIdx == this->vbsTilingData->needCoreNum - 1) {
+    sortCoreLoops = this->vbsTilingData->lastCoreLoops;
+    sortCoreLoopElements = this->vbsTilingData->lastCorePerLoopElements;
+    sortCoreLastLoopElements = this->vbsTilingData->lastCoreLastLoopElements;
+  } else {
+    sortCoreLoops = this->vbsTilingData->perCoreLoops;
+    sortCoreLoopElements = this->vbsTilingData->perCorePerLoopElements;
+    sortCoreLastLoopElements = this->vbsTilingData->perCoreLastLoopElements;
+  }
+
+  this->pipe = tPipe;
+  expertIdxGm.SetGlobalBuffer(
+      (__gm__ int32_t*)expertIdx + this->blockIdx * tilingData->vbsComputeParamsOp.perCoreElements,
+      this->sortTotalLength);
+  sortedexpertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(workspace),
+                                    Align(this->totalLength, sizeof(int32_t)));
+  expandDstToSrcRowGm.SetGlobalBuffer(
+      reinterpret_cast<__gm__ int32_t*>(workspace) + Align(this->totalLength, sizeof(int32_t)),
+      Align(this->totalLength, sizeof(int32_t)));
+
+  this->perCoreExpert = Align((this->expertNum + this->coreNum - 1) / this->coreNum, sizeof(int32_t));
+  this->needInitExpertCore = (this->expertNum + this->perCoreExpert - 1) / this->perCoreExpert;
+  this->currentCoreExpert = this->perCoreExpert;
+  if (this->blockIdx == needInitExpertCore - 1) {
+    this->currentCoreExpert = this->expertNum - (this->needInitExpertCore - 1) * this->perCoreExpert;
+  }
+  if (this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) {
+    expertTokensCountOrCumsumGm.SetGlobalBuffer(
+        (__gm__ int32_t*)expertTokensCountOrCumsum + this->blockIdx * this->perCoreExpert, this->currentCoreExpert);
+  }
+  if (this->expertTokensBeforeCapacityFlag == EXERPT_TOKENS_BEFORE_CAPACITY) {
+    expertTokensBeforeCapacityGm.SetGlobalBuffer(
+        (__gm__ int32_t*)expertTokensBeforeCapacity + this->blockIdx * this->perCoreExpert, this->currentCoreExpert);
+  }
+  // key and value
+  int64_t kvFactor = 2;
+  workspaceGms[0].SetGlobalBuffer((__gm__ float*)workspace + Align(this->totalLength, sizeof(int32_t)) * 2,
+                                  Align(this->totalLength, sizeof(int32_t)) * kvFactor);
+  workspaceGms[1].SetGlobalBuffer((__gm__ float*)workspace + Align(this->totalLength, sizeof(int32_t)) * (kvFactor + 2),
+                                  Align(this->totalLength, sizeof(int32_t)) * kvFactor);
+
+  int64_t bufferSize = Ceil(Max(this->sortOutTilingData->oneLoopMaxElements * MAX_MRGSORT_LIST, sortCoreLoopElements),
+                            ONE_REPEAT_SORT_NUM) *
+                       ONE_REPEAT_SORT_NUM * sizeof(int32_t) * kvFactor;
+  pipe->InitBuffer(sortDataCopyInQueue, bufferNum, bufferSize);
+  pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, bufferSize);
+  pipe->InitBuffer(sortedBuffer, bufferSize);
+}
+
+__aicore__ inline void MoeV2SortMultiCore::Process() {
+  InitExpertTokensGlobalMemory();
+  VBSProcess();
+  VMSProcess();
+  SortOutProcess();
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // INNER_MOE_V2_VBS_ONE_CORE_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_sort_one_core.h

@@ -0,0 +1,162 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_sort_one_core.h
+ * \brief
+ */
+#ifndef INNER_MOE_V2_SORT_ONE_CORE_H
+#define INNER_MOE_V2_SORT_ONE_CORE_H
+
+#include "moe_v2_mrgsort.h"
+#include "moe_v2_sort_base.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+class MoeV2SortOneCore : public MoeV2SortBase {
+ public:
+  __aicore__ inline MoeV2SortOneCore(){};
+  template <typename TilingData>
+  __aicore__ inline void Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity,
+                              GM_ADDR workspace, const TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void CopyIn();
+  __aicore__ inline void SortCompute();
+  __aicore__ inline void CopyOut();
+
+ private:
+  int64_t sortNum;
+  int64_t blockIdx;
+};
+
+__aicore__ inline void MoeV2SortOneCore::CopyIn() {
+  LocalTensor<int32_t> inLocal = sortDataCopyInQueue.AllocTensor<int32_t>();
+  DataCopyExtParams dataCopyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(this->totalLength * sizeof(int32_t)),
+                                   0, 0, 0};
+  DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+  DataCopyPad(inLocal[0], expertIdxGm, dataCopyParams, dataCopyPadParams);
+
+  LocalTensor<int32_t> rowIdxLocal = inLocal[this->sortNum];
+  ArithProgression<int32_t>(rowIdxLocal, 0, 1, this->sortNum);
+  sortDataCopyInQueue.EnQue(inLocal);
+}
+
+__aicore__ inline void MoeV2SortOneCore::SortCompute() {
+  LocalTensor<int32_t> inLocal = sortDataCopyInQueue.DeQue<int32_t>();
+  LocalTensor<int32_t> expertForSourceRowLocal = inLocal[0];
+  LocalTensor<float> expertForSourceRowLocalFp32 = expertForSourceRowLocal.ReinterpretCast<float>();
+  Cast(expertForSourceRowLocalFp32, expertForSourceRowLocal, RoundMode::CAST_ROUND, this->tileLength);
+  AscendC::PipeBarrier<PIPE_V>();
+  Muls(expertForSourceRowLocalFp32, expertForSourceRowLocalFp32, (float)-1, this->tileLength);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM;
+  if (duplicateNum > 0) {
+    int duplicateIndex = this->totalLength - duplicateNum;
+    uint64_t mask0 = UINT64_MAX;
+    mask0 = mask0 << duplicateNum;
+    mask0 = mask0 & (UINT64_MAX >> ONE_REPEAT_SORT_NUM);
+    uint64_t mask[2] = {mask0, 0};
+    Duplicate(expertForSourceRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  LocalTensor<float> concatLocal;
+  LocalTensor<float> tempTensor = tempBuffer.Get<float>(GetSortLen<float>(this->sortNum));
+  Concat(concatLocal, expertForSourceRowLocalFp32, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  LocalTensor<float> sortedLocal = sortedBuffer.Get<float>(GetSortLen<float>(this->sortNum));
+  LocalTensor<uint32_t> sourceRowLocal;
+  sourceRowLocal = inLocal[this->sortNum].ReinterpretCast<uint32_t>();
+  Sort<float, true>(sortedLocal, concatLocal, sourceRowLocal, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
+  LocalTensor<float> sortedExpertForSourceRowLocal = outLocal[0];
+  LocalTensor<uint32_t> expandDstToSrcRowLocal;
+  expandDstToSrcRowLocal = outLocal[this->sortNum].ReinterpretCast<uint32_t>();
+  Extract(sortedExpertForSourceRowLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum / ONE_REPEAT_SORT_NUM);
+  AscendC::PipeBarrier<PIPE_V>();
+  Muls(sortedExpertForSourceRowLocal, sortedExpertForSourceRowLocal, (float)-1, this->tileLength);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  LocalTensor<int32_t> expertForSourceRowLocalInt32;
+  expertForSourceRowLocalInt32 = sortedExpertForSourceRowLocal.ReinterpretCast<int32_t>();
+  Cast(expertForSourceRowLocalInt32, sortedExpertForSourceRowLocal, RoundMode::CAST_ROUND, this->tileLength);
+  sortDataCopyOutQueue.EnQue<float>(outLocal);
+  sortDataCopyInQueue.FreeTensor(inLocal);
+}
+
+__aicore__ inline void MoeV2SortOneCore::CopyOut() {
+  LocalTensor<int32_t> outLocal = sortDataCopyOutQueue.DeQue<int32_t>();
+  DataCopyParams intriParams;
+  intriParams.blockCount = 1;
+  intriParams.blockLen = this->totalLength * sizeof(int32_t);
+  DataCopyPad(sortedexpertIdxGm, outLocal[0], intriParams);
+  DataCopyPad(expandDstToSrcRowGm, outLocal[this->sortNum], intriParams);
+  sortDataCopyOutQueue.FreeTensor(outLocal);
+}
+
+template <typename TilingData>
+__aicore__ inline void MoeV2SortOneCore::Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum,
+                                              GM_ADDR expertTokensBeforeCapacity, GM_ADDR workspace,
+                                              const TilingData* tilingData, TPipe* tPipe) {
+  this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+  this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t));
+  this->sortNum = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
+  this->totalLength = tilingData->n * tilingData->k;
+  this->coreNum = tilingData->coreNum;
+  this->pipe = tPipe;
+  this->n = tilingData->n;
+  this->k = tilingData->k;
+  this->expertNum = tilingData->expertNum;
+  this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag;
+  this->expertTokensBeforeCapacityFlag = tilingData->expertTokensBeforeCapacityFlag;
+
+  expertIdxGm.SetGlobalBuffer((__gm__ int32_t*)expertIdx, this->tileLength);
+  sortedexpertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(workspace), this->tileLength);
+  expandDstToSrcRowGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(workspace) + this->tileLength,
+                                      this->tileLength);
+
+  if (this->blockIdx == this->coreNum - 1) {
+    if (this->expertTokensCountOrCumsumFlag > 0) {
+      expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensCountOrCumsum,
+                                                  Align(this->expertNum, sizeof(int32_t)));
+      InitGlobalMemory(expertTokensCountOrCumsumGm, this->expertNum, 0);
+    }
+    if (this->expertTokensBeforeCapacityFlag == 1) {
+      expertTokensBeforeCapacityGm.SetGlobalBuffer((__gm__ int32_t*)expertTokensBeforeCapacity,
+                                                   Align(this->expertNum, sizeof(int32_t)));
+      InitGlobalMemory(expertTokensBeforeCapacityGm, this->expertNum, 0);
+    }
+  }
+  // key and value
+  int64_t kvFactor = 2;
+  int64_t buffSize = this->sortNum * sizeof(int32_t) * kvFactor;
+  pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize);
+  pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, buffSize);
+  pipe->InitBuffer(tempBuffer, buffSize);
+  pipe->InitBuffer(sortedBuffer, buffSize);
+}
+
+__aicore__ inline void MoeV2SortOneCore::Process() {
+  if (get_block_idx() + get_subblockid() * get_block_num() < 1) {
+    CopyIn();
+    SortCompute();
+    CopyOut();
+  }
+  this->SyncAll();
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // INNER_MOE_V2_SORT_ONE_CORE_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_and_gather.h

@@ -0,0 +1,560 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_src_to_dst_and_gather.h
+ * \brief
+ */
+#ifndef MOE_V2_SRC_TO_DST_AND_GATHER_H
+#define MOE_V2_SRC_TO_DST_AND_GATHER_H
+
+#include "moe_v2_common.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+template <typename T, typename TilingData>
+class MoeV2SrcToDstAndGather {
+ public:
+  __aicore__ inline MoeV2SrcToDstAndGather(){};
+  __aicore__ inline void Init(GM_ADDR x, GM_ADDR scale, GM_ADDR expandedRowIdx, GM_ADDR expandedX,
+                              GM_ADDR dynamicQuantScale, GM_ADDR workspace, const TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void CopyIn(int64_t progress);
+  __aicore__ inline void CopyOut(int64_t progress);
+  __aicore__ inline void CopyOutLoops(int64_t progress);
+  __aicore__ inline void Compute(int32_t srcIdx, int32_t dstIdx, int32_t expertIdx);
+  __aicore__ inline float ComputeMax(LocalTensor<float>& inLocal, LocalTensor<float>& tempLocal,
+                                     LocalTensor<float>& dynamicQuantLocal, int32_t srcIdx, int32_t expertIdx,
+                                     int64_t j);
+  __aicore__ inline void ComputeScale(LocalTensor<float>& inLocal, LocalTensor<float>& tempLocal, float scaleTemp,
+                                      int64_t dstIndex, int64_t j);
+  __aicore__ inline void ComputeLoops(int32_t srcIdx, int32_t dstIdx, int32_t expertIdx);
+
+  __aicore__ inline void CopyOutRemain();
+  __aicore__ inline void SyncAll();
+  __aicore__ inline void AssistInit();
+
+ private:
+  TPipe* pipe;
+  TQue<QuePosition::VECIN, 1> copyInQueue;
+  TQue<QuePosition::VECOUT, 1> copyOutQueue;
+  TQue<QuePosition::VECOUT, 1> copyOutZeroQueue;
+
+  TQue<QuePosition::VECIN, 1> inputXInQueue;
+  TQue<QuePosition::VECIN, 1> smoothInQueue;
+  TQue<QuePosition::VECOUT, 1> calcQueue;
+  TQue<QuePosition::VECOUT, 1> inputXOutQueue;
+  TQue<QuePosition::VECOUT, 1> scaleOutQueue;
+  TQue<QuePosition::VECOUT, 1> scaleOutZeroQueue;
+
+  GlobalTensor<int32_t> expandDstToSrcRowGm;
+  GlobalTensor<int32_t> expandedRowIdxGm;
+  GlobalTensor<int32_t> expertIdxValueGm;
+  GlobalTensor<int32_t> expandedExpertIdxGm;
+  GlobalTensor<int8_t> expandedXGm;
+
+  GlobalTensor<T> inputXGm;
+  GlobalTensor<float> quantSmoothGm;
+  GlobalTensor<float> dynamicQuantScaleGm;
+  GlobalTensor<float> quantSrcGm;
+
+  LocalTensor<int8_t> outTmpLocal;
+  LocalTensor<float> scaleOutTmpLocal;
+  LocalTensor<float> smoothLocal;
+
+  const InnerMoeV2GatherOutComputeTilingData* srcToDstTilingData;
+
+  int64_t coreNum;
+  int64_t blockIdx;
+  int64_t totalLength;
+  int64_t currentLoopRows;
+  int64_t coreRows;
+  int64_t perLoopRows;
+  int64_t lastLoopRows;
+  int64_t rowLoops;
+  int64_t expertCapacity;
+  int64_t expertNum;
+  int64_t cols;
+  int64_t perLoopCols;
+  int64_t lastLoopCols;
+  int64_t colLoops;
+  int64_t perLoopColsAlign;
+  int64_t k;
+  int64_t colsTileLength;
+  int64_t smoothType;
+
+  int64_t tokenCount = 0;
+  int32_t lastExpertId = -1;
+  int32_t lastCoreExpertId = 0;
+  int32_t lastCoreExpertIdNum = 0;
+};
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstAndGather<T, TilingData>::AssistInit() {
+  LocalTensor<int16_t> outLocal = copyOutZeroQueue.AllocTensor<int16_t>();
+  Duplicate<int16_t>(outLocal, static_cast<int16_t>(0), this->perLoopCols);
+  copyOutZeroQueue.EnQue<int16_t>(outLocal);
+  LocalTensor<float> scaleOutLocal = scaleOutZeroQueue.AllocTensor<float>();
+  Duplicate<float>(scaleOutLocal, 0.0f, 8);
+  scaleOutZeroQueue.EnQue<float>(scaleOutLocal);
+
+  if (this->blockIdx != 0) {
+    this->lastCoreExpertId = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2);
+    this->lastCoreExpertIdNum = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2 + 1);
+    for (int64_t i = this->blockIdx - 2; i >= 0; i--) {
+      int32_t lastExpertIdx = expertIdxValueGm.GetValue(i * 2);
+      if (lastExpertIdx < this->lastCoreExpertId) {
+        break;
+      }
+      int32_t lastExpertNum = expertIdxValueGm.GetValue(i * 2 + 1);
+      this->lastCoreExpertIdNum += lastExpertNum;
+    }
+  }
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstAndGather<T, TilingData>::CopyIn(int64_t progress) {
+  LocalTensor<int32_t> inLocal = copyInQueue.AllocTensor<int32_t>();
+  int64_t length = Align(currentLoopRows, sizeof(int32_t));
+  DataCopy(inLocal, expandDstToSrcRowGm[progress * perLoopRows], length);
+  DataCopy(inLocal[length], expandedExpertIdxGm[progress * perLoopRows], length);
+
+  copyInQueue.EnQue<int32_t>(inLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstAndGather<T, TilingData>::Compute(int32_t srcIdx, int32_t dstIdx,
+                                                                      int32_t expertIdx) {
+  DataCopyExtParams copyInParams{1, static_cast<uint32_t>(this->cols * sizeof(T)), 0, 0, 0};
+  DataCopyExtParams smoothParams{1, static_cast<uint32_t>(this->cols * sizeof(float)), 0, 0, 0};
+  DataCopyExtParams copyOutParams{1, static_cast<uint32_t>(this->cols * sizeof(int8_t)), 0, 0, 0};
+
+  LocalTensor<float> inLocal = inputXInQueue.AllocTensor<float>();
+
+  if constexpr (IsSameType<T, float>::value) {
+    DataCopyPad(inLocal, inputXGm[srcIdx / this->k * this->cols], copyInParams, {false, 0, 0, 0});
+  } else {
+    DataCopyPad(inLocal.template ReinterpretCast<T>()[perLoopColsAlign], inputXGm[srcIdx / this->k * this->cols],
+                copyInParams, {false, 0, 0, 0});
+  }
+
+  if (smoothType == 2) {
+    DataCopyPad(smoothLocal, quantSmoothGm[expertIdx * this->cols], smoothParams, {false, 0, 0, 0});
+  }
+
+  inputXInQueue.EnQue<float>(inLocal);
+  smoothInQueue.EnQue(smoothLocal);
+  smoothLocal = smoothInQueue.DeQue<float>();
+
+  inLocal = inputXInQueue.DeQue<float>();
+
+  LocalTensor<float> tempLocal = calcQueue.AllocTensor<float>();
+  LocalTensor<int8_t> outLocal = inputXOutQueue.AllocTensor<int8_t>();
+  LocalTensor<float> dynamicQuantLocal = scaleOutQueue.AllocTensor<float>();
+
+  if constexpr (!IsSameType<T, float>::value) {
+    Cast(inLocal, inLocal.template ReinterpretCast<T>()[perLoopColsAlign], RoundMode::CAST_NONE, this->cols);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  if (smoothType != 0) {
+    Mul(inLocal, inLocal, smoothLocal, this->cols);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  Abs(tempLocal, inLocal, this->cols);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  ReduceMax(dynamicQuantLocal, tempLocal, tempLocal, this->cols);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  float maxValue = dynamicQuantLocal.GetValue(0) / 127.0f;
+
+  Duplicate<float>(dynamicQuantLocal, maxValue, 8);
+  Duplicate<float>(tempLocal, maxValue, this->cols);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Div(tempLocal, inLocal, tempLocal, this->cols);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Cast(tempLocal.ReinterpretCast<half>(), tempLocal, RoundMode::CAST_TRUNC, this->cols);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Cast(outLocal, tempLocal.ReinterpretCast<half>(), RoundMode::CAST_ROUND, this->cols);
+
+  calcQueue.FreeTensor(tempLocal);
+  inputXOutQueue.EnQue(outLocal);
+  scaleOutQueue.EnQue(dynamicQuantLocal);
+
+  LocalTensor<float> quantScaleLocal = scaleOutQueue.DeQue<float>();
+  DataCopyPad(dynamicQuantScaleGm[dstIdx], quantScaleLocal, {1, 4, 0, 0, 0});
+
+  outLocal = inputXOutQueue.DeQue<int8_t>();
+#ifndef __CCE_KT_TEST__
+  DataCopyPad(expandedXGm[dstIdx * this->cols], outLocal, copyOutParams);
+#endif
+  inputXInQueue.FreeTensor(inLocal);
+  inputXOutQueue.FreeTensor(outLocal);
+  scaleOutQueue.FreeTensor(quantScaleLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstAndGather<T, TilingData>::CopyOut(int64_t progress) {
+  LocalTensor<int32_t> inLocal = copyInQueue.DeQue<int32_t>();
+  LocalTensor<int32_t> outLocal = copyOutQueue.AllocTensor<int32_t>();
+  int64_t length = Align(currentLoopRows, sizeof(int32_t));
+  DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(int32_t)), 0, 0, 0};
+  DataCopyExtParams copyParams1{static_cast<uint16_t>(1), static_cast<uint32_t>(this->cols * sizeof(int8_t)), 0, 0, 0};
+
+  SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+  if (this->lastExpertId == -1) {
+    this->lastExpertId = this->lastCoreExpertId;
+    this->tokenCount = this->lastCoreExpertIdNum;
+  }
+  for (int64_t idx = 0; idx < currentLoopRows; idx++) {
+    int32_t expertIdx = inLocal[length].GetValue(idx);
+    SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+    int32_t index = 0;
+    while (this->lastExpertId < expertIdx) {
+      while (this->tokenCount < this->expertCapacity) {
+        index = this->lastExpertId * this->expertCapacity + this->tokenCount;
+        DataCopyPad(expandedXGm[index * this->cols], this->outTmpLocal, copyParams1);
+        DataCopyPad(dynamicQuantScaleGm[index], this->scaleOutTmpLocal, {1, 4, 0, 0, 0});
+        this->tokenCount++;
+      }
+      this->tokenCount = 0;
+      this->lastExpertId++;
+    }
+
+    if (this->tokenCount < this->expertCapacity) {
+      int32_t outOffset = inLocal.GetValue(idx);
+      index = expertIdx * this->expertCapacity + this->tokenCount;
+      outLocal.SetValue(0, index);
+      SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+      DataCopyPad(expandedRowIdxGm[outOffset], outLocal, copyParams);
+      Compute(outOffset, index, expertIdx);
+      SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+      this->tokenCount++;
+    }
+  }
+  copyInQueue.FreeTensor(inLocal);
+  copyOutQueue.FreeTensor(outLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline float MoeV2SrcToDstAndGather<T, TilingData>::ComputeMax(LocalTensor<float>& inLocal,
+                                                                          LocalTensor<float>& tempLocal,
+                                                                          LocalTensor<float>& dynamicQuantLocal,
+                                                                          int32_t srcIdx, int32_t expertIdx,
+                                                                          int64_t j) {
+  LocalTensor<float> smoothLocal = smoothInQueue.AllocTensor<float>();
+
+  DataCopyExtParams intriParamsT{1, static_cast<uint32_t>(colsTileLength * sizeof(T)), 0, 0, 0};
+  DataCopyExtParams intriParamsFp32{1, static_cast<uint32_t>(colsTileLength * sizeof(float)), 0, 0, 0};
+
+  if constexpr (!IsSameType<T, float>::value) {
+    DataCopyPad(inLocal.ReinterpretCast<T>()[perLoopColsAlign], inputXGm[srcIdx * this->cols + j * this->perLoopCols],
+                intriParamsT, {false, 0, 0, 0});
+  } else {
+    DataCopyPad(inLocal, inputXGm[srcIdx * this->cols + j * this->perLoopCols], intriParamsT, {false, 0, 0, 0});
+  }
+
+  inputXInQueue.EnQue<float>(inLocal);
+  inLocal = inputXInQueue.DeQue<float>();
+
+  if constexpr (!IsSameType<T, float>::value) {
+    Cast(inLocal, inLocal.ReinterpretCast<T>()[perLoopColsAlign], RoundMode::CAST_NONE, colsTileLength);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  if (smoothType != 0) {
+    DataCopyPad(smoothLocal, quantSmoothGm[expertIdx * this->cols + j * this->perLoopCols], intriParamsFp32,
+                {false, 0, 0, 0});
+    smoothInQueue.EnQue(smoothLocal);
+    smoothLocal = smoothInQueue.DeQue<float>();
+
+    Mul(inLocal, inLocal, smoothLocal, colsTileLength);
+    AscendC::PipeBarrier<PIPE_V>();
+  }
+
+  Abs(tempLocal, inLocal, colsTileLength);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  ReduceMax(dynamicQuantLocal[8], tempLocal, tempLocal, colsTileLength);
+
+  DataCopyPad(quantSrcGm[j * this->perLoopCols], inLocal, intriParamsFp32);
+  smoothInQueue.FreeTensor(smoothLocal);
+  SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
+
+  return dynamicQuantLocal.GetValue(8);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstAndGather<T, TilingData>::ComputeScale(LocalTensor<float>& inLocal,
+                                                                           LocalTensor<float>& tempLocal,
+                                                                           float scaleTemp, int64_t dstIndex,
+                                                                           int64_t j) {
+  DataCopyExtParams copyInParams{1, static_cast<uint32_t>(colsTileLength * sizeof(float)), 0, 0, 0};
+  DataCopyExtParams copyOutParams{1, static_cast<uint32_t>(colsTileLength * sizeof(int8_t)), 0, 0, 0};
+
+  LocalTensor<int8_t> outLocal = inputXOutQueue.AllocTensor<int8_t>();
+
+  DataCopyPad(inLocal, quantSrcGm[j * this->perLoopCols], copyInParams, {false, 0, 0, 0});
+  inputXInQueue.EnQue<float>(inLocal);
+  inLocal = inputXInQueue.DeQue<float>();
+
+  Duplicate<float>(tempLocal, scaleTemp, colsTileLength);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Div(tempLocal, inLocal, tempLocal, colsTileLength);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Cast(tempLocal.ReinterpretCast<half>(), tempLocal, RoundMode::CAST_TRUNC, colsTileLength);
+  AscendC::PipeBarrier<PIPE_V>();
+
+  Cast(outLocal, tempLocal.ReinterpretCast<half>(), RoundMode::CAST_ROUND, colsTileLength);
+
+  inputXOutQueue.EnQue(outLocal);
+  outLocal = inputXOutQueue.DeQue<int8_t>();
+  DataCopyPad(expandedXGm[dstIndex * this->cols + j * this->perLoopCols], outLocal, copyOutParams);
+
+  inputXOutQueue.FreeTensor(outLocal);
+  SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstAndGather<T, TilingData>::ComputeLoops(int32_t srcIdx, int32_t dstIdx,
+                                                                           int32_t expertIdx) {
+  LocalTensor<float> inLocal = inputXInQueue.AllocTensor<float>();
+  LocalTensor<float> tempLocal = calcQueue.AllocTensor<float>();
+  LocalTensor<float> quantScaleLocal = scaleOutQueue.AllocTensor<float>();
+
+  uint32_t tmp = 0xFF7FFFFF;
+  float reduceMax = *((float*)&tmp);
+  for (int64_t j = 0; j < this->colLoops; j++) {
+    colsTileLength = this->perLoopCols;
+    if (j == this->colLoops - 1) {
+      colsTileLength = this->lastLoopCols;
+    }
+    float tileMax = ComputeMax(inLocal, tempLocal, quantScaleLocal, srcIdx / this->k, expertIdx, j);
+    reduceMax = (reduceMax > tileMax) ? reduceMax : tileMax;
+  }
+
+  float scaleTemp = reduceMax / 127.0f;
+  Duplicate<float>(quantScaleLocal, scaleTemp, 8);
+  scaleOutQueue.EnQue(quantScaleLocal);
+  quantScaleLocal = scaleOutQueue.DeQue<float>();
+
+  DataCopyPad(dynamicQuantScaleGm[dstIdx], quantScaleLocal, {1, 4, 0, 0, 0});
+
+  for (int64_t j = 0; j < this->colLoops; j++) {
+    colsTileLength = this->perLoopCols;
+    if (j == this->colLoops - 1) {
+      colsTileLength = this->lastLoopCols;
+    }
+    ComputeScale(inLocal, tempLocal, scaleTemp, dstIdx, j);
+  }
+
+  inputXInQueue.FreeTensor(inLocal);
+  calcQueue.FreeTensor(tempLocal);
+  scaleOutQueue.FreeTensor(quantScaleLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstAndGather<T, TilingData>::CopyOutLoops(int64_t progress) {
+  LocalTensor<int32_t> inLocal = copyInQueue.DeQue<int32_t>();
+  LocalTensor<int32_t> outLocal = copyOutQueue.AllocTensor<int32_t>();
+  int64_t length = Align(currentLoopRows, sizeof(int32_t));
+  DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(int32_t)), 0, 0, 0};
+
+  SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+  if (this->lastExpertId == -1) {
+    this->lastExpertId = this->lastCoreExpertId;
+    this->tokenCount = this->lastCoreExpertIdNum;
+  }
+  for (int64_t idx = 0; idx < currentLoopRows; idx++) {
+    int32_t expertIdx = inLocal[length].GetValue(idx);
+    SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+    int32_t index = 0;
+    while (this->lastExpertId < expertIdx) {
+      while (this->tokenCount < this->expertCapacity) {
+        index = this->lastExpertId * this->expertCapacity + this->tokenCount;
+        int64_t col = this->perLoopCols;
+        DataCopyPad(dynamicQuantScaleGm[index], this->scaleOutTmpLocal, {1, 4, 0, 0, 0});
+        for (int64_t i = 0; i < this->colLoops; i++) {
+          if (i == this->colLoops - 1) {
+            col = this->lastLoopCols;
+          }
+          DataCopyExtParams copyParams1{static_cast<uint16_t>(1), static_cast<uint32_t>(col * sizeof(int8_t)), 0, 0, 0};
+          DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams1);
+          SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+        }
+        this->tokenCount++;
+      }
+      this->tokenCount = 0;
+      this->lastExpertId++;
+    }
+
+    if (this->tokenCount < this->expertCapacity) {
+      int32_t outOffset = inLocal.GetValue(idx);
+      index = expertIdx * this->expertCapacity + this->tokenCount;
+      outLocal.SetValue(0, index);
+      SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+      DataCopyPad(expandedRowIdxGm[outOffset], outLocal, copyParams);
+      if (smoothType == 2) {
+        ComputeLoops(outOffset, index, expertIdx);
+      } else {
+        ComputeLoops(outOffset, index, 0);
+      }
+      SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+      this->tokenCount++;
+    }
+  }
+  copyInQueue.FreeTensor(inLocal);
+  copyOutQueue.FreeTensor(outLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstAndGather<T, TilingData>::CopyOutRemain() {
+  if (this->blockIdx != this->srcToDstTilingData->needCoreNum - 1) {
+    copyOutZeroQueue.FreeTensor(this->outTmpLocal);
+    scaleOutZeroQueue.FreeTensor(this->scaleOutTmpLocal);
+    return;
+  }
+  while (this->lastExpertId < this->expertNum) {
+    while (this->tokenCount < this->expertCapacity) {
+      int32_t index = this->lastExpertId * this->expertCapacity + this->tokenCount;
+      int64_t col = this->perLoopCols;
+      DataCopyPad(dynamicQuantScaleGm[index], this->scaleOutTmpLocal, {1, 4, 0, 0, 0});
+      for (int64_t i = 0; i < this->colLoops; i++) {
+        if (i == this->colLoops - 1) {
+          col = this->lastLoopCols;
+        }
+        DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(col * sizeof(int8_t)), 0, 0, 0};
+        DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams);
+        SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+      }
+      this->tokenCount++;
+    }
+    this->tokenCount = 0;
+    this->lastExpertId++;
+  }
+  copyOutZeroQueue.FreeTensor(this->outTmpLocal);
+  scaleOutZeroQueue.FreeTensor(this->scaleOutTmpLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstAndGather<T, TilingData>::Init(GM_ADDR x, GM_ADDR scale, GM_ADDR expandedRowIdx,
+                                                                   GM_ADDR expandedX, GM_ADDR dynamicQuantScale,
+                                                                   GM_ADDR workspace, const TilingData* tilingData,
+                                                                   TPipe* tPipe) {
+  int64_t blockNum = GetBlockNum();
+  this->pipe = tPipe;
+  this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+
+  this->coreNum = tilingData->coreNum;
+  this->totalLength = tilingData->n * tilingData->k;
+  this->srcToDstTilingData = &(tilingData->srcToDstCapacityComputeParamsOp);
+  this->expertNum = tilingData->expertNum;
+  this->expertCapacity = tilingData->expertCapacity;
+  this->cols = tilingData->cols;
+  this->k = tilingData->k;
+  this->smoothType = tilingData->smoothType;
+
+  if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) {
+    this->coreRows = this->srcToDstTilingData->lastCoreRows;
+    this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows;
+    this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows;
+    this->rowLoops = this->srcToDstTilingData->lastCoreLoops;
+  } else {
+    this->coreRows = this->srcToDstTilingData->perCoreRows;
+    this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows;
+    this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows;
+    this->rowLoops = this->srcToDstTilingData->perCoreLoops;
+  }
+  this->perLoopCols = this->srcToDstTilingData->perLoopCols;
+  this->lastLoopCols = this->srcToDstTilingData->lastLoopCols;
+  this->colLoops = this->srcToDstTilingData->colLoops;
+  this->perLoopColsAlign = Align(this->perLoopCols, sizeof(T));
+
+  inputXGm.SetGlobalBuffer((__gm__ T*)x);
+  quantSmoothGm.SetGlobalBuffer((__gm__ float*)scale);
+  dynamicQuantScaleGm.SetGlobalBuffer((__gm__ float*)dynamicQuantScale);
+
+  int64_t length = Align(this->totalLength, sizeof(int32_t));
+  expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, length);
+  expandedXGm.SetGlobalBuffer((__gm__ int8_t*)expandedX, this->expertNum * this->expertCapacity * this->cols);
+
+  expandedExpertIdxGm.SetGlobalBuffer(
+      (__gm__ int32_t*)workspace + this->blockIdx * this->srcToDstTilingData->perCoreRows,
+      Align(this->coreRows, sizeof(int32_t)));
+  expandDstToSrcRowGm.SetGlobalBuffer(
+      (__gm__ int32_t*)workspace + length + this->blockIdx * this->srcToDstTilingData->perCoreRows,
+      Align(this->coreRows, sizeof(int32_t)));
+  expertIdxValueGm.SetGlobalBuffer((__gm__ int32_t*)workspace + length * 2, this->coreNum * 2);
+  if (this->colLoops > 1) {
+    quantSrcGm.SetGlobalBuffer((__gm__ float*)workspace + length * 2 + this->coreNum * 2 + this->blockIdx * this->cols,
+                               this->cols * sizeof(float));
+  }
+
+  pipe->InitBuffer(copyInQueue, 1, AlignBytes(this->perLoopRows, sizeof(int32_t)) * 2);
+  pipe->InitBuffer(copyOutQueue, 1, AlignBytes(INT32_ONE_BLOCK_NUM, sizeof(int32_t)));
+  pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(int16_t)));
+
+  int64_t perLoopColsAlignBytes = AlignBytes(this->perLoopCols, sizeof(T));
+  perLoopColsAlignBytes =
+      Max(int64_t(perLoopColsAlignBytes * sizeof(float) / sizeof(T)), int64_t(BLOCK_BYTES + BLOCK_BYTES));
+
+  pipe->InitBuffer(inputXInQueue, 1, perLoopColsAlignBytes);
+  pipe->InitBuffer(smoothInQueue, 1, AlignBytes(this->perLoopCols, sizeof(float)));
+  pipe->InitBuffer(calcQueue, 1, AlignBytes(this->perLoopCols, sizeof(float)));
+  pipe->InitBuffer(inputXOutQueue, 1, AlignBytes(this->perLoopCols, sizeof(int8_t)));
+  pipe->InitBuffer(scaleOutQueue, 1, BLOCK_BYTES + BLOCK_BYTES);
+  pipe->InitBuffer(scaleOutZeroQueue, 1, BLOCK_BYTES);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstAndGather<T, TilingData>::Process() {
+  if (this->blockIdx < this->srcToDstTilingData->needCoreNum) {
+    AssistInit();
+    this->outTmpLocal = copyOutZeroQueue.DeQue<int8_t>();
+    this->scaleOutTmpLocal = scaleOutZeroQueue.DeQue<float>();
+    currentLoopRows = perLoopRows;
+    if (colLoops > 1) {
+      for (int64_t loop = 0; loop < this->rowLoops; loop++) {
+        if (loop == this->rowLoops - 1) {
+          currentLoopRows = lastLoopRows;
+        }
+        CopyIn(loop);
+        CopyOutLoops(loop);
+      }
+    } else {
+      smoothLocal = smoothInQueue.AllocTensor<float>();
+      if (smoothType == 1) {
+        DataCopyExtParams smoothParams{1, static_cast<uint32_t>(this->cols * sizeof(float)), 0, 0, 0};
+        DataCopyPad(smoothLocal, quantSmoothGm, smoothParams, {false, 0, 0, 0});
+      }
+      for (int64_t loop = 0; loop < this->rowLoops; loop++) {
+        if (loop == this->rowLoops - 1) {
+          currentLoopRows = lastLoopRows;
+        }
+        CopyIn(loop);
+        CopyOut(loop);
+      }
+      smoothInQueue.FreeTensor(smoothLocal);
+    }
+    CopyOutRemain();
+  }
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // MOE_V2_SRC_TO_DST_AND_GATHER_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_op.h

@@ -0,0 +1,164 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_src_to_dst_op.h
+ * \brief
+ */
+#ifndef INNER_MOE_V2_SRC_TO_DST_H
+#define INNER_MOE_V2_SRC_TO_DST_H
+
+#include "moe_v2_common.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+class MoeV2SrcToDstOp {
+ public:
+  __aicore__ inline MoeV2SrcToDstOp(){};
+  template <typename TilingData>
+  __aicore__ inline void Init(GM_ADDR expandSrcToDstRow, GM_ADDR workspace, const TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void CopyIn(int64_t progress);
+  __aicore__ inline void Compute(int64_t progress);
+  __aicore__ inline void CopyOut();
+  __aicore__ inline void SyncAll();
+  __aicore__ inline void AssistInit();
+
+ private:
+  TPipe* pipe;
+  TQue<QuePosition::VECIN, 1> copyInQueue;
+  TQue<QuePosition::VECOUT, 1> copyOutQueue;
+  TBuf<TPosition::VECCALC> assistBuffer;
+
+  GlobalTensor<int32_t> expandDstToSrcRowGm;
+  GlobalTensor<int32_t> expandSrcToDstRowGm;
+  GlobalTensor<int32_t> assistGm;
+
+  const InnerMoeV2GatherOutComputeTilingData* srcToDstTilingData;
+
+  int64_t coreNum;
+  int64_t blockIdx;
+  int64_t totalLength;
+  int64_t currentLoopRows;
+  int64_t coreRows;
+  int64_t perLoopRows;
+  int64_t lastLoopRows;
+};
+
+__aicore__ inline void MoeV2SrcToDstOp::AssistInit() {
+#if defined(ASCENDC_OOM) && ASCENDC_OOM == 1
+  OOMCheckAddrRange(assistGm.GetPhyAddr(), ASSIST_NUM * sizeof(int32_t));
+#endif
+  LocalTensor<int32_t> assistTensor = assistBuffer.Get<int32_t>(ASSIST_NUM);
+  DataCopy(assistTensor, assistGm, ASSIST_NUM);
+  SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
+  Adds(assistTensor, assistTensor, (int32_t)(this->blockIdx * this->srcToDstTilingData->perCoreRows), ASSIST_NUM);
+}
+
+__aicore__ inline void MoeV2SrcToDstOp::CopyIn(int64_t progress) {
+  LocalTensor<int32_t> inLocal = copyInQueue.AllocTensor<int32_t>();
+  DataCopy(inLocal, expandDstToSrcRowGm[progress * perLoopRows], Align(currentLoopRows, sizeof(int32_t)));
+  copyInQueue.EnQue<int32_t>(inLocal);
+}
+
+__aicore__ inline void MoeV2SrcToDstOp::Compute(int64_t progress) {
+  LocalTensor<int32_t> outLocal = copyOutQueue.AllocTensor<int32_t>();
+  LocalTensor<int32_t> assistTensor = assistBuffer.Get<int32_t>(ASSIST_NUM);
+
+  AscendC::PipeBarrier<PIPE_V>();
+  int64_t loops = Ceil(currentLoopRows, ASSIST_INDEX_NUM);
+  for (int64_t i = 0; i < loops; i++) {
+    Adds(outLocal[i * ASSIST_NUM], assistTensor,
+         static_cast<int32_t>(this->perLoopRows * progress + i * ASSIST_INDEX_NUM), ASSIST_NUM);
+  }
+  AscendC::PipeBarrier<PIPE_V>();
+  copyOutQueue.EnQue<int32_t>(outLocal);
+}
+
+__aicore__ inline void MoeV2SrcToDstOp::CopyOut() {
+  LocalTensor<int32_t> inLocal = copyInQueue.DeQue<int32_t>();
+  LocalTensor<int32_t> outLocal = copyOutQueue.DeQue<int32_t>();
+  SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+  DataCopyParams intriParams;
+  intriParams.blockCount = 1;
+  intriParams.blockLen = sizeof(int32_t);
+  uint32_t outOffset;
+  for (int64_t idx = 0; idx < currentLoopRows; idx++) {
+    outOffset = inLocal.GetValue(idx);
+    DataCopyPad(expandSrcToDstRowGm[outOffset], outLocal[idx * INT32_ONE_BLOCK_NUM], intriParams);
+  }
+
+  copyInQueue.FreeTensor(inLocal);
+  copyOutQueue.FreeTensor(outLocal);
+}
+
+__aicore__ inline void MoeV2SrcToDstOp::SyncAll() {
+  if (coreNum == 1) {
+    return;
+  }
+#ifndef __CCE_KT_TEST__
+  AscendC::SyncAll();
+#endif
+}
+
+template <typename TilingData>
+__aicore__ inline void MoeV2SrcToDstOp::Init(GM_ADDR expandSrcToDstRow, GM_ADDR workspace, const TilingData* tilingData,
+                                             TPipe* tPipe) {
+  int64_t blockNum = GetBlockNum();
+  this->pipe = tPipe;
+  this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+
+  this->coreNum = tilingData->coreNum;
+  this->totalLength = tilingData->n * tilingData->k;
+  this->srcToDstTilingData = &(tilingData->srcToDstComputeParamsOp);
+
+  if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) {
+    this->coreRows = this->srcToDstTilingData->lastCoreRows;
+    this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows;
+    this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows;
+  } else {
+    this->coreRows = this->srcToDstTilingData->perCoreRows;
+    this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows;
+    this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows;
+  }
+
+  expandSrcToDstRowGm.SetGlobalBuffer((__gm__ int32_t*)expandSrcToDstRow, Align(this->totalLength, sizeof(int32_t)));
+  expandDstToSrcRowGm.SetGlobalBuffer((__gm__ int32_t*)workspace + Align(this->totalLength, sizeof(int32_t)) +
+                                          this->blockIdx * this->srcToDstTilingData->perCoreRows,
+                                      Align(this->coreRows, sizeof(int32_t)));
+  assistGm.SetGlobalBuffer((__gm__ int32_t*)assist, ASSIST_NUM);
+
+  pipe->InitBuffer(copyInQueue, 1, this->perLoopRows * BLOCK_BYTES);
+  pipe->InitBuffer(copyOutQueue, 1, Ceil(this->perLoopRows, ASSIST_NUM) * ASSIST_NUM * BLOCK_BYTES);
+  pipe->InitBuffer(assistBuffer, ASSIST_NUM * sizeof(int32_t));
+}
+
+__aicore__ inline void MoeV2SrcToDstOp::Process() {
+  if (this->blockIdx < this->srcToDstTilingData->needCoreNum) {
+    int64_t loops = (coreRows + perLoopRows - 1) / perLoopRows;
+    currentLoopRows = perLoopRows;
+    AssistInit();
+    for (int64_t loop = 0; loop < loops - 1; loop++) {
+      CopyIn(loop);
+      Compute(loop);
+      CopyOut();
+    }
+    currentLoopRows = lastLoopRows;
+    CopyIn(loops - 1);
+    Compute(loops - 1);
+    CopyOut();
+  }
+  this->SyncAll();
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // INNER_MOE_V2_SRC_TO_DST_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/moe_v2_src_to_dst_with_capacity.h

@@ -0,0 +1,269 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_src_to_dst_with_capacity.h
+ * \brief
+ */
+#ifndef INNER_MOE_V2_SRC_TO_DST_WITH_CAPACITY_H
+#define INNER_MOE_V2_SRC_TO_DST_WITH_CAPACITY_H
+
+#include "moe_v2_common.h"
+
+namespace MoeInitRoutingQuantV2 {
+using namespace AscendC;
+using namespace optiling;
+template <typename T, typename TilingData>
+class MoeV2SrcToDstWithCapacity {
+ public:
+  __aicore__ inline MoeV2SrcToDstWithCapacity(){};
+  __aicore__ inline void Init(GM_ADDR expandedRowIdx, GM_ADDR expandedX, GM_ADDR workspace,
+                              const TilingData* tilingData, TPipe* tPipe);
+  __aicore__ inline void Process();
+
+ private:
+  __aicore__ inline void CopyIn(int64_t progress);
+  __aicore__ inline void CopyOut(int64_t progress);
+  __aicore__ inline void CopyOutRemain();
+  __aicore__ inline void SyncAll();
+  __aicore__ inline void AssistInit();
+
+ private:
+  TPipe* pipe;
+  TQue<QuePosition::VECIN, 1> copyInQueue;
+  TQue<QuePosition::VECOUT, 1> copyOutQueue;
+  TQue<QuePosition::VECOUT, 1> copyOutZeroQueue;
+
+  GlobalTensor<int32_t> expandDstToSrcRowGm;
+  GlobalTensor<int32_t> expandedRowIdxGm;
+  GlobalTensor<int32_t> expertIdxValueGm;
+  GlobalTensor<int32_t> expandedExpertIdxGm;
+  GlobalTensor<T> expandedXGm;
+
+  LocalTensor<T> outTmpLocal;
+
+  const InnerMoeV2GatherOutComputeTilingData* srcToDstTilingData;
+
+  int64_t coreNum;
+  int64_t blockIdx;
+  int64_t totalLength;
+  int64_t currentLoopRows;
+  int64_t coreRows;
+  int64_t perLoopRows;
+  int64_t lastLoopRows;
+  int64_t rowLoops;
+  int64_t expertCapacity;
+  int64_t expertNum;
+  int64_t cols;
+  int64_t perLoopCols;
+  int64_t lastLoopCols;
+  int64_t colLoops;
+
+  int64_t tokenCount = 0;
+  int32_t lastExpertId = -1;
+  int32_t lastCoreExpertId = 0;
+  int32_t lastCoreExpertIdNum = 0;
+};
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::AssistInit() {
+  if constexpr (IsSameType<T, int8_t>::value) {
+    LocalTensor<int16_t> outLocal = copyOutZeroQueue.AllocTensor<int16_t>();
+    Duplicate<int16_t>(outLocal, static_cast<int16_t>(0), this->perLoopCols);
+    copyOutZeroQueue.EnQue<int16_t>(outLocal);
+  } else {
+    LocalTensor<T> outLocal = copyOutZeroQueue.AllocTensor<T>();
+    Duplicate<T>(outLocal, static_cast<T>(0), this->perLoopCols);
+    copyOutZeroQueue.EnQue<T>(outLocal);
+  }
+
+  if (this->blockIdx != 0) {
+    this->lastCoreExpertId = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2);
+    this->lastCoreExpertIdNum = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2 + 1);
+    for (int64_t i = this->blockIdx - 2; i >= 0; i--) {
+      int32_t lastExpertIdx = expertIdxValueGm.GetValue(i * 2);
+      if (lastExpertIdx < this->lastCoreExpertId) {
+        break;
+      }
+      int32_t lastExpertNum = expertIdxValueGm.GetValue(i * 2 + 1);
+      this->lastCoreExpertIdNum += lastExpertNum;
+    }
+  }
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::CopyIn(int64_t progress) {
+  LocalTensor<int32_t> inLocal = copyInQueue.AllocTensor<int32_t>();
+  int64_t length = Align(currentLoopRows, sizeof(int32_t));
+  DataCopy(inLocal, expandDstToSrcRowGm[progress * perLoopRows], length);
+  DataCopy(inLocal[length], expandedExpertIdxGm[progress * perLoopRows], length);
+  copyInQueue.EnQue<int32_t>(inLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::CopyOut(int64_t progress) {
+  LocalTensor<int32_t> inLocal = copyInQueue.DeQue<int32_t>();
+  LocalTensor<int32_t> outLocal = copyOutQueue.AllocTensor<int32_t>();
+  int64_t length = Align(currentLoopRows, sizeof(int32_t));
+  DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(int32_t)), 0, 0, 0};
+
+  SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+  if (this->lastExpertId == -1) {
+    this->lastExpertId = this->lastCoreExpertId;
+    this->tokenCount = this->lastCoreExpertIdNum;
+  }
+  for (int64_t idx = 0; idx < currentLoopRows; idx++) {
+    int32_t expertIdx = inLocal[length].GetValue(idx);
+    SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+    int32_t index = 0;
+    while (this->lastExpertId < expertIdx) {
+      while (this->tokenCount < this->expertCapacity) {
+        index = this->lastExpertId * this->expertCapacity + this->tokenCount;
+        int64_t col = this->perLoopCols;
+        for (int64_t i = 0; i < this->colLoops; i++) {
+          if (i == this->colLoops - 1) {
+            col = this->lastLoopCols;
+          }
+#ifdef __CCE_KT_TEST__
+          // CPU twin debugging cannot use multi-core sync, so index may contain uninitialized dirty data; handle specially
+          if (index * this->cols + i * this->perLoopCols + col * sizeof(T) > expandedXGm.GetSize()) {
+              continue;
+          }
+#endif
+          DataCopyExtParams copyParams1{static_cast<uint16_t>(1), static_cast<uint32_t>(col * sizeof(T)), 0, 0, 0};
+          DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams1);
+          SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+        }
+        this->tokenCount++;
+      }
+      this->tokenCount = 0;
+      this->lastExpertId++;
+    }
+
+    if (this->tokenCount < this->expertCapacity) {
+      int32_t outOffset = inLocal.GetValue(idx);
+      index = expertIdx * this->expertCapacity + this->tokenCount;
+      outLocal.SetValue(0, index);
+      SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+      DataCopyPad(expandedRowIdxGm[outOffset], outLocal, copyParams);
+      SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+      this->tokenCount++;
+    }
+  }
+  copyInQueue.FreeTensor(inLocal);
+  copyOutQueue.FreeTensor(outLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::CopyOutRemain() {
+  if (this->blockIdx != this->srcToDstTilingData->needCoreNum - 1) {
+    copyOutZeroQueue.FreeTensor(this->outTmpLocal);
+    return;
+  }
+  while (this->lastExpertId < this->expertNum) {
+    while (this->tokenCount < this->expertCapacity) {
+      int32_t index = this->lastExpertId * this->expertCapacity + this->tokenCount;
+      int64_t col = this->perLoopCols;
+      for (int64_t i = 0; i < this->colLoops; i++) {
+        if (i == this->colLoops - 1) {
+          col = this->lastLoopCols;
+        }
+        DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(col * sizeof(T)), 0, 0, 0};
+        DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams);
+        SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+      }
+      this->tokenCount++;
+    }
+    this->tokenCount = 0;
+    this->lastExpertId++;
+  }
+  copyOutZeroQueue.FreeTensor(this->outTmpLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::SyncAll() {
+  if (coreNum == 1) {
+    return;
+  }
+#ifndef __CCE_KT_TEST__
+  AscendC::SyncAll();
+#endif
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::Init(GM_ADDR expandedRowIdx, GM_ADDR expandedX,
+                                                                      GM_ADDR workspace, const TilingData* tilingData,
+                                                                      TPipe* tPipe) {
+  int64_t blockNum = GetBlockNum();
+  this->pipe = tPipe;
+  this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+
+  this->coreNum = tilingData->coreNum;
+  this->totalLength = tilingData->n * tilingData->k;
+  this->srcToDstTilingData = &(tilingData->srcToDstCapacityComputeParamsOp);
+  this->expertNum = tilingData->expertNum;
+  this->expertCapacity = tilingData->expertCapacity;
+  this->cols = tilingData->cols;
+
+  if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) {
+    this->coreRows = this->srcToDstTilingData->lastCoreRows;
+    this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows;
+    this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows;
+    this->rowLoops = this->srcToDstTilingData->lastCoreLoops;
+  } else {
+    this->coreRows = this->srcToDstTilingData->perCoreRows;
+    this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows;
+    this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows;
+    this->rowLoops = this->srcToDstTilingData->perCoreLoops;
+  }
+  this->perLoopCols = this->srcToDstTilingData->perLoopCols;
+  this->lastLoopCols = this->srcToDstTilingData->lastLoopCols;
+  this->colLoops = this->srcToDstTilingData->colLoops;
+
+  int64_t length = Align(this->totalLength, sizeof(int32_t));
+  expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t*)expandedRowIdx, length);
+  expandedXGm.SetGlobalBuffer((__gm__ T*)expandedX, this->expertNum * this->expertCapacity * this->cols);
+
+  expandedExpertIdxGm.SetGlobalBuffer(
+      (__gm__ int32_t*)workspace + this->blockIdx * this->srcToDstTilingData->perCoreRows,
+      Align(this->coreRows, sizeof(int32_t)));
+  expandDstToSrcRowGm.SetGlobalBuffer(
+      (__gm__ int32_t*)workspace + length + this->blockIdx * this->srcToDstTilingData->perCoreRows,
+      Align(this->coreRows, sizeof(int32_t)));
+  expertIdxValueGm.SetGlobalBuffer((__gm__ int32_t*)workspace + length * 2, this->coreNum * 2);
+
+  pipe->InitBuffer(copyInQueue, 1, AlignBytes(this->perLoopRows, sizeof(int32_t)) * 2);
+  pipe->InitBuffer(copyOutQueue, 1, AlignBytes(INT32_ONE_BLOCK_NUM, sizeof(int32_t)));
+  if constexpr (IsSameType<T, int8_t>::value) {
+    pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(int16_t)));
+  } else {
+    pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(T)));
+  }
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::Process() {
+  if (this->blockIdx < this->srcToDstTilingData->needCoreNum) {
+    AssistInit();
+    this->outTmpLocal = copyOutZeroQueue.DeQue<T>();
+    currentLoopRows = perLoopRows;
+    for (int64_t loop = 0; loop < this->rowLoops; loop++) {
+      if (loop == this->rowLoops - 1) {
+        currentLoopRows = lastLoopRows;
+      }
+      CopyIn(loop);
+      CopyOut(loop);
+    }
+    CopyOutRemain();
+  }
+  this->SyncAll();
+}
+}  // namespace MoeInitRoutingQuantV2
+#endif  // INNER_MOE_V2_SRC_TO_DST_WITH_CAPACITY_H

csrc/mc2/dispatch_ffn_combine/op_kernel/moe_init_routing_quant_v2/tiling_base.h

@@ -0,0 +1,69 @@
+#pragma once
+#ifndef ASCENDC_DISPATCH_FFN_COMBINE_TILING_BASE_H
+#define ASCENDC_DISPATCH_FFN_COMBINE_TILING_BASE_H
+namespace optiling {
+struct AiCoreParams {
+    uint64_t ubSize;
+    uint64_t blockDim;
+    uint64_t aicNum;
+    
+    uint64_t l1Size;
+    uint64_t l0aSize;
+    uint64_t l0bSize;
+    uint64_t l0cSize;
+};
+
+class TilingBaseClass {
+public:
+    bool DoTiling(
+            int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, 
+            int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag,
+            bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0, 
+            int64_t aivCoreNum, int64_t ubSizePlatForm) 
+    {
+        bool ret = GetShapeAttrsInfo(m, cols, topK, expertCapacity, expertNum, activeNum, dropPadMode, expertTokensCountOrCumsumFlag,
+            expertTokensBeforeCapacityFlag, inuptXDtypeSize, quantMode, scaleDim0);
+
+        if (!ret){
+            return ret;
+        }
+        ret = GetPlatformInfo(aivCoreNum, ubSizePlatForm);
+        if (!ret){
+            return ret;
+        }
+        ret = DoOpTiling();
+        if (!ret){
+            return ret;
+        }
+        ret = GetWorkspaceSize();
+        if (!ret){
+            return ret;
+        }
+        ret = PostTiling();
+        if (!ret){
+            return ret;
+        }
+        tilingKey_ = GetTilingKey();
+        
+        return true;
+    }
+
+//protected:
+    virtual bool GetPlatformInfo(int64_t aivCoreNum, int64_t ubSizePlatForm) = 0;
+    virtual bool GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, 
+        int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag,
+        bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) = 0;
+
+    virtual bool DoOpTiling() = 0;
+    virtual bool GetWorkspaceSize() = 0;
+    virtual bool PostTiling() = 0;
+    virtual uint64_t GetTilingKey() const = 0;
+//protected:
+    uint32_t blockDim_{0};
+    uint64_t workspaceSize_{0};
+    uint64_t tilingKey_{0};
+    AiCoreParams aicoreParams_{0, 0, 0, 0, 0, 0, 0};
+};
+
+}
+#endif

csrc/mc2/dispatch_ffn_combine/op_kernel/unpermute/moe_token_unpermute.h

@@ -0,0 +1,375 @@
+/**
+ * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*!
+ * \file moe_token_unpermute.h
+ * \brief
+ */
+
+#ifndef MOE_TOKEN_UNPERMUTE
+#define MOE_TOKEN_UNPERMUTE
+
+#include "kernel_operator.h"
+#include "moe_token_unpermute_tiling.h"
+using namespace AscendC;
+
+
+template <typename T1, typename T2, typename T3, bool PROBS> class KernelMoeTokenUnpermute {
+public:
+    __aicore__ inline KernelMoeTokenUnpermute()
+    {
+    }
+
+    __aicore__ inline void Init(GM_ADDR permuted_tokens, GM_ADDR sorted_indices, GM_ADDR probs,
+                                GM_ADDR unpermuted_tokens, const MoeTokenUnpermuteTilingData *__restrict tiling_data);
+    __aicore__ inline void Process();
+
+protected:
+    __aicore__ inline void CalMultiOutToken(const int64_t out_offset, const int64_t out_tokens_number);
+    __aicore__ inline void CalSingleOutToken(const int64_t start_token, const int64_t out_token_idx);
+    __aicore__ inline void CalPartOutToken(const int64_t start_token, const int64_t h_index, const int64_t h_length,
+                                           const int64_t out_token_index);
+    __aicore__ inline void CopyTokenIn(const T2 in_token_index, const int64_t h_index, const int64_t h_length);
+    __aicore__ inline void CalFirstToken(const float prob_value, const int64_t h_length);
+    __aicore__ inline void CalToken(const float prob_value, const int64_t h_length);
+    __aicore__ inline void CopyOut(const int64_t out_token_index, const int64_t h_index, const int64_t h_length);
+
+    TPipe pipe;
+    TQue<QuePosition::VECIN, 1> tokens_inque, indices_inque, probs_inque;
+    TBuf<TPosition::VECCALC> temp_buffer0, temp_buffer1, temp_buffer2;
+    TQue<QuePosition::VECOUT, 1> outque;
+    GlobalTensor<T1> tokensGM, outGM;
+    GlobalTensor<T2> indicesGM;
+    GlobalTensor<T3> probsGM;
+    LocalTensor<T2> indicesLocal;
+    LocalTensor<float> token_tensor0, token_tensor1, probs_tensor;
+    DataCopyPadExtParams<T1> extParams1{false, 0, 0, 0};
+    DataCopyPadExtParams<T2> extParams2{false, 0, 0, 0};
+    DataCopyPadExtParams<T3> extParams3{false, 0, 0, 0};
+    DataCopyExtParams copyParams{1, 0, 0, 0, 0};
+
+    constexpr static uint32_t BLOCK_SIZE = 32;
+    constexpr static uint32_t ALIGN_512 = 512;
+
+    int64_t hidden_size;
+    int64_t top_k;
+    int64_t num_out_tokens;
+    int64_t hidden_splited_length;
+    int64_t hidden_splited_num;
+    int64_t hidden_splited_remain;
+    int64_t tokens_core_length;
+    int64_t tokens_core_remain;
+    int64_t tokens_splited_length;
+    int64_t tokens_splited_num;
+    int64_t tokens_splited_remain;
+    int32_t blockIdx;
+    int32_t blockNum;
+};
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void
+KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::Init(GM_ADDR permuted_tokens, GM_ADDR sorted_indices, GM_ADDR probs,
+                                                 GM_ADDR unpermuted_tokens,
+                                                 const MoeTokenUnpermuteTilingData *__restrict tiling_data)
+{
+    this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+    this->blockNum = get_block_num() * get_subblockdim();
+    if (blockIdx >= blockNum) {
+        return;
+    }
+    ASSERT(blockNum != 0 && "block dim can not be zero!");
+    // row_input
+    this->hidden_size = tiling_data->hidden_size;
+    this->top_k = tiling_data->top_k;
+    this->num_out_tokens = tiling_data->num_out_tokens;
+    // hidden_tiling
+    this->hidden_splited_length = tiling_data->hidden_splited_length;
+    this->hidden_splited_num = tiling_data->hidden_splited_num;
+    this->hidden_splited_remain = tiling_data->hidden_splited_remain;
+    // token_tiling
+    this->tokens_core_length = tiling_data->tokens_core_length;
+    this->tokens_core_remain = tiling_data->tokens_core_remain;
+    this->tokens_splited_length = tiling_data->tokens_splited_length;
+    this->tokens_splited_num = tiling_data->tokens_splited_num;
+    this->tokens_splited_remain = tiling_data->tokens_splited_remain;
+
+    // Handle the tail block for token_by_core
+    if (this->tokens_core_remain > 0 && blockIdx < this->tokens_core_remain) {
+        this->tokens_core_length += 1;
+        this->tokens_splited_remain += 1;
+    }
+
+    int64_t hidden_splited_length_align512 = (this->hidden_splited_length + ALIGN_512 - 1) & ~(ALIGN_512 - 1);
+
+    int64_t block_length = this->tokens_core_length * this->top_k;
+    int64_t block_splited_length = this->tokens_splited_length * this->top_k;
+
+    int64_t block_offset;
+    if (this->tokens_core_remain > 0) {
+        if (blockIdx < this->tokens_core_remain) {
+            block_offset = block_length * blockIdx;
+        } else {
+            block_offset = (block_length + this->top_k) * this->tokens_core_remain +
+                           block_length * (blockIdx - this->tokens_core_remain);
+        }
+    } else {
+        block_offset = block_length * blockIdx;
+    }
+
+    this->tokensGM.SetGlobalBuffer((__gm__ T1 *)permuted_tokens);
+    this->indicesGM.SetGlobalBuffer((__gm__ T2 *)sorted_indices + block_offset, block_length);
+
+
+    int64_t out_block_offset;
+    if (this->tokens_core_remain > 0) {
+        if (blockIdx < this->tokens_core_remain) {
+            out_block_offset = this->tokens_core_length * blockIdx * hidden_size;
+        } else {
+            out_block_offset = (this->tokens_core_length + 1) * this->tokens_core_remain +
+                               this->tokens_core_length * (blockIdx - this->tokens_core_remain);
+            out_block_offset *= this->hidden_size;
+        }
+    } else {
+        out_block_offset = this->tokens_core_length * blockIdx * hidden_size;
+    }
+
+    this->outGM.SetGlobalBuffer((__gm__ T1 *)unpermuted_tokens + out_block_offset,
+                                this->tokens_core_length * this->hidden_size);
+
+    this->pipe.InitBuffer(tokens_inque, tiling_data->buffer_num, hidden_splited_length_align512 * sizeof(T1));
+    this->pipe.InitBuffer(indices_inque, 1, block_splited_length * (sizeof(T2)));
+    this->pipe.InitBuffer(outque, 1, hidden_splited_length_align512 * sizeof(T1));
+
+    if constexpr (!IsSameType<T1, float>::value) {
+        this->pipe.InitBuffer(temp_buffer0, hidden_splited_length_align512 * sizeof(float) + 256);
+        this->pipe.InitBuffer(temp_buffer1, hidden_splited_length_align512 * sizeof(float));
+        this->token_tensor0 = this->temp_buffer0.template Get<float>();
+        this->token_tensor1 = this->temp_buffer1.template Get<float>();
+    }
+
+    if constexpr (PROBS) {
+        this->probsGM.SetGlobalBuffer((__gm__ T3 *)probs + block_offset, block_length);
+        this->pipe.InitBuffer(probs_inque, 1, block_splited_length * (sizeof(T3)));
+        if constexpr (!IsSameType<T3, float>::value) {
+            this->pipe.InitBuffer(temp_buffer2, block_splited_length * sizeof(float));
+            this->probs_tensor = this->temp_buffer2.template Get<float>();
+        }
+    }
+};
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::Process()
+{
+
+    if (blockIdx >= blockNum) {
+        return;
+    }
+    for (int64_t i = 0; i < this->tokens_splited_num; ++i) {
+        CalMultiOutToken(i * this->tokens_splited_length, this->tokens_splited_length);
+    }
+    // Handle the tail block when tokens_num is not evenly divisible by core count
+    if (this->tokens_splited_remain > 0) {
+        CalMultiOutToken(this->tokens_splited_num * this->tokens_splited_length, this->tokens_splited_remain);
+    }
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CalMultiOutToken(const int64_t out_offset,
+                                                                                    const int64_t out_tokens_number)
+{
+    this->indicesLocal = this->indices_inque.template AllocTensor<T2>();
+    int64_t in_offset = out_offset * this->top_k;
+    this->copyParams.blockLen = out_tokens_number * this->top_k * sizeof(T2);
+    DataCopyPad(this->indicesLocal, this->indicesGM[in_offset], this->copyParams, this->extParams2);
+    this->indices_inque.template EnQue(this->indicesLocal);
+
+    if constexpr (PROBS) {
+        LocalTensor<T3> temp_probs_tensor = this->probs_inque.template AllocTensor<T3>();
+        this->copyParams.blockLen = out_tokens_number * this->top_k * sizeof(T3);
+        DataCopyPad(temp_probs_tensor, this->probsGM[in_offset], this->copyParams, this->extParams3);
+        this->probs_inque.template EnQue(temp_probs_tensor);
+        temp_probs_tensor = this->probs_inque.template DeQue<T3>();
+        if constexpr (!IsSameType<T3, float>::value) {
+            Cast(this->probs_tensor, temp_probs_tensor, RoundMode::CAST_NONE, out_tokens_number * this->top_k);
+            this->probs_inque.FreeTensor(temp_probs_tensor);
+            PipeBarrier<PIPE_V>();
+        } else {
+            this->probs_tensor = temp_probs_tensor;
+        }
+    }
+    this->indicesLocal = this->indices_inque.template DeQue<T2>();
+
+    
+    for (int64_t out_token_idx = 0; out_token_idx < out_tokens_number; ++out_token_idx) {
+        CalSingleOutToken(out_token_idx * this->top_k, out_offset + out_token_idx);
+    }
+    // Free Tensor
+    this->indices_inque.FreeTensor(this->indicesLocal);
+    if constexpr (PROBS && IsSameType<T3, float>::value) {
+        this->probs_inque.FreeTensor(this->probs_tensor);
+    }
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CalSingleOutToken(const int64_t start_token,
+                                                                                     const int64_t out_token_idx)
+{
+    for (int64_t h_index = 0; h_index < this->hidden_splited_num; ++h_index) {
+        CalPartOutToken(start_token, h_index, this->hidden_splited_length, out_token_idx);
+    }
+    // Handle the tail block when a full hidden_size does not fit in one pass
+    if (this->hidden_splited_remain > 0) {
+        CalPartOutToken(start_token, this->hidden_splited_num, this->hidden_splited_remain, out_token_idx);
+    }
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void
+KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CalPartOutToken(const int64_t start_token, const int64_t h_index,
+                                                            const int64_t h_length, const int64_t out_token_index)
+{
+    if constexpr (IsSameType<T1, float>::value) {
+        this->token_tensor0 = this->outque.template AllocTensor<T1>();
+    }
+    int64_t end_token = start_token + this->top_k;
+    T2 cal_token_idx = this->indicesLocal.GetValue(start_token);
+
+    // Handle the first token
+    if (cal_token_idx < this->num_out_tokens) {
+        float probsValue = 0;
+        if constexpr (PROBS) {
+            probsValue = this->probs_tensor.GetValue(start_token);
+        }
+
+        CopyTokenIn(cal_token_idx, h_index, h_length);
+        PipeBarrier<PIPE_V>();
+        CalFirstToken(probsValue, h_length);
+    } else {
+        PipeBarrier<PIPE_V>();
+        Duplicate(this->token_tensor0, static_cast<float>(0), h_length);
+    }
+
+    // Handle the remaining tokens
+    for (int64_t token_index = start_token + 1; token_index < end_token; ++token_index) {
+        cal_token_idx = this->indicesLocal.GetValue(token_index);
+        if (cal_token_idx < this->num_out_tokens) {
+            float probsValue = 0;
+            if constexpr (PROBS) {
+                probsValue = this->probs_tensor.GetValue(token_index);
+            }
+        
+            CopyTokenIn(cal_token_idx, h_index, h_length);
+            PipeBarrier<PIPE_V>();
+            CalToken(probsValue, h_length);
+        }
+    }
+
+    // Write out the computed result
+    CopyOut(out_token_index, h_index, h_length);
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CopyTokenIn(const T2 in_token_index,
+                                                                               const int64_t h_index,
+                                                                               const int64_t h_length)
+{
+    LocalTensor<T1> tokensLocal = this->tokens_inque.template AllocTensor<T1>();
+    int64_t offset = in_token_index * this->hidden_size + h_index * this->hidden_splited_length;
+
+    if (likely((h_length * sizeof(T1)) % BLOCK_SIZE == 0)) {
+        DataCopy(tokensLocal, this->tokensGM[offset], h_length);
+    } else {
+        this->copyParams.blockLen = h_length * sizeof(T1);
+        DataCopyPad(tokensLocal, this->tokensGM[offset], this->copyParams, this->extParams1);
+    }
+
+    this->tokens_inque.template EnQue(tokensLocal);
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CalFirstToken(const float prob_value,
+                                                                                 const int64_t h_length)
+{
+    LocalTensor<T1> tokensLocal = this->tokens_inque.template DeQue<T1>();
+
+    if constexpr (!IsSameType<T1, float>::value) {
+        Cast(this->token_tensor0, tokensLocal, RoundMode::CAST_NONE, h_length);
+    } else {
+        uint64_t byteAlign32 = (h_length * sizeof(float) + BLOCK_SIZE - 1) & ~(BLOCK_SIZE - 1);
+        DataCopy(this->token_tensor0, tokensLocal, byteAlign32 / sizeof(float));
+    }
+
+    this->tokens_inque.FreeTensor(tokensLocal);
+
+    if constexpr (PROBS) {
+        PipeBarrier<PIPE_V>();
+        Muls(this->token_tensor0, this->token_tensor0, prob_value, h_length);
+    }
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CalToken(const float prob_value,
+                                                                            const int64_t h_length)
+{
+    LocalTensor<T1> tokensLocal = this->tokens_inque.template DeQue<T1>();
+
+    if constexpr (!IsSameType<T1, float>::value) {
+        Cast(this->token_tensor1, tokensLocal, RoundMode::CAST_NONE, h_length);
+        this->tokens_inque.FreeTensor(tokensLocal);
+        if constexpr (PROBS) {
+            PipeBarrier<PIPE_V>();
+            Muls(this->token_tensor1, this->token_tensor1, prob_value, h_length);
+        }
+        PipeBarrier<PIPE_V>();
+        Add(this->token_tensor0, this->token_tensor0, this->token_tensor1, h_length);
+    } else {
+        if constexpr (PROBS) {
+            Muls(tokensLocal, tokensLocal, prob_value, h_length);
+            PipeBarrier<PIPE_V>();
+        }
+        Add(this->token_tensor0, this->token_tensor0, tokensLocal, h_length);
+        this->tokens_inque.FreeTensor(tokensLocal);
+    }
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CopyOut(const int64_t out_token_index,
+                                                                           const int64_t h_index,
+                                                                           const int64_t h_length)
+{
+    LocalTensor<T1> temp_out_tensors;
+    if constexpr (!IsSameType<T1, float>::value) {
+        temp_out_tensors = this->outque.template AllocTensor<T1>();
+        PipeBarrier<PIPE_V>();
+        Cast(temp_out_tensors, this->token_tensor0, RoundMode::CAST_RINT, h_length);
+    } else {
+        temp_out_tensors = this->token_tensor0;
+    }
+
+    this->outque.template EnQue<T1>(temp_out_tensors);
+    temp_out_tensors = this->outque.template DeQue<T1>();
+
+    int64_t offset = out_token_index * this->hidden_size + h_index * this->hidden_splited_length;
+    if (likely((h_length * sizeof(T1)) % BLOCK_SIZE == 0)) {
+        DataCopy(this->outGM[offset], temp_out_tensors, h_length);
+    } else {
+        this->copyParams.blockLen = h_length * sizeof(T1);
+        DataCopyPad(this->outGM[offset], temp_out_tensors, this->copyParams);
+    }
+
+    this->outque.FreeTensor(temp_out_tensors);
+}
+#endif // MOE_TOKEN_UNPERMUTE

csrc/mc2/dispatch_ffn_combine/op_kernel/unpermute/moe_token_unpermute_tiling.h

@@ -0,0 +1,38 @@
+#ifndef MOE_TOKEN_UNPERMUTE_TILING
+#define MOE_TOKEN_UNPERMUTE_TILING
+
+struct MoeTokenUnpermuteTilingData {
+    int64_t hidden_size;
+    int64_t top_k;
+    int64_t num_out_tokens;
+    int64_t hidden_splited_length;
+    int64_t hidden_splited_num;
+    int64_t hidden_splited_remain;
+    int64_t tokens_core_length;
+    int64_t tokens_core_remain;
+    int64_t tokens_splited_length;
+    int64_t tokens_splited_num;
+    int64_t tokens_splited_remain;
+    int64_t buffer_num;
+};
+
+__forceinline__ [host, aicore] void
+MoeTokenUnpermuteTiling(int32_t m, int32_t n, int32_t topK, MoeTokenUnpermuteTilingData &tilingData, uint32_t coreNum)
+{
+    #define I64(x) static_cast<int64_t>(x)
+    tilingData.hidden_size = I64(n);
+    tilingData.top_k = I64(topK);
+    tilingData.num_out_tokens = I64(m);
+    tilingData.hidden_splited_length = tilingData.hidden_size;
+    tilingData.hidden_splited_num = 1;
+    tilingData.hidden_splited_remain = 0;
+    uint32_t outTokens = m / topK;
+    tilingData.tokens_core_length = I64(outTokens / coreNum);
+    tilingData.tokens_core_remain = I64(outTokens % coreNum);
+    tilingData.tokens_splited_length = I64(min(tilingData.tokens_core_length, 600));
+    tilingData.tokens_splited_num = I64(tilingData.tokens_core_length / tilingData.tokens_splited_length);
+    tilingData.tokens_splited_remain = I64(tilingData.tokens_core_length % tilingData.tokens_splited_length);
+    tilingData.buffer_num = 4;
+}
+
+#endif

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_row.hpp

@@ -0,0 +1,216 @@
+/*
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_ROW_HPP
+#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_ROW_HPP
+
+#include "catlass/catlass.hpp"
+#include "catlass/arch/resource.hpp"
+#include "catlass/epilogue/dispatch_policy.hpp"
+#include "catlass/gemm_coord.hpp"
+#include "catlass/matrix_coord.hpp"
+#include "catlass/layout/layout.hpp"
+#include "catlass/detail/callback.hpp"
+#include "catlass/epilogue/block/block_epilogue.hpp"
+
+namespace Catlass::Epilogue::Block {
+
+// float scale, dequant per expert
+template <
+    uint32_t UB_STAGES_,
+    class CType_,
+    class LayoutPerTokenScale_,
+    class DType_,
+    class TileCopy_
+>
+class BlockEpilogue <
+    EpilogueAtlasA2PerTokenDequant<UB_STAGES_>,
+    CType_,
+    Gemm::GemmType<float, LayoutPerTokenScale_>,
+    DType_,
+    TileCopy_
+> {
+public:
+    using DispatchPolicy = EpilogueAtlasA2PerTokenDequant<UB_STAGES_>;
+    using ArchTag = typename DispatchPolicy::ArchTag;
+    static constexpr uint32_t UB_STAGES = UB_STAGES_;
+
+    // Data infos
+    using ElementC = typename CType_::Element;
+    using LayoutC = typename CType_::Layout;
+    using ElementPerTokenScale = float;
+    using LayoutPerTokenScale = LayoutPerTokenScale_;
+    using ElementD = typename DType_::Element;
+    using LayoutD = typename DType_::Layout;
+
+    // Check data infos
+    static_assert(
+        std::is_same_v<ElementC, half> && (std::is_same_v<ElementD, half> || std::is_same_v<ElementD, bfloat16_t>),
+        "The element type template parameters of BlockEpilogue are wrong"
+    );
+    static_assert(
+        std::is_same_v<LayoutC, layout::RowMajor> && 
+            std::is_same_v<LayoutPerTokenScale, layout::VectorLayout> && std::is_same_v<LayoutD, layout::RowMajor>,
+        "The layout template parameters of BlockEpilogue are wrong"
+    );
+
+
+    // Tile copy
+    using CopyGmToUbC = typename TileCopy_::CopyGmToUbC;
+    using CopyUbToGmD = typename TileCopy_::CopyUbToGmD;
+
+    struct Params {
+        __gm__ int32_t *ptrTokenPerExpert{nullptr};
+        int32_t EP;
+        int32_t expertPerRank;
+        int32_t n2;
+
+        CATLASS_DEVICE
+        Params() {};
+
+        CATLASS_DEVICE
+        Params(int32_t EP_, int32_t expertPerRank_, __gm__ int32_t *ptrTokenPerExpert_, int32_t n2_) : ptrTokenPerExpert(ptrTokenPerExpert_), EP(EP_), expertPerRank(expertPerRank_), n2(n2_) {}
+    };
+
+    CATLASS_DEVICE
+    BlockEpilogue(Arch::Resource<ArchTag> const &resource, Params const &params = Params{}) : params(params)
+    {
+        size_t ubOffset = 0;
+        int32_t eventVMTE2 = 0;
+        int32_t eventMTE2V = 0;
+        int32_t eventMTE3V = 0;
+        int32_t eventVMTE3 = 0;
+        int32_t blockN = params.n2;
+        for (uint32_t i = 0; i < UB_STAGES; ++i) {
+            ubCList[i] = resource.ubBuf.template GetBufferByByte<ElementC>(ubOffset);
+            ubOffset += blockN * sizeof(ElementC);
+            ubDList[i] = resource.ubBuf.template GetBufferByByte<ElementD>(ubOffset);
+            ubOffset += blockN * sizeof(ElementD);
+
+            eventUbCVMTE2List[i] = eventVMTE2++;
+            eventUbCMTE2VList[i] = eventMTE2V++;
+            eventUbDMTE3VList[i] = eventMTE3V++;
+            eventUbDVMTE3List[i] = eventVMTE3++;
+
+
+            ubCFp32List[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += blockN * sizeof(float);
+        }
+    }
+    CATLASS_DEVICE
+    void SetFlag() 
+    {
+        for (uint32_t i = 0; i < UB_STAGES; ++i) {
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[i]);
+            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[i]);
+        }
+    }
+    
+    CATLASS_DEVICE
+    void Finalize()
+    {
+        for (uint32_t i = 0; i < UB_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[i]);
+            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[i]);
+        }
+    }
+    CATLASS_DEVICE
+    ~BlockEpilogue()
+    {
+        
+    }
+
+    CATLASS_DEVICE
+    void UpdateParams(Params const &params_)
+    {
+        params = params_;
+    }
+
+    CATLASS_DEVICE
+    void operator() (
+        AscendC::GlobalTensor<ElementC> const &gmC,
+        MatrixCoord const &shapeC,
+        AscendC::GlobalTensor<ElementPerTokenScale> const &gmPerTokenScale,
+        AscendC::GlobalTensor<ElementD> const &gmD
+    )
+    {
+        uint32_t blockM = shapeC.row();
+        uint32_t blockN = shapeC.column();
+
+        uint32_t tileLoops = blockM;
+
+        for (uint32_t loopIdx = 0; loopIdx < tileLoops; loopIdx ++) {
+            auto gmTileC = gmC[loopIdx * blockN];
+            auto &ubC = ubCList[ubListId];
+            auto &ubCFp32 = ubCFp32List[ubListId];
+            auto &ubMul = ubMulList[ubListId];
+            auto &ubD = ubDList[ubListId];
+            auto gmTileD = gmD[loopIdx * blockN];
+            LayoutC layoutUbC{1, blockN};
+
+            // Move C from GM workspace to UB
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[ubListId]);
+            copyGmToUbC(ubC, gmTileC, layoutUbC, layoutUbC);
+            AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(eventUbCMTE2VList[ubListId]);
+
+            // Cast C to FP32 in UB
+            AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(eventUbCMTE2VList[ubListId]);
+            AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_NONE, blockN);
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[ubListId]);
+
+            // Get per-token scale from row loopIdx of gmPerTokenScale
+            ElementPerTokenScale perTokenScale = gmPerTokenScale(loopIdx);
+
+            AscendC::SetFlag<AscendC::HardEvent::S_V>(0);
+            AscendC::WaitFlag<AscendC::HardEvent::S_V>(0);
+            // Multiply FP32 C by the per-token scale
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Muls(ubCFp32, ubCFp32, perTokenScale, blockN);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            // Cast the muls result back to fp16/bf16
+            LayoutD layoutUbD{1, blockN};
+            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[ubListId]);
+
+            AscendC::Cast(ubD, ubCFp32, AscendC::RoundMode::CAST_RINT, blockN);
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(eventUbDVMTE3List[ubListId]);
+            
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(eventUbDVMTE3List[ubListId]);
+            copyUbToGmD(gmTileD, ubD, layoutUbD, layoutUbD);
+            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[ubListId]);
+
+            ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0;
+        }
+    }
+
+private:
+    Params params;
+
+    AscendC::LocalTensor<ElementC> ubCList[UB_STAGES];
+    AscendC::LocalTensor<ElementD> ubDList[UB_STAGES];
+
+    int32_t eventUbCVMTE2List[UB_STAGES];
+    int32_t eventUbCMTE2VList[UB_STAGES];
+    int32_t eventUbDMTE3VList[UB_STAGES];
+    int32_t eventUbDVMTE3List[UB_STAGES];
+
+    uint32_t ubListId{0};
+
+    AscendC::LocalTensor<float> ubCFp32List[UB_STAGES];
+    AscendC::LocalTensor<float> ubMulList[UB_STAGES];
+
+
+    CopyGmToUbC copyGmToUbC;
+    CopyUbToGmD copyUbToGmD;
+};
+
+}  // namespace Catlass::Epilogue::Block
+
+#endif  // CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_ROW_HPP

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_swiglu.hpp

@@ -0,0 +1,330 @@
+/*
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_SWIGLU_HPP
+#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_SWIGLU_HPP
+
+#include "catlass/catlass.hpp"
+#include "catlass/arch/resource.hpp"
+#include "catlass/epilogue/dispatch_policy.hpp"
+#include "catlass/gemm_coord.hpp"
+#include "catlass/matrix_coord.hpp"
+#include "catlass/layout/layout.hpp"
+#include "catlass/detail/callback.hpp"
+
+namespace Catlass::Epilogue::Block {
+
+// float scale, dequant per expert
+template <
+    uint32_t UB_STAGES_,
+    class CType_,
+    class LayoutPerTokenScale_,
+    class DType_,
+    class TileElemWiseMuls_,
+    class TileCopy_
+>
+class BlockEpilogue <
+    EpilogueAtlasA2PerTokenDequantSwigluQuant<UB_STAGES_>,
+    CType_,
+    Gemm::GemmType<float, LayoutPerTokenScale_>,
+    DType_,
+    TileElemWiseMuls_,
+    TileCopy_
+> {
+public:
+    using DispatchPolicy = EpilogueAtlasA2PerTokenDequantSwigluQuant<UB_STAGES_>;
+    using ArchTag = typename DispatchPolicy::ArchTag;
+    static constexpr uint32_t UB_STAGES = UB_STAGES_;
+
+    // Data infos
+    using ElementC = typename CType_::Element;
+    using LayoutC = typename CType_::Layout;
+    using ElementPerTokenScale = float;
+    using LayoutPerTokenScale = LayoutPerTokenScale_;
+    using ElementD = typename DType_::Element;
+    using LayoutD = typename DType_::Layout;
+
+    // Check data infos
+    static_assert(
+        std::is_same_v<ElementC, half> && (std::is_same_v<ElementD, float> || std::is_same_v<ElementD, int8_t>),
+        "The element type template parameters of BlockEpilogue are wrong"
+    );
+    static_assert(
+        std::is_same_v<LayoutC, layout::RowMajor> && 
+            std::is_same_v<LayoutPerTokenScale, layout::VectorLayout> && std::is_same_v<LayoutD, layout::RowMajor>,
+        "The layout template parameters of BlockEpilogue are wrong"
+    );
+
+    // Tile copy
+    using CopyGmToUbC = typename TileCopy_::CopyGmToUbC;
+    using CopyUbToGmD = typename TileCopy_::CopyUbToGmD;
+    using CopyUbToGmDequantScale = Epilogue::Tile::CopyUb2Gm<ArchTag, Gemm::GemmType<ElementPerTokenScale, LayoutPerTokenScale>>;
+
+    struct Params {
+        __gm__ ElementPerTokenScale *ptrPerTokenScale{nullptr};
+        LayoutPerTokenScale layoutPerTokenScale{};
+        __gm__ ElementD *ptrD{nullptr};
+        LayoutD layoutD{};
+
+        CATLASS_DEVICE
+        Params() {};
+
+        CATLASS_DEVICE
+        Params(__gm__ ElementPerTokenScale *ptrPerTokenScale_, LayoutPerTokenScale const &layoutPerTokenScale_,
+            __gm__ ElementD *ptrD_, LayoutD const &layoutD_
+        ) : ptrPerTokenScale(ptrPerTokenScale_), layoutPerTokenScale(layoutPerTokenScale_),
+            ptrD(ptrD_), layoutD(layoutD_) {}
+    };
+
+    CATLASS_DEVICE
+    BlockEpilogue(Arch::Resource<ArchTag> const &resource, int32_t n, Params const &params = Params{}) : params(params)
+    {
+        int32_t eventVMTE2 = 0;
+        int32_t eventMTE2V = 0;
+        int32_t eventMTE3V = 0;
+        int32_t eventVMTE3 = 0;
+        uint32_t blockN = n;
+        uint32_t ChunkTileLen = blockN / 2;
+        uint32_t HalfChunkTileLen = ChunkTileLen / 2;
+
+        for (uint32_t i = 0; i < UB_STAGES; ++i) {
+            ubCList[i] = resource.ubBuf.template GetBufferByByte<ElementC>(ubOffset);
+            ubOffset += blockN * sizeof(ElementC);
+            ubDList[i] = resource.ubBuf.template GetBufferByByte<ElementD>(ubOffset);
+            ubOffset += blockN * sizeof(ElementD);
+            ubCFp32List[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += blockN * sizeof(float);
+            ubCFp32ChunkNList[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += ChunkTileLen * sizeof(float);
+            ubCFp32ChunkNAbsList[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += ChunkTileLen * sizeof(float);
+            ubCFp32ChunkNMaxList[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += HalfChunkTileLen * sizeof(float);
+            ubQuantS32List[i] = ubCFp32ChunkNAbsList[i].template ReinterpretCast<int32_t>();
+            ubQuantF16List[i] = ubCFp32ChunkNAbsList[i].template ReinterpretCast<half>();
+
+            eventUbCVMTE2List[i] = eventVMTE2++;
+            eventUbCMTE2VList[i] = eventMTE2V++;
+            eventUbDMTE3VList[i] = eventMTE3V++;
+            eventUbDVMTE3List[i] = eventVMTE3++;
+
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[i]);
+            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[i]);  
+        }
+        ubPerTokenScaleOutput = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+    }
+    CATLASS_DEVICE
+    void Finalize()
+    {
+        for (uint32_t i = 0; i < UB_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[i]);
+            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[i]);
+        }
+    }
+    CATLASS_DEVICE
+    ~BlockEpilogue()
+    {
+    }
+
+    CATLASS_DEVICE
+    void UpdateParams(Params const &params_)
+    {
+        params = params_;
+    }
+    // Each tile is 1x7168, and each block covers all tokens for one expert = [group[i], 7168]
+    CATLASS_DEVICE
+    void operator() (
+        AscendC::GlobalTensor<ElementC> const &gmC,
+        MatrixCoord const &shapeC,
+        AscendC::GlobalTensor<ElementPerTokenScale> const &gmPerTokenScale1,
+        AscendC::GlobalTensor<ElementD> const &gmD,
+        AscendC::GlobalTensor<ElementPerTokenScale> const &gmPerTokenScale2,
+        Arch::Resource<ArchTag> const &resource,
+
+        uint32_t epilogueCoreNum = 40,
+        float swigluLimit = 0.0f,
+        uint32_t blockK = 1,
+        Callback &&callback = Callback{}
+    )
+    {
+        callback();
+        uint32_t blockM = shapeC.row();
+        uint32_t blockN = shapeC.column();
+
+        uint32_t tileLoops = blockM;
+        uint32_t subblockIdx = get_block_idx() + get_subblockid() * get_block_num();
+
+        uint32_t subblockNum = get_block_num() * 2;
+        uint32_t moveDataCoreNum = subblockNum - epilogueCoreNum;
+
+        if (subblockIdx < moveDataCoreNum) {
+            return;
+        }
+        uint32_t epilogueCoreIdx = subblockIdx - moveDataCoreNum;
+
+        uint32_t perCoreData =  blockM / epilogueCoreNum;
+        uint32_t remainderData = blockM % epilogueCoreNum;
+
+        uint32_t scaleBlock = (perCoreData * sizeof(ElementPerTokenScale) + 31) / 32 * 32;
+        sharedTmpBuffer = resource.ubBuf.template GetBufferByByte<uint8_t>(ubOffset + scaleBlock);
+
+        uint32_t tasksForIdx  = epilogueCoreIdx < remainderData ? perCoreData + 1 : perCoreData;
+        uint32_t loopStartIdx = epilogueCoreIdx * perCoreData + (epilogueCoreIdx < remainderData? epilogueCoreIdx : remainderData);
+
+        uint32_t alignedPerCoreData = RoundUp<BYTE_PER_BLK / sizeof(ElementPerTokenScale)>(perCoreData + 1);
+
+        uint32_t ChunkTileLen = blockN / 2;
+        uint32_t HalfChunkTileLen = ChunkTileLen / 2;
+
+
+        for (uint32_t loopIdx = loopStartIdx; loopIdx < loopStartIdx + tasksForIdx; ++loopIdx) {
+
+            auto gmTileC = gmC[loopIdx * blockK];
+
+            auto &ubC = ubCList[ubListId];
+            auto &ubD = ubDList[ubListId];
+
+            auto &ubCFp32 = ubCFp32List[ubListId];
+            auto &ubCFp32ChunkN = ubCFp32ChunkNList[ubListId]; 
+            auto &ubAbs = ubCFp32ChunkNAbsList[ubListId];
+            // auto &ubMax = ubCFp32ChunkNMaxList[ubListId];
+            auto &ubReduceMax = ubCFp32ChunkNMaxList[ubListId];
+            auto &ubOutputTmp = ubAbs;
+            auto &sharedUbTmpBuffer = ubReduceMax;
+            auto &ubQuantS32 = ubQuantS32List[ubListId];
+            auto &ubQuantF16 = ubQuantF16List[ubListId];
+
+            auto gmTileD = gmD[loopIdx * ChunkTileLen];
+            LayoutC layoutUbC{1, blockN};
+
+            // Move C from GM workspace to UB
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[ubListId]);
+            copyGmToUbC(ubC, gmTileC, layoutUbC, layoutUbC);
+            AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(eventUbCMTE2VList[ubListId]);
+
+            // Cast C to FP32 in UB
+            AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(eventUbCMTE2VList[ubListId]);
+            AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_NONE, blockN);
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[ubListId]);
+
+            // Get per-token scale from row loopIdx of gmPerTokenScale
+            ElementPerTokenScale perTokenScale = gmPerTokenScale1(loopIdx);
+
+            AscendC::SetFlag<AscendC::HardEvent::S_V>(0);
+            AscendC::WaitFlag<AscendC::HardEvent::S_V>(0);
+            // Multiply FP32 C by the per-token scale
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Muls(ubCFp32, ubCFp32, perTokenScale, blockN);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            // swiglu limit clamp
+            if (swigluLimit > 0.0f) {
+                AscendC::ClampMax(ubCFp32, ubCFp32, sharedTmpBuffer, swigluLimit, blockN);
+                AscendC::PipeBarrier<PIPE_V>();
+                AscendC::ClampMin(ubCFp32[ChunkTileLen], ubCFp32[ChunkTileLen], sharedTmpBuffer, -1.0f * swigluLimit, ChunkTileLen);
+                //AscendC::ClampMin(ubCFp32, ubCFp32, sharedTmpBuffer, -1.0f * swigluLimit, ChunkTileLen);
+                AscendC::PipeBarrier<PIPE_V>();
+            }
+            
+            // Swiglu computation process
+            AscendC::Muls(ubCFp32ChunkN, ubCFp32, -1.0f, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Exp(ubCFp32ChunkN, ubCFp32ChunkN, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Adds(ubCFp32ChunkN, ubCFp32ChunkN, 1.0f, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            // TODO: confirm whether the division impacts subsequent data
+            AscendC::Div(ubCFp32ChunkN, ubCFp32, ubCFp32ChunkN, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Mul(ubCFp32ChunkN, ubCFp32ChunkN, ubCFp32[ChunkTileLen], ChunkTileLen);
+            
+            // Quantization process; difference between the two approaches
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Abs(ubAbs, ubCFp32ChunkN, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            AscendC::ReduceMax<float>(ubReduceMax, ubAbs, sharedUbTmpBuffer, ChunkTileLen, false);
+            AscendC::PipeBarrier<PIPE_V>();
+            
+            AscendC::SetFlag<AscendC::HardEvent::V_S>(0);
+            AscendC::WaitFlag<AscendC::HardEvent::V_S>(0);
+
+            // TODO: compare the efficiency of the two calculation methods
+            ElementPerTokenScale GMubDequantScale = ubReduceMax.GetValue(0);
+            AscendC::SetFlag<AscendC::HardEvent::S_V>(0);
+
+            auto ubPerTokenScaleOutputOffset = loopIdx - loopStartIdx;
+            ubPerTokenScaleOutput.SetValue(ubPerTokenScaleOutputOffset, GMubDequantScale / 127.f);
+
+            AscendC::WaitFlag<AscendC::HardEvent::S_V>(0);
+            AscendC::Muls(ubOutputTmp, ubCFp32ChunkN, 127.f / GMubDequantScale, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            AscendC::Cast(ubQuantS32, ubOutputTmp, AscendC::RoundMode::CAST_RINT, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::SetDeqScale(static_cast<half>(1.0));
+            AscendC::Cast(ubQuantF16, ubQuantS32, AscendC::RoundMode::CAST_RINT, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventUbDVMTE3List[ubListId]);
+            AscendC::Cast(ubD, ubQuantF16, AscendC::RoundMode::CAST_RINT, ChunkTileLen);
+            // AscendC::Muls(ubD, ubCFp32ChunkN, 127.f / GMubDequantScale, ChunkTileLen);
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(eventUbDMTE3VList[ubListId]);         
+
+            LayoutD layoutUbD{1, ChunkTileLen};
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(eventUbDVMTE3List[ubListId]);
+            copyUbToGmD(gmTileD, ubD, layoutUbD, layoutUbD);
+            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[ubListId]);
+            ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0;
+        }
+
+        if(tasksForIdx > 0){
+            LayoutPerTokenScale layoutGmPerTokenScale2{tasksForIdx};
+
+            AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
+            AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
+
+            copyUbToGmDequantScale(gmPerTokenScale2[loopStartIdx], ubPerTokenScaleOutput[0], layoutGmPerTokenScale2, layoutGmPerTokenScale2);
+        }
+        
+
+    }
+
+private:
+    Params params;
+
+    AscendC::LocalTensor<ElementC> ubCList[UB_STAGES];
+    AscendC::LocalTensor<ElementD> ubDList[UB_STAGES];
+
+    int32_t eventUbCVMTE2List[UB_STAGES];
+    int32_t eventUbCMTE2VList[UB_STAGES];
+    int32_t eventUbDMTE3VList[UB_STAGES];
+    int32_t eventUbDVMTE3List[UB_STAGES];
+
+    uint32_t ubListId{0};
+    size_t ubOffset = 0;
+
+    AscendC::LocalTensor<float> ubCFp32List[UB_STAGES];
+    AscendC::LocalTensor<float> ubCFp32ChunkNList[UB_STAGES];
+    AscendC::LocalTensor<float> ubCFp32ChunkNAbsList[UB_STAGES];
+    AscendC::LocalTensor<float> ubCFp32ChunkNMaxList[UB_STAGES];
+    AscendC::LocalTensor<int32_t> ubQuantS32List[UB_STAGES];
+    AscendC::LocalTensor<half> ubQuantF16List[UB_STAGES];
+    AscendC::LocalTensor<float> ubPerTokenScaleOutput;
+    AscendC::LocalTensor<uint8_t> sharedTmpBuffer;
+
+    CopyGmToUbC copyGmToUbC;
+    CopyUbToGmD copyUbToGmD;
+    CopyUbToGmDequantScale copyUbToGmDequantScale;
+};
+
+}  // namespace Catlass::Epilogue::Block
+
+#endif  // CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_SWIGLU_HPP

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/block_epilogue_pertoken_v2.hpp

@@ -0,0 +1,255 @@
+#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_V2_ONLY_HPP
+#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_V2_ONLY_HPP
+
+#include "catlass/catlass.hpp"
+#include "catlass/arch/resource.hpp"
+#include "catlass/epilogue/dispatch_policy.hpp"
+#include "catlass/gemm_coord.hpp"
+#include "catlass/matrix_coord.hpp"
+#include "catlass/layout/layout.hpp"
+#include "catlass/detail/callback.hpp"
+
+#include "hccl_shmem.hpp"
+#include "layout3d.hpp"
+
+namespace Catlass::Epilogue::Block {
+template <
+    uint32_t UB_STAGES_,
+    class CType_,
+    class LayoutPerTokenScale_,
+    class DType_,
+    class TileCopy_
+>
+class BlockEpilogue <
+    EpilogueAtlasA2PerTokenDequantV2<UB_STAGES_>,
+    CType_,
+    Gemm::GemmType<float, LayoutPerTokenScale_>,
+    DType_,
+    TileCopy_
+> {
+public:
+    using DispatchPolicy = EpilogueAtlasA2PerTokenDequantV2<UB_STAGES_>;
+    using ArchTag = typename DispatchPolicy::ArchTag;
+    static constexpr uint32_t UB_STAGES = UB_STAGES_;
+
+    // Data infos
+    using ElementC = typename CType_::Element;
+    using LayoutC = typename CType_::Layout;
+    using ElementPerTokenScale = float;
+    using LayoutPerTokenScale = LayoutPerTokenScale_;
+    using ElementD = typename DType_::Element;
+    using LayoutD = typename DType_::Layout;
+
+    using CopyScaleGmToUb = Epilogue::Tile::CopyGm2Ub<ArchTag, Gemm::GemmType<float, layout::VectorLayout>>;
+    using CopyGmToUbC = typename TileCopy_::CopyGmToUbC;
+    using CopyUbToGmD = typename TileCopy_::CopyUbToGmD;
+
+    struct Params {
+        __gm__ int32_t *ptrTokenPerExpert{nullptr};
+        int32_t EP;
+        int32_t expertPerRank;
+        int32_t n2;
+        LayoutC layoutC;
+        int32_t n0;
+        int32_t rank;
+        HcclShmem shmem;
+        int32_t offsetD;
+        Layout3D tokenPerExpertLayout;
+        CATLASS_DEVICE
+        Params() {};
+        CATLASS_DEVICE
+        Params(int32_t EP_, int32_t expertPerRank_, int32_t rank_, __gm__ int32_t *ptrTokenPerExpert_, 
+        LayoutC layoutC_, int32_t n2_, int32_t n0_, HcclShmem& shmem_, int32_t offsetD_, Layout3D tokenPerExpertLayout_) : 
+        ptrTokenPerExpert(ptrTokenPerExpert_), EP(EP_), 
+        expertPerRank(expertPerRank_),rank(rank_), layoutC(layoutC_), n2(n2_), n0(n0_),
+        shmem(shmem_), offsetD(offsetD_), tokenPerExpertLayout(tokenPerExpertLayout_)
+         {}
+    };
+
+
+    CATLASS_DEVICE
+    BlockEpilogue(Arch::Resource<ArchTag> const &resource, Params const &params = Params{}) : params(params)
+    {
+        //ub:192KB
+        n0 = params.n0;
+        size_t ubOffset = 0;
+        for(int32_t i = 0; i < 2; i++) {
+            ubCList[i] = resource.ubBuf.template GetBufferByByte<ElementC>(ubOffset);
+            ubOffset += max_len * sizeof(ElementC);
+            ubDList[i] = resource.ubBuf.template GetBufferByByte<ElementD>(ubOffset);
+            ubOffset += max_len * sizeof(ElementD);
+            ubFp32List[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += max_len * sizeof(float);
+            scaleUbList[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += (max_len / n0) * sizeof(float);
+            source_scale_offset[i] = -1;
+        }
+        tokenPerExpert.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(params.ptrTokenPerExpert));
+        tokenPerExpertLayout = params.tokenPerExpertLayout;
+        is_ping = true;
+    }
+    CATLASS_DEVICE
+    void SetFlag()
+    {
+        AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
+        AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID1);
+        AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
+        AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
+        AscendC::SetFlag<AscendC::HardEvent::S_MTE2>(EVENT_ID2);
+        AscendC::SetFlag<AscendC::HardEvent::S_MTE2>(EVENT_ID3);
+        AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
+        AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID1);
+    }
+
+    CATLASS_DEVICE
+    void Finalize()
+    {
+        AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
+        AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID1);
+        AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
+        AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
+        AscendC::WaitFlag<AscendC::HardEvent::S_MTE2>(EVENT_ID2);
+        AscendC::WaitFlag<AscendC::HardEvent::S_MTE2>(EVENT_ID3);
+        AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
+        AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID1);
+
+    }
+    CATLASS_DEVICE
+    ~BlockEpilogue()
+    {
+        
+    }
+    CATLASS_DEVICE
+    void operator() (
+        AscendC::GlobalTensor<ElementC> const &gmC,
+        AscendC::GlobalTensor<ElementPerTokenScale> const &gmPerTokenScale,
+        GemmCoord& blockCoord,
+        GemmCoord& actualBlockShape,
+        int32_t groupIdx,
+        int32_t preSrcExpertSum,
+        AscendC::GlobalTensor<int32_t> preSumBeforeRank
+    ){
+        is_ping = !is_ping;
+        auto event_id = is_ping ? EVENT_ID0 : EVENT_ID1;
+        auto event_id_2 = is_ping ? EVENT_ID2 : EVENT_ID3;
+
+        auto &ubC = ubCList[is_ping];
+        auto &ubD = ubDList[is_ping];
+        int32_t gmCOffset = preSrcExpertSum * params.n2 + blockCoord.m() * params.n2 + blockCoord.n();
+        auto gmTileC = gmC[gmCOffset];
+        auto &ubCFp32 = ubFp32List[is_ping];
+        auto &scaleUb = scaleUbList[is_ping];
+
+        LayoutC layoutGM{actualBlockShape.m(), actualBlockShape.n(), params.n2};
+        LayoutC layoutUB{actualBlockShape.m(), actualBlockShape.n(), n0};
+
+
+        AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(event_id);
+        copyGmToUbC(ubC, gmTileC, layoutUB, layoutGM);
+        AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(event_id);
+
+        AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(event_id);
+        AscendC::Cast<float, ElementC, false>(ubCFp32, ubC, AscendC::RoundMode::CAST_NONE, -1, repeat, {1, 1, 8, 4});
+        AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(event_id);
+
+
+        AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(event_id_2);
+        AscendC::WaitFlag<AscendC::HardEvent::S_MTE2>(event_id_2);
+
+        int32_t gmScaleOffset = preSrcExpertSum + blockCoord.m();
+        layout::VectorLayout scaleLauout{actualBlockShape.m()};
+        if (source_scale_offset[event_id] != gmScaleOffset) {
+                source_scale_offset[event_id] = gmScaleOffset;
+                copyScaleGmToUb(scaleUb, gmPerTokenScale[gmScaleOffset],  scaleLauout, scaleLauout);
+        }
+
+        AscendC::SetFlag<AscendC::HardEvent::MTE2_S>(event_id_2);
+        AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(event_id_2);
+
+        
+
+        
+        AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(event_id_2);
+        AscendC::WaitFlag<AscendC::HardEvent::MTE2_S>(event_id_2); // Note that the value must be MTE2_S instead of MTE2_V.
+                                                                   // Otherwise, 0 will be read, causing garbled characters.
+        AscendC::PipeBarrier<PIPE_V>();
+        for (int32_t row = 0; row < actualBlockShape.m(); ++row) {
+                float scale = scaleUb(row);
+                Muls<float, false>(ubCFp32[n0* row], ubCFp32[n0 * row] , scale, -1, (actualBlockShape.n() + 127) / 128 * 2, {1, 1, 8, 8});
+        }
+        AscendC::PipeBarrier<PIPE_V>();
+        AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(event_id);
+        AscendC::Cast<ElementD, float, false>(ubD, ubCFp32, AscendC::RoundMode::CAST_RINT, -1, repeat, {1, 1, 4, 8});
+        AscendC::SetFlag<AscendC::HardEvent::S_MTE2>(event_id_2);
+        AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(event_id_2);
+        AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(event_id);
+
+        int32_t lenTile = actualBlockShape.m();
+        int32_t stTile = blockCoord.m();
+        int32_t edTile = stTile + lenTile;
+        int32_t preSumRankInExpert = 0;
+        int32_t tileOffset = 0;
+
+        AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(event_id);
+        for (int32_t dstEpIdx = 0; dstEpIdx < params.EP; dstEpIdx ++) {
+            int32_t lenRankInExpert = tokenPerExpert(tokenPerExpertLayout(dstEpIdx, params.rank, groupIdx));
+            int32_t dstExpertOffset = preSumBeforeRank(dstEpIdx * params.expertPerRank + groupIdx);
+            int32_t stRankInExpert = preSumRankInExpert;
+            int32_t edRankInExpert = stRankInExpert + lenRankInExpert;
+            preSumRankInExpert += lenRankInExpert;
+            if (stRankInExpert >= edTile) {
+                break;
+            }
+            else if (edRankInExpert <= stTile) {
+                continue;
+            }
+            int32_t stData = max(stRankInExpert, stTile);
+            int32_t edData = min(edRankInExpert, edTile);
+            uint32_t lenData = edData - stData;
+            if (lenData <= 0){
+                continue;
+            }
+            
+            uint32_t dstOffsetInExpert = 0;
+            if (stTile > stRankInExpert) {
+                dstOffsetInExpert = stTile - stRankInExpert;
+            }
+            AscendC::GlobalTensor<ElementD> gmRemotePeer;
+            __gm__ void* dstPeermemPtr = params.shmem(params.offsetD, dstEpIdx);
+            gmRemotePeer.SetGlobalBuffer(reinterpret_cast<__gm__ ElementD*>(dstPeermemPtr));
+            MatrixCoord dstOffset{dstOffsetInExpert + dstExpertOffset, blockCoord.n()};
+            int64_t gmDstOffset = params.layoutC.GetOffset(dstOffset);
+            auto gmTileD = gmRemotePeer[gmDstOffset];
+            LayoutC layoutGM2{lenData, actualBlockShape.n(), params.n2};
+            LayoutC layoutUB2{lenData, actualBlockShape.n(), n0};
+            copyUbToGmD(gmTileD, ubD[tileOffset *  n0], layoutGM2, layoutUB2);
+            tileOffset += lenData;
+        }
+        AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(event_id);
+
+    }
+private:
+
+    Params params;
+    AscendC::LocalTensor<ElementC> ubCList[UB_STAGES];
+    AscendC::LocalTensor<ElementD> ubDList[UB_STAGES];
+    AscendC::LocalTensor<float> ubFp32List[UB_STAGES];
+    AscendC::LocalTensor<float> scaleUbList[UB_STAGES];
+    int32_t source_scale_offset[UB_STAGES];
+
+    int32_t max_len = 8 * 32 / 4 * 128;
+    int32_t n0;
+    bool is_ping = false;
+
+    
+    int32_t repeat = 128;
+
+    CopyGmToUbC copyGmToUbC;
+    CopyUbToGmD copyUbToGmD;
+    
+    CopyScaleGmToUb copyScaleGmToUb;
+    AscendC::GlobalTensor<int32_t> tokenPerExpert;
+    Layout3D tokenPerExpertLayout;
+};
+}
+#endif

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/block_mmad_preload_async_fixpipe_quant.hpp

@@ -0,0 +1,554 @@
+/*
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_FIXPIPE_QUANT_HPP
+#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_FIXPIPE_QUANT_HPP
+
+#include "catlass/catlass.hpp"
+#include "catlass/arch/resource.hpp"
+#include "catlass/coord.hpp"
+#include "catlass/gemm_coord.hpp"
+#include "catlass/gemm/dispatch_policy.hpp"
+#include "catlass/gemm/helper.hpp"
+#include "dispatch_policy_custom.hpp"
+
+
+namespace Catlass::Gemm::Block {
+
+template<AscendC::HardEvent event>
+__aicore__ inline void SyncFlagFunc(int32_t eventID)
+{
+    AscendC::SetFlag<event>(eventID);
+    AscendC::WaitFlag<event>(eventID);
+}
+
+template <
+    uint32_t PRELOAD_STAGES_,
+    uint32_t L1_STAGES_,
+    uint32_t L0A_STAGES_,
+    uint32_t L0B_STAGES_,
+    uint32_t L0C_STAGES_,
+    bool ENABLE_UNIT_FLAG_,
+    bool ENABLE_SHUFFLE_K_,
+    class L1TileShape_,
+    class L0TileShape_,
+    class AType_,
+    class BType_,
+    class CType_,
+    class BiasType_,
+    class TileCopy_,
+    class TileMmad_
+>
+struct BlockMmad <
+    MmadAtlasA2PreloadAsyncFixpipe<
+        PRELOAD_STAGES_,
+        L1_STAGES_,
+        L0A_STAGES_,
+        L0B_STAGES_,
+        L0C_STAGES_,
+        ENABLE_UNIT_FLAG_,
+        ENABLE_SHUFFLE_K_
+    >,
+    L1TileShape_,
+    L0TileShape_,
+    AType_,
+    BType_,
+    CType_,
+    BiasType_,
+    TileCopy_,
+    TileMmad_
+> {
+public:
+    // Type Aliases
+    using DispatchPolicy = MmadAtlasA2PreloadAsyncFixpipe<
+        PRELOAD_STAGES_,
+        L1_STAGES_,
+        L0A_STAGES_,
+        L0B_STAGES_,
+        L0C_STAGES_,
+        ENABLE_UNIT_FLAG_,
+        ENABLE_SHUFFLE_K_
+    >;
+    using ArchTag = typename DispatchPolicy::ArchTag;
+    using L1TileShape = L1TileShape_;
+    using L0TileShape = L0TileShape_;
+    using ElementA = typename AType_::Element;
+    using LayoutA = typename AType_::Layout;
+    using ElementB = typename BType_::Element;
+    using LayoutB = typename BType_::Layout;
+    using ElementC = typename CType_::Element;
+    using LayoutC = typename CType_::Layout;
+    using TileMmad = TileMmad_;
+    using CopyGmToL1A = typename TileCopy_::CopyGmToL1A;
+    using CopyGmToL1B = typename TileCopy_::CopyGmToL1B;
+    using CopyGmToL1S = Gemm::Tile::CopyGmToL1<ArchTag, Gemm::GemmType<uint64_t, layout::VectorLayout>>;
+    using CopyL1ToFP = typename Gemm::Tile::QuantTileCopy<ArchTag, AType_, BType_, CType_, void, Catlass::Gemm::Tile::ScaleGranularity::PER_CHANNEL>::CopyL1ToFP;
+    using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A;
+    using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B;
+    
+    using ElementAccumulator =
+        typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
+    using CopyL0CToGm = typename std::conditional<
+        std::is_same_v<ElementA, int8_t>,
+        Gemm::Tile::CopyL0CToGm<ArchTag, ElementAccumulator, CType_, Gemm::Tile::ScaleGranularity::PER_CHANNEL>,
+        typename TileCopy_::CopyL0CToGm
+    >::type;
+    using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc;
+    using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc;
+    using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst;
+    using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst;
+    using LayoutCInL0 = layout::zN;
+
+    using L1AAlignHelper = Gemm::helper::L1AlignHelper<ElementA, LayoutA>;
+    using L1BAlignHelper = Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
+
+    static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES;
+    static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES;
+    static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES;
+    static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES;
+    static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES;
+
+    static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG;
+    static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K;
+
+    // L1 tile size
+    static constexpr uint32_t L1A_TILE_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA);
+    static constexpr uint32_t L1B_TILE_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB);
+    static constexpr uint32_t L1S_TILE_SIZE = L1TileShape::N * sizeof(int64_t);
+    // L0 tile size
+    static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA);
+    static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB);
+    static constexpr uint32_t L0C_TILE_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator);
+
+    // Check LayoutC
+    static_assert(std::is_same_v<LayoutC, layout::RowMajor>, "LayoutC only support RowMajor yet!");
+
+    // Check L1TileShape
+    static_assert(
+        (std::is_same_v<ElementA, int8_t> 
+            ? (L1A_TILE_SIZE + L1B_TILE_SIZE + L1S_TILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE
+            : (L1A_TILE_SIZE + L1B_TILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE),
+        "L1TileShape exceeding the L1 space for the given data type"
+    );
+
+    // Check L0TileShape
+    static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!");
+    static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!");
+    static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!");
+
+    static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N,
+        "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet");
+
+    static constexpr auto L1A_LAYOUT = LayoutAInL1::template MakeLayout<ElementA>(
+        L1TileShape::M, L1TileShape::K);
+    static constexpr auto L1B_LAYOUT = LayoutBInL1::template MakeLayout<ElementB>(
+        L1TileShape::K, L1TileShape::N);
+
+    CATLASS_DEVICE
+    BlockMmad(Arch::Resource<ArchTag> &resource, __gm__ int32_t* flagPtr = nullptr, int32_t expertPerRank = 0, 
+                uint32_t l1BufAddrStart = 0, uint32_t FpAddrStart = 0)
+    {
+        syncGroupIdx = 0;
+        ptrSoftFlagBase_ = flagPtr;
+        expertPerRank_ = expertPerRank;
+        InitL1(resource, l1BufAddrStart);
+        InitFpBuf(resource, FpAddrStart);
+        InitL0A(resource);
+        InitL0B(resource);
+        InitL0C(resource);
+    }
+
+    CATLASS_DEVICE
+    ~BlockMmad()
+    {
+        SynchronizeBlock();
+        for (uint32_t i = 0; i < L1_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[i]);
+            AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[i]);
+        }
+        for (uint32_t i = 0; i < L0A_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[i]);
+        }
+        for (uint32_t i = 0; i < L0B_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[i]);
+        }
+        for (uint32_t i = 0; i < L0C_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList[i]);
+        }
+        if constexpr (std::is_same_v<ElementA, int8_t>) {
+            AscendC::WaitFlag<AscendC::HardEvent::FIX_MTE2>(0);
+        }
+    }
+
+    CATLASS_DEVICE
+    void operator()(
+        AscendC::GlobalTensor<ElementA> const &gmBlockA, LayoutA const &layoutA,
+        AscendC::GlobalTensor<ElementB> const &gmBlockB, LayoutB const &layoutB,
+        AscendC::GlobalTensor<ElementC> const &gmBlockC, LayoutC const &layoutC,
+        AscendC::GlobalTensor<uint64_t> const &gmBlockS, layout::VectorLayout const &layoutScale,
+        GemmCoord const &actualShape, int32_t syncLoopIdx = -1, int32_t flag = 0
+    )
+    {
+        uint32_t kTileCount = CeilDiv<L1TileShape::K>(actualShape.k());
+
+        uint32_t mRound = RoundUp<L1AAlignHelper::M_ALIGNED>(actualShape.m());
+        uint32_t nRound = RoundUp<L1BAlignHelper::N_ALIGNED>(actualShape.n());
+
+        uint32_t startTileIdx = 0;
+        if constexpr (ENABLE_SHUFFLE_K) {
+            startTileIdx = AscendC::GetBlockIdx() % kTileCount;
+        }
+
+        for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) {
+            uint32_t kTileIdx = (startTileIdx + kLoopIdx < kTileCount) ?
+                (startTileIdx + kLoopIdx) : (startTileIdx + kLoopIdx - kTileCount);
+
+            uint32_t kActual = (kTileIdx < kTileCount - 1) ?
+                L1TileShape::K : (actualShape.k() - kTileIdx * L1TileShape::K);
+
+            // Emission load instruction from GM to L1
+            MatrixCoord gmTileAOffset{0, kTileIdx * L1TileShape::K};
+            MatrixCoord gmTileBOffset{kTileIdx * L1TileShape::K, 0};
+            auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)];
+            auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)];
+            // Load first matrix A tile from GM to L1
+            AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[l1ListId]);
+            auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual));
+            copyGmToL1A(l1ATensorList[l1ListId], gmTileA, L1A_LAYOUT, layoutTileA);
+            AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList[l1ListId]);
+            // Load first matrix B tile from GM to L1
+            AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[l1ListId]);
+            auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n()));
+            copyGmToL1B(l1BTensorList[l1ListId], gmTileB, L1B_LAYOUT, layoutTileB);
+            AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList[l1ListId]);
+
+            // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile
+            if (preloadCount == PRELOAD_STAGES) {
+                L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]);
+            }
+
+            // Store the current load status
+            uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ?
+                (l1TileMmadParamsId + preloadCount) : (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES);
+            auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId];
+            l1TileMmadParams.l1ListId = l1ListId;
+            l1TileMmadParams.mRound = mRound;
+            l1TileMmadParams.nRound = nRound;
+            l1TileMmadParams.kActual = kActual;
+            l1TileMmadParams.isKLoopFirst = (kLoopIdx == 0);
+            l1TileMmadParams.isKLoopLast = (kLoopIdx == kTileCount - 1);
+            l1TileMmadParams.flag = flag;
+            if (kLoopIdx == kTileCount - 1) {
+                l1TileMmadParams.gmBlockC = gmBlockC;
+                l1TileMmadParams.gmBlockS = gmBlockS;
+                l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN());
+                l1TileMmadParams.layoutScale = layoutScale;
+                l1TileMmadParams.syncLoopIdx = syncLoopIdx;
+            }
+
+            if (preloadCount < PRELOAD_STAGES) {
+                ++preloadCount;
+            } else {
+                l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0;
+            }
+            l1ListId = (l1ListId + 1 < L1_STAGES) ? (l1ListId + 1) : 0;
+        }
+    }
+
+    CATLASS_DEVICE
+    void SynchronizeBlock()
+    {
+        while (preloadCount > 0) {
+            L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]);
+            l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0;
+            --preloadCount;
+        }
+    }
+
+    CATLASS_DEVICE
+    void Finalize(int32_t target, int32_t flag = 0)
+    {
+        if (ptrSoftFlagBase_ != nullptr) {
+            if (target < 0) {
+                return;
+            }
+            AscendC::SetFlag<AscendC::HardEvent::FIX_MTE3>(EVENT_ID0);
+            AscendC::WaitFlag<AscendC::HardEvent::FIX_MTE3>(EVENT_ID0);
+            AscendC::GlobalTensor<int32_t> flagGlobal;
+            flagGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(ptrSoftFlagBase_) + (expertPerRank_ + AscendC::GetBlockIdx()) * FLAGSTRIDE);
+            AscendC::DataCopy(flagGlobal, l1FTensor[target * 16], FLAGSTRIDE);
+        }
+        else {
+            for(;syncGroupIdx <= target; syncGroupIdx++) {
+                int32_t flagId = syncGroupIdx / 15 + flag;
+                AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(flagId);
+            }
+        }
+    }
+private:
+    struct L1TileMmadParams {
+        uint32_t l1ListId;
+        uint32_t mRound;
+        uint32_t nRound;
+        uint32_t kActual;
+        bool isKLoopFirst;
+        bool isKLoopLast;
+        AscendC::GlobalTensor<ElementC> gmBlockC;
+        AscendC::GlobalTensor<uint64_t> gmBlockS;
+        LayoutC layoutCInGm;
+        layout::VectorLayout layoutScale;
+        int32_t syncLoopIdx;
+        int32_t flag;
+        CATLASS_DEVICE
+        L1TileMmadParams() = default;
+    };
+
+    CATLASS_DEVICE
+    void InitL1(Arch::Resource<ArchTag> &resource, uint32_t l1BufAddrStart)
+    {
+        uint32_t l1AOffset = l1BufAddrStart;
+        uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1_STAGES;
+
+        for (uint32_t i = 0; i < L1_STAGES; ++i) {
+            l1ATensorList[i] = resource.l1Buf.template GetBufferByByte<ElementA>(l1AOffset + L1A_TILE_SIZE * i);
+            l1BTensorList[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1BOffset + L1B_TILE_SIZE * i);
+            l1AEventList[i] = i;
+            l1BEventList[i] = i + L1_STAGES;
+            AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[i]);
+            AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[i]);
+        }
+        uint32_t l1SOffset = l1BOffset + L1B_TILE_SIZE * L1_STAGES;
+        if constexpr (std::is_same_v<ElementA, int8_t>) {
+            l1STensor = resource.l1Buf.template GetBufferByByte<uint64_t>(l1SOffset);
+            AscendC::SetFlag<AscendC::HardEvent::FIX_MTE2>(0);
+        }
+        if (ptrSoftFlagBase_ != nullptr) {
+            // Initialize the flag matrix (structure as below):
+            // 1 0 0 0 0 0 0 0
+            // 2 0 0 0 0 0 0 0
+            // ...
+            // 16 0 0 0 0 0 0 0
+            // Then move it to L1
+            uint32_t l1FOffset = l1SOffset + L1S_TILE_SIZE;
+            l1FTensor = resource.l1Buf.template GetBufferByByte<int32_t>(l1FOffset);
+            AscendC::GlobalTensor<int32_t> flagBase;
+            flagBase.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(ptrSoftFlagBase_));
+            AscendC::DataCopy(l1FTensor, flagBase, expertPerRank_ * FLAGSTRIDE);
+        }
+    }
+
+    CATLASS_DEVICE
+    void InitFpBuf(Arch::Resource<ArchTag> &resource, uint32_t FpAddrStart)
+    {
+        uint32_t FpOffset = FpAddrStart;
+        fixpipeBuf = resource.fpBuf.template GetBufferByByte<uint64_t>(FpOffset);
+    }
+
+    CATLASS_DEVICE
+    void InitL0A(Arch::Resource<ArchTag> &resource)
+    {
+        for (uint32_t i = 0; i < L0A_STAGES; ++i) {
+            l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_TILE_SIZE * i);
+            l0AEventList[i] = i;
+            AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[i]);
+        }
+    }
+
+    CATLASS_DEVICE
+    void InitL0B(Arch::Resource<ArchTag> &resource)
+    {
+        for (uint32_t i = 0; i < L0B_STAGES; ++i) {
+            l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_TILE_SIZE * i);
+            l0BEventList[i] = i + L0A_STAGES;
+            AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[i]);
+        }
+    }
+
+    CATLASS_DEVICE
+    void InitL0C(Arch::Resource<ArchTag> &resource)
+    {
+        for (uint32_t i = 0; i < L0C_STAGES; ++i) {
+            l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(L0C_TILE_SIZE * i);
+            l0CEventList[i] = i;
+            AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList[i]);
+        }
+    }
+
+    CATLASS_DEVICE
+    void L1TileMmad(L1TileMmadParams const &params)
+    {
+        uint32_t mPartLoop = CeilDiv<L0TileShape::M>(params.mRound);
+        uint32_t nPartLoop = CeilDiv<L0TileShape::N>(params.nRound);
+        uint32_t kPartLoop = CeilDiv<L0TileShape::K>(params.kActual);
+        auto &l1ATensor = l1ATensorList[params.l1ListId];
+        auto &l1BTensor = l1BTensorList[params.l1ListId];
+
+        auto &l0CTensor = l0CTensorList[l0CListId];
+        LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound));
+
+        if constexpr (!ENABLE_UNIT_FLAG) {
+            if (params.isKLoopFirst) {
+                AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList[l0CListId]);
+            }
+        }
+
+        for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) {
+            uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ?
+                L0TileShape::M : (params.mRound - mPartIdx * L0TileShape::M);
+
+            for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) {
+                uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ?
+                    L0TileShape::K : (params.kActual - kPartIdx * L0TileShape::K);
+
+                auto &l0ATile = l0ATensorList[l0AListId];
+                auto layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mPartActual, kPartActual);
+                auto l1AOffset = MakeCoord(mPartIdx, kPartIdx) * L0TileShape::ToCoordMK();
+                auto l1ATile = l1ATensor[L1A_LAYOUT.GetOffset(l1AOffset)];
+
+                AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]);
+                if ((mPartIdx == 0) && (kPartIdx == 0)) {
+                    AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList[params.l1ListId]);
+                }
+                copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT);
+                if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) {
+                    AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[params.l1ListId]);
+                }
+
+                for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) {
+                    uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ?
+                        L0TileShape::N : (params.nRound - nPartIdx * L0TileShape::N);
+
+                    auto &l0BTile = l0BTensorList[l0BListId];
+                    auto layoutBInL0 = LayoutBInL0::template MakeLayout<ElementB>(kPartActual, nPartActual);
+                    auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * L0TileShape::ToCoordKN();
+                    auto l1BTile = l1BTensor[L1B_LAYOUT.GetOffset(l1BOffset)];
+
+                    AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[l0BListId]);
+                    if ((kPartIdx == 0) && (nPartIdx == 0)) {
+                        AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList[params.l1ListId]);
+                    }
+                    copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, L1B_LAYOUT);
+                    if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) {
+                        AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[params.l1ListId]);
+                    }
+
+                    AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
+
+                    auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * L0TileShape::ToCoordMN();
+                    auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)];
+
+                    AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
+                    // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0
+                    bool initC = (params.isKLoopFirst && (kPartIdx == 0));
+                    // If the unit flag is enabled, the unit flag is set according to the calculation progress
+                    uint8_t unitFlag = 0b00;
+                    if constexpr (ENABLE_UNIT_FLAG) {
+                        if (params.isKLoopLast &&
+                            (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) {
+                            unitFlag = 0b11;
+                        } else {
+                            unitFlag = 0b10;
+                        }
+                    }
+                    tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag);
+
+                    AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[l0BListId]);
+                    l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0;
+                }
+                AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]);
+                l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0;
+            }
+        }
+
+        if (params.isKLoopLast) {
+            auto layoutCInGm = params.layoutCInGm;
+            if constexpr (std::is_same_v<ElementA, int8_t>) {
+                auto layoutScale = params.layoutScale;
+                auto layoutTileS = layoutScale.GetTileLayout(MakeCoord(layoutCInGm.shape(1)));
+                AscendC::WaitFlag<AscendC::HardEvent::FIX_MTE2>(0);
+                copyGmToL1S(l1STensor, params.gmBlockS, layoutTileS, layoutTileS);
+                AscendC::SetFlag<AscendC::HardEvent::MTE2_FIX>(0);
+                AscendC::WaitFlag<AscendC::HardEvent::MTE2_FIX>(0);
+
+                copyL1ToFP(fixpipeBuf, l1STensor, layoutTileS, layoutTileS);
+                AscendC::SetFlag<AscendC::HardEvent::FIX_MTE2>(0);
+                AscendC::PipeBarrier<PIPE_FIX>();
+            }
+            if constexpr (!ENABLE_UNIT_FLAG) {
+                AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventList[l0CListId]);
+                AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventList[l0CListId]);
+                if constexpr (std::is_same_v<ElementA, int8_t>) {
+                    copyL0CToGm(params.gmBlockC, l0CTensor, l1STensor, layoutCInGm, layoutCInL0);
+                } else {
+                    copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0);
+                }
+                AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList[l0CListId]);
+            } else {
+                if constexpr (std::is_same_v<ElementA, int8_t>) {
+                    copyL0CToGm(params.gmBlockC, l0CTensor, l1STensor, layoutCInGm, layoutCInL0, 0b11);
+                } else {
+                    copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0, 0b11);
+                }
+            }
+            l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0;
+            #ifdef __TILE_SYNC__
+            if (params.flag > 0) {
+                int32_t flagId = params.flag + params.syncLoopIdx / 8;
+                AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(flagId);
+            }
+            #else
+            Finalize(params.syncLoopIdx, params.flag);
+            #endif
+        }
+    }
+
+    AscendC::LocalTensor<uint64_t> fixpipeBuf;
+
+    AscendC::LocalTensor<ElementA> l1ATensorList[L1_STAGES];
+    AscendC::LocalTensor<ElementB> l1BTensorList[L1_STAGES];
+    AscendC::LocalTensor<uint64_t> l1STensor;
+    AscendC::LocalTensor<int32_t> l1FTensor;
+    int32_t syncGroupIdx;
+    int32_t l1AEventList[L1_STAGES];
+    int32_t l1BEventList[L1_STAGES];
+    uint32_t l1ListId{0};
+
+    AscendC::LocalTensor<ElementA> l0ATensorList[L0A_STAGES];
+    int32_t l0AEventList[L0A_STAGES];
+    uint32_t l0AListId{0};
+
+    AscendC::LocalTensor<ElementB> l0BTensorList[L0B_STAGES];
+    int32_t l0BEventList[L0B_STAGES];
+    uint32_t l0BListId{0};
+
+    AscendC::LocalTensor<ElementAccumulator> l0CTensorList[L0C_STAGES_];
+    int32_t l0CEventList[L0C_STAGES_];
+    uint32_t l0CListId{0};
+
+    L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES];
+    uint32_t l1TileMmadParamsId{0};
+    uint32_t preloadCount{0};
+
+    TileMmad tileMmad;
+    CopyGmToL1A copyGmToL1A;
+    CopyGmToL1B copyGmToL1B;
+    CopyGmToL1S copyGmToL1S;
+    CopyL1ToL0A copyL1ToL0A;
+    CopyL1ToL0B copyL1ToL0B;
+    CopyL0CToGm copyL0CToGm;
+    CopyL1ToFP copyL1ToFP;
+
+    __gm__ int32_t* ptrSoftFlagBase_ = nullptr;
+    int32_t expertPerRank_;
+};
+
+}  // namespace Catlass::Gemm::Block
+
+#endif  // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_FIXPIPE_QUANT_HPP

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/const_args.hpp

@@ -0,0 +1,11 @@
+
+#ifndef CONST_ARGS_HPP
+#define CONST_ARGS_HPP
+constexpr static uint64_t MB_SIZE = 1024 * 1024UL;
+constexpr static int32_t NUMS_PER_FLAG = 16;
+constexpr static int32_t CACHE_LINE = 512;
+constexpr static int32_t RESET_VAL = 0xffff;
+constexpr static int32_t FLAGSTRIDE = 16;
+constexpr static int32_t UB_ALIGN = 32;
+constexpr uint16_t CROSS_CORE_FLAG_MAX_SET_COUNT = 15;
+#endif

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/copy_gm_to_l1_custom.hpp

@@ -0,0 +1,40 @@
+#ifndef COPY_GM_TO_L1_CUSTOM_HPP
+#define COPY_GM_TO_L1_CUSTOM_HPP
+
+namespace Catlass::Gemm::Tile {
+    /// Partial specialization for nZ in and nZ out.
+    template <
+        class ArchTag,
+        class Element
+    >
+    struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::VectorLayout>> {
+        using LayoutDst = layout::VectorLayout;
+        using LayoutSrc = layout::VectorLayout;
+
+        static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);   // int64, 32/8=4
+
+        // Methods  
+
+        CATLASS_DEVICE
+        CopyGmToL1() {};
+
+        CATLASS_DEVICE
+        void operator()(
+            AscendC::LocalTensor<Element> const &dstTensor,
+            AscendC::GlobalTensor<Element> const &srcTensor,
+            LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
+        {
+            uint32_t blockCount = 1;
+            uint32_t blockLen = CeilDiv<ELE_NUM_PER_C0>(layoutSrc.shape(0));
+
+            AscendC::DataCopyParams repeatParams;
+
+            repeatParams.blockCount = blockCount;
+            repeatParams.blockLen = blockLen;
+            repeatParams.srcStride = 0;
+            repeatParams.dstStride = 0;
+            AscendC::DataCopy(dstTensor, srcTensor, repeatParams);
+        }
+    };
+}
+#endif // COPY_GM_TO_L1_CUSTOM_HPP

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/copy_l0c_to_gm_custom.hpp

@@ -0,0 +1,47 @@
+#ifndef COPY_L0C_TO_GM_CUSTOM_HPP
+#define COPY_L0C_TO_GM_CUSTOM_HPP
+
+namespace Catlass::Gemm::Tile {
+    template <
+        class ElementAccumulator_,
+        class ElementDst_,
+        bool ReluEnable_
+    >
+    struct CopyL0CToGm<Catlass::Arch::AtlasA2,
+                    ElementAccumulator_,
+                    Gemm::GemmType<ElementDst_, layout::RowMajor>,
+                    ScaleGranularity::PER_CHANNEL,
+                    ReluEnable_>
+    {
+        using ArchTag = Catlass::Arch::AtlasA2;
+        using ElementDst = ElementDst_;
+        using ElementSrc = ElementAccumulator_;
+        using LayoutSrc = Catlass::layout::zN;
+        using LayoutDst = Catlass::layout::RowMajor;
+        static constexpr auto quantPre = CopyL0CToGmQuantMode<ArchTag, ElementSrc, ElementDst,
+            ScaleGranularity::PER_CHANNEL>::VALUE;
+        static constexpr auto reluEn = ReluEnable_;
+
+        CATLASS_DEVICE
+        void operator()(AscendC::GlobalTensor<ElementDst> const &dst, AscendC::LocalTensor<ElementSrc> const &src, AscendC::LocalTensor<uint64_t> cbufWorkspace,
+            LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0)
+        {
+            AscendC::FixpipeParamsV220 intriParams;
+
+            // Fixpipe layout information
+            intriParams.nSize = dstLayout.shape(1);
+            intriParams.mSize = dstLayout.shape(0);
+            intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0);
+            intriParams.dstStride = dstLayout.stride(0);
+
+            // Fixpipe auxiliary arguments
+            intriParams.quantPre = quantPre;
+            intriParams.reluEn = reluEn;
+            intriParams.unitFlag = unitFlag;
+
+            // Call AscendC Fixpipe
+            AscendC::Fixpipe<ElementDst, ElementSrc, AscendC::CFG_ROW_MAJOR>(dst, src, cbufWorkspace, intriParams);
+        }
+    };
+}
+#endif // COPY_L0C_TO_GM_CUSTOM_HPP

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/dispatch_policy_custom.hpp

@@ -0,0 +1,47 @@
+#ifndef DISPATH_POLICY_CUSTOM_HPP
+#define DISPATH_POLICY_CUSTOM_HPP
+
+namespace Catlass::Gemm {
+    template <bool ENABLE_UNIT_FLAG_ = false, bool ENABLE_SHUFFLE_K_ = false>
+    struct MmadAtlasA2PreloadFixpipeQuant : public MmadAtlasA2 {
+        static constexpr uint32_t STAGES = 2;
+        static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
+        static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_;
+    };
+
+    template <uint32_t PRELOAD_STAGES_, uint32_t L1_STAGES_, uint32_t L0A_STAGES_, uint32_t L0B_STAGES_,
+        uint32_t L0C_STAGES_, bool ENABLE_UNIT_FLAG_, bool ENABLE_SHUFFLE_K_>
+    struct MmadAtlasA2PreloadAsyncFixpipe :
+        public MmadAtlasA2PreloadAsync<
+            PRELOAD_STAGES_,
+            L1_STAGES_, 
+            L0A_STAGES_,
+            L0B_STAGES_,
+            L0C_STAGES_,
+            ENABLE_UNIT_FLAG_,
+            ENABLE_SHUFFLE_K_
+        > {
+    };
+}
+
+namespace Catlass::Epilogue {
+        
+    template <uint32_t UB_STAGES_>
+    struct EpilogueAtlasA2UnQuant {
+        using ArchTag = Arch::AtlasA2;
+        static constexpr uint32_t UB_STAGES = UB_STAGES_;
+    };
+
+    template <uint32_t UB_STAGES_>
+    struct EpilogueAtlasA2PerTokenDequantSwigluQuant {
+        using ArchTag = Arch::AtlasA2;
+        static constexpr uint32_t UB_STAGES = UB_STAGES_;
+    };
+    
+    template <uint32_t UB_STAGES_>
+    struct EpilogueAtlasA2PerTokenDequantV2  {
+        using ArchTag = Arch::AtlasA2;
+        static constexpr uint32_t UB_STAGES = UB_STAGES_;
+    };
+}
+#endif // DISPATH_POLICY_CUSTOM_HPP

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/get_tensor_addr.hpp

@@ -0,0 +1,16 @@
+#ifndef GET_TENSOR_ADDR_HPP
+#define GET_TENSOR_ADDR_HPP
+#include "kernel_operator.h"
+
+#define FORCE_INLINE_AICORE inline __attribute__((always_inline)) __aicore__
+
+template <typename T>
+FORCE_INLINE_AICORE __gm__ T* GetTensorAddr(uint32_t index, GM_ADDR tensorPtr) {
+    __gm__ uint64_t* dataAddr = reinterpret_cast<__gm__ uint64_t*>(tensorPtr);
+    uint64_t tensorPtrOffset = *dataAddr;  // The offset of the data address from the first address.
+    // Moving 3 bits to the right means dividing by sizeof(uint64 t).
+    __gm__ uint64_t* retPtr = dataAddr + (tensorPtrOffset >> 3);
+    return reinterpret_cast<__gm__ T*>(*(retPtr + index));
+}
+
+#endif // GET_TENSOR_ADDR_HPP

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/hccl_shmem.hpp

@@ -0,0 +1,327 @@
+#ifndef SYNC_UTIL_HPP
+#define SYNC_UTIL_HPP
+
+
+#include "kernel_operator.h"
+#include "const_args.hpp"
+
+#ifdef HCCL_COMM
+#include "moe_distribute_base.h"
+using namespace AscendC::HcclContextDef;
+
+#else
+#include "shmem_api.h"
+#endif
+
+#define FORCE_INLINE_AICORE inline __attribute__((always_inline)) __aicore__
+constexpr int32_t MAX_RANK_SIZE = 32;
+constexpr int32_t SHMEM_MEM = 700 * MB_SIZE;
+
+constexpr uint16_t SEND_SYNC_EVENT_ID = 9;
+constexpr uint16_t RECV_SYNC_EVENT_ID = 10;
+
+constexpr uint32_t SELF_STATE_OFFSET = 256 * 1024;
+constexpr uint32_t STATE_OFFSET = 512;
+
+FORCE_INLINE_AICORE void AicSyncAll() {
+    AscendC::CrossCoreSetFlag<0x0, PIPE_FIX>(8);
+    AscendC::CrossCoreWaitFlag<0x0>(8);
+}
+
+template<typename T>
+FORCE_INLINE_AICORE void gm_store(__gm__ T *addr, T val) {
+    *((__gm__ T *)addr) = val;
+}
+
+template<typename T>
+FORCE_INLINE_AICORE T gm_load(__gm__ T *cache) {
+    return *((__gm__ T *)cache);
+}
+
+template<typename T>
+FORCE_INLINE_AICORE void gm_dcci(__gm__ T * addr) {
+    using namespace AscendC;
+    GlobalTensor<uint8_t> global;
+    global.SetGlobalBuffer(reinterpret_cast<GM_ADDR>(addr));
+
+    // Important: add hint to avoid dcci being optimized by compiler
+    __asm__ __volatile__("");
+    DataCacheCleanAndInvalid<uint8_t, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(global);
+    __asm__ __volatile__("");
+}
+
+FORCE_INLINE_AICORE int32_t gm_signal_wait_until_eq_for_barrier(__gm__ int32_t *sig_addr, int32_t cmp_val) {
+    do {
+        gm_dcci((__gm__ uint8_t *)sig_addr);
+        if (*sig_addr == cmp_val) {
+            return *sig_addr;
+        }
+        if (*sig_addr == cmp_val + 1) {
+            return *sig_addr;
+        }
+    } while (true);
+    return -1;
+}
+
+FORCE_INLINE_AICORE void gm_signal_wait_until_ne(__gm__ int32_t *sig_addr, int32_t cmp_val) {
+    do {
+        AscendC::LocalTensor<int32_t> ub;
+        ub.address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::VECIN);
+        ub.address_.bufferAddr = 0;
+        AscendC::GlobalTensor<int32_t> sig;
+        sig.SetGlobalBuffer(sig_addr);
+        AscendC::DataCopy(ub, sig, 8);
+        AscendC::SetFlag<AscendC::HardEvent::MTE2_S>(EVENT_ID0);
+        AscendC::WaitFlag<AscendC::HardEvent::MTE2_S>(EVENT_ID0);
+        if (ub(0) != cmp_val) {
+            return;
+        }
+    } while (true);
+    return;
+}
+
+
+class HcclShmem {
+public:
+    #ifdef HCCL_COMM    // HCCL needs to initialize the HCCL context
+        __gm__ HcclOpResParamCustom *WinContext_{nullptr};
+        Hccl<HCCL_SERVER_TYPE_AICPU> hccl_;
+        AscendC::LocalTensor<int32_t> ub;
+        FORCE_INLINE_AICORE
+        HcclShmem(){
+            auto contextGM0 = AscendC::GetHcclContext<HCCL_GROUP_ID_0>();
+            WinContext_ = (__gm__ HcclOpResParamCustom *)contextGM0;
+
+            m_rank = WinContext_->localUsrRankId;
+            m_rankSize = WinContext_->rankSize;
+            m_segmentSize = WinContext_->winSize;
+        }
+    #else
+        FORCE_INLINE_AICORE
+        HcclShmem(){
+            m_segmentSize = SHMEM_MEM;
+        }
+        FORCE_INLINE_AICORE 
+        void initShmem(GM_ADDR symmetricPtr_, size_t rank, size_t rankSize) {
+            symmetricPtr = symmetricPtr_;
+            m_rank = rank;
+            m_rankSize = rankSize;
+        }
+    #endif
+
+    FORCE_INLINE_AICORE
+    GM_ADDR operator() () const {   // No parameters: return pointer to local peermem
+        #ifdef HCCL_COMM
+            return (GM_ADDR)(WinContext_->localWindowsIn);
+        #else
+            return reinterpret_cast<GM_ADDR>(shmem_ptr(symmetricPtr, m_rank));
+        #endif
+    }
+
+    FORCE_INLINE_AICORE
+    GM_ADDR operator() (int32_t index) const {  // With index parameter: return pointer to the base address of remote peermem
+        #ifdef HCCL_COMM
+            return (GM_ADDR)((index == m_rank) ? WinContext_->localWindowsIn :
+                                    ((HcclRankRelationResV2Custom *)(WinContext_->remoteRes[index].nextDevicePtr))->windowsIn);
+        #else
+            return reinterpret_cast<GM_ADDR>(shmem_ptr(symmetricPtr, index));
+        #endif
+    }
+
+    FORCE_INLINE_AICORE
+    GM_ADDR operator () (int64_t offset, int32_t rankId) const  {  
+        #ifdef HCCL_COMM
+            if (offset < 0 || offset >= m_segmentSize) {
+                return nullptr;
+            }
+            if (rankId < 0 || rankId >= m_rankSize) {
+                return nullptr;
+            }
+            return (GM_ADDR)((rankId == m_rank) ? WinContext_->localWindowsIn :
+                                    ((HcclRankRelationResV2Custom *)(WinContext_->remoteRes[rankId].nextDevicePtr))->windowsIn) + offset;
+        #else
+            return reinterpret_cast<GM_ADDR>(shmem_ptr((symmetricPtr + offset), rankId));
+        #endif
+    }
+
+
+
+    FORCE_INLINE_AICORE
+    size_t SegmentSize() const {
+        return m_segmentSize;
+    }
+
+    FORCE_INLINE_AICORE
+    int32_t RankSize() const {
+        return m_rankSize;
+    }
+
+
+    FORCE_INLINE_AICORE
+    ~HcclShmem() {
+    }
+
+
+    FORCE_INLINE_AICORE
+    void CrossRankSync() {
+        uint64_t flag_offset = (m_segmentSize - MB_SIZE) / sizeof(int32_t);
+        __gm__ int32_t* sync_counter = (__gm__ int32_t*)(*this)() + flag_offset;
+        __gm__ int32_t* sync_base = (__gm__ int32_t*)(*this)() + flag_offset + 2048;
+        int count = gm_load(sync_base) + 1;
+        int vec_id = AscendC::GetBlockIdx();
+        int vec_size = AscendC::GetBlockNum() * AscendC::GetTaskRation();
+        for(int i = vec_id; i < m_rankSize; i += vec_size) {
+            __gm__ int32_t* sync_remote = (__gm__ int32_t*)((*this)(i)) + flag_offset + m_rank * 16;
+            gm_store(sync_remote, count);
+            gm_dcci((__gm__ uint8_t*)sync_remote);
+            auto sync_check = sync_counter + i * 16;
+            gm_signal_wait_until_eq_for_barrier(sync_check, count);
+        }
+
+        AscendC::SyncAll<true>();
+        gm_store(sync_base, count);
+    }
+
+
+    FORCE_INLINE_AICORE
+    void InitStatusTargetSum()
+    {
+        using namespace AscendC;
+        uint64_t flag_offset = (m_segmentSize - MB_SIZE) + SELF_STATE_OFFSET;
+        //uint64_t self_state_offset = (m_segmentSize - 2 * MB_SIZE);
+        // ep state
+        //uint32_t coreIdx = get_block_idx();;
+        uint32_t coreIdx = GetBlockIdx();
+        GlobalTensor<int32_t> selfStatusTensor;
+        selfStatusTensor.SetGlobalBuffer((__gm__ int32_t *)((*this)() + flag_offset));
+        __asm__ __volatile__("");
+        DataCacheCleanAndInvalid<int32_t, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(selfStatusTensor[coreIdx * UB_ALIGN]);
+        __asm__ __volatile__("");
+        int32_t state = selfStatusTensor(coreIdx * UB_ALIGN);
+        if (state == 0) {
+            sumTarget_ = static_cast<float>(1.0);
+            selfStatusTensor(coreIdx * UB_ALIGN) = 0x3F800000;  // 1.0f
+            epStateValue_ = 0x3F800000;                          // 1.0f
+        } else {
+            sumTarget_ = static_cast<float>(0.0);
+            selfStatusTensor(coreIdx * UB_ALIGN) = 0;
+            epStateValue_ = 0;
+        }
+        __asm__ __volatile__("");
+        DataCacheCleanAndInvalid<int32_t, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(selfStatusTensor[coreIdx * UB_ALIGN]);
+        __asm__ __volatile__("");
+    }
+
+    FORCE_INLINE_AICORE
+    void CrossRankSyncV2Set(AscendC::LocalTensor<int32_t> ctrBuffer) {
+        //subblockid = 0
+        uint32_t stateOffset_ =  STATE_OFFSET;
+        // uint32_t epStateOffsetOnWin_ = m_rank * stateOffset_;
+        
+        uint64_t flag_offset = (m_segmentSize - MB_SIZE) + m_rank * stateOffset_;
+        //uint64_t flag_offset = (m_segmentSize - MB_SIZE);
+        int vec_size = get_block_num();
+        int vec_id = get_block_idx();
+ 
+        AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(RECV_SYNC_EVENT_ID);
+        AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(SEND_SYNC_EVENT_ID);
+        AscendC::CrossCoreWaitFlag(SEND_SYNC_EVENT_ID);
+        AscendC::PipeBarrier<PIPE_ALL>();
+ 
+        ctrBuffer.SetValue(0, epStateValue_);
+        AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
+        AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
+        for (uint32_t dstEpIdx = vec_id; dstEpIdx < m_rankSize; dstEpIdx += vec_size) {
+            AscendC::GlobalTensor<int32_t> gmDstStates;
+            gmDstStates.SetGlobalBuffer((__gm__ int32_t*)((*this)(flag_offset, dstEpIdx)));
+            DataCopy(gmDstStates, ctrBuffer, 8);
+        }
+        AscendC::CrossCoreWaitFlag(RECV_SYNC_EVENT_ID);
+    }
+
+    FORCE_INLINE_AICORE
+    void CrossRankSyncV2Wait(AscendC::LocalTensor<float> statusTensor, AscendC::LocalTensor<float> gatherMaskOutTensor,
+        AscendC::LocalTensor<uint32_t> gatherTmpTensor, AscendC::LocalTensor<float> statusSumOutTensor) {
+
+        uint64_t flag_offset = (m_segmentSize - MB_SIZE);
+        int vec_size = get_block_num();
+        int vec_id = get_block_idx();
+        uint32_t stateOffset_ =  STATE_OFFSET;
+
+        uint32_t sendRankNum_ = m_rankSize / vec_size;
+        uint32_t remainderRankNum = m_rankSize % vec_size;
+        uint32_t startRankId_ = sendRankNum_ * vec_id;
+        if (vec_id < remainderRankNum) {
+            sendRankNum_++;
+            startRankId_ += vec_id;
+        } else {
+            startRankId_ += remainderRankNum;
+        }
+        uint32_t endRankId_ = startRankId_ + sendRankNum_;
+        AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(SEND_SYNC_EVENT_ID);
+
+        AscendC::GlobalTensor<float> epStatusSpaceGlobalTensor_;
+        epStatusSpaceGlobalTensor_.SetGlobalBuffer((__gm__ float *)((*this)() + flag_offset));
+
+        if (startRankId_ < m_rankSize) {
+            AscendC::PipeBarrier<PIPE_ALL>();
+            gatherTmpTensor.SetValue(0, 1);
+            uint32_t mask = 1;  // gatherMask + sum
+            uint64_t rsvdCnt = 0;
+            // DataCopyParams intriParams{static_cast<uint16_t>(sendRankNum_), 1,
+            //                         static_cast<uint16_t>((moeSendNum_ > 512) ? 7 : 15), 0}; 
+            AscendC::DataCopyParams intriParams{static_cast<uint16_t>(sendRankNum_), 1,
+                                    static_cast<uint16_t>(15), 0}; 
+
+            float sumOfFlag = static_cast<float>(-1.0);
+            float minTarget = (sumTarget_ * sendRankNum_) - (float)0.5;
+            float maxTarget = (sumTarget_ * sendRankNum_) + (float)0.5;
+            AscendC::SumParams sumParams{1, sendRankNum_, sendRankNum_};
+
+            AscendC::SetFlag<AscendC::HardEvent::S_V>(EVENT_ID0);
+            AscendC::WaitFlag<AscendC::HardEvent::S_V>(EVENT_ID0);
+
+            while ((sumOfFlag < minTarget) || (sumOfFlag > maxTarget)) {
+                AscendC::DataCopy<float>(statusTensor, epStatusSpaceGlobalTensor_[startRankId_ * stateOffset_ / sizeof(float)],
+                                intriParams);
+                AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
+                AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
+
+                GatherMask(gatherMaskOutTensor, statusTensor, gatherTmpTensor, true, mask,
+                        {1, (uint16_t)sendRankNum_, 1, 0}, rsvdCnt);
+
+                AscendC::PipeBarrier<PIPE_V>();
+                AscendC::Sum(statusSumOutTensor, gatherMaskOutTensor, sumParams);
+                AscendC::SetFlag<AscendC::HardEvent::V_S>(EVENT_ID0);
+                AscendC::WaitFlag<AscendC::HardEvent::V_S>(EVENT_ID0);
+                sumOfFlag = statusSumOutTensor.GetValue(0);
+            }
+        }
+
+        AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(RECV_SYNC_EVENT_ID);
+        AscendC::CrossCoreWaitFlag(RECV_SYNC_EVENT_ID);
+
+        //unpermute
+        AscendC::CrossCoreWaitFlag(SEND_SYNC_EVENT_ID);
+    }
+
+
+    FORCE_INLINE_AICORE
+    __gm__ int32_t* SyncBaseAddr() {
+        uint64_t flag_offset = (m_segmentSize - MB_SIZE) / sizeof(int32_t);
+        return (__gm__ int32_t*)(*this)() + flag_offset + 2048;
+    }
+
+private:
+    GM_ADDR symmetricPtr;
+    int32_t m_rank;
+    int32_t m_rankSize;
+    size_t m_segmentSize;
+    float sumTarget_{0.0};
+    int32_t epStateValue_;
+};
+
+
+
+
+#endif

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/layout3d.hpp

@@ -0,0 +1,20 @@
+#ifndef LAYOUT_3D_HPP
+#define LAYOUT_3D_HPP
+#include "kernel_operator.h"
+#include "catlass/catlass.hpp"
+class Layout3D {
+    int64_t strides[2];
+    public:
+    CATLASS_DEVICE
+    Layout3D() {}
+    CATLASS_DEVICE
+    Layout3D(int64_t stride0, int64_t stride1) {
+        strides[0] = stride0;
+        strides[1] = stride1;
+    }
+    CATLASS_DEVICE
+    int64_t operator() (int64_t dim0, int64_t dim1, int64_t dim2) {
+        return dim0 * strides[0] + dim1 * strides[1] + dim2;
+    }
+};
+#endif // LAYOUT_3D_HPP

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/moe_distribute_base.h

@@ -0,0 +1,365 @@
+/**
+ * This program is free software, you can redistribute it and/or modify.
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_distribute_base.h
+ * \brief
+ */
+
+#ifndef MOE_DISTRIBUTE_BASE_H
+#define MOE_DISTRIBUTE_BASE_H
+
+#include "kernel_operator.h"
+
+constexpr uint32_t LOCAL_NOTIFY_MAX_NUM = 64;
+constexpr uint32_t LOCAL_STREAM_MAX_NUM = 19U;
+constexpr uint32_t AICPU_OP_NOTIFY_MAX_NUM = 2;
+constexpr uint32_t AICPU_MAX_RANK_NUM = 128 * 1024;
+constexpr uint32_t TIME_CYCLE = 50; // 系统cycle数转换成时间的基准单位，固定为50
+
+struct HcclSignalInfo {
+    uint64_t resId; // 在代表event时为eventid，notify时为notifyid
+    uint64_t addr;
+    uint32_t devId;
+    uint32_t tsId;
+    uint32_t rankId;
+    uint32_t flag;
+};
+
+struct ListCommon {
+    uint64_t nextHost;
+    uint64_t preHost;
+    uint64_t nextDevice;
+    uint64_t preDevice;
+};
+
+struct HcclStreamInfo {
+    int32_t streamIds;
+    uint32_t sqIds;
+    uint32_t cqIds;      // 记录物理cqId
+    uint32_t logicCqids; // 记录逻辑cqId
+};
+
+struct LocalResInfoV2 {
+    uint32_t streamNum;
+    uint32_t signalNum;
+    HcclSignalInfo localSignals[LOCAL_NOTIFY_MAX_NUM];
+    HcclStreamInfo streamInfo[LOCAL_STREAM_MAX_NUM];
+    HcclStreamInfo mainStreamInfo;
+    HcclSignalInfo aicpuOpNotify[AICPU_OP_NOTIFY_MAX_NUM];  // 集合通信AICPU展开资源
+    ListCommon nextTagRes;                                  // HccltagLocalResV2
+};
+
+enum class rtFloatOverflowMode_t {
+    RT_OVERFLOW_MODE_SATURATION = 0,
+    RT_OVERFLOW_MODE_INFNAN,
+    RT_OVERFLOW_MODE_UNDEF,
+};
+
+struct AlgoTopoInfo {
+    uint32_t userRank;     // 通信域 RankID
+    uint32_t userRankSize; // 通信域的Rank数量
+    int32_t deviceLogicId;
+    bool isSingleMeshAggregation;
+    uint32_t deviceNumPerAggregation;  // 每个Module中的Device数量
+    uint32_t superPodNum;              // 集群中总的超节点数
+    uint32_t devicePhyId;
+    uint32_t topoType; // TopoType
+    uint32_t deviceType;
+    uint32_t serverNum;
+    uint32_t meshAggregationRankSize;
+    uint32_t multiModuleDiffDeviceNumMode;
+    uint32_t multiSuperPodDiffServerNumMode;
+    uint32_t realUserRank;
+    bool isDiffDeviceModule;
+    bool isDiffDeviceType;
+    uint32_t gcdDeviceNumPerAggregation;
+    uint32_t moduleNum;
+    uint32_t isUsedRdmaRankPairNum;
+    uint64_t isUsedRdmaRankPair;
+    uint32_t pairLinkCounterNum;
+    uint64_t pairLinkCounter;
+    uint32_t nicNum;
+    uint64_t nicList;           // niclist数组指针
+    uint64_t complanRankLength; // complanRank占用的字节数
+    uint64_t complanRank;       // 指针
+    uint64_t bridgeRankNum;     // bridgeRank占用的个数
+    uint64_t bridgeRank;        // 指针
+    uint64_t serverAndsuperPodRankLength; // serverAndsuperPodRank占用的字节数
+    uint64_t serverAndsuperPodRank; // 指针
+};
+
+struct HcclOpConfig {
+    uint8_t deterministic; //确定性计算开关
+    uint8_t retryEnable;   // 是否重执行
+    uint8_t highPerfEnable;
+    uint8_t padding[5];    // 大小需要64By对齐，未来添加参数时减小padding
+    uint8_t linkTimeOut[8]; // 发送超时时长
+    uint64_t notifyWaitTime; // 超时时长，同HCCL_EXEC_TIMEOUT
+    uint32_t retryHoldTime;
+    uint32_t retryIntervalTime;
+    bool interHccsDisable = false; //使能rdma开关
+    rtFloatOverflowMode_t floatOverflowMode = rtFloatOverflowMode_t::RT_OVERFLOW_MODE_UNDEF;
+    uint32_t multiQpThreshold = 512;  // 多QP每个QP分担数据量最小阈值
+};
+
+struct HcclMC2WorkSpace {
+    uint64_t workSpace;
+    uint64_t workSpaceSize;
+};
+
+struct RemoteResPtr {
+    uint64_t nextHostPtr;
+    uint64_t nextDevicePtr;
+};
+
+struct HDCommunicateParams {
+    uint64_t hostAddr { 0 };
+    uint64_t deviceAddr { 0 };
+    uint64_t readCacheAddr { 0 };
+    uint32_t devMemSize{ 0 };
+    uint32_t buffLen{ 0 };
+    uint32_t flag{ 0 };
+};
+
+struct HcclRankRelationResV2Custom {
+    uint32_t remoteUsrRankId;
+    uint32_t remoteWorldRank;
+    uint64_t windowsIn;
+    uint64_t windowsOut;
+    uint64_t windowsExp;
+    ListCommon nextTagRes;
+};
+
+struct HcclOpResParamCustom {
+    // 本地资源
+    HcclMC2WorkSpace mc2WorkSpace;
+    uint32_t localUsrRankId; // usrrankid
+    uint32_t rankSize;       // 通信域内total rank个数
+    uint64_t winSize; // 每个win大小，静态图时，可能是0，如果通信域内也有动态图，则可能为非0
+    uint64_t localWindowsIn; // 全F为无效值
+    uint64_t localWindowsOut; // 全F为无效值
+    char hcomId[128];
+    // aicore识别remote window
+    uint64_t winExpSize;
+    uint64_t localWindowsExp;
+    uint32_t rWinStart; // 为HcclRankRelationRes起始位置
+    uint32_t rWinOffset; // 为HcclRemoteRes的大小
+    uint64_t version;
+    LocalResInfoV2 localRes;
+    AlgoTopoInfo topoInfo;
+
+    // 外部配置参数
+    HcclOpConfig config;
+    uint64_t hostStateInfo;
+    uint64_t aicpuStateInfo;
+    uint64_t lockAddr;
+    uint32_t rsv[16];
+    uint32_t notifysize;                         // RDMA场景使用，910B/910_93为4B，其余芯片为8B
+    uint32_t remoteResNum;                       // 有效的remoteResNum
+    RemoteResPtr remoteRes[AICPU_MAX_RANK_NUM];  //数组指针，指向HcclRankRelationResV2，下标为remoteUserRankId
+
+    // communicate retry
+    HDCommunicateParams kfcControlTransferH2DParams;
+    HDCommunicateParams kfcStatusTransferD2HParams;
+    uint64_t tinyMem;  // for all2all
+    uint64_t tinyMemSize;
+    // 零拷贝场景使用
+    uint64_t zeroCopyHeadPtr;
+    uint64_t zeroCopyTailPtr;
+    uint64_t zeroCopyRingBuffer;
+    uint64_t zeroCopyIpcPtrs[16];                // 保存集合通信时每个对端的输入输出内存地址
+    uint32_t zeroCopyDevicePhyId[16];            // 保存每个rank对应的物理卡Id
+
+    bool utraceStatusFlag;
+};
+
+// Transport 内存类型
+enum class HcclAiRMAMemType : uint32_t {
+    LOCAL_INPUT = 0,
+    REMOTE_INPUT,
+
+    LOCAL_OUTPUT,
+    REMOTE_OUTPUT,
+
+    // 可透传更多的内存，可在MAX_NUM之前追加，例如：
+    // LOCAL_EXP,
+    // REMOTE_EXP,
+    MAX_NUM
+};
+
+// Transport 内存信息
+struct HcclAiRMAMemInfo {
+    uint32_t memMaxNum{0};  // 最大内存数量，等于 HcclAiRMAMemType::MAX_NUM
+    uint32_t sizeOfMemDetails{0};  // sizeof(MemDetails)，用于内存校验和偏移计算
+    uint64_t memDetailPtr{0};  // MemDetails数组首地址, 个数: HcclAiRMAMemType::MAX_NUM
+    // 可往后追加字段
+};
+
+// 全部 Transport QP/Mem 信息
+struct HcclAiRMAInfo {
+    uint32_t curRankId{0};  // 当前rankId
+    uint32_t rankNum{0};  // rank数量
+    uint32_t qpNum{0};  // 单个Transport的QP数量
+
+    uint32_t sizeOfAiRMAWQ{0};  // sizeof(HcclAiRMAWQ)
+    uint32_t sizeOfAiRMACQ{0};  // sizeof(HcclAiRMACQ)
+    uint32_t sizeOfAiRMAMem{0};  // sizeof(HcclAiRMAMemInfo)
+
+    // HcclAiRMAWQ二维数组首地址
+    // QP个数: rankNum * qpNum
+    // 计算偏移获取SQ指针：sqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMAWQ
+    // 0 <= qpIndex < qpNum
+    uint64_t sqPtr{0};
+
+    // HcclAiRMACQ二维数组首地址
+    // QP个数: rankNum * qpNum
+    // 计算偏移获取SCQ指针：scqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMACQ
+    // 0 <= qpIndex < qpNum
+    uint64_t scqPtr{0};
+
+    // HcclAiRMAWQ二维数组首地址
+    // QP个数: rankNum * qpNum
+    // 计算偏移获取RQ指针：rqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMAWQ
+    // 0 <= qpIndex < qpNum
+    uint64_t rqPtr{0};
+
+    // HcclAiRMACQ二维数组首地址
+    // QP个数: rankNum * qpNum
+    // 计算偏移获取RCQ指针: rcqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMACQ
+    // 0 <= qpIndex < qpNum
+    uint64_t rcqPtr{0};
+
+    // HcclAivMemInfo一维数组
+    // 内存信息个数: rankNum
+    // 计算偏移获取内存信息指针: memPtr + rankId * sizeOfAiRMAMem
+    // srcRankId 获取自身内存信息，dstRankId 获取 Transport 内存信息
+    uint64_t memPtr{0};
+    // 可往后追加字段
+};
+struct CombinedCapability {
+    uint64_t dataplaneModeBitmap;
+};
+
+struct HcclA2CombineOpParam {
+    uint64_t workSpace;                         // Address for communication between client and server,
+                                                // hccl requests and clears
+    uint64_t workSpaceSize;                     // Space for communication between client and server
+    uint32_t rankId;                            // id of this rank
+    uint32_t rankNum;                           // num of ranks in this comm group
+    uint64_t winSize;                           // size of each windows memory
+    uint64_t windowsIn[AscendC::HCCL_MAX_RANK_NUM];      // windows address for input, windowsIn[rankId] corresponds
+                                                // to the local card address,
+                                                // and others are cross-card mapping addresses.
+    uint64_t windowsOut[AscendC::HCCL_MAX_RANK_NUM];     // windows address for output, windowsOut[rankId] corresponds
+                                                // to the local card address,
+                                                // and others are cross-card mapping addresses.
+    uint8_t res[8328];
+    uint8_t multiFlag;
+    __gm__ AscendC::IbVerbsData *data;
+    uint64_t dataSize;
+    // 追加字段
+    uint64_t sizeOfAiRMAInfo; // sizeof(HcclAiRMAInfo)
+    uint64_t aiRMAInfo; // HcclAiRMAInfo* 单个结构体指针
+
+    CombinedCapability* capability;             // address of the communication capability information structure on the Device
+    uint64_t capabilitySize;                    // size of the communication capability information structure
+};
+enum class DataplaneMode : uint32_t {
+    HOST = 0,
+    AICPU = 1,
+    AIV = 2,
+};
+
+enum class DBMode : int32_t {
+    INVALID_DB = -1,
+    HW_DB = 0,
+    SW_DB
+};
+
+struct HcclAiRMAWQ {
+    uint32_t wqn{0};
+    uint64_t bufAddr{0};
+    uint32_t wqeSize{0};
+    uint32_t depth{0};
+    uint64_t headAddr{0};
+    uint64_t tailAddr{0};
+    DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw
+    uint64_t dbAddr{0};
+    uint32_t sl{0};
+};
+
+struct HcclAiRMACQ {
+    uint32_t cqn{0};
+    uint64_t bufAddr{0};
+    uint32_t cqeSize{0};
+    uint32_t depth{0};
+    uint64_t headAddr{0};
+    uint64_t tailAddr{0};
+    DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw
+    uint64_t dbAddr{0};
+};
+
+struct hns_roce_rc_sq_wqe {
+    uint32_t byte_4;
+    uint32_t msg_len;
+    uint32_t immtdata;
+    uint32_t byte_16;
+    uint32_t byte_20;
+    uint32_t rkey;
+    uint64_t remoteVA;
+};
+
+
+struct hns_roce_lite_wqe_data_seg {
+    uint32_t len;
+    uint32_t lkey;
+    uint64_t localVA;
+};
+
+__aicore__ inline void cacheWriteThrough(__gm__ uint8_t* sourceAddr, uint64_t length) {
+    __gm__ uint8_t* start =
+        (__gm__ uint8_t*)((uint64_t)sourceAddr / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE);
+    __gm__ uint8_t* end =
+        (__gm__ uint8_t*)(((uint64_t)sourceAddr + length) / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE);
+    AscendC::GlobalTensor<uint8_t> global;
+    global.SetGlobalBuffer(start);
+    for (uint32_t i = 0; i <= end - start; i += AscendC::CACHE_LINE_SIZE) {
+        AscendC::DataCacheCleanAndInvalid<uint8_t, AscendC::CacheLine::SINGLE_CACHE_LINE,
+            AscendC::DcciDst::CACHELINE_OUT>(global[i]);
+    }
+}
+__aicore__ inline DataplaneMode GetDataplaneMode(GM_ADDR contextGM0) {
+    __gm__ HcclA2CombineOpParam *winContext_ = (__gm__ HcclA2CombineOpParam *)contextGM0;
+    CombinedCapability* capability = winContext_->capability;
+    uint64_t capabilitySize = winContext_->capabilitySize;
+    DataplaneMode dataplaneMode = DataplaneMode::AICPU;
+    if (capability == 0) {
+        return dataplaneMode;
+    }
+    uint64_t dataplaneModeBitmap = capability->dataplaneModeBitmap;
+    if ((dataplaneModeBitmap & 0x04) == 0x04) {
+        dataplaneMode = DataplaneMode::AIV;
+    }
+    return dataplaneMode;
+}
+
+__aicore__ inline int64_t GetCurrentTimestampUs()
+{
+    return AscendC::GetSystemCycle() / TIME_CYCLE;
+}
+
+__aicore__ inline void RecordRankCommDuration(AscendC::LocalTensor<int32_t> performanceInfoU32Tensor, uint32_t rankId, int64_t startTime)
+{
+    int64_t endTime = GetCurrentTimestampUs();
+    int32_t duration = static_cast<int32_t>(endTime - startTime); // int32_t可以表示2^31(us)，约35min在实际场景下满足需要
+    performanceInfoU32Tensor.SetValue(rankId * sizeof(int64_t) / sizeof(int32_t), duration); // 使用int32_t是因为atomicAdd不支持int64_t类型，这里只赋值到int64_t的低32位。
+}
+#endif // MOE_DISTRIBUTE_BASE_H

csrc/mc2/dispatch_ffn_combine/op_kernel/utils/select_helper.hpp

@@ -0,0 +1,25 @@
+#ifndef SELECT_HELPER_HPP
+#define SELECT_HELPER_HPP
+
+#include "catlass/layout/layout.hpp"
+using namespace AscendC;
+using namespace Catlass;
+
+template <typename Layout, typename ElementType, typename = void>
+struct LayoutBInitializer {
+    CATLASS_DEVICE
+    static Layout create(uint32_t k, uint32_t n) {
+        return Layout{k, n};
+    }
+};
+
+template <typename Layout, typename ElementType>
+struct LayoutBInitializer<Layout, ElementType,
+    std::enable_if_t<std::is_same_v<Layout, layout::zN>>
+> {
+    CATLASS_DEVICE
+    static Layout create(uint32_t k, uint32_t n) {
+        return Layout::template MakeLayout<ElementType>(k, n);
+    }
+};
+#endif // SELECT_HELPER_HPP

csrc/mc2/dispatch_ffn_combine_bf16/CMakeLists.txt

@@ -0,0 +1,19 @@
+# -----------------------------------------------------------------------------------------------------------
+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+# CANN Open Software License Agreement Version 2.0 (the "License").
+# Please refer to the License for details. You may not use this file except in compliance with the License.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+# See LICENSE in the root of the software repository for the full text of the License.
+# -----------------------------------------------------------------------------------------------------------
+
+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
+if(NOT ENABLE_TEST AND NOT BENCHMARK)
+    list(REMOVE_ITEM CURRENT_DIRS tests)
+endif()
+foreach(SUB_DIR ${CURRENT_DIRS})
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
+        add_subdirectory(${SUB_DIR})
+    endif()
+endforeach()

csrc/mc2/dispatch_ffn_combine_bf16/op_host/CMakeLists.txt

@@ -0,0 +1,38 @@
+set(_DISPATCH_FFN_INC_OPTS)
+if (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/aarch64-linux/ascendc/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/aarch64-linux/ascendc/include)
+elseif (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/arm64-linux/ascendc/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/arm64-linux/ascendc/include)
+elseif (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/ascendc/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/ascendc/include)
+endif()
+if (EXISTS ${CMAKE_SOURCE_DIR}/third_party/catlass/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${CMAKE_SOURCE_DIR}/third_party/catlass/include)
+endif()
+
+add_op_to_compiled_list()
+
+if (BUILD_OPEN_PROJECT)
+    target_sources(op_host_aclnnInner PRIVATE
+        dispatch_ffn_combine_bf16_def.cpp
+    )
+endif()
+
+add_ops_compile_options(
+    OP_NAME DispatchFFNCombineBF16
+    OPTIONS
+        --cce-auto-sync=on
+        -Wno-deprecated-declarations
+        -DHCCL_COMM
+        -DCATLASS_ARCH=2201
+        ${_DISPATCH_FFN_INC_OPTS}
+)
+
+if (NOT BUILD_OPS_RTY_KERNEL)
+    add_modules_sources(OPTYPE dispatch_ffn_combine_bf16 ACLNNTYPE aclnn_inner)
+    target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
+        ${CMAKE_CURRENT_SOURCE_DIR}
+        ${CMAKE_CURRENT_SOURCE_DIR}/../op_kernel
+    )
+endif()
+

csrc/mc2/dispatch_ffn_combine_bf16/op_host/dispatch_ffn_combine_bf16_def.cpp

@@ -0,0 +1,93 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file dispatch_ffn_combine_bf16_def.cpp
+ * \brief
+ */
+#include "register/op_def_registry.h"
+
+namespace ops {
+class DispatchFFNCombineBF16 : public OpDef {
+ public:
+  explicit DispatchFFNCombineBF16(const char *name) : OpDef(name) {
+    this->Input("a")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("w1")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ})
+        .IgnoreContiguous();
+    this->Input("w2")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ})
+        .IgnoreContiguous();
+    this->Input("expertIdx")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("scale1")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("scale2")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("probs")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+
+    // 输出
+    this->Output("out")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16})
+        .Format({ ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Output("expert_token_nums")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+
+    this->Attr("group").AttrType(REQUIRED).String();
+    this->Attr("M").AttrType(OPTIONAL).Int();
+    this->Attr("transB").AttrType(OPTIONAL).Bool(false);
+    this->Attr("weightNz").AttrType(OPTIONAL).Bool(false);
+
+    OpAICoreConfig aicore_config;
+    aicore_config.DynamicCompileStaticFlag(true)
+        .DynamicFormatFlag(true)
+        .DynamicRankSupportFlag(true)
+        .DynamicShapeSupportFlag(true)
+        .NeedCheckSupportFlag(false)
+        .PrecisionReduceFlag(true)
+        .ExtendCfgInfo("aclnnSupport.value", "support_aclnn")
+        .ExtendCfgInfo("jitCompile.flag", "static_false")
+        .ExtendCfgInfo("multiKernelSupportDynamicGraph.value", "multi_kernel");
+    this->AICore().AddConfig("ascend910_93", aicore_config);
+    // this->AICore().AddConfig("ascend910b", aicore_config);
+    this->MC2().HcclGroup("group");
+  }
+};
+
+OP_ADD(DispatchFFNCombineBF16);
+}  // namespace ops

csrc/mc2/dispatch_ffn_combine_bf16/op_host/dispatch_ffn_combine_bf16_proto.cpp

@@ -0,0 +1,42 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file dispatch_ffn_proto.cpp
+ * \brief
+ */
+#include <graph/utils/type_utils.h>
+#include <register/op_impl_registry.h>
+// #include "../../common/ophost/op_util.h"
+// #include "../../common/ophost/hcom_topo_info.h"
+// #include "log/ops_log.h"
+
+using namespace ge;
+namespace ops {
+const size_t ATTR_GROUP = 0;
+const size_t ATTR_RANK_SIZE = 1;
+const size_t SUPPORT_DIM_SIZE = 2;
+
+static ge::graphStatus InferShapeDispatchFFNCombineBF16(gert::InferShapeContext* context) {
+  (void) context;
+  return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus InferDataTypeDispatchFFNCombineBF16(gert::InferDataTypeContext* context) {
+  (void) context;
+  // auto d_type = context->GetInputDataType(0);
+  // context->SetOutputDataType(0, d_type);
+  return ge::GRAPH_SUCCESS;
+}
+
+IMPL_OP_INFERSHAPE(DispatchFFNCombineBF16)
+  .InferShape(InferShapeDispatchFFNCombineBF16)
+  .InferDataType(InferDataTypeDispatchFFNCombineBF16);
+}  // namespace ops

csrc/mc2/dispatch_ffn_combine_bf16/op_host/dispatch_ffn_combine_bf16_tiling.cpp

@@ -0,0 +1,280 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+/*!
+ * \file dispatch_ffn_tiling.cpp
+ * \brief
+ */
+#include "vector"
+#include "register/tilingdata_base.h"
+#include "tiling/tiling_api.h"
+#include "tiling_base/error_log.h"
+#include "hcom_topo_info.h"
+#include "register/op_def_registry.h"
+#include "dispatch_ffn_combine_bf16_tiling.h"
+#include <vector>
+#include <map>
+#include <algorithm>
+#include "moe_init_routing_v2/moe_init_routing_v2_tiling.h"
+
+using namespace AscendC;
+using namespace ge;
+
+namespace {
+    const char *K_INNER_DEBUG = "DispatchFFNCombineBF16 Tiling Debug";
+    constexpr uint32_t ATTR_GROUP_INDEX = 0;
+    constexpr uint32_t ATTR_MAX_OUTPUT_SIZE_INDEX = 1;
+    constexpr uint32_t ATTR_IS_TRANS_B = 2;
+    constexpr uint32_t ATTR_WEIGHT_NZ = 3;
+    constexpr uint64_t INIT_TILINGKEY = 1000000;
+    constexpr uint64_t TILINGKEY_TRANS_B = 1U;
+    constexpr uint64_t TILINGKEY_WEIGHT_NZ = 10;
+    constexpr uint32_t X_INDEX = 0;
+    constexpr uint32_t WEIGHT_INDEX = 1;
+    constexpr uint32_t WEIGHT2_INDEX = 2;
+    constexpr uint32_t EXPERTID_INDEX = 3;
+    constexpr uint32_t BLOCK_NUM = 20;
+    constexpr uint32_t SYSTEM_NEED_WORKSPACE = 16 * 1024 * 1024;
+    constexpr uint64_t MB_SIZE = 1024 * 1024UL;
+}
+
+namespace optiling {
+
+static int32_t CeilDev(int32_t num, int32_t div)
+{
+    if (div == 0) {
+        return 0;
+    }
+    return (num + div - 1) / div;
+}
+
+static ge::graphStatus DispatchFFNCombineBF16CheckAttrAndSetTiling(gert::TilingContext *context, DispatchFFNCombineBF16Info& info)
+{
+    auto attrs = context->GetAttrs();
+    OP_TILING_CHECK(attrs == nullptr, OP_LOGE(K_INNER_DEBUG, "attrs is null."), return ge::GRAPH_FAILED);
+
+    auto groupPtr = attrs->GetAttrPointer<char>(static_cast<int>(ATTR_GROUP_INDEX));
+    auto maxOutputSizePtr = attrs->GetAttrPointer<int>(ATTR_MAX_OUTPUT_SIZE_INDEX);
+    auto is_trans_b = attrs->GetAttrPointer<bool>(ATTR_IS_TRANS_B);
+    auto weight_nz = attrs->GetAttrPointer<bool>(ATTR_WEIGHT_NZ);
+    OP_TILING_CHECK(groupPtr == nullptr || strlen(groupPtr) == 0,
+    OP_LOGE(K_INNER_DEBUG, "group is invalid."), return GRAPH_FAILED);
+
+    OP_TILING_CHECK(is_trans_b == nullptr,
+        OP_LOGE(K_INNER_DEBUG, "is_trans_b is invalid."), return GRAPH_FAILED);
+    OP_TILING_CHECK(weight_nz == nullptr,
+        OP_LOGE(K_INNER_DEBUG, "weight_nz is invalid."), return GRAPH_FAILED);
+
+    info.maxOutputSize = *maxOutputSizePtr;
+    info.isTransposeB = *is_trans_b;
+    info.isWeightNz = *weight_nz;
+
+    int64_t rankSize;
+    (void)ge::HcomTopoInfo::Instance().GetGroupRankSize(groupPtr, rankSize);
+    info.worldSize = rankSize;
+
+    OP_LOGD(K_INNER_DEBUG, "maxOutputSize=%d ", info.maxOutputSize);
+    OP_LOGD(K_INNER_DEBUG, "rankSize=%d ", info.worldSize);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus DispatchFFNCombineBF16CheckShapeAndSetTiling(gert::TilingContext *context, DispatchFFNCombineBF16Info &info)
+{
+    const char *nodeName = context->GetNodeName();
+
+    const gert::StorageShape *aStorageShape = context->GetInputShape(X_INDEX);
+    auto expertIdxTensor = context->GetDynamicInputTensor(EXPERTID_INDEX, 0);
+    uint32_t M = aStorageShape->GetStorageShape().GetDim(0);
+    uint32_t K = aStorageShape->GetStorageShape().GetDim(1);
+
+    auto wTensor = context->GetDynamicInputTensor(WEIGHT_INDEX, 0);
+    uint32_t wTensorDims = wTensor->GetOriginShape().GetDimNum();
+    uint32_t N = wTensor->GetStorageShape().GetDim(wTensorDims - 1);
+
+    uint32_t topK = expertIdxTensor->GetStorageShape().GetDim(1);
+    uint32_t listLen = 0;
+    while (true) {
+        auto wTensorT = context->GetDynamicInputTensor(WEIGHT_INDEX, ++listLen);
+        if (wTensorT == nullptr) {break;}
+    }
+
+    uint32_t expertPerRank;
+    if (listLen == 1) {
+        expertPerRank = wTensor->GetStorageShape().GetDim(0);
+    } else {
+        expertPerRank = listLen;
+    }
+
+    info.M = M;
+    info.N = N;
+    info.K = K;
+    info.expertPerRank = expertPerRank;
+    info.topK = topK;
+    info.listLen = listLen;
+    OP_LOGD(K_INNER_DEBUG, "M=%d ", info.M);
+    OP_LOGD(K_INNER_DEBUG, "K=%d ", info.K);
+    OP_LOGD(K_INNER_DEBUG, "N=%d ", info.N);
+    OP_LOGD(K_INNER_DEBUG, "expertPerRank=%d ", info.expertPerRank);
+    OP_LOGD(K_INNER_DEBUG, "topK=%d ", info.topK);
+    OP_LOGD(K_INNER_DEBUG, "listLen=%d ", info.listLen);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus DispatchFFNCombineBF16GetPlatformInfoAndSetTiling(gert::TilingContext *context, DispatchFFNCombineBF16Info& info)
+{
+    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
+    uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
+    uint64_t ubSize = 0U;
+    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
+    info.aivNum = aivNum;
+    info.totalUbSize = ubSize;
+
+    OP_LOGD(K_INNER_DEBUG, "aivNum=%d", info.aivNum);
+    OP_LOGD(K_INNER_DEBUG, "ubSize=%lu", info.totalUbSize);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+void SetTilingData(CoCTiling &cocTilingData, DispatchFFNCombineBF16Info &info)
+{
+    cocTilingData.m0 = 128;
+    cocTilingData.k0 = 256;
+    cocTilingData.n0 = 256;
+    cocTilingData.swizzleDirect = 1;
+    cocTilingData.swizzleOffset = 7;
+    cocTilingData.ubMoveNum = 16 * 1024;
+    cocTilingData.pValue = 1;
+    cocTilingData.commNpuSplit = info.worldSize;
+    cocTilingData.commDataSplit = 1;
+    cocTilingData.lenPerLoop = cocTilingData.m0 * cocTilingData.n0 / 2;
+}
+
+static ge::graphStatus DispatchFFNCombineBF16TilingFuncImpl(gert::TilingContext *context)
+{
+    const char *nodeName = context->GetNodeName();
+    OP_LOGI(nodeName, "Enter DispatchFFNCombineBF16 tiling func.");
+
+    // 1. tilingData
+    DispatchFFNCombineBF16TilingData *tilingData = context->GetTilingData<DispatchFFNCombineBF16TilingData>();
+    OP_TILING_CHECK(tilingData == nullptr, OP_LOGE(nodeName, "tilingData is nullptr."),
+        return ge::GRAPH_FAILED);
+    OP_LOGI(nodeName, "DispatchFFNCombineBF16 get tilingData.");
+    DispatchFFNCombineBF16Info& info = tilingData->dispatchFFNCombineBF16Info;
+    OP_LOGI(nodeName, "DispatchFFNCombineBF16 get tilingData info.");
+
+    OP_TILING_CHECK(DispatchFFNCombineBF16CheckAttrAndSetTiling(context, info) != ge::GRAPH_SUCCESS,
+        OP_LOGE(context->GetNodeName(), "DispatchFFNCombineBF16 CheckAttrAndSetTiling Failed"),
+        return ge::GRAPH_FAILED);
+    OP_TILING_CHECK(DispatchFFNCombineBF16CheckShapeAndSetTiling(context, info) != ge::GRAPH_SUCCESS,
+        OP_LOGE(context->GetNodeName(), "DispatchFFNCombineBF16 CheckShapeAndSetTiling Failed"),
+        return ge::GRAPH_FAILED);
+    OP_TILING_CHECK(DispatchFFNCombineBF16GetPlatformInfoAndSetTiling(context, info) != ge::GRAPH_SUCCESS,
+        OP_LOGE(context->GetNodeName(), "DispatchFFNCombineBF16 GetPlatformInfoAndSetTiling Failed"),
+        return ge::GRAPH_FAILED);
+
+    SetTilingData(tilingData->cocTiling, info);
+
+    // 2. set blockDim
+    uint32_t blockDim = 1U;
+    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
+    auto aicNum = ascendcPlatform.GetCoreNumAic();
+    auto aivNum = ascendcPlatform.GetCoreNumAiv();
+    blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum);
+    context->SetBlockDim(blockDim);
+
+    // 3. set tilingKey
+    uint64_t tilingKey = INIT_TILINGKEY;
+    tilingKey += info.isTransposeB ? TILINGKEY_TRANS_B : 0;
+    tilingKey += info.isWeightNz ? TILINGKEY_WEIGHT_NZ : 0;
+    context->SetTilingKey(tilingKey);
+
+    OP_LOGD(K_INNER_DEBUG, "tilingKey=%d", tilingKey);
+
+    optiling::MoeInitRoutingV2TilingBase moeInitRoutingQuantV2TilingBase;
+    int64_t inuptXDtypeSize = sizeof(int16_t);
+    int64_t scaleDim0 = 0;
+    int64_t ubSize = 196352;
+    int64_t expertCapacity = 0;
+    int64_t expertNum = info.expertPerRank * info.worldSize;
+    int64_t activeNum = info.M * info.topK;
+    int64_t dropPadMode = 0;
+     int64_t expertTokensCountOrCumsumFlag = 2;
+     bool expertTokensBeforeCapacityFlag = false;
+     int64_t quantMode = 1;
+     uint32_t aivNumInitRouting = 2 * BLOCK_NUM;
+    moeInitRoutingQuantV2TilingBase.DoTiling(info.M, info.K, info.topK, expertCapacity, expertNum, activeNum, dropPadMode,
+        expertTokensCountOrCumsumFlag, expertTokensBeforeCapacityFlag, inuptXDtypeSize, quantMode, scaleDim0, aivNumInitRouting, ubSize);
+    uint64_t initRoutingQuantTilingKey = moeInitRoutingQuantV2TilingBase.tilingKey_;
+    size_t initRoutingWorkspace = moeInitRoutingQuantV2TilingBase.workspaceSize_;
+
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData = moeInitRoutingQuantV2TilingBase.moeInitRoutingTilingData;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.vbsComputeParamsOp = moeInitRoutingQuantV2TilingBase.moeInitRoutingTilingData.vbsComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.vmsMiddleComputeParamsOp = moeInitRoutingQuantV2TilingBase.moeInitRoutingTilingData.vmsMiddleComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.sortOutComputeParamsOp = moeInitRoutingQuantV2TilingBase.moeInitRoutingTilingData.sortOutComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.srcToDstComputeParamsOp = moeInitRoutingQuantV2TilingBase.moeInitRoutingTilingData.srcToDstComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.srcToDstCapacityComputeParamsOp = moeInitRoutingQuantV2TilingBase.moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.gatherOutComputeParamsOp = moeInitRoutingQuantV2TilingBase.moeInitRoutingTilingData.gatherOutComputeParamsOp;
+    tilingData->cocTiling.initRoutingQuantTilingKey = initRoutingQuantTilingKey;
+    // OP_LOGE(initRoutingTilingKey, " initRoutingTilingKey.");
+    OP_LOGD(K_INNER_DEBUG, "tilingKey=%ld", initRoutingQuantTilingKey);
+
+    // 4. workspace
+    size_t *workSpaces = context->GetWorkspaceSizes(1);
+    OP_TILING_CHECK(workSpaces == nullptr, OP_LOGE(nodeName, "workSpaces is nullptr."),
+        return ge::GRAPH_FAILED);
+
+    uint32_t n2 = info.K;
+    uint32_t k2 = info.N / 2;
+
+    uint64_t cocWorkspace = (info.M + 256 - 1) / 256 * 256 * info.topK *sizeof(int32_t) +
+                            info.worldSize * info.worldSize * info.expertPerRank * sizeof(int32_t) * 3 +
+                            info.maxOutputSize * sizeof(float) * 2 +
+                            info.maxOutputSize * info.N * sizeof(int16_t) +
+                            info.maxOutputSize * n2 * sizeof(int16_t) +
+                            info.maxOutputSize * info.K * sizeof(int16_t) +
+                            info.maxOutputSize * k2 * sizeof(int16_t) +
+                            info.worldSize * sizeof(int32_t) * 16 +
+                            (info.expertPerRank + info.worldSize) * sizeof(int32_t) * 16;
+                            // std::max(info.maxOutputSize * info.N * sizeof(int16_t), info.maxOutputSize * n2 * sizeof(int16_t)) +
+                            // std::max(info.maxOutputSize * info.K * sizeof(int8_t), info.maxOutputSize * k2 * sizeof(int8_t));
+
+    workSpaces[0] = SYSTEM_NEED_WORKSPACE + std::max(cocWorkspace, initRoutingWorkspace);
+
+
+    // 5. communication
+    auto attrs = context->GetAttrs();
+    auto group = attrs->GetAttrPointer<char>(static_cast<int>(ATTR_GROUP_INDEX));
+    uint32_t opType = 8U;
+    std::string algConfig = "AlltoAll=level0:fullmesh;level1:pairwise";
+    AscendC::Mc2CcTilingConfig mc2CcTilingConfig(group, opType, algConfig);
+    mc2CcTilingConfig.GetTiling(tilingData->mc2InitTiling);
+    mc2CcTilingConfig.GetTiling(tilingData->mc2CcTiling);
+
+    OP_LOGI(nodeName, "Leave DispatchFFNCombineBF16 tiling func.");
+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus DispatchFFNCombineBF16TilingFunc(gert::TilingContext* context)
+{
+    return DispatchFFNCombineBF16TilingFuncImpl(context);
+}
+
+struct DispatchFFNCombineBF16CompileInfo {};
+ge::graphStatus TilingParseForDispatchFFNCombineBF16(gert::TilingParseContext *context)
+{
+    (void)context;
+    return ge::GRAPH_SUCCESS;
+}
+
+IMPL_OP_OPTILING(DispatchFFNCombineBF16)
+    .Tiling(DispatchFFNCombineBF16TilingFunc)
+    .TilingParse<DispatchFFNCombineBF16CompileInfo>(TilingParseForDispatchFFNCombineBF16);
+} // namespace optiling

csrc/mc2/dispatch_ffn_combine_bf16/op_host/hcom_topo_info.h

@@ -0,0 +1,72 @@
+/* Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ * ===================================================================================================================*/
+
+#ifndef METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_
+#define METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_
+
+#include <unordered_map>
+#include <mutex>
+
+using Status = int32_t;
+
+namespace ge {
+static constexpr uint32_t COMM_MESH = 0b1U;
+static constexpr uint32_t COMM_SWITCH = (COMM_MESH << 1U);
+static constexpr uint32_t COMM_RING = (COMM_MESH << 2U);
+static constexpr uint32_t COMM_PAIRWISE = (COMM_MESH << 3U);
+class HcomTopoInfo {
+ public:
+  enum class TopoLevel {
+    L0 = 0,
+    L1,
+    MAX,
+  };
+  struct TopoLevelDesc {
+    uint32_t comm_sets;
+    uint32_t rank_size;
+  };
+  using TopoDescs = TopoLevelDesc[static_cast<int32_t>(TopoLevel::MAX)];
+  struct TopoInfo {
+    int64_t rank_size;
+    void *notify_handle;
+    TopoDescs topo_level_descs;
+  };
+  static HcomTopoInfo &Instance();
+  bool TopoInfoHasBeenSet(const char_t *group);
+  bool TryGetGroupTopoInfo(const char_t *group, TopoInfo &info);
+  Status SetGroupTopoInfo(const char_t *group, const TopoInfo &info);
+  Status GetGroupRankSize(const char_t *group, int64_t &rank_size);
+  TopoDescs *GetGroupTopoDesc(const char_t *group);
+  Status GetGroupNotifyHandle(const char_t *group, void *&notify_handle);
+  void UnsetGroupTopoInfo(const char_t *group) {
+    const std::lock_guard<std::mutex> lock(mutex_);
+    (void) rank_info_.erase(group);
+  }
+
+  Status SetGroupOrderedStream(const char_t *group, void *stream);
+  Status GetGroupOrderedStream(const char_t *group, void *&stream);
+  void UnsetGroupOrderedStream(const char_t *group) {
+    const std::lock_guard<std::mutex> lock(mutex_);
+    (void) group_to_ordered_stream_.erase(group);
+  };
+
+  Status SetGroupOrderedStream(const int32_t device_id, const char_t *group, void *stream);
+  Status GetGroupOrderedStream(const int32_t device_id, const char_t *group, void *&stream);
+  void UnsetGroupOrderedStream(const int32_t device_id, const char_t *group);
+ private:
+  HcomTopoInfo() = default;
+  ~HcomTopoInfo() = default;
+  std::unordered_map<std::string, TopoInfo> rank_info_;
+  std::mutex mutex_;
+  std::unordered_map<std::string, void*> group_to_ordered_stream_; // Ordered stream for the communication domain
+  std::unordered_map<int32_t, std::unordered_map<std::string, void*>> device_id_to_group_to_ordered_stream_; // Ordered stream for the communication domain
+};
+}
+
+#endif // METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_

csrc/mc2/dispatch_ffn_combine_bf16/op_host/op_api/aclnn_dispatch_ffn_combine_bf16.cpp

@@ -0,0 +1,83 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+#include "aclnn_dispatch_ffn_combine_bf16.h"
+#include <algorithm>
+// #include "aclnn_kernels/common/op_error_check.h"
+// #include "opdev/op_log.h"
+// #include "opdev/common_types.h"
+// #include "opdev/platform.h"
+// #include "ophost/matmul_util.h"
+#include <unistd.h>
+#include <vector>
+#include <string>
+#include <iostream>
+#include <fcntl.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <sys/file.h>
+#include <climits>
+// using namespace op;
+
+// using namespace op;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+static constexpr size_t TWO_DIMS = 2;
+static constexpr int64_t KVALUE_MIN = 256;
+static constexpr int64_t KVALUE_MAX = 65535;
+static constexpr size_t HCCL_GROUP_NAME_MAX = 128U;
+enum NnopbaseHcclServerType {
+    NNOPBASE_HCCL_SERVER_TYPE_AICPU = 0,
+    NNOPBASE_HCCL_SERVER_TYPE_MTE,
+    NNOPBASE_HCCL_SERVER_TYPE_END
+};
+
+extern aclnnStatus aclnnInnerDispatchFFNCombineBF16GetWorkspaceSize(const aclTensor* x, const aclTensorList* weight1, const aclTensorList* weight2,
+                                                         const aclTensor* expertId, const aclTensorList* scale1, const aclTensorList* scale2,
+                                                         const aclTensor* probs,
+                                                         const char* group, int64_t maxOutputSize,
+                                                         bool transB, bool weightNz,
+                                                         const aclTensor* out, const aclTensor* expertTokenNums,
+                                                         uint64_t* workspaceSize, aclOpExecutor** executor);
+extern aclnnStatus aclnnInnerDispatchFFNCombineBF16(void *workspace, uint64_t workspaceSize,
+                                            aclOpExecutor *executor, aclrtStream stream);
+extern "C" void __attribute__((weak)) NnopbaseSetHcclServerType(void *executor, NnopbaseHcclServerType sType);
+
+
+
+aclnnStatus aclnnDispatchFFNCombineBF16GetWorkspaceSize(const aclTensor* x, const aclTensorList* weight1, const aclTensorList* weight2,
+                                                    const aclTensor* expertId, const aclTensorList* scale1, const aclTensorList* scale2,
+                                                    const aclTensor* probs,
+                                                    const char* group, int64_t maxOutputSize,
+                                                    const aclTensor* out, const aclTensor* expertTokenNums,
+                                                    uint64_t* workspaceSize, aclOpExecutor** executor)
+{
+    bool transB = false;
+    bool weightNz = true;
+
+    aclnnStatus ret = aclnnInnerDispatchFFNCombineBF16GetWorkspaceSize(x, weight1, weight2, expertId, scale1, scale2, probs, group,
+                                                                    maxOutputSize, transB, weightNz,
+                                                                    out, expertTokenNums, workspaceSize, executor);
+    return ret;
+}
+
+aclnnStatus aclnnDispatchFFNCombineBF16(void* workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)
+{
+    if (NnopbaseSetHcclServerType) {
+        NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);
+    }
+    aclnnStatus ret = aclnnInnerDispatchFFNCombineBF16(workspace, workspaceSize, executor, stream);
+    return ret;
+}
+#ifdef __cplusplus
+}
+#endif

csrc/mc2/dispatch_ffn_combine_bf16/op_host/op_api/aclnn_dispatch_ffn_combine_bf16.h

@@ -0,0 +1,39 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+#ifndef OP_API_INC_DISPATCH_FFN_COMBINE_BF16_
+#define OP_API_INC_DISPATCH_FFN_COMBINE_BF16_
+
+#include <string>
+
+#include "aclnn/aclnn_base.h"
+#include "hccl/hccl.h"
+#include "hccl/hccl_types.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+__attribute__((visibility("default"))) aclnnStatus aclnnDispatchFFNCombineBF16GetWorkspaceSize(const aclTensor* x, const aclTensorList* weight1, const aclTensorList* weight2,
+                                                                                        const aclTensor* expertId, const aclTensorList* scale1, const aclTensorList* scale2,
+                                                                                        const aclTensor* probs,
+                                                                                        const char* group, int64_t maxOutputSize,
+                                                                                        const aclTensor* out, const aclTensor* expertTokenNums,
+                                                                                        uint64_t* workspaceSize, aclOpExecutor** executor);
+
+
+__attribute__((visibility("default"))) aclnnStatus aclnnDispatchFFNCombineBF16(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor,
+                                           aclrtStream stream);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // OP_API_INC_DISPATCH_FFN_COMBINE_BF16_

csrc/mc2/dispatch_ffn_combine_bf16/op_host/tiling_args.h

@@ -0,0 +1,9 @@
+#ifndef TILING_ARGS_H
+#define TILING_ARGS_H
+#include <cstdint>
+
+namespace Moe {
+constexpr uint64_t COMBINE_STATE_WIN_OFFSET = 3U * 1024UL * 1024UL;
+constexpr uint64_t NOTIFY_DISPATCH_WIN_OFFSET = 204U * 1024UL * 1024UL;
+}  // namespace Moe
+#endif  // TILING_ARGS_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/dispatch_ffn_combine_bf16.cpp

@@ -0,0 +1,51 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/* !
+ * \file dispatch_ffn_combine.cpp
+ * \brief
+ */
+#include "kernel_operator.h"
+#include "lib/matmul_intf.h"
+#include "dispatch_ffn_combine_bf16_tiling.h"
+#include "dispatch_ffn_combine_bf16.h"
+
+using namespace AscendC;
+using namespace DispatchFFNCombineBF16Impl;
+extern "C" __global__ __aicore__ void dispatch_ffn_combine_bf16(GM_ADDR x, GM_ADDR w1, GM_ADDR w2,  GM_ADDR expertId, GM_ADDR scale1, GM_ADDR scale2, GM_ADDR probs,
+    GM_ADDR c, GM_ADDR expertTokenNums, GM_ADDR workspaceGM,  GM_ADDR tilingGM)
+{
+    REGISTER_TILING_DEFAULT(DispatchFFNCombineBF16TilingData);
+    if (TILING_KEY_IS(1000000)) {
+        KERNEL_TASK_TYPE(1000000, KERNEL_TYPE_MIX_AIC_1_2);
+        GET_TILING_DATA_WITH_STRUCT(DispatchFFNCombineBF16TilingData, tilingData, tilingGM);
+        DispatchFFNCombineBF16<DTYPE_A, DTYPE_W1, DTYPE_OUT, false, false> op;
+        op.Init(x, w1, w2, expertId, scale1, scale2, probs, c, expertTokenNums, workspaceGM, tilingGM);
+        op.Process();
+    } else if (TILING_KEY_IS(1000001)) {
+        KERNEL_TASK_TYPE(1000001, KERNEL_TYPE_MIX_AIC_1_2);
+        GET_TILING_DATA_WITH_STRUCT(DispatchFFNCombineBF16TilingData, tilingData, tilingGM);
+        DispatchFFNCombineBF16<DTYPE_A, DTYPE_W1, DTYPE_OUT, true, false> op;
+        op.Init(x, w1, w2, expertId, scale1, scale2, probs, c, expertTokenNums, workspaceGM, tilingGM);
+        op.Process();
+    } else if (TILING_KEY_IS(1000010)) {
+        KERNEL_TASK_TYPE(1000010, KERNEL_TYPE_MIX_AIC_1_2);
+        GET_TILING_DATA_WITH_STRUCT(DispatchFFNCombineBF16TilingData, tilingData, tilingGM);
+        DispatchFFNCombineBF16<DTYPE_A, DTYPE_W1, DTYPE_OUT, false, true> op;
+        op.Init(x, w1, w2, expertId, scale1, scale2, probs, c, expertTokenNums, workspaceGM, tilingGM);
+        op.Process();
+    } else if (TILING_KEY_IS(1000011)) {
+        KERNEL_TASK_TYPE(1000011, KERNEL_TYPE_MIX_AIC_1_2);
+        GET_TILING_DATA_WITH_STRUCT(DispatchFFNCombineBF16TilingData, tilingData, tilingGM);
+        DispatchFFNCombineBF16<DTYPE_A, DTYPE_W1, DTYPE_OUT, true, true> op;
+        op.Init(x, w1, w2, expertId, scale1, scale2, probs, c, expertTokenNums, workspaceGM, tilingGM);
+        op.Process();
+    }
+}

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/dispatch_ffn_combine_bf16.h

@@ -0,0 +1,291 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file dispatch_ffn_combine.h
+ * \brief
+ */
+
+#ifndef DISPATCH_FFN_COMBINE_BF16_H
+#define DISPATCH_FFN_COMBINE_BF16_H
+
+using namespace AscendC;
+
+#include "kernel_operator.h"
+
+#include "utils/moe_distribute_base.h"
+
+#include "dispatch_ffn_combine_bf16_tiling.h"
+
+#include "catlass/catlass.hpp"
+#include "catlass/arch/arch.hpp"
+#include "catlass/epilogue/dispatch_policy.hpp"
+#include "catlass/epilogue/block/block_epilogue.hpp"
+#include "catlass/epilogue/tile/tile_copy.hpp"
+#include "catlass/epilogue/tile/tile_elemwise_add.hpp"
+#include "catlass/epilogue/tile/tile_elemwise_muls.hpp"
+#include "catlass/gemm/block/block_mmad.hpp"
+#include "catlass/gemm/block/block_swizzle.hpp"
+#include "catlass/gemm/dispatch_policy.hpp"
+#include "catlass/gemm/kernel/matmul_epilogue.hpp"
+#include "catlass/gemm/gemm_type.hpp"
+#include "catlass/layout/layout.hpp"
+// #include "catlass/status.hpp"
+
+#include "utils/select_helper.hpp"
+#include "utils/const_args.hpp"
+#include "dispatch_ffn_combine_bf16_kernel.hpp"
+#include "moe_init_routing_v2/moe_init_routing_v2_tiling.h"
+
+using namespace Catlass;
+using namespace AscendC::HcclContextDef;
+
+namespace DispatchFFNCombineBF16Impl {
+#define TemplateMMA2AClass typename AType_, typename BType_, typename CType_, bool TB_, bool Nz_
+#define TemplateMMA2ACFunc AType_, BType_, CType_, TB_, Nz_
+
+using namespace AscendC;
+template <TemplateMMA2AClass>
+class DispatchFFNCombineBF16 {
+public:
+    __aicore__ inline DispatchFFNCombineBF16() {};
+    __aicore__ inline void Init(GM_ADDR xGM, GM_ADDR weight1GM, GM_ADDR weight2GM, GM_ADDR expertIdGM, GM_ADDR scale1GM, GM_ADDR scale2GM,
+                                GM_ADDR probs, GM_ADDR outGM, GM_ADDR expertTokenNums, GM_ADDR workspaceGM, GM_ADDR tilingGM);
+    __aicore__ inline void Process();
+
+
+private:
+    GM_ADDR xGM_;
+    GM_ADDR weight1GM_;
+    GM_ADDR weight2GM_;
+    GM_ADDR expertIdGM_;
+    GM_ADDR scale1GM_;
+    GM_ADDR scale2GM_;
+    GM_ADDR probs_;
+    GM_ADDR outGM_;
+    GM_ADDR gmExpertTokenNums_;
+    GM_ADDR workspaceGM_;
+
+    GM_ADDR moeInitRoutingQuantV2Scale = nullptr;
+    GM_ADDR moeInitRoutingQuantV2Offset = nullptr;
+    GM_ADDR expertTokensBeforeCapacity = nullptr;
+
+
+    TBuf<AscendC::TPosition::VECCALC> uBuf_;
+
+    int32_t rank;
+    int32_t rankSize;
+    int32_t aivNum;
+
+    int32_t m0;
+    int32_t k0;
+    int32_t n0;
+    int32_t swizzlOffset;
+    int32_t swizzlDirect;
+    int32_t ubMoveNum;
+    int32_t pValue;
+
+    int32_t commNpuSplit;
+    int32_t commDataSplit;
+    int32_t lenPerLoop;
+
+    int32_t m;
+    int32_t k;
+    int32_t n;
+    int32_t topK;
+    int32_t expertPerRank;
+    int32_t maxOutputSize;
+    int32_t EP;
+    int32_t listLen;
+
+    optiling::MoeInitRoutingV2TilingData moeInitRoutingQuantV2TilingData;
+    uint64_t initRoutingQuantTilingKey;
+
+    // Hccl<HCCL_SERVER_TYPE_AICPU> hccl_;
+
+};
+
+
+template <TemplateMMA2AClass>
+__aicore__ inline void DispatchFFNCombineBF16<TemplateMMA2ACFunc>::Init(GM_ADDR xGM, GM_ADDR weight1GM, GM_ADDR weight2GM, GM_ADDR expertIdGM, GM_ADDR scale1GM, GM_ADDR scale2GM,
+                                                                    GM_ADDR probs, GM_ADDR outGM, GM_ADDR expertTokenNums, GM_ADDR workspaceGM, GM_ADDR tilingGM)
+{
+    REGISTER_TILING_DEFAULT(DispatchFFNCombineBF16TilingData);
+    auto tiling = (__gm__ DispatchFFNCombineBF16TilingData*)tilingGM;
+    GET_TILING_DATA(tilingData, tilingGM);
+
+    xGM_ = xGM;
+    weight1GM_ = weight1GM;
+    weight2GM_ = weight2GM;
+    expertIdGM_ = expertIdGM;
+    scale1GM_ = scale1GM;
+    scale2GM_ = scale2GM;
+    probs_ = probs;
+
+    outGM_ = outGM;
+    gmExpertTokenNums_ = expertTokenNums;
+
+    workspaceGM_ = workspaceGM;
+
+    aivNum = tilingData.dispatchFFNCombineBF16Info.aivNum;
+
+    m = tilingData.dispatchFFNCombineBF16Info.M;
+    k = tilingData.dispatchFFNCombineBF16Info.K;
+    n = tilingData.dispatchFFNCombineBF16Info.N;
+    EP =  tilingData.dispatchFFNCombineBF16Info.worldSize;
+    topK = tilingData.dispatchFFNCombineBF16Info.topK;
+    expertPerRank = tilingData.dispatchFFNCombineBF16Info.expertPerRank;
+    maxOutputSize = tilingData.dispatchFFNCombineBF16Info.maxOutputSize;
+    listLen = tilingData.dispatchFFNCombineBF16Info.listLen;
+
+    m0 = tilingData.cocTiling.m0;
+    k0 = tilingData.cocTiling.k0;
+    n0 = tilingData.cocTiling.n0;
+    swizzlDirect = tilingData.cocTiling.swizzleDirect;
+    swizzlOffset = tilingData.cocTiling.swizzleOffset;
+    ubMoveNum = tilingData.cocTiling.ubMoveNum;
+    pValue = tilingData.cocTiling.pValue;
+    commNpuSplit = tilingData.cocTiling.commNpuSplit;
+    commDataSplit = tilingData.cocTiling.commDataSplit;
+    lenPerLoop = tilingData.cocTiling.lenPerLoop;
+    moeInitRoutingQuantV2TilingData = tilingData.cocTiling.moeInitRoutingQuantV2TilingData;
+    initRoutingQuantTilingKey = tilingData.cocTiling.initRoutingQuantTilingKey;
+
+    auto contextGM0 = AscendC::GetHcclContext<HCCL_GROUP_ID_0>();
+    __gm__ HcclOpResParamCustom *WinContext_{nullptr};
+    WinContext_ = (__gm__ HcclOpResParamCustom *)contextGM0;
+
+    rank = WinContext_->localUsrRankId;
+    rankSize = WinContext_->rankSize;
+}
+
+template <TemplateMMA2AClass>
+__aicore__ inline void DispatchFFNCombineBF16<TemplateMMA2ACFunc>::Process()
+{
+    // Define ArchTag
+    using ArchTag = Arch::AtlasA2;
+    constexpr bool enableUnitFlag = false;
+    constexpr bool enableShuffleK = true;
+
+    uint32_t k2 = n/2;
+    uint32_t n2 = k;
+
+    int64_t activeNum = 0;
+    int64_t expertCapacity = 0;
+    int64_t expertNum = expertPerRank * EP;
+    int64_t dropPadMode = 0;
+    int64_t expertTokensCountOrCumsumFlag = 2;
+    bool expertTokensBeforeCapacityFlag = false;
+    int64_t quantMode = 1;
+
+    using LayoutA = layout::RowMajor;
+    using LayoutB = typename std::conditional<
+        Nz_,
+        layout::zN,
+        typename std::conditional<TB_, layout::ColumnMajor, layout::RowMajor>::type
+    >::type;
+
+    LayoutB layoutB1 = LayoutBInitializer<LayoutB, BType_>::create(k, n);
+    LayoutB layoutB2 = LayoutBInitializer<LayoutB, BType_>::create(k2, n2);
+    using LayoutC = layout::RowMajor;
+    using L1TileShape = typename std::conditional<
+            std::is_same_v<AType_, int8_t>,
+            GemmShape<128, 256, 512>,
+            GemmShape<128, 256, 256>
+            >::type;
+
+    constexpr uint32_t workspaceStages = 2;
+    constexpr uint32_t preloadStages = 1;
+    constexpr uint32_t l1Stages = 2;
+    constexpr uint32_t l0AStages = 2;
+    constexpr uint32_t l0BStages = 2;
+    constexpr uint32_t l0CStages = 1;
+
+    using DispatchPolicy = Gemm::MmadAtlasA2PreloadAsyncFixpipe<
+        preloadStages,
+        l1Stages, l0AStages, l0BStages, l0CStages,
+        enableUnitFlag, enableShuffleK
+    >;
+
+    using L0TileShape = typename std::conditional<
+            std::is_same_v<AType_, int8_t>,
+            GemmShape<128, 256, 128>,
+            GemmShape<128, 256, 64>
+            >::type;
+    using AType = Gemm::GemmType<AType_, layout::RowMajor>;
+    using BType = Gemm::GemmType<BType_, LayoutB>;
+    using CType = Gemm::GemmType<CType_, layout::RowMajor>;
+    using D1Type = Gemm::GemmType<CType_, layout::RowMajor>;
+
+    using D2Type = typename std::conditional<
+        std::is_same_v<CType_, bfloat16_t>,
+        Gemm::GemmType<bfloat16_t, layout::RowMajor>,
+        Gemm::GemmType<CType_, layout::RowMajor>
+        >::type;
+
+    using BlockMmad = Gemm::Block::BlockMmad<DispatchPolicy, L1TileShape, L0TileShape, AType, BType, CType>;
+    constexpr uint32_t ubStages = 2;
+
+    using EpilogueDispatchPolicy1 = Epilogue::EpilogueAtlasA2PerTokenDequantSwigluQuant<ubStages>;
+
+    using ScaleType = Gemm::GemmType<uint64_t, layout::VectorLayout>;
+    using PerTokenScaleType = Gemm::GemmType<float, layout::VectorLayout>;
+    using ElementMulType = Gemm::GemmType<float, layout::RowMajor>;
+    using TileElemWiseMuls = Epilogue::Tile::TileElemWiseMuls<ArchTag, ElementMulType, 0>;
+
+    using TileCopy1 = Epilogue::Tile::TileCopy<ArchTag, CType, ScaleType, PerTokenScaleType, D1Type>;
+    using BlockEpilogue1 = Epilogue::Block::BlockEpilogue<EpilogueDispatchPolicy1, CType, PerTokenScaleType,
+        D1Type, TileElemWiseMuls, TileCopy1>;
+
+    using EpilogueDispatchPolicy2 = Epilogue::EpilogueAtlasA2PerTokenDequantV2<ubStages>;
+    using TileCopy2 = Epilogue::Tile::TileCopy<ArchTag, CType, ScaleType, PerTokenScaleType, D2Type>;
+    using BlockEpilogue2 = Epilogue::Block::BlockEpilogue<EpilogueDispatchPolicy2, CType,PerTokenScaleType,
+        D2Type, TileCopy2>;
+
+    using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<9, 1>;
+    using ElementGroupList = int64_t;
+    using MatmulKernel = Gemm::Kernel::DispatchFFNCombineBF16Kernel<BlockMmad,
+        BlockScheduler, ElementGroupList, BlockEpilogue1, BlockEpilogue2>;
+
+    LayoutA layoutA1{static_cast<uint32_t>(m), static_cast<uint32_t>(k)};
+    LayoutA layoutA2{static_cast<uint32_t>(m), static_cast<uint32_t>(k2)};
+    layout::VectorLayout layoutScale1{static_cast<uint32_t>(n)};
+    layout::VectorLayout layoutScale2{static_cast<uint32_t>(n2)};
+    layout::RowMajor layoutD1{static_cast<uint32_t>(maxOutputSize), static_cast<uint32_t>(k2)};
+    layout::RowMajor layoutD2{static_cast<uint32_t>(m*topK), static_cast<uint32_t>(n2)};
+    // Prepare params
+
+    GemmCoord problemShape{static_cast<uint32_t>(m), static_cast<uint32_t>(n), static_cast<uint32_t>(k)};
+
+    uint32_t epilogueCoreNum = aivNum;
+    uint32_t epilogueGranularity = expertPerRank - 2;
+
+    typename MatmulKernel::Params params{
+        problemShape, static_cast<uint32_t>(EP), static_cast<uint32_t>(listLen), static_cast<uint32_t>(expertPerRank), static_cast<uint32_t>(maxOutputSize),
+        static_cast<uint32_t>(rank), static_cast<uint32_t>(rankSize),
+        static_cast<uint32_t>(topK), initRoutingQuantTilingKey,
+        epilogueCoreNum, epilogueGranularity,
+        xGM_, layoutA1, layoutA2,
+        weight1GM_, layoutB1,
+        weight2GM_, layoutB2,
+        scale1GM_, layoutScale1,
+        scale2GM_, layoutScale2,
+        outGM_, layoutD1, layoutD2,
+        expertIdGM_, moeInitRoutingQuantV2Scale, moeInitRoutingQuantV2Offset,
+        expertTokensBeforeCapacity, probs_,
+        workspaceGM_, gmExpertTokenNums_, ubMoveNum, moeInitRoutingQuantV2TilingData};
+    //Call kernel
+    MatmulKernel kernel(params);
+    kernel(params);
+}
+
+} // DispatchFFNCombineBF16Impl
+#endif // DISPATCH_FFN_COMBINE_H
+

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/dispatch_ffn_combine_bf16_tiling.h

@@ -0,0 +1,56 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file dispatch_ffn_combine_tiling.h
+ * \brief
+ */
+
+#include "moe_init_routing_v2/moe_init_routing_v2_tiling.h"
+
+#ifndef ASCENDC_DISPATCH_FFN_COMBINE_BF16_TILING_H
+#define ASCENDC_DISPATCH_FFN_COMBINE_BF16_TILING_H
+struct DispatchFFNCombineBF16Info {
+    uint32_t M;
+    uint32_t K;
+    uint32_t N;
+    uint32_t expertPerRank;
+    uint32_t maxOutputSize;
+    uint32_t isTransposeB;
+    uint32_t isWeightNz;
+    uint32_t aivNum;
+    uint32_t totalUbSize;
+    uint32_t topK;
+    uint32_t worldSize;
+    uint32_t listLen;
+};
+
+struct CoCTiling {
+    int32_t m0 = -1;
+    int32_t k0 = -1;
+    int32_t n0 = -1;
+    int32_t swizzleDirect = -1;
+    int32_t swizzleOffset = -1;
+    int32_t ubMoveNum = -1;
+    int32_t pValue = -1;
+    int32_t commNpuSplit = -1;
+    int32_t commDataSplit = -1;
+    int32_t lenPerLoop = -1;
+    uint64_t initRoutingQuantTilingKey;
+    optiling::MoeInitRoutingV2TilingData moeInitRoutingQuantV2TilingData;
+};
+
+struct DispatchFFNCombineBF16TilingData {
+    Mc2InitTiling mc2InitTiling;
+    Mc2CcTiling mc2CcTiling;
+    DispatchFFNCombineBF16Info dispatchFFNCombineBF16Info;
+    CoCTiling cocTiling;
+};
+#endif

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_init_routing_v2.cpp

@@ -0,0 +1,125 @@
+/**
+ * This program is free software, you can redistribute it and/or modify.
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/* !
+ * \file moe_init_routing_v2.cpp
+ * \brief
+ */
+
+
+#ifdef __DAV_C310__
+#include "arch35/moe_v2_mrgsort_out.h"
+#include "arch35/moe_v2_mrgsort.h"
+#include "arch35/moe_v2_sort_multi_core.h"
+#include "arch35/moe_v2_sort_one_core.h"
+#include "arch35/moe_v2_expert_token_out_regbase.h"
+#include "arch35/moe_v2_expert_token_out_simt.h"
+#include "arch35/moe_v2_src_to_dst_op_simt.h"
+#include "arch35/moe_v2_src_to_dst_with_capacity_simt.h"
+#include "arch35/moe_v2_gather_out_for_simt.h"
+#else
+#include "moe_v2_mrgsort_out.h"
+#include "moe_v2_mrgsort.h"
+#include "moe_v2_sort_multi_core.h"
+#include "moe_v2_sort_one_core.h"
+#include "moe_v2_expert_token_out.h"
+#include "moe_v2_src_to_dst_op.h"
+#include "moe_v2_src_to_dst_with_capacity.h"
+#include "moe_v2_gather_out.h"
+#include "moe_v2_init_routing_fullload.h"
+#endif
+
+using namespace AscendC;
+using namespace MoeInitRoutingV2;
+using namespace optiling;
+
+template <class DTYPE_X = bfloat16_t>
+__aicore__ inline  void moe_init_routing_v2(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX,
+                                                          GM_ADDR expandedRowIdx, GM_ADDR expertTokensCountOrCumsum,
+                                                          GM_ADDR expertTokensBeforeCapacity, GM_ADDR workspace,
+                                                          const MoeInitRoutingV2TilingData* tilingData, uint64_t tilingKey)
+{
+
+    if (g_coreType == AIC) {
+        return;
+    }
+
+    // GET_TILING_DATA(tilingData, tiling);
+    if (workspace == nullptr) {
+        return;
+    }
+
+    GM_ADDR userWS = workspace;
+    if (userWS == nullptr) {
+        return;
+    }
+    // auto t = tilingData;
+    if (tilingKey == 20000) {
+        TPipe sortPipe;
+        MoeV2FullLoad<DTYPE_X> op;
+        op.Init(x, expertIdx, expandedX, expandedRowIdx, expertTokensCountOrCumsum, userWS, tilingData, &sortPipe);
+        op.Process();
+        sortPipe.Destroy();
+        // trap();
+        return;
+    }
+
+    if (tilingKey == 10001 || tilingKey == 10011) {
+        TPipe sortPipe;
+        MoeV2SortOneCore op;
+        op.Init<MoeInitRoutingV2TilingData>(expertIdx, expertTokensCountOrCumsum, expertTokensBeforeCapacity, userWS, tilingData,
+                                            &sortPipe);
+        op.Process();
+        sortPipe.Destroy();
+    } else if (tilingKey == 10002 || tilingKey == 10012) {
+        TPipe sortPipe;
+        MoeV2SortMultiCore op;
+        op.Init<MoeInitRoutingV2TilingData>(expertIdx, expertTokensCountOrCumsum, expertTokensBeforeCapacity, userWS, tilingData,
+                                            &sortPipe);
+        op.Process();
+        sortPipe.Destroy();
+    }
+
+    if (tilingKey == 10001 || tilingKey == 10002) {
+        if (tilingData->expertTokensCountOrCumsumFlag != EXERPT_TOKENS_NONE) {
+            TPipe expertTokenOutPipe;
+            MoeV2ExpertTokenOut expertTokenOutOp;
+            expertTokenOutOp.Init<MoeInitRoutingV2TilingData>(expertTokensCountOrCumsum, expertTokensBeforeCapacity,
+                                                              expandedRowIdx, userWS, tilingData, &expertTokenOutPipe);
+            expertTokenOutOp.Process();
+            expertTokenOutPipe.Destroy();
+        }
+        TPipe srcToDstPipe;
+        MoeV2SrcToDstOp srcToDstOp;
+        srcToDstOp.Init<MoeInitRoutingV2TilingData>(expandedRowIdx, userWS, tilingData, &srcToDstPipe);
+        srcToDstOp.Process();
+        srcToDstPipe.Destroy();
+    } else if (tilingKey == 10011 || tilingKey == 10012) {
+        TPipe expertTokenOutPipe;
+        MoeV2ExpertTokenOut expertTokenOutOp;
+        expertTokenOutOp.Init<MoeInitRoutingV2TilingData>(expertTokensCountOrCumsum, expertTokensBeforeCapacity,
+                                                          expandedRowIdx, userWS, tilingData, &expertTokenOutPipe);
+        expertTokenOutOp.Process();
+        expertTokenOutPipe.Destroy();
+
+        TPipe srcToDstPipe;
+        MoeV2SrcToDstWithCapacity<DTYPE_X, MoeInitRoutingV2TilingData> srcToDstWithCapacityOp;
+        srcToDstWithCapacityOp.Init(expandedRowIdx, expandedX, userWS, tilingData, &srcToDstPipe);
+        srcToDstWithCapacityOp.Process();
+        srcToDstPipe.Destroy();
+    }
+
+    TPipe gatherPipe;
+    MoeV2GatherOut<DTYPE_X> gatherOp;
+    gatherOp.Init(x, expandedRowIdx, expandedX, userWS, tilingData, &gatherPipe);
+    gatherOp.Process();
+    gatherPipe.Destroy();
+
+}

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_init_routing_v2_tiling.h

@@ -0,0 +1,557 @@
+#pragma once
+
+#include "tiling_base.h"
+
+
+namespace optiling {
+const static int64_t TILING_KEY_DROPLESS_SORT_ONE_CORE = 10001;
+const static int64_t TILING_KEY_DROPLESS_SORT_MULTI_CORE = 10002;
+const static int64_t TILING_KEY_DROP_PAD_MODE_SORT_ONE_CORE = 10011;
+const static int64_t TILING_KEY_DROP_PAD_MODE_SORT_MULTI_CORE = 10012;
+const static int64_t TILING_KEY_HIGH_PERFORMANCE = 20000;
+const static int64_t NUM_TWO = 2;
+const static int64_t NUM_THREE = 3;
+const static int64_t NUM_FOUR = 4;
+const static int64_t MRG_LIST_NUM = 4;
+const static int64_t SORT32_ALIGN_ELEMENT = 32;
+const static int64_t ONE_BLOCK_BYTE = 32;
+const static size_t DIM_ONE = 1;
+const static size_t DIM_TWO = 2;
+const static size_t DIM_THREE = 3;
+const static int32_t SIZE_16 = 16;
+const static int32_t LENGTH_1024 = 1024;
+const static int64_t MAX_COLS_ONE_LOOP = 16376;
+const static int64_t ASSIST_NUM = 256;
+const static int64_t INDEX_INPUT_X = 0;
+const static int64_t INDEX_INPUT_EXPERT_IDX = 1;
+const static int64_t ATTR_ACTIVE_ROWS = 0;
+const static int64_t ATTR_EXPERT_CAPACITY = 1;
+const static int64_t ATTR_EXPERT_NUM = 2;
+const static int64_t ATTR_DROP_PAD_MODE = 3;
+const static int64_t ATTR_EXPERT_TOKENS_COUNT_OR_CUMSUM_FLAG = 4;
+const static int64_t ATTR_EXPERT_TOKENS_BEFORE_CAPACITY_FLAG = 5;
+const static int64_t OUTOUT_EXPANDED_X = 0;
+const static int64_t OUTOUT_EXPANDED_ROW_IDX = 1;
+const static int64_t OUTOUT_EXPERT_TOKENS_COUNT_OR_CUMSUM = 2;
+const static int64_t OUTOUT_EXPERT_TOKENS_BEFORE_CAPACITY = 3;
+const static int64_t KV_FACTOR = 2;
+const static int64_t ONE_CORE_SORT_BUFFER = 6;
+const static int64_t EXPERT_TOKENS_COUNT = 2;
+const static int64_t ONE_CORE_SORT_BUFFER_310P = 24;
+
+
+inline static int64_t CeilLog4(int64_t x) {
+  return static_cast<int64_t>(std::ceil(std::log(x) / std::log(NUM_FOUR)));
+}
+
+inline static int64_t GetPerOrLastValue(int64_t x, int64_t y) {
+  if (y == 0) {
+    return 0;
+  }
+  return x <= y ? x : x % y;
+}
+
+template <class T>
+constexpr T CeilDiv(const T dividend, const T divisor)
+{
+    return (dividend + divisor - 1) / divisor;
+}
+
+
+struct MoeV2VBSComputeTilingData {
+    int64_t needCoreNum = 0;
+    int64_t perCoreElements = 0;
+    int64_t perCoreLoops = 0;
+    int64_t perCorePerLoopElements = 0;
+    int64_t perCoreLastLoopElements = 0;
+    int64_t lastCoreElements = 0;
+    int64_t lastCoreLoops = 0;
+    int64_t lastCorePerLoopElements = 0;
+    int64_t lastCoreLastLoopElements = 0;
+    int64_t oneLoopMaxElements = 0;
+};
+
+struct MoeV2VMSMiddleComputeTilingData {
+    int64_t needCoreNum = 0;
+};
+
+struct MoeV2SortOutComputeTilingData {
+    int64_t oneLoopMaxElements = 0;
+};
+
+struct MoeV2GatherOutComputeTilingData {
+    int64_t needCoreNum = 0;
+    int64_t activateRows = 0;
+    int64_t perCoreRows = 0;
+    int64_t perCorePerLoopRows = 0;
+    int64_t perCoreLastLoopRows = 0;
+    int64_t lastCoreRows = 0;
+    int64_t lastCorePerLoopRows = 0;
+    int64_t lastCoreLastLoopRows = 0;
+    int64_t perCoreLoops = 0;
+    int64_t lastCoreLoops = 0;
+    int64_t perLoopCols = 0;
+    int64_t lastLoopCols = 0;
+    int64_t colLoops = 0;
+};
+
+struct MoeInitRoutingV2TilingData {
+    int64_t coreNum;
+    int64_t n;
+    int64_t cols;
+    int64_t k;
+    int64_t expertCapacity;
+    int64_t expertNum;
+    int64_t dropPadMode;
+    int64_t expertTokensCountOrCumsumFlag;
+    int64_t expertTokensBeforeCapacityFlag;
+    MoeV2VBSComputeTilingData vbsComputeParamsOp;
+    MoeV2VMSMiddleComputeTilingData vmsMiddleComputeParamsOp;
+    MoeV2SortOutComputeTilingData sortOutComputeParamsOp;
+    MoeV2GatherOutComputeTilingData srcToDstComputeParamsOp;
+    MoeV2GatherOutComputeTilingData srcToDstCapacityComputeParamsOp;
+    MoeV2GatherOutComputeTilingData gatherOutComputeParamsOp;
+};
+
+
+class MoeInitRoutingV2TilingBase : public TilingBaseClass {
+
+protected:
+  bool GetPlatformInfo(int64_t aivCoreNum, int64_t ubSizePlatForm) override;
+  bool GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity,
+  int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag,
+  bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) override;
+
+  bool DoOpTiling() override;
+  uint64_t GetTilingKey() const override;
+  bool GetWorkspaceSize() override;
+
+
+protected:
+  bool CheckTokenCount(int64_t num, const char* tag);
+
+  void Tiling4GatherOutCompute();
+  void Tiling4SrcToDstCompute();
+  void Tiling4SrcToDstCapacityCompute();
+  void Tiling4SortOutCompute();
+  void Tiling4VMSMiddleCompute();
+  void Tiling4VBSCompute();
+  void ShowTilingData();
+  void Tiling4VBSMultiCoreCompute(MoeV2VBSComputeTilingData* tilingData);
+  void Tiling4VBSOneCoreCompute(MoeV2VBSComputeTilingData* tilingData);
+  bool IsFullLoad();
+
+
+
+
+  int64_t aivNum = 0;
+  int64_t sortLoopMaxElement = 0;
+  int64_t mrgSortListMaxElement = 2040;
+  int64_t totalLength = 0;
+  int64_t activateNum = 0;
+  int64_t expertCapacity = 0;
+  int64_t expertNum = 0;
+  int64_t dropPadMode = 0;
+  int64_t expertTokensCountOrCumsumFlag = 0;
+  bool expertTokensBeforeCapacityFlag = false;
+  int64_t inuptXDtypeSize_ = 0;
+  bool isFullLoad = false;
+  const char *opName = "DispatchFFNCombine Tiling Debug";
+
+public:
+  MoeInitRoutingV2TilingData moeInitRoutingTilingData;
+};
+
+
+bool MoeInitRoutingV2TilingBase::DoOpTiling() {
+  sortLoopMaxElement =
+      (aicoreParams_.ubSize) / (sizeof(int32_t) * NUM_TWO * NUM_FOUR) / SORT32_ALIGN_ELEMENT * SORT32_ALIGN_ELEMENT;
+  isFullLoad = IsFullLoad();
+  Tiling4VBSCompute();
+  Tiling4VMSMiddleCompute();
+  Tiling4SortOutCompute();
+  Tiling4SrcToDstCompute();
+  Tiling4SrcToDstCapacityCompute();
+  Tiling4GatherOutCompute();
+  ShowTilingData();
+  return true;
+};
+
+uint64_t MoeInitRoutingV2TilingBase::GetTilingKey() const {
+  if (isFullLoad) {
+    return TILING_KEY_HIGH_PERFORMANCE;
+  }
+  if (dropPadMode == 0) {
+    if (totalLength <= sortLoopMaxElement) {  
+      return TILING_KEY_DROPLESS_SORT_ONE_CORE;
+    } else {
+      return TILING_KEY_DROPLESS_SORT_MULTI_CORE;
+    }
+  } else {
+    if (totalLength <= sortLoopMaxElement) {
+      return TILING_KEY_DROP_PAD_MODE_SORT_ONE_CORE;
+    } else {
+      return TILING_KEY_DROP_PAD_MODE_SORT_MULTI_CORE;
+    }
+  }
+  return tilingKey_;
+}
+
+
+void MoeInitRoutingV2TilingBase::ShowTilingData()
+{
+    OP_LOGD(opName,
+              "moeInitRoutingTilingData is coreNum:%ld, n:%ld, cols:%ld, k:%ld, expertCapacity:%ld, expertNum:%ld, "
+              "dropPadMode:%ld, expertTokensCountOrCumsumFlag:%ld, expertTokensBeforeCapacityFlag:%ld",
+              moeInitRoutingTilingData.coreNum, moeInitRoutingTilingData.n,
+              moeInitRoutingTilingData.cols, moeInitRoutingTilingData.k,
+              moeInitRoutingTilingData.expertCapacity, moeInitRoutingTilingData.expertNum,
+              moeInitRoutingTilingData.dropPadMode, moeInitRoutingTilingData.expertTokensCountOrCumsumFlag,
+              moeInitRoutingTilingData.expertTokensBeforeCapacityFlag);
+    OP_LOGD(opName,
+              "MoeV2VBSComputeTilingData is needCoreNum:%ld, perCoreElements:%ld, perCoreLoops:%ld, "
+              "perCorePerLoopElements:%ld, "
+              "perCoreLastLoopElements:%ld, lastCoreElements:%ld, lastCoreLoops:%ld, lastCorePerLoopElements:%ld, "
+              "lastCoreLastLoopElements:%ld, oneLoopMaxElements:%ld",
+              moeInitRoutingTilingData.vbsComputeParamsOp.needCoreNum,
+              moeInitRoutingTilingData.vbsComputeParamsOp.perCoreElements,
+              moeInitRoutingTilingData.vbsComputeParamsOp.perCoreLoops,
+              moeInitRoutingTilingData.vbsComputeParamsOp.perCorePerLoopElements,
+              moeInitRoutingTilingData.vbsComputeParamsOp.perCoreLastLoopElements,
+              moeInitRoutingTilingData.vbsComputeParamsOp.lastCoreElements,
+              moeInitRoutingTilingData.vbsComputeParamsOp.lastCoreLoops,
+              moeInitRoutingTilingData.vbsComputeParamsOp.lastCorePerLoopElements,
+              moeInitRoutingTilingData.vbsComputeParamsOp.lastCoreLastLoopElements,
+              moeInitRoutingTilingData.vbsComputeParamsOp.oneLoopMaxElements);
+    OP_LOGD(opName, "VMSMiddleComputeTilingData is needCoreNum:%ld",
+              moeInitRoutingTilingData.vmsMiddleComputeParamsOp.needCoreNum);
+    OP_LOGD(opName, "SortOutComputeTilingData is oneLoopMaxElements:%ld",
+              moeInitRoutingTilingData.sortOutComputeParamsOp.oneLoopMaxElements);
+    OP_LOGD(
+        opName,
+        "SrcToDstComputeTilingData is needCoreNum:%ld, activateRows:%ld, perCoreRows:%ld, perCorePerLoopRows:%ld, "
+        "perCoreLastLoopRows:%ld, lastCoreRows:%ld, lastCorePerLoopRows:%ld, lastCoreLastLoopRows:%ld,",
+        moeInitRoutingTilingData.srcToDstComputeParamsOp.needCoreNum,
+        moeInitRoutingTilingData.srcToDstComputeParamsOp.activateRows,
+        moeInitRoutingTilingData.srcToDstComputeParamsOp.perCoreRows,
+        moeInitRoutingTilingData.srcToDstComputeParamsOp.perCorePerLoopRows,
+        moeInitRoutingTilingData.srcToDstComputeParamsOp.perCoreLastLoopRows,
+        moeInitRoutingTilingData.srcToDstComputeParamsOp.lastCoreRows,
+        moeInitRoutingTilingData.srcToDstComputeParamsOp.lastCorePerLoopRows,
+        moeInitRoutingTilingData.srcToDstComputeParamsOp.lastCoreLastLoopRows);
+    OP_LOGD(opName,
+              "SrcToDstComputeCapacityTilingData is needCoreNum:%ld, perCoreRows:%ld, perCorePerLoopRows:%ld, "
+              "perCoreLastLoopRows:%ld, lastCoreRows:%ld, lastCorePerLoopRows:%ld, lastCoreLastLoopRows:%ld,",
+              moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp.needCoreNum,
+              moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp.perCoreRows,
+              moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp.perCorePerLoopRows,
+              moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp.perCoreLastLoopRows,
+              moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp.lastCoreRows,
+              moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp.lastCorePerLoopRows,
+              moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp.lastCoreLastLoopRows);
+    OP_LOGD(
+        opName,
+        "GatherOutComputeTilingData is needCoreNum:%ld, activateRows:%ld, perCoreRows:%ld, perCorePerLoopRows:%ld, "
+        "perCoreLastLoopRows:%ld, lastCoreRows:%ld, lastCorePerLoopRows:%ld, lastCoreLastLoopRows:%ld,",
+        moeInitRoutingTilingData.gatherOutComputeParamsOp.needCoreNum,
+        moeInitRoutingTilingData.gatherOutComputeParamsOp.activateRows,
+        moeInitRoutingTilingData.gatherOutComputeParamsOp.perCoreRows,
+        moeInitRoutingTilingData.gatherOutComputeParamsOp.perCorePerLoopRows,
+        moeInitRoutingTilingData.gatherOutComputeParamsOp.perCoreLastLoopRows,
+        moeInitRoutingTilingData.gatherOutComputeParamsOp.lastCoreRows,
+        moeInitRoutingTilingData.gatherOutComputeParamsOp.lastCorePerLoopRows,
+        moeInitRoutingTilingData.gatherOutComputeParamsOp.lastCoreLastLoopRows);
+}
+
+
+
+bool MoeInitRoutingV2TilingBase::GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity,
+  int64_t expertNum, int64_t activateNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag,
+  bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) {
+
+  this->activateNum = activateNum;
+  this->expertCapacity = expertCapacity;
+  this->expertNum = expertNum;
+  this->dropPadMode = dropPadMode;
+  this->expertTokensCountOrCumsumFlag = expertTokensCountOrCumsumFlag;
+  this->expertTokensBeforeCapacityFlag = expertTokensBeforeCapacityFlag;
+  if (dropPadMode == 1) {
+    expertTokensCountOrCumsumFlag = 0;
+  } else {
+    expertTokensBeforeCapacityFlag = false;
+  }
+  moeInitRoutingTilingData.cols = cols;
+  moeInitRoutingTilingData.n = m;
+  moeInitRoutingTilingData.k = topK;
+  moeInitRoutingTilingData.expertCapacity = expertCapacity;
+  moeInitRoutingTilingData.expertNum = expertNum;
+  moeInitRoutingTilingData.dropPadMode = dropPadMode;
+  moeInitRoutingTilingData.expertTokensCountOrCumsumFlag = expertTokensCountOrCumsumFlag;
+  moeInitRoutingTilingData.expertTokensBeforeCapacityFlag = expertTokensBeforeCapacityFlag;
+  totalLength = moeInitRoutingTilingData.n * moeInitRoutingTilingData.k;
+  inuptXDtypeSize_ = inuptXDtypeSize;
+  return true;
+}
+
+bool MoeInitRoutingV2TilingBase::GetPlatformInfo(int64_t aivCoreNum, int64_t ubSizePlatForm) {
+  aivNum = aivCoreNum;
+  aicoreParams_.blockDim = aivCoreNum;
+  aicoreParams_.ubSize = ubSizePlatForm;
+  moeInitRoutingTilingData.coreNum = aivCoreNum;
+  return true;
+}
+
+bool MoeInitRoutingV2TilingBase::GetWorkspaceSize() {
+  size_t sortWorkspaceSize = totalLength * sizeof(float) * NUM_TWO * NUM_THREE;
+  size_t scatterWorkspaceSize = totalLength * sizeof(int32_t) * NUM_TWO;
+  size_t expertTokenFlagSize = aivNum * 2 * sizeof(int32_t);
+  workspaceSize_ = sortWorkspaceSize + scatterWorkspaceSize + expertTokenFlagSize + SIZE_16 * LENGTH_1024 * LENGTH_1024;
+  return true;
+}
+
+
+void MoeInitRoutingV2TilingBase::Tiling4VBSOneCoreCompute(MoeV2VBSComputeTilingData* tilingData) {
+    tilingData->needCoreNum = 1;
+    tilingData->perCoreElements = totalLength;
+    tilingData->perCoreLoops = 1;
+    tilingData->perCorePerLoopElements = tilingData->perCoreElements;
+    tilingData->perCoreLastLoopElements = tilingData->perCoreElements;
+    tilingData->lastCoreElements = tilingData->perCoreElements;
+    tilingData->lastCoreLoops = 1;
+    tilingData->lastCorePerLoopElements = tilingData->perCoreElements;
+    tilingData->lastCoreLastLoopElements = tilingData->perCoreElements;
+}
+
+void MoeInitRoutingV2TilingBase::Tiling4VBSMultiCoreCompute(MoeV2VBSComputeTilingData* tilingData) {
+      //Tiling4VBSMultiCoreCompute
+      int64_t needCoreNum = CeilDiv(totalLength, sortLoopMaxElement);
+      needCoreNum = static_cast<int64_t>(std::pow(4, CeilLog4(needCoreNum)));
+      needCoreNum = std::min(needCoreNum, aivNum);
+      if (needCoreNum > 0) {
+          int64_t perCoreElements = totalLength / needCoreNum;
+          int64_t alineFloorPerCoreElements = perCoreElements - perCoreElements % SORT32_ALIGN_ELEMENT;
+          int64_t lastCoreElement = totalLength - (needCoreNum - 1) * alineFloorPerCoreElements;
+          int64_t alineCeilPerCoreElements = perCoreElements + SORT32_ALIGN_ELEMENT - perCoreElements % SORT32_ALIGN_ELEMENT;
+          if (lastCoreElement > alineCeilPerCoreElements) {
+            perCoreElements = alineCeilPerCoreElements;
+            needCoreNum = CeilDiv(totalLength, perCoreElements);
+          } else {
+            perCoreElements = alineFloorPerCoreElements;
+          }
+          tilingData->needCoreNum = needCoreNum;
+          do {
+              tilingData->perCoreElements = perCoreElements;
+              tilingData->perCoreLoops = CeilDiv(tilingData->perCoreElements, sortLoopMaxElement);  // 每个核处理的loop数
+              tilingData->perCorePerLoopElements = std::min(tilingData->perCoreElements, sortLoopMaxElement);
+              tilingData->perCoreLastLoopElements = tilingData->perCoreElements - (tilingData->perCoreLoops - 1) * tilingData->perCorePerLoopElements;
+              tilingData->lastCoreElements = totalLength - (tilingData->needCoreNum - 1) * tilingData->perCoreElements;
+              tilingData->lastCoreLoops = tilingData->perCoreLoops;
+              int64_t tmp = CeilDiv(tilingData->lastCoreElements, tilingData->lastCoreLoops);
+              int64_t lastCorePerLoopElements =
+                  CeilDiv(CeilDiv(tilingData->lastCoreElements, tilingData->lastCoreLoops), SORT32_ALIGN_ELEMENT) *
+                  SORT32_ALIGN_ELEMENT;
+              tilingData->lastCorePerLoopElements = lastCorePerLoopElements;
+              tilingData->lastCoreLastLoopElements = tilingData-> lastCoreElements - (tilingData->lastCoreLoops - 1) * tilingData->lastCorePerLoopElements;
+              perCoreElements -= SORT32_ALIGN_ELEMENT;
+          } while (tilingData->lastCoreLastLoopElements <= 0 && perCoreElements > 0);
+      }
+}
+
+void MoeInitRoutingV2TilingBase::Tiling4VBSCompute() {
+  auto tilingData = &moeInitRoutingTilingData.vbsComputeParamsOp;
+  tilingData->oneLoopMaxElements = sortLoopMaxElement;
+  if (totalLength <= sortLoopMaxElement) {
+    Tiling4VBSOneCoreCompute(tilingData);
+    return;
+  }
+  Tiling4VBSMultiCoreCompute(tilingData);
+}
+
+void MoeInitRoutingV2TilingBase::Tiling4VMSMiddleCompute() {
+  auto vbsComputeTilingData = &moeInitRoutingTilingData.vbsComputeParamsOp;
+  auto tilingData = &moeInitRoutingTilingData.vmsMiddleComputeParamsOp;
+  if (vbsComputeTilingData->needCoreNum <= MRG_LIST_NUM) {
+      tilingData->needCoreNum = 0;
+  } else {
+      int64_t needCoreNum = CeilDiv(vbsComputeTilingData->needCoreNum, MRG_LIST_NUM);
+      tilingData->needCoreNum = needCoreNum;
+  }
+}
+
+void MoeInitRoutingV2TilingBase::Tiling4SortOutCompute() {
+  auto tilingData = &moeInitRoutingTilingData.sortOutComputeParamsOp;
+  tilingData->oneLoopMaxElements = mrgSortListMaxElement;
+}
+
+
+void MoeInitRoutingV2TilingBase::Tiling4SrcToDstCompute() {
+  auto tilingData = &moeInitRoutingTilingData.srcToDstComputeParamsOp;
+
+  int64_t perLoopMaxRows = (aicoreParams_.ubSize - ASSIST_NUM * sizeof(float) - aivNum * SORT32_ALIGN_ELEMENT) /
+                           (SORT32_ALIGN_ELEMENT * NUM_TWO) / NUM_TWO;
+  int64_t perCoreRows = CeilDiv(totalLength, aivNum);
+  if (perCoreRows <= 0) {
+    tilingData->needCoreNum = 0;
+    return;
+  }
+
+  int64_t needCoreNum = CeilDiv(totalLength, perCoreRows);
+  tilingData->needCoreNum = needCoreNum;
+  int64_t lastCoreNum = totalLength - perCoreRows * (tilingData->needCoreNum - 1);
+  tilingData->perCoreRows = perCoreRows;
+  if (perLoopMaxRows >= tilingData->perCoreRows) {
+      tilingData->perCorePerLoopRows = tilingData->perCoreRows;
+      tilingData->perCoreLastLoopRows = tilingData->perCoreRows;
+  } else {
+      tilingData->perCorePerLoopRows = perLoopMaxRows;
+      tilingData->perCoreLastLoopRows = tilingData->perCoreRows - (CeilDiv(tilingData->perCoreRows, perLoopMaxRows) - 1) * perLoopMaxRows;
+  }
+  tilingData->lastCoreRows = lastCoreNum;
+  if (perLoopMaxRows >= tilingData->lastCoreRows) {
+      tilingData->lastCorePerLoopRows = tilingData->lastCoreRows;
+      tilingData->lastCoreLastLoopRows = tilingData->lastCoreRows;
+  } else {
+      tilingData->lastCorePerLoopRows = perLoopMaxRows;
+      tilingData->lastCoreLastLoopRows = tilingData->lastCoreRows - (CeilDiv(tilingData->lastCoreRows, perLoopMaxRows) - 1) * perLoopMaxRows;
+  }
+}
+
+void MoeInitRoutingV2TilingBase::Tiling4SrcToDstCapacityCompute() {
+  auto tilingData = &moeInitRoutingTilingData.srcToDstCapacityComputeParamsOp;
+  int64_t perCoreRows = CeilDiv(totalLength, aivNum);
+
+  if (perCoreRows <= 0) {
+    tilingData->needCoreNum = 0;
+    return;
+  }
+
+  int64_t needCoreNum = CeilDiv(totalLength, perCoreRows);
+  tilingData->needCoreNum = needCoreNum;
+  int64_t cols = moeInitRoutingTilingData.cols;
+  tilingData->perCoreRows = perCoreRows;
+  int64_t lastCoreRows = totalLength - perCoreRows * (needCoreNum - 1);
+  tilingData->lastCoreRows = lastCoreRows;
+
+
+  int64_t rowSize =
+      (perCoreRows * sizeof(int32_t) * 2 + ONE_BLOCK_BYTE + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+  int64_t colSize = (cols * inuptXDtypeSize_ + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+
+  if (rowSize + colSize < static_cast<int64_t>(aicoreParams_.ubSize)) {
+    tilingData->perCorePerLoopRows = perCoreRows;
+    tilingData->perCoreLastLoopRows = perCoreRows;
+    tilingData->lastCorePerLoopRows = lastCoreRows;
+    tilingData->lastCoreLastLoopRows = lastCoreRows;
+    tilingData->perCoreLoops = 1;
+    tilingData->lastCoreLoops = 1;
+    tilingData->perLoopCols = cols;
+    tilingData->lastLoopCols = cols;
+    tilingData->colLoops = 1;
+
+  } else {
+    int64_t baseMaxCols = MAX_COLS_ONE_LOOP;
+    int64_t baseMaxColsSize = (baseMaxCols * inuptXDtypeSize_ + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+    int64_t basePerLoopMaxRows = (static_cast<int64_t>(aicoreParams_.ubSize) - baseMaxColsSize - ONE_BLOCK_BYTE) /
+                                 static_cast<int64_t>(sizeof(int32_t)) / NUM_TWO / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+    if (cols < MAX_COLS_ONE_LOOP) {
+      basePerLoopMaxRows = (static_cast<int64_t>(aicoreParams_.ubSize) - colSize - ONE_BLOCK_BYTE) /
+                           static_cast<int64_t>(sizeof(int32_t)) / NUM_TWO / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+    } else if (perCoreRows < basePerLoopMaxRows) {
+      baseMaxCols =
+          (static_cast<int64_t>(aicoreParams_.ubSize) - rowSize) / inuptXDtypeSize_ / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+    }
+    tilingData->perLoopCols = (std::min(baseMaxCols, cols));
+    tilingData->lastLoopCols = (GetPerOrLastValue(cols, baseMaxCols));
+    tilingData->colLoops = ((cols + baseMaxCols - 1) / baseMaxCols);
+    tilingData->perCorePerLoopRows = (std::min(perCoreRows, basePerLoopMaxRows));
+    tilingData->perCoreLastLoopRows = (GetPerOrLastValue(perCoreRows, basePerLoopMaxRows));
+    tilingData->perCoreLoops = ((perCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows);
+    tilingData->lastCorePerLoopRows = (std::min(lastCoreRows, basePerLoopMaxRows));
+    tilingData->lastCoreLastLoopRows = (GetPerOrLastValue(lastCoreRows, basePerLoopMaxRows));
+    tilingData->lastCoreLoops = ((lastCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows);
+  }
+}
+
+void MoeInitRoutingV2TilingBase::Tiling4GatherOutCompute()
+{
+    auto tilingData = &moeInitRoutingTilingData.gatherOutComputeParamsOp;
+    tilingData->activateRows = totalLength;
+    if (dropPadMode == 0) {
+        tilingData->activateRows = activateNum;
+    }
+    int64_t perCoreRows = CeilDiv(totalLength, aivNum);
+    if (perCoreRows <= 0 || moeInitRoutingTilingData.cols <= 0) {
+        tilingData->needCoreNum = 0;
+        return;
+    }
+    tilingData->needCoreNum = CeilDiv(totalLength, perCoreRows);
+    int64_t cols = moeInitRoutingTilingData.cols;
+    tilingData->perCoreRows = perCoreRows;
+    int64_t lastCoreRows = totalLength - perCoreRows * (tilingData->needCoreNum - 1);
+    tilingData->lastCoreRows = lastCoreRows;
+
+    int64_t rowSize = (perCoreRows * sizeof(int32_t) + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+    int64_t colSize = (cols * inuptXDtypeSize_ + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+
+    if (rowSize + colSize < static_cast<int64_t>(aicoreParams_.ubSize) / NUM_TWO) {
+        tilingData->perCorePerLoopRows = perCoreRows;
+        tilingData->perCoreLastLoopRows = perCoreRows;
+        tilingData->lastCorePerLoopRows = lastCoreRows;
+        tilingData->lastCoreLastLoopRows = lastCoreRows;
+        tilingData->perCoreLoops = 1;
+        tilingData->lastCoreLoops = 1;
+        tilingData->perLoopCols = cols;
+        tilingData->lastLoopCols = cols;
+        tilingData->colLoops = 1;
+    } else {
+        int64_t baseMaxCols = MAX_COLS_ONE_LOOP;
+        int64_t baseMaxColsSize =
+            (baseMaxCols * inuptXDtypeSize_ + ONE_BLOCK_BYTE - 1) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+        int64_t basePerLoopMaxRows = (static_cast<int64_t>(aicoreParams_.ubSize) / NUM_TWO - baseMaxColsSize) /
+                                     static_cast<int64_t>(sizeof(int32_t)) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+        if (cols < MAX_COLS_ONE_LOOP) {
+            basePerLoopMaxRows = (static_cast<int64_t>(aicoreParams_.ubSize) / NUM_TWO - colSize) /
+                                 static_cast<int64_t>(sizeof(int32_t)) / ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+        } else if (perCoreRows < basePerLoopMaxRows) {
+            baseMaxCols = (static_cast<int64_t>(aicoreParams_.ubSize) / NUM_TWO - rowSize) / inuptXDtypeSize_ /
+                          ONE_BLOCK_BYTE * ONE_BLOCK_BYTE;
+        }
+        tilingData->perLoopCols = (std::min(baseMaxCols, cols));
+        tilingData->lastLoopCols = (GetPerOrLastValue(cols, baseMaxCols));
+        tilingData->colLoops = ((cols + baseMaxCols - 1) / baseMaxCols);
+
+        tilingData->perCorePerLoopRows = (std::min(perCoreRows, basePerLoopMaxRows));
+        tilingData->perCoreLastLoopRows = (GetPerOrLastValue(perCoreRows, basePerLoopMaxRows));
+        tilingData->perCoreLoops = ((perCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows);
+
+        tilingData->lastCorePerLoopRows = (std::min(lastCoreRows, basePerLoopMaxRows));
+        tilingData->lastCoreLastLoopRows = (GetPerOrLastValue(lastCoreRows, basePerLoopMaxRows));
+        tilingData->lastCoreLoops = ((lastCoreRows + basePerLoopMaxRows - 1) / basePerLoopMaxRows);
+    }
+}
+
+
+bool MoeInitRoutingV2TilingBase::IsFullLoad()
+{
+    if (totalLength > sortLoopMaxElement || moeInitRoutingTilingData.cols > MAX_COLS_ONE_LOOP ||
+        this->dropPadMode == 1) {
+        return false;
+    }
+    int64_t sortBufferNum = ONE_CORE_SORT_BUFFER;
+
+    int64_t sortSpace =
+        CeilDiv(this->totalLength, SORT32_ALIGN_ELEMENT) * SORT32_ALIGN_ELEMENT * sizeof(int32_t) * sortBufferNum;
+    int64_t otherSpace =
+        CeilDiv(this->totalLength, SORT32_ALIGN_ELEMENT) * SORT32_ALIGN_ELEMENT * sizeof(int32_t) * NUM_THREE;
+    int64_t expertSpace = CeilDiv(this->expertNum * int64_t(sizeof(int32_t)), ONE_BLOCK_BYTE) * ONE_BLOCK_BYTE;
+    int64_t perCoreXRows = moeInitRoutingTilingData.n / aivNum;
+    int64_t remainder = moeInitRoutingTilingData.n % aivNum;
+    // NUM_TWO is Max xRows need add 2 becauseof the left and right row may be another row.
+    perCoreXRows = remainder <= 1 ? perCoreXRows + 1 : perCoreXRows + NUM_TWO;
+    int64_t gatherSpace =
+        CeilDiv(moeInitRoutingTilingData.cols * inuptXDtypeSize_, ONE_BLOCK_BYTE) * ONE_BLOCK_BYTE * perCoreXRows;
+    int64_t remainUbAfterSort = aicoreParams_.ubSize - sortSpace - otherSpace - expertSpace - gatherSpace;
+    return remainUbAfterSort > 0;
+}
+
+}

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_common.h

@@ -0,0 +1,201 @@
+/**
+ * This program is free software, you can redistribute it and/or modify.
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_common.h
+ * \brief
+ */
+#ifndef MOE_V2_COMMON_H
+#define MOE_V2_COMMON_H
+
+#include "kernel_operator.h"
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+using namespace optiling;
+constexpr int64_t SPLIT_N = 0;
+constexpr int64_t SPLIT_K = 1;
+constexpr float MIN_FP32 = -3.4e38f;
+#if __CCE_AICORE__ == 200
+constexpr int64_t ONE_REPEAT_SORT_NUM = 16;
+constexpr int64_t REGIONP_ROPOSAL_KV_RATIO = 4; // 8 / 2
+constexpr int64_t SYNC_LEN = 8 * 8 * 32;
+#else
+constexpr int64_t ONE_REPEAT_SORT_NUM = 32;
+#endif
+constexpr int64_t BLOCK_BYTES = 32;
+constexpr int64_t INT32_ONE_BLOCK_NUM = 8;
+constexpr int64_t FP32_ONE_REPEAT_NUM = 64;
+
+constexpr int64_t ASSIST_NUM = 256;
+constexpr int64_t ASSIST_INDEX_NUM = 32;
+
+constexpr int64_t MERGE_LIST_TWO = 2;
+constexpr int64_t MERGE_LIST_THREE = 3;
+constexpr int64_t MERGE_LIST_FOUR = 4;
+
+constexpr int64_t MERGE_LIST_IDX_TWO = 2;
+constexpr int64_t MERGE_LIST_IDX_THREE = 3;
+
+constexpr int64_t MAX_EXPERT_NUM = 5120;
+constexpr int64_t DROPLESS_MODE = 0;
+constexpr int64_t DROP_PAD_MODE = 1;
+constexpr int64_t EXERPT_TOKENS_COUNT = 2;
+constexpr int64_t EXERPT_TOKENS_CUMSUM = 1;
+constexpr int64_t EXERPT_TOKENS_NONE = 0;
+constexpr int64_t EXERPT_TOKENS_BEFORE_CAPACITY = 1;
+
+const __gm__ int32_t assist[256] = {
+    0,  0, 0, 0, 0, 0, 0, 0, 1,  0, 0, 0, 0, 0, 0, 0, 2,  0, 0, 0, 0, 0, 0, 0, 3,  0, 0, 0, 0, 0, 0, 0,
+    4,  0, 0, 0, 0, 0, 0, 0, 5,  0, 0, 0, 0, 0, 0, 0, 6,  0, 0, 0, 0, 0, 0, 0, 7,  0, 0, 0, 0, 0, 0, 0,
+    8,  0, 0, 0, 0, 0, 0, 0, 9,  0, 0, 0, 0, 0, 0, 0, 10, 0, 0, 0, 0, 0, 0, 0, 11, 0, 0, 0, 0, 0, 0, 0,
+    12, 0, 0, 0, 0, 0, 0, 0, 13, 0, 0, 0, 0, 0, 0, 0, 14, 0, 0, 0, 0, 0, 0, 0, 15, 0, 0, 0, 0, 0, 0, 0,
+    16, 0, 0, 0, 0, 0, 0, 0, 17, 0, 0, 0, 0, 0, 0, 0, 18, 0, 0, 0, 0, 0, 0, 0, 19, 0, 0, 0, 0, 0, 0, 0,
+    20, 0, 0, 0, 0, 0, 0, 0, 21, 0, 0, 0, 0, 0, 0, 0, 22, 0, 0, 0, 0, 0, 0, 0, 23, 0, 0, 0, 0, 0, 0, 0,
+    24, 0, 0, 0, 0, 0, 0, 0, 25, 0, 0, 0, 0, 0, 0, 0, 26, 0, 0, 0, 0, 0, 0, 0, 27, 0, 0, 0, 0, 0, 0, 0,
+    28, 0, 0, 0, 0, 0, 0, 0, 29, 0, 0, 0, 0, 0, 0, 0, 30, 0, 0, 0, 0, 0, 0, 0, 31, 0, 0, 0, 0, 0, 0, 0};
+
+__aicore__ inline int64_t Ceil(int64_t a, int64_t b)
+{
+    if (b == 0) {
+        return 0;
+    }
+    return (a + b - 1) / b;
+}
+
+__aicore__ inline int64_t Align(int64_t elementNum, int64_t bytes)
+{
+    if (bytes == 0) {
+        return 0;
+    }
+    return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES / bytes;
+}
+
+__aicore__ inline int64_t AlignBytes(int64_t elementNum, int64_t bytes)
+{
+    return (elementNum * bytes + BLOCK_BYTES - 1) / BLOCK_BYTES * BLOCK_BYTES;
+}
+
+template <typename T>
+__aicore__ inline T Min(T a, T b)
+{
+    return a > b ? b : a;
+}
+
+template <typename T>
+__aicore__ inline T Max(T a, T b)
+{
+    return a < b ? b : a;
+}
+
+template <HardEvent event>
+__aicore__ inline void SetWaitFlag(HardEvent evt)
+{
+    event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(evt));
+    SetFlag<event>(eventId);
+    WaitFlag<event>(eventId);
+}
+
+template <typename T>
+__aicore__ inline void DataCopyPadCustom(LocalTensor<T> inLocal, GlobalTensor<T> srcGm, DataCopyExtParams tokenCopyParams, DataCopyPadExtParams<T> padParams)
+{
+#if __CCE_AICORE__ == 220
+  DataCopyPad(inLocal, srcGm, tokenCopyParams, padParams);
+#else
+  int64_t elem =  tokenCopyParams.blockLen / sizeof(T);
+  int64_t numPerBlock = BLOCK_BYTES / sizeof(T);
+  int64_t alignElem = AlignUp(elem, numPerBlock);
+
+  if (likely(alignElem == elem)) {
+
+    DataCopyParams copyParams = {tokenCopyParams.blockCount, static_cast<uint16_t>(alignElem / numPerBlock) , 0, 0};
+    DataCopy(inLocal, srcGm, copyParams);
+  } else {
+    DataCopyParams copyParams = {1, static_cast<uint16_t>(alignElem / numPerBlock) , 0, 0};
+    for (uint32_t i = 0; i < tokenCopyParams.blockCount; i++) {
+        DataCopy(inLocal[i * alignElem], srcGm[i * elem], copyParams);
+    }
+  }
+#endif
+}
+
+template <typename T, bool needBack = false, bool isAtomic = false>
+__aicore__ inline void DataCopyCustom(GlobalTensor<T> dstGm, LocalTensor<T> inLocal, int64_t blockCount, int64_t blockLen)
+{
+    int64_t elem =  blockLen / sizeof(T);
+    int64_t numPerBlock = sizeof(T) == 0 ? 1 : BLOCK_BYTES / sizeof(T);
+    int64_t alignElem = AlignUp(elem, numPerBlock);
+
+    if (likely(alignElem == elem)) {
+        DataCopyParams copyParams = {static_cast<uint16_t>(blockCount), static_cast<uint16_t>(alignElem / numPerBlock) , 0, 0};
+        DataCopy(dstGm, inLocal, copyParams);
+    } else {
+        if (blockCount == 1) {
+            if constexpr (needBack) {
+                int64_t elemAlignDown = numPerBlock == 0 ? 0 :  elem / numPerBlock * numPerBlock;
+                if (elemAlignDown != 0) {
+                    DataCopyParams copyParams = {static_cast<uint16_t>(blockCount), static_cast<uint16_t>(elemAlignDown / numPerBlock) , 0, 0};
+                    DataCopy(dstGm, inLocal, copyParams);
+                    SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+                    SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
+
+                    for (uint32_t i = 0; i < numPerBlock; i++) {
+                        inLocal.SetValue(alignElem-1-i, inLocal.GetValue(elem - 1 - i));
+                    }
+                    SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+
+                    DataCopyParams copyParamslast = {1, 1, 0, 0};
+
+                    DataCopy(dstGm[elem-numPerBlock], inLocal[elemAlignDown], copyParamslast);
+                } else {
+                    T tmp[BLOCK_BYTES];
+                    SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+                    SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
+                    for (uint32_t i = 0; i < elem; i++) {
+                        tmp[i] = inLocal.GetValue(elem - 1 - i);
+                    }
+                    DataCopyParams copyParamslast = {1, 1, 0, 0};
+                    SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::S_MTE2);
+                    SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
+                    DataCopy(inLocal, dstGm[elem-numPerBlock], copyParamslast);
+                    SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+                    for (uint32_t i = 0; i < elem; i++) {
+                        inLocal.SetValue(numPerBlock-1-i, tmp[i]);
+                    }
+                    SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+                    DataCopy(dstGm[elem-numPerBlock], inLocal, copyParamslast);
+                }
+
+            } else if constexpr (isAtomic) {
+                SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+                SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
+                for (uint32_t i = 0; i < alignElem - elem; i++) {
+                    inLocal.SetValue(alignElem-1-i, T(0));
+                }
+                SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+
+                DataCopyParams copyParams = {static_cast<uint16_t>(blockCount), static_cast<uint16_t>(alignElem / numPerBlock) , 0, 0};
+                DataCopy(dstGm, inLocal, copyParams);
+            } else {
+                DataCopyParams copyParams = {static_cast<uint16_t>(blockCount), static_cast<uint16_t>(alignElem / numPerBlock) , 0, 0};
+                DataCopy(dstGm, inLocal, copyParams);
+            }
+        } else {
+            DataCopyParams copyParams = {1, static_cast<uint16_t>(alignElem / numPerBlock) , 0, 0};
+            for (uint32_t i = 0; i < blockCount; i++) {
+                DataCopy(dstGm[i * elem], inLocal[i * alignElem], copyParams);
+                PipeBarrier<PIPE_MTE3>();
+            }
+        }
+    }
+}
+
+} // namespace MoeInitRoutingV2
+#endif // MOE_V2_COMMON_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_expert_token_out.h

@@ -0,0 +1,380 @@
+/**
+ * This program is free software, you can redistribute it and/or modify.
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_expert_token_out.h
+ * \brief
+ */
+#ifndef MOE_V2_EXPERT_TOKEN_OUT_H
+#define MOE_V2_EXPERT_TOKEN_OUT_H
+
+#include "moe_v2_common.h"
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+using namespace optiling;
+
+constexpr int64_t EXPERT_ID_VALUE_NUM = 2;
+
+class MoeV2ExpertTokenOut {
+public:
+    __aicore__ inline MoeV2ExpertTokenOut(){};
+    template <typename TilingData>
+    __aicore__ inline void Init(GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity,
+                                GM_ADDR expandedRowIdx, GM_ADDR workspace, const TilingData *tilingData, TPipe *tPipe);
+    __aicore__ inline void Process();
+
+private:
+    __aicore__ inline void CopyIn(int64_t progress);
+    __aicore__ inline void Compute(int64_t progress);
+    __aicore__ inline void SyncAll();
+    __aicore__ inline void InitLocal();
+    __aicore__ inline void GetExpertTokenCount(int32_t curExpertId);
+    __aicore__ inline void CopyOutTokenGm();
+    __aicore__ inline void CopyOutExpertTokensCumsum(bool isTail);
+    __aicore__ inline void CopyOutExpertTokensCount(bool isTail);
+
+private:
+    TPipe *pipe;
+    TQue<QuePosition::VECIN, 1> copyInQueue;
+    TQue<QuePosition::VECIN, 1> expertTokenIdxCopyInQueue;
+    TQue<QuePosition::VECOUT, 1> expertTokenIdxCopyOutQueue;
+
+    GlobalTensor<int32_t> expertTokensCountOrCumsumGm;
+    GlobalTensor<int32_t> expertTokensBeforeCapacityGm;
+    GlobalTensor<int32_t> expandedExpertIdxGm;
+    GlobalTensor<int32_t> expertIdxValueGm;
+    GlobalTensor<int32_t> expandedRowIdxGm;
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    GlobalTensor<int32_t> syncTmpSpaceGm_;
+    TBuf<TPosition::VECCALC> workBuffer;
+#endif
+    LocalTensor<int32_t> expertTokenIdxOutLocal;
+
+    const MoeV2GatherOutComputeTilingData *srcToDstTilingData;
+
+    int64_t coreNum;
+    int64_t blockIdx;
+    int64_t totalLength;
+    int64_t currentLoopRows;
+    int64_t coreRows;
+    int64_t perLoopRows;
+    int64_t lastLoopRows;
+    int64_t expertNum;
+    int64_t expertNumUbAlign;
+    int64_t dropPadMode = 0;
+    int64_t expertTokensCountOrCumsumFlag = 0;
+    int64_t expertTokensBeforeCapacityFlag = 0;
+
+    int64_t tokenCount = 0;
+    int64_t expertIdx = 0;
+    int32_t lastExpertId = -1;
+    int32_t firstExpertId = -1;
+
+    int32_t expertTokenValue = 0;
+};
+
+__aicore__ inline void MoeV2ExpertTokenOut::InitLocal()
+{
+    LocalTensor<int32_t> tokenIdxLocal = expertTokenIdxCopyOutQueue.AllocTensor<int32_t>();
+    Duplicate<int32_t>(tokenIdxLocal, 0, this->expertNumUbAlign);
+    expertTokenIdxCopyOutQueue.EnQue<int32_t>(tokenIdxLocal);
+
+    // expandedRowIdx initialized to -1, which is used in the src_to_dst_with_capacity step.
+    // use this step SyncAll to synchronize every core data
+    if (this->dropPadMode == 0) {
+        return;
+    }
+    LocalTensor<int32_t> outLocal = copyInQueue.AllocTensor<int32_t>();
+    int64_t loops = (coreRows + perLoopRows - 1) / perLoopRows;
+    Duplicate<int32_t>(outLocal, -1, perLoopRows);
+    SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+    for (int64_t loop = 0; loop < loops; loop++) {
+        int64_t copyLength = perLoopRows;
+        if (loop == loops - 1) {
+            copyLength = lastLoopRows;
+        }
+        DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(copyLength * sizeof(int32_t)), 0,
+                                     0, 0};
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+        DataCopyCustom<int32_t, false, false>(
+            expandedRowIdxGm[this->blockIdx * this->srcToDstTilingData->perCoreRows + loop * perLoopRows], outLocal,
+            copyParams.blockCount, copyParams.blockLen);
+#else
+        DataCopyPad(expandedRowIdxGm[this->blockIdx * this->srcToDstTilingData->perCoreRows + loop * perLoopRows],
+                    outLocal, copyParams);
+#endif
+    }
+    SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
+    copyInQueue.FreeTensor(outLocal);
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::CopyIn(int64_t progress)
+{
+    LocalTensor<int32_t> inLocal = copyInQueue.AllocTensor<int32_t>();
+    DataCopy(inLocal, expandedExpertIdxGm[progress * perLoopRows], Align(currentLoopRows, sizeof(int32_t)));
+    copyInQueue.EnQue<int32_t>(inLocal);
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::GetExpertTokenCount(int32_t curExpertId)
+{
+    this->tokenCount++;
+    if (this->lastExpertId < curExpertId) {
+        this->expertTokenIdxOutLocal.SetValue(this->expertIdx, this->tokenCount - 1);
+        this->tokenCount = 1;
+        this->expertIdx += (curExpertId - this->lastExpertId);
+        while (curExpertId - this->firstExpertId + 1 > this->expertNumUbAlign) {
+            SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+            CopyOutExpertTokensCumsum(false);
+            CopyOutExpertTokensCount(false);
+            SetWaitFlag<HardEvent::MTE3_V>(HardEvent::MTE3_V);
+            Duplicate<int32_t>(this->expertTokenIdxOutLocal, 0, this->expertNumUbAlign);
+            SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
+            this->firstExpertId += this->expertNumUbAlign;
+            this->expertIdx = curExpertId - this->firstExpertId;
+        }
+        this->lastExpertId = curExpertId;
+    }
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::Compute(int64_t progress)
+{
+    LocalTensor<int32_t> inLocal = copyInQueue.DeQue<int32_t>();
+    SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+    if (this->lastExpertId == -1) {
+        this->lastExpertId = inLocal.GetValue(0);
+        this->firstExpertId = this->lastExpertId;
+    }
+    for (int64_t i = 0; i < currentLoopRows; i++) {
+        int32_t expertId = inLocal.GetValue(i);
+        GetExpertTokenCount(expertId);
+    }
+    this->expertTokenIdxOutLocal.SetValue(this->expertIdx, this->tokenCount);
+    copyInQueue.FreeTensor(inLocal);
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::CopyOutExpertTokensCumsum(bool isTail)
+{
+    if (this->dropPadMode != DROPLESS_MODE || expertTokensCountOrCumsumFlag != EXERPT_TOKENS_CUMSUM) {
+        return;
+    }
+#ifdef __CCE_KT_TEST__
+    // CPU孪生调试无法使用多核同步，可能导致index为未初始化的脏数据，因此需要特殊处理
+    if (this->firstExpertId > expertTokensCountOrCumsumGm.GetSize()) {
+        return;
+    }
+#endif
+    int64_t copyLength = isTail ? this->lastExpertId - this->firstExpertId + 1 : this->expertNumUbAlign;
+    int64_t end = this->expertNum - this->firstExpertId;
+    for (int64_t i = 0; i < copyLength; i++) {
+        this->expertTokenValue += this->expertTokenIdxOutLocal.GetValue(i);
+        this->expertTokenIdxOutLocal.SetValue(i, this->expertTokenValue);
+    }
+    // if the remaining UB is sufficient, use the UB space to copy
+    // otherwise, copy the calculated data first, and then copy the last tokenValue to remaining expert position
+    if (isTail && end <= this->expertNumUbAlign) {
+        int64_t startAlign = Min(Align(copyLength, sizeof(int32_t)), end);
+        for (int64_t i = copyLength; i < startAlign; i++) {
+            this->expertTokenIdxOutLocal.SetValue(i, this->expertTokenValue);
+        }
+        if (startAlign < end) {
+            Duplicate<int32_t>(this->expertTokenIdxOutLocal[startAlign], this->expertTokenValue, end - startAlign);
+        }
+        copyLength = end;
+        SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+    }
+    DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(copyLength * sizeof(int32_t)), 0, 0,
+                                 0};
+    SetAtomicAdd<int32_t>();
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyCustom<int32_t,false,false>(
+        expertTokensCountOrCumsumGm[this->firstExpertId], this->expertTokenIdxOutLocal,
+        copyParams.blockCount, copyParams.blockLen);
+#else
+    DataCopyPad(expertTokensCountOrCumsumGm[this->firstExpertId], this->expertTokenIdxOutLocal, copyParams);
+#endif
+    SetAtomicNone();
+    if (isTail && end > this->expertNumUbAlign) {
+        int64_t remainderLength = end - copyLength;
+        SetWaitFlag<HardEvent::MTE3_V>(HardEvent::MTE3_V);
+        Duplicate<int32_t>(this->expertTokenIdxOutLocal, this->expertTokenValue, this->expertNumUbAlign);
+        SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+        int64_t loopTimes = remainderLength / this->expertNumUbAlign + 1;
+        for (int64_t i = 0; i < loopTimes; i++) {
+            copyLength = i == loopTimes - 1 ? remainderLength - this->expertNumUbAlign * i : this->expertNumUbAlign;
+            DataCopyExtParams params{static_cast<uint16_t>(1), static_cast<uint32_t>(copyLength * sizeof(int32_t)), 0,
+                                     0, 0};
+            SetAtomicAdd<int32_t>();
+
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+            DataCopyCustom<int32_t,false,false>(
+                expertTokensCountOrCumsumGm[this->lastExpertId + 1 + this->expertNumUbAlign * i],
+                this->expertTokenIdxOutLocal,
+                params.blockCount, params.blockLen);
+
+#else
+            DataCopyPad(expertTokensCountOrCumsumGm[this->lastExpertId + 1 + this->expertNumUbAlign * i],
+                        this->expertTokenIdxOutLocal, params);
+#endif
+            SetAtomicNone();
+        }
+    }
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::CopyOutExpertTokensCount(bool isTail)
+{
+    int64_t copyLength = isTail ? this->lastExpertId - this->firstExpertId + 1 : this->expertNumUbAlign;
+    DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(copyLength * sizeof(int32_t)), 0, 0,
+                                 0};
+#ifdef __CCE_KT_TEST__
+    // CPU孪生调试不进行输出拷贝
+    return;
+#endif
+    SetAtomicAdd<int32_t>();
+    if (this->dropPadMode == DROP_PAD_MODE && expertTokensBeforeCapacityFlag > EXERPT_TOKENS_NONE) {
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+        DataCopyCustom<int32_t,false,false>(expertTokensBeforeCapacityGm[this->firstExpertId],
+            this->expertTokenIdxOutLocal, copyParams.blockCount, copyParams.blockLen);
+ #else
+        DataCopyPad(expertTokensBeforeCapacityGm[this->firstExpertId], this->expertTokenIdxOutLocal, copyParams);
+#endif
+    }
+    if (this->dropPadMode == DROPLESS_MODE && expertTokensCountOrCumsumFlag == EXERPT_TOKENS_COUNT) {
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+        DataCopyCustom<int32_t,false,false>(expertTokensCountOrCumsumGm[this->firstExpertId], this->expertTokenIdxOutLocal,
+            copyParams.blockCount, copyParams.blockLen);
+#else
+        DataCopyPad(expertTokensCountOrCumsumGm[this->firstExpertId], this->expertTokenIdxOutLocal, copyParams);
+#endif
+    }
+    SetAtomicNone();
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::CopyOutTokenGm()
+{
+    if (this->dropPadMode == DROPLESS_MODE) {
+        SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+        CopyOutExpertTokensCumsum(true);
+        CopyOutExpertTokensCount(true);
+        return;
+    }
+    this->expertTokenIdxOutLocal.SetValue(this->expertNumUbAlign, this->lastExpertId);
+    this->expertTokenIdxOutLocal.SetValue(this->expertNumUbAlign + 1, this->tokenCount);
+    DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(EXPERT_ID_VALUE_NUM * sizeof(int32_t)),
+                                 0, 0, 0};
+    SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyCustom<int32_t,false,false>(expertIdxValueGm[this->blockIdx * BLOCK_BYTES / sizeof(int32_t)],
+             this->expertTokenIdxOutLocal[this->expertNumUbAlign], copyParams.blockCount, copyParams.blockLen);
+#else
+    DataCopyPad(expertIdxValueGm[this->blockIdx * EXPERT_ID_VALUE_NUM],
+                this->expertTokenIdxOutLocal[this->expertNumUbAlign], copyParams);
+#endif
+    CopyOutExpertTokensCount(true);
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::SyncAll()
+{
+    if (coreNum == 1) {
+        return;
+    }
+#ifndef __CCE_KT_TEST__
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    LocalTensor<int32_t> syncLocal = workBuffer.Get<int32_t>();
+    AscendC::SyncAll(syncTmpSpaceGm_, syncLocal, GetBlockNum());
+#else
+    AscendC::SyncAll();
+#endif
+#endif
+}
+
+template <typename TilingData>
+__aicore__ inline void MoeV2ExpertTokenOut::Init(GM_ADDR expertTokensCountOrCumsum, GM_ADDR expertTokensBeforeCapacity,
+                                                 GM_ADDR expandedRowIdx, GM_ADDR workspace,
+                                                 const TilingData *tilingData, TPipe *tPipe)
+{
+    int64_t blockNum = GetBlockNum() * 2;
+    this->pipe = tPipe;
+    this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+    // this->blockIdx = GetBlockIdx();
+
+    this->coreNum = tilingData->coreNum;
+    this->totalLength = tilingData->n * tilingData->k;
+    this->srcToDstTilingData = &(tilingData->srcToDstComputeParamsOp);
+    this->expertNum = tilingData->expertNum;
+    this->dropPadMode = tilingData->dropPadMode;
+    this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag;
+    this->expertTokensBeforeCapacityFlag = tilingData->expertTokensBeforeCapacityFlag;
+
+    if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) {
+        this->coreRows = this->srcToDstTilingData->lastCoreRows;
+        this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows;
+        this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows;
+    } else {
+        this->coreRows = this->srcToDstTilingData->perCoreRows;
+        this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows;
+        this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows;
+    }
+
+    expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx, Align(this->totalLength, sizeof(int32_t)));
+    if (this->dropPadMode == DROPLESS_MODE && this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) {
+        expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t *)expertTokensCountOrCumsum, this->expertNum);
+    }
+    if (this->dropPadMode == DROP_PAD_MODE && this->expertTokensBeforeCapacityFlag == EXERPT_TOKENS_BEFORE_CAPACITY) {
+        expertTokensBeforeCapacityGm.SetGlobalBuffer((__gm__ int32_t *)expertTokensBeforeCapacity, this->expertNum);
+    }
+
+    expandedExpertIdxGm.SetGlobalBuffer((__gm__ int32_t *)workspace +
+                                            this->blockIdx * this->srcToDstTilingData->perCoreRows,
+                                        Align(this->coreRows, sizeof(int32_t)));
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    syncTmpSpaceGm_.SetGlobalBuffer((__gm__ int32_t *)workspace +
+                                        Align(this->totalLength, sizeof(int32_t)) * EXERPT_TOKENS_COUNT +
+                                        this->coreNum * BLOCK_BYTES / sizeof(int32_t),
+                                    INT32_ONE_BLOCK_NUM * GetBlockNum() * 2 * BLOCK_BYTES);
+    pipe->InitBuffer(workBuffer, INT32_ONE_BLOCK_NUM * GetBlockNum() * 2 * BLOCK_BYTES);
+    LocalTensor<int32_t> syncLocal = workBuffer.Get<int32_t>();
+    Duplicate<int32_t>(syncLocal, 0, SYNC_LEN);
+    SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+    DataCopy(syncTmpSpaceGm_, syncLocal, SYNC_LEN);
+    expertIdxValueGm.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(this->totalLength, sizeof(int32_t)) * 2,
+                                     this->coreNum * BLOCK_BYTES / sizeof(int32_t));
+#else
+    expertIdxValueGm.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(this->totalLength, sizeof(int32_t)) * 2,
+                                     this->coreNum * 2);
+#endif
+    this->expertNumUbAlign = Min(Align(this->expertNum, sizeof(int32_t)), MAX_EXPERT_NUM);
+    pipe->InitBuffer(copyInQueue, 1, this->perLoopRows * BLOCK_BYTES);
+    pipe->InitBuffer(expertTokenIdxCopyInQueue, 1, this->expertNumUbAlign * sizeof(int32_t));
+    pipe->InitBuffer(expertTokenIdxCopyOutQueue, 1, (this->expertNumUbAlign + EXPERT_ID_VALUE_NUM) * sizeof(int32_t));
+}
+
+__aicore__ inline void MoeV2ExpertTokenOut::Process()
+{
+    if (this->blockIdx < this->srcToDstTilingData->needCoreNum) {
+        int64_t loops = (coreRows + perLoopRows - 1) / perLoopRows;
+        currentLoopRows = perLoopRows;
+        InitLocal();
+        this->expertTokenIdxOutLocal = expertTokenIdxCopyOutQueue.DeQue<int32_t>();
+        for (int64_t loop = 0; loop < loops - 1; loop++) {
+            CopyIn(loop);
+            Compute(loop);
+        }
+        currentLoopRows = lastLoopRows;
+        CopyIn(loops - 1);
+        Compute(loops - 1);
+        CopyOutTokenGm();
+        expertTokenIdxCopyOutQueue.FreeTensor(this->expertTokenIdxOutLocal);
+    }
+    this->SyncAll();
+}
+
+} // namespace MoeInitRoutingV2
+#endif // MOE_V2_EXPERT_TOKEN_OUT_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_gather_out.h

@@ -0,0 +1,198 @@
+/**
+ * This program is free software, you can redistribute it and/or modify.
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_gather_out.h
+ * \brief
+ */
+#ifndef MOE_V2_GATHER_OUT_H
+#define MOE_V2_GATHER_OUT_H
+
+#include "moe_v2_common.h"
+#include "kernel_operator.h"
+using namespace optiling;
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+
+constexpr int64_t BUFFER_NUM = 2;
+
+template <typename T>
+class MoeV2GatherOut {
+public:
+    __aicore__ inline MoeV2GatherOut(){};
+    __aicore__ inline void Init(GM_ADDR inputX, GM_ADDR expandedRowIdx, GM_ADDR expandedX, GM_ADDR workspace,
+                                const MoeInitRoutingV2TilingData *tilingData, TPipe *tPipe);
+    __aicore__ inline void Process();
+
+private:
+    __aicore__ inline void CopyInIndices(int64_t progress);
+    __aicore__ inline void CopyOut(int64_t progress);
+
+private:
+    TPipe *pipe;
+    TQueBind<QuePosition::VECIN, QuePosition::VECOUT, BUFFER_NUM> inputActivationsCopyInQueue;
+    TQue<QuePosition::VECIN, BUFFER_NUM> expandDstToSrcRowCopyInQueue;
+
+    GlobalTensor<T> inputXGm;
+    GlobalTensor<T> expandedXGm;
+    GlobalTensor<int32_t> expandedRowIdxGm;
+
+    const MoeV2GatherOutComputeTilingData *gatherOutTilingData;
+
+    int64_t needCoreNum;
+    int64_t blockIdx;
+    int64_t cols;
+    int64_t n;
+    int64_t k;
+    int64_t activateRows;
+    int64_t currentLoopRows;
+    int64_t coreRows;
+    int64_t perLoopRows;
+    int64_t lastLoopRows;
+    int64_t rowLoops;
+    int64_t colsTileLength;
+    int64_t perLoopCols;
+    int64_t lastLoopCols;
+    int64_t colLoops;
+    int64_t dropPadMode;
+
+    int64_t indicesOffset;
+    int64_t inputOffset;
+    int64_t outOffset;
+};
+
+template <typename T>
+__aicore__ inline void MoeV2GatherOut<T>::CopyInIndices(int64_t progress)
+{
+    this->indicesOffset = progress * this->perLoopRows;
+    LocalTensor<int32_t> indicesLocal = expandDstToSrcRowCopyInQueue.AllocTensor<int32_t>();
+    DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(this->currentLoopRows * sizeof(int32_t)), 0, 0, 0};
+    DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyPadCustom(indicesLocal, expandedRowIdxGm[indicesOffset], dataCopyParams, dataCopyPadParams);
+#else
+    DataCopyPad(indicesLocal, expandedRowIdxGm[indicesOffset], dataCopyParams, dataCopyPadParams);
+#endif
+    expandDstToSrcRowCopyInQueue.EnQue<int32_t>(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherOut<T>::CopyOut(int64_t progress)
+{
+    LocalTensor<int32_t> indicesLocal = expandDstToSrcRowCopyInQueue.DeQue<int32_t>();
+    SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+    colsTileLength = this->perLoopCols;
+    for (int64_t colsLoop = 0; colsLoop < this->colLoops; colsLoop++) {
+        int64_t initialRow = this->gatherOutTilingData->perCoreRows * this->blockIdx + this->perLoopRows * progress;
+        int64_t curLoopRow = 0;
+        if (colsLoop == this->colLoops - 1) {
+            colsTileLength = this->lastLoopCols;
+        }
+        int64_t currentLoopStartRow = initialRow / this->k;
+        int64_t currentLoopLastRow = (initialRow + this->currentLoopRows - 1) / this->k;
+        for (int64_t row = currentLoopStartRow; row <= currentLoopLastRow; row++) {
+            LocalTensor<T> inLocal = inputActivationsCopyInQueue.AllocTensor<T>();
+            // input row position
+            inputOffset = row * this->cols + colsLoop * this->perLoopCols;
+            DataCopyExtParams dataCopyParams{1, static_cast<uint32_t>(this->colsTileLength * sizeof(T)), 0, 0, 0};
+            DataCopyPadExtParams<T> dataCopyPadParams{false, 0, 0, 0};
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+            DataCopyPadCustom(inLocal, inputXGm[inputOffset], dataCopyParams, dataCopyPadParams);
+#else
+            DataCopyPad(inLocal, inputXGm[inputOffset], dataCopyParams, dataCopyPadParams);
+#endif
+            SetWaitFlag<HardEvent::MTE2_MTE3>(HardEvent::MTE2_MTE3);
+
+            DataCopyExtParams intriParams{1, static_cast<uint32_t>(this->colsTileLength * sizeof(T)), 0, 0, 0};
+            while (curLoopRow < this->currentLoopRows && initialRow / this->k == row) {
+                int32_t outIndex = indicesLocal.GetValue(curLoopRow);
+                curLoopRow++;
+                initialRow++;
+                if (outIndex == -1 || (this->dropPadMode == DROPLESS_MODE && outIndex >= this->activateRows)) {
+                    continue;
+                }
+                outOffset = outIndex * cols + colsLoop * this->perLoopCols;
+#ifdef __CCE_KT_TEST__
+                // CPU孪生调试无法使用多核同步，可能导致index为未初始化的脏数据，因此需要特殊处理
+                if (outOffset > expandedXGm.GetSize()) {
+                    continue;
+                }
+#endif
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+                DataCopyCustom(expandedXGm[outOffset], inLocal, intriParams.blockCount, intriParams.blockLen);
+#else
+                DataCopyPad(expandedXGm[outOffset], inLocal, intriParams);
+#endif
+            }
+            inputActivationsCopyInQueue.FreeTensor(inLocal);
+        }
+    }
+    expandDstToSrcRowCopyInQueue.FreeTensor(indicesLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherOut<T>::Init(GM_ADDR inputX, GM_ADDR expandedRowIdx, GM_ADDR expandedX,
+                                               GM_ADDR workspace, const MoeInitRoutingV2TilingData *tilingData,
+                                               TPipe *tPipe)
+{
+    this->pipe = tPipe;
+    // this->blockIdx = GetBlockIdx();
+    this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+    this->gatherOutTilingData = &(tilingData->gatherOutComputeParamsOp);
+
+    this->needCoreNum = this->gatherOutTilingData->needCoreNum;
+    this->activateRows = this->gatherOutTilingData->activateRows;
+    this->cols = tilingData->cols;
+    this->n = tilingData->n;
+    this->k = tilingData->k;
+    this->dropPadMode = tilingData->dropPadMode;
+
+    if (this->blockIdx == this->gatherOutTilingData->needCoreNum - 1) {
+        this->coreRows = this->gatherOutTilingData->lastCoreRows;
+        this->perLoopRows = this->gatherOutTilingData->lastCorePerLoopRows;
+        this->lastLoopRows = this->gatherOutTilingData->lastCoreLastLoopRows;
+        this->rowLoops = this->gatherOutTilingData->lastCoreLoops;
+    } else {
+        this->coreRows = this->gatherOutTilingData->perCoreRows;
+        this->perLoopRows = this->gatherOutTilingData->perCorePerLoopRows;
+        this->lastLoopRows = this->gatherOutTilingData->perCoreLastLoopRows;
+        this->rowLoops = this->gatherOutTilingData->perCoreLoops;
+    }
+    this->perLoopCols = this->gatherOutTilingData->perLoopCols;
+    this->lastLoopCols = this->gatherOutTilingData->lastLoopCols;
+    this->colLoops = this->gatherOutTilingData->colLoops;
+
+    inputXGm.SetGlobalBuffer((__gm__ T *)inputX, this->coreRows * this->cols);
+    expandedXGm.SetGlobalBuffer((__gm__ T *)expandedX, tilingData->n * tilingData->k * this->cols);
+    expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx +
+                                         this->blockIdx * this->gatherOutTilingData->perCoreRows,
+                                     Align(this->coreRows, sizeof(int32_t)));
+
+    pipe->InitBuffer(inputActivationsCopyInQueue, BUFFER_NUM, AlignBytes(this->perLoopCols, sizeof(T)));
+    pipe->InitBuffer(expandDstToSrcRowCopyInQueue, BUFFER_NUM, AlignBytes(this->perLoopRows, sizeof(int32_t)));
+}
+
+template <typename T>
+__aicore__ inline void MoeV2GatherOut<T>::Process()
+{
+    if (this->blockIdx < this->needCoreNum) {
+        currentLoopRows = perLoopRows;
+        for (int64_t loop = 0; loop < this->rowLoops; loop++) {
+            if (loop == this->rowLoops - 1) {
+                currentLoopRows = lastLoopRows;
+            }
+            CopyInIndices(loop);
+            CopyOut(loop);
+        }
+    }
+}
+} // namespace MoeInitRoutingV2
+#endif // MOE_V2_GATHER_OUT_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_init_routing_fullload.h

@@ -0,0 +1,388 @@
+/**
+ * This program is free software, you can redistribute it and/or modify.
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/* !
+ * \file moe_v2_init_routing_fullload.h
+ * \brief
+ */
+#ifndef MOE_V2_FULL_LOAD_H
+#define MOE_V2_FULL_LOAD_H
+
+#include "moe_v2_mrgsort.h"
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+using namespace optiling;
+
+template <typename T>
+class MoeV2FullLoad : public MoeV2SortBase {
+public:
+    __aicore__ inline MoeV2FullLoad(){};
+    __aicore__ inline void Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, GM_ADDR expandedRowIdx,
+                                GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace,
+                                const MoeInitRoutingV2TilingData *tilingData, TPipe *tPipe);
+    __aicore__ inline void Process();
+
+private:
+    __aicore__ inline void CopyIn();
+    __aicore__ inline void SortCompute();
+    __aicore__ inline void CopyOutIdx();
+    __aicore__ inline void CopyOutEmpty();
+    __aicore__ inline void CopyOutX();
+    __aicore__ inline void ComputeExpertTokenCountOrCumsum();
+
+private:
+    int64_t sortNum_;
+    const MoeV2GatherOutComputeTilingData *gatherOutTilingData_;
+    int64_t blockIdx_;
+    int64_t needCoreNum_;
+    int64_t sortNeedCoreNum_;
+    int64_t coreRows_;
+    int64_t perCoreRows_;
+    int64_t k_;
+    int64_t n_;
+    int64_t cols_;
+    int64_t activateRows_;
+    int64_t expertNum;
+    int64_t expertCapacity;
+
+    TQue<QuePosition::VECIN, 1> xCopyInQueue_;
+    TQue<QuePosition::VECOUT, 1> expandedRowIdxCopyOutQueue_;
+    TQue<QuePosition::VECOUT, 1> expandedExpertIdxCopyOutQueue_;
+    TQue<QuePosition::VECOUT, 1> expandDstToSrcRowQueue_;
+    TQue<QuePosition::VECOUT, 1> expertTokensCopyOutQueue_;
+
+    GlobalTensor<T> xGm_;
+    GlobalTensor<int32_t> expertIdxGm_;
+
+    GlobalTensor<T> expandedXGm_;
+    GlobalTensor<int32_t> expandedRowIdxGm_;
+    GlobalTensor<int32_t> expandedExpertIdxGm_;
+    GlobalTensor<int32_t> expertTokensCountOrCumsumGm;
+    GlobalTensor<int32_t> expertTokensBeforeCapacityGm;
+
+    int64_t expertTokensCountOrCumsumFlag = 0;
+    int64_t expertTokensBeforeCapacityFlag = 0;
+    int64_t dropPadMode = 0;
+};
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::CopyIn()
+{
+    LocalTensor<int32_t> inLocal = sortDataCopyInQueue.AllocTensor<int32_t>();
+    DataCopyExtParams dataCopyParams{static_cast<uint16_t>(1),
+                                     static_cast<uint32_t>(this->totalLength * sizeof(int32_t)), 0, 0, 0};
+    DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyPadCustom(inLocal[0], expertIdxGm_, dataCopyParams, dataCopyPadParams);
+#else
+    DataCopyPad(inLocal[0], expertIdxGm_, dataCopyParams, dataCopyPadParams);
+#endif
+    ArithProgression<int32_t>(inLocal[this->sortNum_], 0, 1, this->totalLength);
+    sortDataCopyInQueue.EnQue(inLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::SortCompute()
+{
+    LocalTensor<int32_t> inLocal = sortDataCopyInQueue.DeQue<int32_t>();
+    LocalTensor<int32_t> expertIdxLocal = inLocal[0];
+    LocalTensor<float> expertIdxLocalFp32 = expertIdxLocal.ReinterpretCast<float>();
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    Cast(expertIdxLocalFp32, expertIdxLocal, RoundMode::CAST_NONE, this->totalLength);
+#else
+    Cast(expertIdxLocalFp32, expertIdxLocal, RoundMode::CAST_ROUND, this->totalLength);
+#endif
+    PipeBarrier<PIPE_V>();
+    Muls(expertIdxLocalFp32, expertIdxLocalFp32, (float)-1, this->totalLength);
+    PipeBarrier<PIPE_V>();
+    int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM;
+    if (duplicateNum > 0) {
+        int duplicateIndex = this->totalLength - duplicateNum;
+        uint64_t mask0 = UINT64_MAX;
+        mask0 = mask0 << duplicateNum;
+        mask0 = mask0 & (UINT64_MAX >> (FP32_ONE_REPEAT_NUM - ONE_REPEAT_SORT_NUM));
+        uint64_t mask[2] = {mask0, 0};
+        Duplicate(expertIdxLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+        PipeBarrier<PIPE_V>();
+    }
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    LocalTensor<float> concatLocal;
+#else
+    LocalTensor<float> concatLocal = expertIdxLocalFp32;
+#endif
+    LocalTensor<float> tempTensor = tempBuffer.Get<float>(GetSortLen<float>(this->sortNum_));
+    Concat(concatLocal, expertIdxLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+    PipeBarrier<PIPE_V>();
+    LocalTensor<uint32_t> rowIdxLocal = inLocal[this->sortNum_].template ReinterpretCast<uint32_t>();
+    LocalTensor<float> sortedLocal = sortedBuffer.Get<float>(GetSortLen<float>(this->sortNum_));
+    Sort<float, true>(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+    PipeBarrier<PIPE_V>();
+    LocalTensor<float> expandedExpertIdxLocal = expandedExpertIdxCopyOutQueue_.AllocTensor<float>();
+    LocalTensor<uint32_t> expandDstToSrcRowLocal = expandDstToSrcRowQueue_.AllocTensor<uint32_t>();
+    LocalTensor<float> expandDstToSrcRowLocalFp32 = expandDstToSrcRowLocal.ReinterpretCast<float>();
+    Extract(expandedExpertIdxLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+    PipeBarrier<PIPE_V>();
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    Cast(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocal.ReinterpretCast<int32_t>(), RoundMode::CAST_NONE,
+         this->totalLength);
+#else
+    Cast(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocal.ReinterpretCast<int32_t>(), RoundMode::CAST_ROUND,
+         this->totalLength);
+#endif
+    PipeBarrier<PIPE_V>();
+    Muls(expandedExpertIdxLocal, expandedExpertIdxLocal, (float)-1, this->totalLength);
+    PipeBarrier<PIPE_V>();
+    LocalTensor<int32_t> expandedExpertIdxLocalInt32;
+    expandedExpertIdxLocalInt32 = expandedExpertIdxLocal.ReinterpretCast<int32_t>();
+    Cast(expandedExpertIdxLocalInt32, expandedExpertIdxLocal, RoundMode::CAST_ROUND, this->totalLength);
+    PipeBarrier<PIPE_V>();
+    expandedExpertIdxCopyOutQueue_.EnQue<int32_t>(expandedExpertIdxLocalInt32);
+
+    LocalTensor<uint32_t> expandedRowIdx = expandedRowIdxCopyOutQueue_.AllocTensor<uint32_t>();
+    LocalTensor<uint32_t> expandedRowIdxU32 = expandedRowIdx.ReinterpretCast<uint32_t>();
+    Muls(expandDstToSrcRowLocalFp32, expandDstToSrcRowLocalFp32, (float)-1, this->totalLength);
+    PipeBarrier<PIPE_V>();
+    ArithProgression<int32_t>(inLocal[this->sortNum_], 0, 1, this->totalLength);
+    PipeBarrier<PIPE_V>();
+    if (duplicateNum > 0) {
+        int duplicateIndex = this->totalLength - duplicateNum;
+        uint64_t mask0 = UINT64_MAX;
+        mask0 = mask0 << duplicateNum;
+        mask0 = mask0 & (UINT64_MAX >> (FP32_ONE_REPEAT_NUM - ONE_REPEAT_SORT_NUM));
+        uint64_t mask[2] = {mask0, 0};
+        Duplicate(expandDstToSrcRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+        PipeBarrier<PIPE_V>();
+    }
+    Concat(concatLocal, expandDstToSrcRowLocalFp32, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+    PipeBarrier<PIPE_V>();
+    Sort<float, true>(sortedLocal, concatLocal, rowIdxLocal, tempTensor, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+    PipeBarrier<PIPE_V>();
+    Extract(tempTensor, expandedRowIdxU32, sortedLocal, this->sortNum_ / ONE_REPEAT_SORT_NUM);
+    PipeBarrier<PIPE_V>();
+    expandedRowIdxCopyOutQueue_.EnQue<uint32_t>(expandedRowIdx);
+    sortDataCopyInQueue.FreeTensor(inLocal);
+
+    expandDstToSrcRowQueue_.FreeTensor(expandDstToSrcRowLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::CopyOutIdx()
+{
+    LocalTensor<int32_t> expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue<int32_t>();
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyExtParams intriParams;
+#else
+    DataCopyParams intriParams;
+#endif
+    intriParams.blockCount = 1;
+    intriParams.blockLen = this->totalLength * sizeof(int32_t);
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyCustom(expandedRowIdxGm_, expandedRowIdx, intriParams.blockCount, intriParams.blockLen);
+#else
+    DataCopyPad(expandedRowIdxGm_, expandedRowIdx, intriParams);
+#endif
+    expandedRowIdxCopyOutQueue_.EnQue(expandedRowIdx);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::ComputeExpertTokenCountOrCumsum()
+{
+    LocalTensor<int32_t> expandedExpertIdx = expandedExpertIdxCopyOutQueue_.DeQue<int32_t>();
+    LocalTensor<int32_t> expertTokensCount = expertTokensCopyOutQueue_.AllocTensor<int32_t>();
+
+    int64_t expertNumAlign = Align(this->expertNum, sizeof(int32_t));
+    Duplicate(expertTokensCount, 0, expertNumAlign);
+    SetWaitFlag<HardEvent::V_S>(HardEvent::V_S);
+
+    int32_t lastExpertId = expandedExpertIdx.GetValue(0);
+#ifdef __CCE_KT_TEST__
+    // CPU孪生调试无法使用多核同步，可能导致lastExpertId为未初始化的脏数据，因此需要特殊处理
+    if (lastExpertId > expertTokensCount.GetSize()) {
+        return;
+    }
+#endif
+    int64_t tokenCount = 0;
+    int64_t lastExpertCount = 0;
+    for (int64_t i = 0; i < this->totalLength; i++) {
+        int32_t curExpertId = expandedExpertIdx.GetValue(i);
+        tokenCount++;
+        while (lastExpertId < curExpertId) {
+            expertTokensCount.SetValue(lastExpertId, tokenCount - 1);
+            if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_COUNT) {
+                tokenCount = 1;
+            }
+            lastExpertId++;
+        }
+    }
+    expertTokensCount.SetValue(lastExpertId, tokenCount);
+    if (this->expertTokensCountOrCumsumFlag == EXERPT_TOKENS_CUMSUM) {
+        lastExpertId++;
+        while (lastExpertId < this->expertNum) {
+            expertTokensCount.SetValue(lastExpertId, tokenCount);
+            lastExpertId++;
+        }
+    }
+    DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(this->expertNum * sizeof(int32_t)), 0,
+                                 0, 0};
+    if (this->expertTokensCountOrCumsumFlag > 0) {
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+        DataCopyCustom(expertTokensCountOrCumsumGm, expertTokensCount, copyParams.blockCount, copyParams.blockLen);
+#else
+        DataCopyPad(expertTokensCountOrCumsumGm, expertTokensCount, copyParams);
+#endif
+    }
+    expertTokensCopyOutQueue_.FreeTensor(expertTokensCount);
+    expandedExpertIdxCopyOutQueue_.FreeTensor(expandedExpertIdx);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::CopyOutX()
+{
+    LocalTensor<T> xLocal = xCopyInQueue_.AllocTensor<T>();
+    LocalTensor<int32_t> expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue<int32_t>();
+    DataCopyParams intriParams;
+    intriParams.blockCount = 1;
+    intriParams.blockLen = this->cols_ * sizeof(T);
+    int64_t inFactor = Align(this->cols_, sizeof(T));
+    int64_t curRowsStart = this->blockIdx_ * this->perCoreRows_;
+    int64_t startXRow = curRowsStart / this->k_;
+    int64_t endXRow = (curRowsStart + this->coreRows_ - 1) / this->k_;
+
+    DataCopyExtParams dataXCopyParams{static_cast<uint16_t>(endXRow - startXRow + 1),
+                                      static_cast<uint32_t>(this->cols_ * sizeof(T)), 0, 0, 0};
+    DataCopyPadExtParams<T> dataXCopyPadParams{false, 0, 0, 0};
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyPadCustom(xLocal, xGm_[startXRow * this->cols_], dataXCopyParams, dataXCopyPadParams);
+#else
+    DataCopyPad(xLocal, xGm_[startXRow * this->cols_], dataXCopyParams, dataXCopyPadParams);
+#endif
+    SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+
+    int64_t k = 0;
+    for (int64_t i = startXRow; i <= endXRow; i++) {
+        for (; k < this->perCoreRows_ && curRowsStart / this->k_ == i; curRowsStart++, k++) {
+            int32_t outIndex = expandedRowIdx.GetValue(curRowsStart);
+            if (outIndex < this->activateRows_) {
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+                DataCopyCustom(expandedXGm_[outIndex * this->cols_], xLocal[(i - startXRow) * inFactor],
+                               intriParams.blockCount, intriParams.blockLen);
+#else
+                DataCopyPad(expandedXGm_[outIndex * this->cols_], xLocal[(i - startXRow) * inFactor], intriParams);
+#endif
+            }
+        }
+    }
+    expandedRowIdxCopyOutQueue_.FreeTensor(expandedRowIdx);
+    xCopyInQueue_.FreeTensor(xLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::CopyOutEmpty()
+{
+    LocalTensor<int32_t> outLocal = expandedExpertIdxCopyOutQueue_.DeQue<int32_t>();
+    expandedExpertIdxCopyOutQueue_.FreeTensor(outLocal);
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::Init(GM_ADDR x, GM_ADDR expertIdx, GM_ADDR expandedX, GM_ADDR expandedRowIdx,
+                                              GM_ADDR expertTokensCountOrCumsum, GM_ADDR workspace,
+                                              const MoeInitRoutingV2TilingData *tilingData, TPipe *tPipe)
+{
+    this->gatherOutTilingData_ = &(tilingData->gatherOutComputeParamsOp);
+    //this->blockIdx_ = GetBlockIdx();
+    this->blockIdx_ = get_block_idx() + get_subblockid() * get_block_num();
+    this->n_ = tilingData->n;
+    this->k_ = tilingData->k;
+    this->cols_ = tilingData->cols;
+    this->sortNeedCoreNum_ = this->gatherOutTilingData_->needCoreNum;
+    this->needCoreNum_ = this->gatherOutTilingData_->needCoreNum;
+    if (needCoreNum_ == 0) {
+        this->sortNeedCoreNum_ = tilingData->vbsComputeParamsOp.needCoreNum;
+    }
+    this->perCoreRows_ = this->gatherOutTilingData_->perCoreRows;
+    this->activateRows_ = this->gatherOutTilingData_->activateRows;
+    if (this->blockIdx_ == this->gatherOutTilingData_->needCoreNum - 1) {
+        this->coreRows_ = this->gatherOutTilingData_->lastCoreRows;
+    } else {
+        this->coreRows_ = this->gatherOutTilingData_->perCoreRows;
+    }
+    this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag;
+    this->dropPadMode = tilingData->dropPadMode;
+    this->expertNum = tilingData->expertNum;
+
+    this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t));
+    this->sortNum_ = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
+    this->totalLength = tilingData->n * tilingData->k;
+    this->pipe = tPipe;
+
+    xGm_.SetGlobalBuffer((__gm__ T *)x);
+    expertIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expertIdx, this->tileLength);
+
+    expandedXGm_.SetGlobalBuffer((__gm__ T *)expandedX);
+    expandedRowIdxGm_.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx, this->tileLength);
+    if (this->expertTokensCountOrCumsumFlag > 0) {
+        // dropless
+        expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t *)expertTokensCountOrCumsum,
+                                                    Align(this->expertNum, sizeof(int32_t)));
+    }
+
+    int64_t kvFactor = 2;
+    int64_t buffSize = this->sortNum_ * sizeof(int32_t);
+
+    int64_t curRowsStart = this->blockIdx_ * this->perCoreRows_;
+    int64_t startXRow = curRowsStart / this->k_;
+    int64_t endXRow = (curRowsStart + this->coreRows_ - 1) / this->k_;
+    pipe->InitBuffer(xCopyInQueue_, bufferNum, AlignBytes(this->cols_, sizeof(T)) * (endXRow - startXRow + 1));
+    pipe->InitBuffer(expandedRowIdxCopyOutQueue_, bufferNum, buffSize);
+    pipe->InitBuffer(expandedExpertIdxCopyOutQueue_, bufferNum, buffSize);
+    pipe->InitBuffer(expertTokensCopyOutQueue_, bufferNum, AlignBytes(this->expertNum, sizeof(int32_t)));
+    pipe->InitBuffer(expandDstToSrcRowQueue_, bufferNum, buffSize);
+    pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize * kvFactor);
+
+// sort310p
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    pipe->InitBuffer(tempBuffer, buffSize * REGIONP_ROPOSAL_KV_RATIO * kvFactor);
+    pipe->InitBuffer(sortedBuffer, buffSize * REGIONP_ROPOSAL_KV_RATIO * kvFactor);
+#else
+    pipe->InitBuffer(tempBuffer, buffSize * kvFactor);
+    pipe->InitBuffer(sortedBuffer, buffSize * kvFactor);
+#endif
+}
+
+template <typename T>
+__aicore__ inline void MoeV2FullLoad<T>::Process()
+{
+    if (this->blockIdx_ < this->sortNeedCoreNum_) {
+        CopyIn();
+        SortCompute();
+        if (this->blockIdx_ == 0) {
+            CopyOutIdx();
+        }
+        if (this->blockIdx_ == this->sortNeedCoreNum_ - 1 && this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) {
+            ComputeExpertTokenCountOrCumsum();
+        } else {
+            CopyOutEmpty();
+        }
+    }
+    if (this->needCoreNum_ != 0) {
+        if (this->blockIdx_ < this->needCoreNum_) {
+            CopyOutX();
+        }
+
+    } else {
+        if (this->blockIdx_ < this->sortNeedCoreNum_) {
+            LocalTensor<int32_t> expandedRowIdx = expandedRowIdxCopyOutQueue_.DeQue<int32_t>();
+            expandedRowIdxCopyOutQueue_.FreeTensor(expandedRowIdx);
+        }
+    }
+}
+} // namespace MoeInitRoutingV2
+#endif // MOE_V2_FULL_LOAD_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_mrgsort.h

@@ -0,0 +1,211 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_mrgsort.h
+ * \brief
+ */
+#ifndef MOE_V2_MRGSORT_H
+#define MOE_V2_MRGSORT_H
+
+#include "moe_v2_common.h"
+#include "kernel_operator.h"
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+using namespace optiling;
+
+struct MoeV2MrgsortParam {
+    int64_t perListElements;
+    int64_t lastListElements;
+    int64_t oneLoopMaxElements;
+};
+
+class MoeV2Mrgsort {
+public:
+    __aicore__ inline MoeV2Mrgsort(){};
+    __aicore__ inline void Init(MoeV2MrgsortParam *param);
+    __aicore__ inline void Process();
+    __aicore__ inline void SetInput(GlobalTensor<float> &gmInput, LocalTensor<float> &ubInput);
+    __aicore__ inline void SetOutput(GlobalTensor<float> &gmOutput, LocalTensor<float> &ubOutput);
+
+private:
+    __aicore__ inline void CopyIn();
+    __aicore__ inline void UpdateMrgParam();
+    __aicore__ inline void MrgsortCompute();
+    __aicore__ inline void UpdateSortInfo();
+    __aicore__ inline void CopyOut();
+    __aicore__ inline void ClearCache();
+
+private:
+    MoeV2MrgsortParam *param = nullptr;
+
+    GlobalTensor<float> gmInputs[4];
+    GlobalTensor<float> gmOutput;
+
+    LocalTensor<float> ubInputs[4];
+    LocalTensor<float> ubOutput;
+
+    int64_t listNum{0};
+    int64_t remainListNum{0};
+    int64_t outOffset{0};
+    int64_t offsets[4];
+    int64_t listRemainElements[4];
+    int64_t lengths[4];
+    int64_t allRemainElements{0};
+    int64_t curLoopSortedNum{0};
+
+    // for MrgSort
+    uint16_t validBitTail{0};
+    uint16_t elementCountListTail[4];
+    uint32_t listSortedNums[4];
+    LocalTensor<float> tmpUbInputs[4];
+};
+
+__aicore__ inline void MoeV2Mrgsort::ClearCache()
+{
+    this->listNum = 0;
+    this->allRemainElements = 0;
+    this->outOffset = 0;
+}
+
+__aicore__ inline void MoeV2Mrgsort::SetInput(GlobalTensor<float> &gmInput, LocalTensor<float> &ubInput)
+{
+    this->gmInputs[listNum] = gmInput;
+    this->ubInputs[listNum] = ubInput;
+    this->listNum += 1;
+}
+
+__aicore__ inline void MoeV2Mrgsort::SetOutput(GlobalTensor<float> &gmOutput, LocalTensor<float> &ubOutput)
+{
+    this->gmOutput = gmOutput;
+    this->ubOutput = ubOutput;
+}
+
+__aicore__ inline void MoeV2Mrgsort::UpdateMrgParam()
+{
+    if (this->remainListNum == MERGE_LIST_TWO) {
+        elementCountListTail[MERGE_LIST_IDX_TWO] = 0;
+        elementCountListTail[MERGE_LIST_IDX_THREE] = 0;
+        validBitTail = 0b0011;
+    } else if (this->remainListNum == MERGE_LIST_THREE) {
+        elementCountListTail[MERGE_LIST_IDX_THREE] = 0;
+        validBitTail = 0b0111;
+    } else if (this->remainListNum == MERGE_LIST_FOUR) {
+        validBitTail = 0b1111;
+    } else {
+        validBitTail = 0b0001;
+    }
+}
+
+__aicore__ inline void MoeV2Mrgsort::CopyIn()
+{
+    this->remainListNum = 0;
+    SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
+    for (int64_t i = 0, j = 0; i < listNum; i++) {
+        lengths[i] = Min(param->oneLoopMaxElements, listRemainElements[i]);
+        if (lengths[i] > 0) {
+            DataCopy(this->ubInputs[i], this->gmInputs[i][offsets[i]],
+                     Align(GetSortLen<float>(lengths[i]), sizeof(float)));
+            tmpUbInputs[j] = this->ubInputs[i];
+            elementCountListTail[j] = lengths[i];
+            this->remainListNum += 1;
+            j++;
+        }
+    }
+}
+
+__aicore__ inline void MoeV2Mrgsort::MrgsortCompute()
+{
+    SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
+    if (this->remainListNum == MERGE_LIST_TWO) {
+        MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]);
+        MrgSort<float, true>(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+    } else if (this->remainListNum == MERGE_LIST_THREE) {
+        MrgSortSrcList sortListTail =
+            MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], tmpUbInputs[0]);
+        MrgSort<float, true>(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+    } else if (this->remainListNum == MERGE_LIST_FOUR) {
+        MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO],
+                                                     tmpUbInputs[MERGE_LIST_IDX_THREE]);
+        MrgSort<float, true>(this->ubOutput, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+    } else {
+        DataCopy(this->ubOutput, this->tmpUbInputs[0],
+                 Align(GetSortLen<float>(elementCountListTail[0]), sizeof(float)));
+        listSortedNums[0] = elementCountListTail[0];
+    }
+}
+
+__aicore__ inline void MoeV2Mrgsort::UpdateSortInfo()
+{
+    curLoopSortedNum = 0;
+    for (int64_t i = 0, j = 0; i < listNum; i++) {
+        if (lengths[i] > 0) {
+            // update remain size
+            listRemainElements[i] -= listSortedNums[j];
+            allRemainElements -= listSortedNums[j];
+            // update offset
+            offsets[i] += GetSortOffset<float>(listSortedNums[j]);
+            // update current loop sorted nums
+            curLoopSortedNum += listSortedNums[j];
+            j += 1;
+        }
+    }
+}
+
+__aicore__ inline void MoeV2Mrgsort::CopyOut()
+{
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyExtParams intriParams;
+#else
+    DataCopyParams intriParams;
+#endif
+    intriParams.blockCount = 1;
+    intriParams.blockLen = GetSortLen<float>(curLoopSortedNum) * sizeof(float);
+    SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyCustom<float, false, false>(this->gmOutput[outOffset], this->ubOutput, intriParams.blockCount,
+                                        intriParams.blockLen);
+#else
+    DataCopyPad(this->gmOutput[outOffset], this->ubOutput, intriParams);
+#endif
+    outOffset += GetSortLen<float>(curLoopSortedNum);
+}
+
+__aicore__ inline void MoeV2Mrgsort::Init(MoeV2MrgsortParam *param)
+{
+    this->param = param;
+    this->remainListNum = listNum;
+
+    for (int64_t i = 0; i < listNum; i++) {
+        offsets[i] = GetSortOffset<float>(param->perListElements * i);
+        if (i == listNum - 1) {
+            listRemainElements[i] = param->lastListElements;
+        } else {
+            listRemainElements[i] = param->perListElements;
+        }
+        allRemainElements += listRemainElements[i];
+    }
+}
+
+__aicore__ inline void MoeV2Mrgsort::Process()
+{
+    for (; allRemainElements > 0;) {
+        CopyIn();
+        UpdateMrgParam();
+        MrgsortCompute();
+        UpdateSortInfo();
+        CopyOut();
+    }
+
+    ClearCache();
+}
+} // namespace MoeInitRoutingV2
+#endif // MOE_V2_MRGSORT_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_mrgsort_out.h

@@ -0,0 +1,245 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_mrgsort_out.h
+ * \brief
+ */
+#ifndef MOE_V2_MRGSORT_OUT_H
+#define MOE_V2_MRGSORT_OUT_H
+
+#include "moe_v2_mrgsort.h"
+#include "kernel_operator.h"
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+using namespace optiling;
+
+class MoeV2MrgsortOut {
+public:
+    __aicore__ inline MoeV2MrgsortOut(){};
+    __aicore__ inline void Init(MoeV2MrgsortParam *param, TPipe *tPipe);
+    __aicore__ inline void Process();
+    __aicore__ inline void SetInput(GlobalTensor<float> &gmInput, LocalTensor<float> &ubInput);
+    __aicore__ inline void SetOutput(GlobalTensor<int32_t> &gmOutput1, GlobalTensor<int32_t> &gmOutput2,
+                                     LocalTensor<float> &ubOutput1, LocalTensor<float> &ubOutput2);
+    __aicore__ inline void SetBuffer(LocalTensor<float> &tempBuffer);
+
+private:
+    __aicore__ inline void CopyIn();
+    __aicore__ inline void UpdateMrgParam();
+    __aicore__ inline void MrgsortCompute();
+    __aicore__ inline void UpdateSortInfo();
+    __aicore__ inline void Extract();
+    __aicore__ inline void CopyOut();
+    __aicore__ inline void ClearCache();
+
+private:
+    MoeV2MrgsortParam *param = nullptr;
+
+    GlobalTensor<float> gmInputs[4];
+    GlobalTensor<int32_t> gmOutput1;
+    GlobalTensor<int32_t> gmOutput2;
+
+    LocalTensor<float> ubInputs[4];
+    LocalTensor<float> tempBuffer;
+
+    // for extract
+    LocalTensor<float> ubOutput1;
+    LocalTensor<uint32_t> ubOutput2;
+
+    // for copy out
+    LocalTensor<int32_t> ubOutputInt1;
+    LocalTensor<int32_t> ubOutputInt2;
+
+    int64_t listNum{0};
+    int64_t remainListNum{0};
+    int64_t outOffset{0};
+    int64_t offsets[4];
+    int64_t listRemainElements[4];
+    int64_t lengths[4];
+    int64_t allRemainElements{0};
+    int64_t curLoopSortedNum{0};
+
+    // for MrgSort
+    uint16_t validBitTail;
+    uint16_t elementCountListTail[4];
+    uint32_t listSortedNums[4];
+    LocalTensor<float> tmpUbInputs[4];
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    bool needCopyOut = true;
+#endif
+};
+
+__aicore__ inline void MoeV2MrgsortOut::ClearCache()
+{
+    this->listNum = 0;
+    this->allRemainElements = 0;
+    this->outOffset = 0;
+}
+
+__aicore__ inline void MoeV2MrgsortOut::SetInput(GlobalTensor<float> &gmInput, LocalTensor<float> &ubInput)
+{
+    this->gmInputs[listNum] = gmInput;
+    this->ubInputs[listNum] = ubInput;
+    this->listNum += 1;
+}
+
+__aicore__ inline void MoeV2MrgsortOut::SetOutput(GlobalTensor<int32_t> &gmOutput1, GlobalTensor<int32_t> &gmOutput2,
+                                                  LocalTensor<float> &ubOutput1, LocalTensor<float> &ubOutput2)
+{
+    this->gmOutput1 = gmOutput1;
+    this->ubOutput1 = ubOutput1;
+    this->ubOutputInt1 = ubOutput1.ReinterpretCast<int32_t>();
+
+    this->gmOutput2 = gmOutput2;
+    this->ubOutput2 = ubOutput2.ReinterpretCast<uint32_t>();
+    this->ubOutputInt2 = ubOutput2.ReinterpretCast<int32_t>();
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    needCopyOut = !(this->gmOutput1.GetPhyAddr() == this->gmOutput2.GetPhyAddr());
+#endif
+}
+
+__aicore__ inline void MoeV2MrgsortOut::SetBuffer(LocalTensor<float> &tempBuffer)
+{
+    this->tempBuffer = tempBuffer;
+}
+
+__aicore__ inline void MoeV2MrgsortOut::UpdateMrgParam()
+{
+    if (this->remainListNum == MERGE_LIST_TWO) {
+        elementCountListTail[MERGE_LIST_IDX_TWO] = 0;
+        elementCountListTail[MERGE_LIST_IDX_THREE] = 0;
+        validBitTail = 0b0011;
+    } else if (this->remainListNum == MERGE_LIST_THREE) {
+        elementCountListTail[MERGE_LIST_IDX_THREE] = 0;
+        validBitTail = 0b0111;
+    } else if (this->remainListNum == MERGE_LIST_FOUR) {
+        validBitTail = 0b1111;
+    } else {
+        validBitTail = 0b0001;
+    }
+}
+
+__aicore__ inline void MoeV2MrgsortOut::CopyIn()
+{
+    this->remainListNum = 0;
+    SetWaitFlag<HardEvent::MTE3_MTE2>(HardEvent::MTE3_MTE2);
+    for (int64_t i = 0, j = 0; i < listNum; i++) {
+        lengths[i] = Min(param->oneLoopMaxElements, listRemainElements[i]);
+        if (lengths[i] > 0) {
+            DataCopy(this->ubInputs[i], this->gmInputs[i][offsets[i]],
+                     Align(GetSortLen<float>(lengths[i]), sizeof(float)));
+            tmpUbInputs[j] = this->ubInputs[i];
+            elementCountListTail[j] = lengths[i];
+            this->remainListNum += 1;
+            j++;
+        }
+    }
+}
+
+__aicore__ inline void MoeV2MrgsortOut::MrgsortCompute()
+{
+    SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
+    if (this->remainListNum == MERGE_LIST_TWO) {
+        MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], tmpUbInputs[0]);
+        MrgSort<float, true>(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+    } else if (this->remainListNum == MERGE_LIST_THREE) {
+        MrgSortSrcList sortListTail =
+            MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO], tmpUbInputs[0]);
+        MrgSort<float, true>(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+    } else if (this->remainListNum == MERGE_LIST_FOUR) {
+        MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[MERGE_LIST_IDX_TWO],
+                                                     tmpUbInputs[MERGE_LIST_IDX_THREE]);
+        MrgSort<float, true>(this->tempBuffer, sortListTail, elementCountListTail, listSortedNums, validBitTail, 1);
+    } else {
+        DataCopy(this->tempBuffer, this->tmpUbInputs[0],
+                 Align(GetSortLen<float>(elementCountListTail[0]), sizeof(float)));
+        listSortedNums[0] = elementCountListTail[0];
+    }
+}
+
+__aicore__ inline void MoeV2MrgsortOut::UpdateSortInfo()
+{
+    curLoopSortedNum = 0;
+    for (int64_t i = 0, j = 0; i < listNum; i++) {
+        if (lengths[i] > 0) {
+            // update remain size
+            listRemainElements[i] -= listSortedNums[j];
+            allRemainElements -= listSortedNums[j];
+            // update offset
+            offsets[i] += GetSortOffset<float>(listSortedNums[j]);
+            // update current loop sorted nums
+            curLoopSortedNum += listSortedNums[j];
+            j += 1;
+        }
+    }
+}
+
+__aicore__ inline void MoeV2MrgsortOut::Extract()
+{
+    AscendC::Extract(this->ubOutput1, this->ubOutput2, this->tempBuffer, Ceil(curLoopSortedNum, ONE_REPEAT_SORT_NUM));
+    PipeBarrier<PIPE_V>();
+    Muls(this->ubOutput1, this->ubOutput1, (float)-1, Align(curLoopSortedNum, sizeof(float)));
+    PipeBarrier<PIPE_V>();
+    Cast(this->ubOutputInt1, this->ubOutput1, RoundMode::CAST_ROUND, Align(curLoopSortedNum, sizeof(float)));
+}
+
+__aicore__ inline void MoeV2MrgsortOut::CopyOut()
+{
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyExtParams intriParams;
+#else
+    DataCopyParams intriParams;
+#endif
+    intriParams.blockCount = 1;
+    intriParams.blockLen = curLoopSortedNum * sizeof(int32_t);
+    SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    if (needCopyOut) {
+        DataCopyCustom<int32_t,true,false>(this->gmOutput1[outOffset], this->ubOutputInt1, intriParams.blockCount, intriParams.blockLen);
+    }
+    DataCopyCustom<int32_t,true,false>(this->gmOutput2[outOffset], this->ubOutputInt2, intriParams.blockCount, intriParams.blockLen);
+#else
+    DataCopyPad(this->gmOutput2[outOffset], this->ubOutputInt2, intriParams);
+    DataCopyPad(this->gmOutput1[outOffset], this->ubOutputInt1, intriParams);
+#endif
+    outOffset += curLoopSortedNum;
+}
+
+__aicore__ inline void MoeV2MrgsortOut::Init(MoeV2MrgsortParam *param, TPipe *tPipe)
+{
+    this->param = param;
+    this->allRemainElements = 0;
+    for (int64_t i = 0; i < listNum; i++) {
+        offsets[i] = GetSortOffset<float>(param->perListElements * i);
+        if (i == listNum - 1) {
+            listRemainElements[i] = param->lastListElements;
+        } else {
+            listRemainElements[i] = param->perListElements;
+        }
+        allRemainElements += listRemainElements[i];
+    }
+}
+
+__aicore__ inline void MoeV2MrgsortOut::Process()
+{
+    for (; allRemainElements > 0;) {
+        CopyIn();
+        UpdateMrgParam();
+        MrgsortCompute();
+        UpdateSortInfo();
+        Extract();
+        CopyOut();
+    }
+    ClearCache();
+}
+} // namespace MoeInitRoutingV2
+#endif // MOE_V2_MRGSORT_OUT_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_sort_base.h

@@ -0,0 +1,74 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_sort_base.h
+ * \brief
+ */
+#ifndef MOE_V2_SORT_BASE_H
+#define MOE_V2_SORT_BASE_H
+
+#include "kernel_operator.h"
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+using namespace optiling;
+
+class MoeV2SortBase {
+public:
+    __aicore__ inline MoeV2SortBase(){};
+
+protected:
+    __aicore__ inline void SyncAll();
+
+protected:
+    TPipe *pipe;
+    TQue<QuePosition::VECIN, 1> sortDataCopyInQueue;
+    TQue<QuePosition::VECOUT, 1> sortDataCopyOutQueue;
+    TBuf<TPosition::VECCALC> tempBuffer;
+    TBuf<TPosition::VECCALC> sortedBuffer;
+
+    GlobalTensor<int32_t> expertIdxGm;
+    GlobalTensor<int32_t> sortedexpertIdxGm;
+    GlobalTensor<int32_t> expandDstToSrcRowGm;
+    GlobalTensor<int32_t> expertTokensCountOrCumsumGm;
+    GlobalTensor<int32_t> expertTokensBeforeCapacityGm;
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    GlobalTensor<int32_t> syncTmpSpaceGm_;
+#endif
+    int64_t tileLength;
+    int64_t bufferNum = 1;
+    int64_t totalLength;
+    int64_t coreNum;
+    int64_t n;
+    int64_t k;
+    int64_t existRowIdx;
+    int64_t expertNum;
+    int64_t expertTokensCountOrCumsumFlag = 0;
+    int64_t expertTokensBeforeCapacityFlag = 0;
+
+    static constexpr int64_t SYNC_GM_NUM = 2;
+    static constexpr int64_t WORK_GM_NUM = 2;
+    static constexpr int64_t DST_BLK_STRIDE = 1;
+    static constexpr int64_t DST_REP_STRIDE = 8;
+};
+
+__aicore__ inline void MoeV2SortBase::SyncAll()
+{
+#ifndef __CCE_KT_TEST__
+    if (coreNum == 1) {
+        return;
+    }
+    AscendC::SyncAll();
+#endif
+}
+
+} // namespace MoeInitRoutingV2
+#endif // MOE_V2_SORT_BASE_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_sort_multi_core.h

@@ -0,0 +1,507 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_sort_multi_core.h
+ * \brief
+ */
+#ifndef MOE_V2_VBS_ONE_CORE_H
+#define MOE_V2_VBS_ONE_CORE_H
+
+#include "moe_v2_sort_base.h"
+#include "moe_v2_mrgsort.h"
+#include "moe_v2_mrgsort_out.h"
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+using namespace optiling;
+
+class MoeV2SortMultiCore : public MoeV2SortBase {
+public:
+    __aicore__ inline MoeV2SortMultiCore(){};
+    template <typename TilingData>
+    __aicore__ inline void Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum,
+                                GM_ADDR expertTokensBeforeCapacity, GM_ADDR workspace, const TilingData *tilingData,
+                                TPipe *tPipe);
+    __aicore__ inline void Process();
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    __aicore__ inline void ResetIO(GM_ADDR expandedRowIdx, GM_ADDR workspace);
+#endif
+private:
+    __aicore__ inline void VBSProcess();
+    __aicore__ inline void UBSortProcess(int64_t progress, int64_t size, int64_t sortNum);
+    __aicore__ inline void OneCoreVMSProcess(int64_t listNum, int64_t perListElements, int64_t lastListElements);
+    __aicore__ inline void VMSProcess();
+    __aicore__ inline void SortOutProcess();
+    __aicore__ inline void VBSCopyIn(int64_t progress, int64_t size, int64_t sortNum);
+    __aicore__ inline void UBSortCompute(int64_t progress, int64_t size, int64_t sortNum);
+    __aicore__ inline void VBSCopyOut(int64_t progress, int64_t size, int64_t sortNum);
+    __aicore__ inline void InitMoeMrgSort(MoeV2Mrgsort *sorter, int64_t listNum, int64_t coreOffset,
+                                          int64_t loopOffset);
+    __aicore__ inline void InitMoeMrgSortOut(MoeV2MrgsortOut *sorter, int64_t listNum, int64_t coreOffset);
+    __aicore__ inline void InitExpertTokensGlobalMemory();
+
+private:
+    GlobalTensor<float> workspaceGms[2];
+
+    const MoeV2VBSComputeTilingData *vbsTilingData;
+    const MoeV2VMSMiddleComputeTilingData *vmsTilingData;
+    const MoeV2SortOutComputeTilingData *sortOutTilingData;
+
+    // for MoeMrgsort
+    MoeV2Mrgsort mrgsorter;
+    MoeV2MrgsortParam mrgsortParam;
+
+    int64_t blockIdx;
+    int64_t srcWsIndex = 0;
+
+    int64_t listNum;
+    int64_t perListElements;
+    int64_t lastListElements;
+
+    int64_t sortTotalLength;
+    int64_t sortCoreLoops;
+    int64_t sortCoreLoopElements;
+    int64_t sortCoreLastLoopElements;
+
+    int64_t perCoreExpert;
+    int64_t needInitExpertCore;
+    int64_t currentCoreExpert;
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    int64_t perCoreOffset;
+#endif
+    static constexpr int64_t MAX_MRGSORT_LIST = 4;
+};
+
+__aicore__ inline void MoeV2SortMultiCore::InitExpertTokensGlobalMemory()
+{
+    if (this->blockIdx < this->needInitExpertCore) {
+        if (this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) {
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+            InitGlobalMemory(expertTokensCountOrCumsumGm, Align(this->currentCoreExpert, sizeof(int32_t)), 0);
+#else
+            InitGlobalMemory(expertTokensCountOrCumsumGm, currentCoreExpert, 0);
+#endif
+        }
+        if (this->expertTokensBeforeCapacityFlag == EXERPT_TOKENS_BEFORE_CAPACITY) {
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+            InitGlobalMemory(expertTokensBeforeCapacityGm, Align(this->currentCoreExpert, sizeof(int32_t)), 0);
+#else
+            InitGlobalMemory(expertTokensBeforeCapacityGm, currentCoreExpert, 0);
+#endif
+        }
+    }
+}
+
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+__aicore__ inline void MoeV2SortMultiCore::ResetIO(GM_ADDR expandedRowIdx, GM_ADDR workspace)
+{
+    sortedexpertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(expandedRowIdx),
+        Align(this->totalLength, sizeof(int32_t)));
+    expandDstToSrcRowGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(expandedRowIdx),
+        Align(this->totalLength, sizeof(int32_t)));
+    expertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace) +
+                                Align(this->totalLength, sizeof(int32_t)) + this->blockIdx * perCoreOffset,
+                                this->sortTotalLength);
+    this->srcWsIndex = 0;
+    this->needInitExpertCore = 0;
+}
+#endif
+
+__aicore__ inline void MoeV2SortMultiCore::VBSCopyIn(int64_t progress, int64_t size, int64_t sortNum)
+{
+    LocalTensor<int32_t> inLocal = sortDataCopyInQueue.AllocTensor<int32_t>();
+    int64_t inOffset = progress * sortCoreLoopElements;
+    DataCopyExtParams dataCopyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(size * sizeof(int32_t)), 0, 0, 0};
+    DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyPadCustom(inLocal[0], expertIdxGm[inOffset], dataCopyParams, dataCopyPadParams);
+#else
+    DataCopyPad(inLocal[0], expertIdxGm[inOffset], dataCopyParams, dataCopyPadParams);
+#endif
+    LocalTensor<int32_t> rowIdxLocal = inLocal[sortNum];
+    int64_t startValue = this->blockIdx * this->vbsTilingData->perCoreElements + inOffset;
+    SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+    ArithProgression<int32_t>(rowIdxLocal, startValue, 1, size);
+    sortDataCopyInQueue.EnQue(inLocal);
+}
+
+__aicore__ inline void MoeV2SortMultiCore::UBSortCompute(int64_t progress, int64_t size, int64_t sortNum)
+{
+    LocalTensor<int32_t> inLocal = sortDataCopyInQueue.DeQue<int32_t>();
+    LocalTensor<int32_t> expertForSourceRowLocal = inLocal[0];
+    LocalTensor<float> expertForSourceRowLocalFp32;
+
+    expertForSourceRowLocalFp32 = expertForSourceRowLocal.ReinterpretCast<float>();
+    #if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    Cast(expertForSourceRowLocalFp32, expertForSourceRowLocal, RoundMode::CAST_NONE, sortNum);
+    #else
+    Cast(expertForSourceRowLocalFp32, expertForSourceRowLocal, RoundMode::CAST_ROUND, sortNum);
+    #endif
+    PipeBarrier<PIPE_V>();
+    Muls(expertForSourceRowLocalFp32, expertForSourceRowLocalFp32, (float)-1, sortNum);
+    PipeBarrier<PIPE_V>();
+
+    int64_t duplicateNum = size % ONE_REPEAT_SORT_NUM;
+    if (duplicateNum > 0) {
+        int duplicateIndex = size - duplicateNum;
+        uint64_t mask0 = UINT64_MAX;
+        mask0 = mask0 << duplicateNum;
+        mask0 = mask0 & (UINT64_MAX >> (FP32_ONE_REPEAT_NUM - ONE_REPEAT_SORT_NUM));
+        uint64_t mask[2] = {mask0, 0};
+        Duplicate(expertForSourceRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+        PipeBarrier<PIPE_V>();
+    }
+
+    LocalTensor<float> sortedLocal = sortedBuffer.Get<float>(GetSortLen<float>(sortNum));
+    LocalTensor<uint32_t> sourceRowLocal;
+    sourceRowLocal = inLocal[sortNum].ReinterpretCast<uint32_t>();
+
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    LocalTensor<float> concatLocal;
+    SetWaitFlag<HardEvent::MTE3_V>(HardEvent::MTE3_V);
+
+    LocalTensor<float> outLocal = tempBuffer.Get<float>(GetSortLen<float>(sortNum));
+
+    Concat(concatLocal, expertForSourceRowLocalFp32, sortedLocal, sortNum / ONE_REPEAT_SORT_NUM);
+
+    PipeBarrier<PIPE_V>();
+
+    Sort<float, true>(outLocal, concatLocal, sourceRowLocal, concatLocal, sortNum / ONE_REPEAT_SORT_NUM);
+
+    SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+#else
+    LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
+
+    LocalTensor<float> concatLocal = expertForSourceRowLocalFp32;
+
+    Sort<float, true>(outLocal, concatLocal, sourceRowLocal, sortedLocal, sortNum / ONE_REPEAT_SORT_NUM);
+
+    sortDataCopyOutQueue.EnQue<float>(outLocal);
+#endif
+    sortDataCopyInQueue.FreeTensor(inLocal);
+}
+
+__aicore__ inline void MoeV2SortMultiCore::VBSCopyOut(int64_t progress, int64_t size, int64_t sortNum)
+{
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    LocalTensor<float> outLocal = tempBuffer.Get<float>(GetSortLen<float>(sortNum));
+#else
+    LocalTensor<float> outLocal = sortDataCopyOutQueue.DeQue<float>();
+#endif
+    DataCopy(workspaceGms[0][this->blockIdx * GetSortLen<float>(this->vbsTilingData->perCoreElements) +
+                             GetSortLen<float>(progress * sortCoreLoopElements)],
+             outLocal, Align(GetSortLen<float>(size), sizeof(float)));
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    SetWaitFlag<HardEvent::MTE3_V>(HardEvent::MTE3_V);
+#else
+    sortDataCopyOutQueue.FreeTensor(outLocal);
+#endif       
+}
+
+__aicore__ inline void MoeV2SortMultiCore::InitMoeMrgSort(MoeV2Mrgsort *sorter, int64_t listNum, int64_t coreOffset,
+                                                          int64_t loopOffset)
+{
+    GlobalTensor<float> srcWsGm = workspaceGms[srcWsIndex][blockIdx * coreOffset + loopOffset];
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    LocalTensor<float> inLocal = sortedBuffer.Get<float>(GetSortLen<float>(this->sortOutTilingData->oneLoopMaxElements) * MAX_MRGSORT_LIST);
+    LocalTensor<float> outLocal = tempBuffer.Get<float>(GetSortLen<float>(this->sortOutTilingData->oneLoopMaxElements) * MAX_MRGSORT_LIST);
+#else
+    LocalTensor<float> inLocal = sortDataCopyInQueue.AllocTensor<float>();
+    LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
+#endif
+    for (int64_t i = 0; i < listNum; i++) {
+        LocalTensor<float> inLocalT = inLocal[GetSortLen<float>(this->sortOutTilingData->oneLoopMaxElements) * i];
+        sorter->SetInput(srcWsGm, inLocalT);
+    }
+    GlobalTensor<float> dstWsGm = workspaceGms[1 - srcWsIndex][blockIdx * coreOffset + loopOffset];
+    sorter->SetOutput(dstWsGm, outLocal);
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ != 200
+    sortDataCopyInQueue.FreeTensor(inLocal);
+    sortDataCopyOutQueue.FreeTensor(outLocal);
+#endif
+}
+
+__aicore__ inline void MoeV2SortMultiCore::InitMoeMrgSortOut(MoeV2MrgsortOut *sorter, int64_t listNum,
+                                                             int64_t coreOffset)
+{
+    GlobalTensor<float> srcWsGm = workspaceGms[srcWsIndex];
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    LocalTensor<float> inLocal = sortedBuffer.Get<float>(GetSortLen<float>(this->sortOutTilingData->oneLoopMaxElements) * MAX_MRGSORT_LIST);
+    LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
+#else
+    LocalTensor<float> inLocal = sortDataCopyInQueue.AllocTensor<float>();
+    LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
+#endif
+    for (int64_t i = 0; i < listNum; i++) {
+        LocalTensor<float> inLocalT = inLocal[GetSortLen<float>(this->sortOutTilingData->oneLoopMaxElements) * i];
+        sorter->SetInput(srcWsGm, inLocalT);
+    }
+
+    LocalTensor<float> outLocalV = outLocal[this->sortOutTilingData->oneLoopMaxElements * MAX_MRGSORT_LIST];
+    sorter->SetOutput(this->sortedexpertIdxGm, this->expandDstToSrcRowGm, outLocal, outLocalV);
+
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    LocalTensor<float> tempLocal = tempBuffer.Get<float>(GetSortLen<float>(this->sortOutTilingData->oneLoopMaxElements) * MAX_MRGSORT_LIST);
+    // buffer for Extract
+    sorter->SetBuffer(tempLocal);
+    sortDataCopyOutQueue.FreeTensor(outLocal);
+#else
+    LocalTensor<float> tempLocal =
+        sortedBuffer.Get<float>(GetSortLen<float>(this->sortOutTilingData->oneLoopMaxElements) * MAX_MRGSORT_LIST);
+    // buffer for Extract
+    sorter->SetBuffer(tempLocal);
+    sortDataCopyInQueue.FreeTensor(inLocal);
+    sortDataCopyOutQueue.FreeTensor(outLocal);
+#endif
+}
+
+__aicore__ inline void MoeV2SortMultiCore::OneCoreVMSProcess(int64_t listNum, int64_t perListElements,
+                                                             int64_t lastListElements)
+{
+    int64_t coreOffset = GetSortLen<float>(this->vbsTilingData->perCoreElements);
+    mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements;
+
+    for (int64_t i = 0; listNum >= 1; i++) {
+        int64_t loops = (listNum + MAX_MRGSORT_LIST - 1) / MAX_MRGSORT_LIST;
+        int64_t remainListNum = listNum - (loops - 1) * MAX_MRGSORT_LIST;
+
+        mrgsortParam.perListElements = perListElements;
+        mrgsortParam.lastListElements = perListElements;
+
+        int64_t loopOffset = GetSortLen<float>(mrgsortParam.perListElements * MAX_MRGSORT_LIST);
+        for (int64_t loop = 0; loop < loops - 1; loop++) {
+            InitMoeMrgSort(&mrgsorter, MAX_MRGSORT_LIST, coreOffset, loop * loopOffset);
+            mrgsorter.Init(&mrgsortParam);
+            mrgsorter.Process();
+        }
+
+        mrgsortParam.perListElements = perListElements;
+        mrgsortParam.lastListElements = lastListElements;
+        InitMoeMrgSort(&mrgsorter, remainListNum, coreOffset, (loops - 1) * loopOffset);
+        mrgsorter.Init(&mrgsortParam);
+        mrgsorter.Process();
+
+        listNum = loops;
+        lastListElements = perListElements * (remainListNum - 1) + lastListElements;
+        perListElements = perListElements * MAX_MRGSORT_LIST;
+        srcWsIndex = (srcWsIndex + 1) % WORK_GM_NUM;
+
+        if (loops == 1) {
+            break;
+        }
+    }
+}
+
+__aicore__ inline void MoeV2SortMultiCore::UBSortProcess(int64_t progress, int64_t size, int64_t sortNum)
+{
+    VBSCopyIn(progress, size, sortNum);
+    UBSortCompute(progress, size, sortNum);
+    VBSCopyOut(progress, size, sortNum);
+}
+
+__aicore__ inline void MoeV2SortMultiCore::VBSProcess()
+{
+    if (this->blockIdx < this->vbsTilingData->needCoreNum) {
+        int64_t sortNum = Ceil(sortCoreLoopElements, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
+        for (int64_t loop = 0; loop < sortCoreLoops - 1; loop++) {
+            UBSortProcess(loop, sortCoreLoopElements, sortNum);
+        }
+
+        sortNum = Ceil(sortCoreLastLoopElements, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
+        UBSortProcess(sortCoreLoops - 1, sortCoreLastLoopElements, sortNum);
+
+        if (sortCoreLoops > 1) {
+            OneCoreVMSProcess(sortCoreLoops, sortCoreLoopElements, sortCoreLastLoopElements);
+        }
+    }
+
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    LocalTensor<int32_t> syncLocal = tempBuffer.Get<int32_t>();
+    AscendC::SyncAll(syncTmpSpaceGm_, syncLocal, GetBlockNum());
+#else
+    SyncAll();
+#endif
+}
+
+__aicore__ inline void MoeV2SortMultiCore::VMSProcess()
+{
+    int64_t currentStageNeedCoreNum = this->vmsTilingData->needCoreNum;
+    perListElements = this->vbsTilingData->perCoreElements;
+    lastListElements = this->vbsTilingData->lastCoreElements;
+    listNum = this->vbsTilingData->needCoreNum;
+
+    for (; listNum > MAX_MRGSORT_LIST;) {
+        currentStageNeedCoreNum = Ceil(listNum, MAX_MRGSORT_LIST);
+        int64_t coreOffset = GetSortLen<float>(perListElements * MAX_MRGSORT_LIST);
+        int64_t remainListNum = listNum - (currentStageNeedCoreNum - 1) * MAX_MRGSORT_LIST;
+
+        if (this->blockIdx < currentStageNeedCoreNum - 1) {
+            mrgsortParam.perListElements = perListElements;
+            mrgsortParam.lastListElements = perListElements;
+            mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements;
+            InitMoeMrgSort(&mrgsorter, MAX_MRGSORT_LIST, coreOffset, 0);
+            mrgsorter.Init(&mrgsortParam);
+            mrgsorter.Process();
+        } else if (this->blockIdx == currentStageNeedCoreNum - 1) {
+            mrgsortParam.perListElements = perListElements;
+            mrgsortParam.lastListElements = lastListElements;
+            mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements;
+            InitMoeMrgSort(&mrgsorter, remainListNum, coreOffset, 0);
+            mrgsorter.Init(&mrgsortParam);
+            mrgsorter.Process();
+        }
+        listNum = currentStageNeedCoreNum;
+        currentStageNeedCoreNum = Ceil(listNum, MAX_MRGSORT_LIST);
+        srcWsIndex = (srcWsIndex + 1) % WORK_GM_NUM;
+
+        lastListElements = perListElements * (remainListNum - 1) + lastListElements;
+        perListElements = perListElements * MAX_MRGSORT_LIST;
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+        LocalTensor<int32_t> syncLocal = tempBuffer.Get<int32_t>();
+        AscendC::SyncAll(syncTmpSpaceGm_, syncLocal, GetBlockNum());
+#else
+        SyncAll();
+#endif
+    }
+}
+
+__aicore__ inline void MoeV2SortMultiCore::SortOutProcess()
+{
+    if (this->blockIdx < 1) {
+        mrgsortParam.perListElements = perListElements;
+        mrgsortParam.lastListElements = lastListElements;
+        mrgsortParam.oneLoopMaxElements = this->sortOutTilingData->oneLoopMaxElements;
+
+        MoeV2MrgsortOut sorter;
+        InitMoeMrgSortOut(&sorter, listNum, GetSortLen<float>(perListElements));
+        sorter.Init(&mrgsortParam, pipe);
+        sorter.Process();
+    }
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    LocalTensor<int32_t> syncLocal = tempBuffer.Get<int32_t>();
+    AscendC::SyncAll(syncTmpSpaceGm_, syncLocal, GetBlockNum());
+#else
+    SyncAll();
+#endif
+}
+
+template <typename TilingData>
+__aicore__ inline void MoeV2SortMultiCore::Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum,
+                                                GM_ADDR expertTokensBeforeCapacity, GM_ADDR workspace,
+                                                const TilingData *tilingData, TPipe *tPipe)
+{
+    this->totalLength = tilingData->n * tilingData->k;
+    this->coreNum = tilingData->coreNum;
+    this->vbsTilingData = &(tilingData->vbsComputeParamsOp);
+    this->vmsTilingData = &(tilingData->vmsMiddleComputeParamsOp);
+    this->sortOutTilingData = &(tilingData->sortOutComputeParamsOp);
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    this->perCoreOffset = this->vbsTilingData->perCoreElements;
+#endif
+    this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+    this->tileLength = this->vbsTilingData->perCorePerLoopElements;
+    this->sortTotalLength = this->vbsTilingData->perCoreElements;
+    if (this->blockIdx == tilingData->vbsComputeParamsOp.needCoreNum - 1) {
+        this->tileLength = this->vbsTilingData->lastCorePerLoopElements;
+        this->sortTotalLength = this->vbsTilingData->lastCoreElements;
+    }
+    this->n = tilingData->n;
+    this->k = tilingData->k;
+    this->expertNum = tilingData->expertNum;
+    this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag;
+    this->expertTokensBeforeCapacityFlag = tilingData->expertTokensBeforeCapacityFlag;
+
+    // VBS param init
+    if (this->blockIdx == this->vbsTilingData->needCoreNum - 1) {
+        sortCoreLoops = this->vbsTilingData->lastCoreLoops;
+        sortCoreLoopElements = this->vbsTilingData->lastCorePerLoopElements;
+        sortCoreLastLoopElements = this->vbsTilingData->lastCoreLastLoopElements;
+    } else {
+        sortCoreLoops = this->vbsTilingData->perCoreLoops;
+        sortCoreLoopElements = this->vbsTilingData->perCorePerLoopElements;
+        sortCoreLastLoopElements = this->vbsTilingData->perCoreLastLoopElements;
+    }
+
+    this->pipe = tPipe;
+    expertIdxGm.SetGlobalBuffer((__gm__ int32_t *)expertIdx +
+                                    this->blockIdx * tilingData->vbsComputeParamsOp.perCoreElements,
+                                this->sortTotalLength);
+    sortedexpertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace),
+                                      Align(this->totalLength, sizeof(int32_t)));
+    expandDstToSrcRowGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace) +
+                                            Align(this->totalLength, sizeof(int32_t)),
+                                        Align(this->totalLength, sizeof(int32_t)));
+
+    this->perCoreExpert = Align((this->expertNum + this->coreNum - 1) / this->coreNum, sizeof(int32_t));
+    this->needInitExpertCore = (this->expertNum + this->perCoreExpert - 1) / this->perCoreExpert;
+    this->currentCoreExpert = this->perCoreExpert;
+    if (this->blockIdx == needInitExpertCore - 1) {
+        this->currentCoreExpert = this->expertNum - (this->needInitExpertCore - 1) * this->perCoreExpert;
+    }
+    if (this->expertTokensCountOrCumsumFlag > EXERPT_TOKENS_NONE) {
+        expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t *)expertTokensCountOrCumsum +
+                                                        this->blockIdx * this->perCoreExpert,
+                                                    this->currentCoreExpert);
+    }
+    if (this->expertTokensBeforeCapacityFlag == EXERPT_TOKENS_BEFORE_CAPACITY) {
+        expertTokensBeforeCapacityGm.SetGlobalBuffer((__gm__ int32_t *)expertTokensBeforeCapacity +
+                                                         this->blockIdx * this->perCoreExpert,
+                                                     this->currentCoreExpert);
+    }
+    // key and value
+    int64_t kvFactor = 2;
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    int64_t workspaceLen = GetSortLen<float>(Align(this->totalLength, sizeof(int32_t)));
+    workspaceGms[0].SetGlobalBuffer((__gm__ float *)workspace + Align(this->totalLength, sizeof(int32_t)) * 2,
+                                    workspaceLen);
+    workspaceGms[1].SetGlobalBuffer((__gm__ float *)workspace +
+                                        Align(this->totalLength, sizeof(int32_t)) * 2 + workspaceLen,
+                                    workspaceLen);
+                                    
+#else
+    workspaceGms[0].SetGlobalBuffer((__gm__ float *)workspace + Align(this->totalLength, sizeof(int32_t)) * 2,
+                                    Align(this->totalLength, sizeof(int32_t)) * kvFactor);
+    workspaceGms[1].SetGlobalBuffer((__gm__ float *)workspace +
+                                        Align(this->totalLength, sizeof(int32_t)) * (kvFactor + 2),
+                                    Align(this->totalLength, sizeof(int32_t)) * kvFactor);
+#endif
+
+    int64_t bufferSize = Ceil(Max(this->sortOutTilingData->oneLoopMaxElements * MAX_MRGSORT_LIST, sortCoreLoopElements),
+                              ONE_REPEAT_SORT_NUM) *
+                         ONE_REPEAT_SORT_NUM * sizeof(int32_t) * kvFactor;
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    syncTmpSpaceGm_.SetGlobalBuffer((__gm__ int32_t *)workspace +
+                                        Align(this->totalLength, sizeof(int32_t)) * 2 + 2 * workspaceLen,
+                                    INT32_ONE_BLOCK_NUM * GetBlockNum() * BLOCK_BYTES);
+    pipe->InitBuffer(sortDataCopyInQueue, bufferNum, bufferSize);
+    pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, bufferSize);
+    pipe->InitBuffer(tempBuffer, bufferSize * REGIONP_ROPOSAL_KV_RATIO);
+    pipe->InitBuffer(sortedBuffer, bufferSize* REGIONP_ROPOSAL_KV_RATIO);
+    LocalTensor<int32_t> syncLocal = tempBuffer.Get<int32_t>();
+    Duplicate<int32_t>(syncLocal, 0, SYNC_LEN);
+    SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+    DataCopy(syncTmpSpaceGm_, syncLocal, SYNC_LEN);
+#else
+    pipe->InitBuffer(sortDataCopyInQueue, bufferNum, bufferSize);
+    pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, bufferSize);
+    pipe->InitBuffer(sortedBuffer, bufferSize);
+#endif
+}
+
+__aicore__ inline void MoeV2SortMultiCore::Process()
+{
+    InitExpertTokensGlobalMemory();
+    VBSProcess();
+    VMSProcess();
+    SortOutProcess();
+}
+} // namespace MoeInitRoutingV2
+#endif // MOE_V2_VBS_ONE_CORE_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_sort_one_core.h

@@ -0,0 +1,226 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_sort_one_core.h
+ * \brief
+ */
+#ifndef MOE_V2_SORT_ONE_CORE_H
+#define MOE_V2_SORT_ONE_CORE_H
+
+#include "moe_v2_mrgsort.h"
+#include "moe_v2_sort_base.h"
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+using namespace optiling;
+
+class MoeV2SortOneCore : public MoeV2SortBase {
+public:
+    __aicore__ inline MoeV2SortOneCore(){};
+    template <typename TilingData>
+    __aicore__ inline void Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum,
+                                GM_ADDR expertTokensBeforeCapacity, GM_ADDR workspace, const TilingData *tilingData,
+                                TPipe *tPipe);
+    __aicore__ inline void Process();
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    __aicore__ inline void ResetIO(GM_ADDR expandedRowIdx, GM_ADDR workspace);
+    bool needCopy = true;
+#endif
+
+private:
+    __aicore__ inline void CopyIn();
+    __aicore__ inline void SortCompute();
+    __aicore__ inline void CopyOut();
+
+private:
+    int64_t sortNum;
+    int64_t blockIdx;
+    int64_t needCoreNum;
+};
+
+__aicore__ inline void MoeV2SortOneCore::CopyIn()
+{
+    LocalTensor<int32_t> inLocal = sortDataCopyInQueue.AllocTensor<int32_t>();
+    DataCopyExtParams dataCopyParams{static_cast<uint16_t>(1),
+                                     static_cast<uint32_t>(this->totalLength * sizeof(int32_t)), 0, 0, 0};
+    DataCopyPadExtParams<int32_t> dataCopyPadParams{false, 0, 0, 0};
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyPadCustom(inLocal[0], expertIdxGm, dataCopyParams, dataCopyPadParams);
+#else
+    DataCopyPad(inLocal[0], expertIdxGm, dataCopyParams, dataCopyPadParams);
+#endif
+    LocalTensor<int32_t> rowIdxLocal = inLocal[this->sortNum];
+    ArithProgression<int32_t>(rowIdxLocal, 0, 1, this->sortNum);
+    sortDataCopyInQueue.EnQue(inLocal);
+}
+
+__aicore__ inline void MoeV2SortOneCore::SortCompute()
+{
+    LocalTensor<int32_t> inLocal = sortDataCopyInQueue.DeQue<int32_t>();
+    LocalTensor<int32_t> expertForSourceRowLocal = inLocal[0];
+    LocalTensor<float> expertForSourceRowLocalFp32 = expertForSourceRowLocal.ReinterpretCast<float>();
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    Cast(expertForSourceRowLocalFp32, expertForSourceRowLocal, RoundMode::CAST_NONE, this->tileLength);
+#else
+    Cast(expertForSourceRowLocalFp32, expertForSourceRowLocal, RoundMode::CAST_ROUND, this->tileLength);
+#endif
+    PipeBarrier<PIPE_V>();
+    Muls(expertForSourceRowLocalFp32, expertForSourceRowLocalFp32, (float)-1, this->tileLength);
+    PipeBarrier<PIPE_V>();
+
+    int64_t duplicateNum = this->totalLength % ONE_REPEAT_SORT_NUM;
+    if (duplicateNum > 0) {
+        int duplicateIndex = this->totalLength - duplicateNum;
+        uint64_t mask0 = UINT64_MAX;
+        mask0 = mask0 << duplicateNum;
+        mask0 = mask0 & (UINT64_MAX >> (FP32_ONE_REPEAT_NUM - ONE_REPEAT_SORT_NUM));
+        uint64_t mask[2] = {mask0, 0};
+        Duplicate(expertForSourceRowLocalFp32[duplicateIndex], MIN_FP32, mask, 1, DST_BLK_STRIDE, DST_REP_STRIDE);
+        PipeBarrier<PIPE_V>();
+    }
+
+    LocalTensor<float> concatLocal = expertForSourceRowLocalFp32;
+    LocalTensor<float> tempTensor = tempBuffer.Get<float>(GetSortLen<float>(this->sortNum));
+    Concat(concatLocal, expertForSourceRowLocalFp32, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM);
+    PipeBarrier<PIPE_V>();
+
+    LocalTensor<float> sortedLocal = sortedBuffer.Get<float>(GetSortLen<float>(this->sortNum));
+    LocalTensor<uint32_t> sourceRowLocal;
+    sourceRowLocal = inLocal[this->sortNum].ReinterpretCast<uint32_t>();
+    Sort<float, true>(sortedLocal, concatLocal, sourceRowLocal, tempTensor, this->sortNum / ONE_REPEAT_SORT_NUM);
+    PipeBarrier<PIPE_V>();
+
+    LocalTensor<float> outLocal = sortDataCopyOutQueue.AllocTensor<float>();
+    LocalTensor<float> sortedExpertForSourceRowLocal = outLocal[0];
+    LocalTensor<uint32_t> expandDstToSrcRowLocal;
+    expandDstToSrcRowLocal = outLocal[this->sortNum].ReinterpretCast<uint32_t>();
+    Extract(sortedExpertForSourceRowLocal, expandDstToSrcRowLocal, sortedLocal, this->sortNum / ONE_REPEAT_SORT_NUM);
+    PipeBarrier<PIPE_V>();
+    Muls(sortedExpertForSourceRowLocal, sortedExpertForSourceRowLocal, (float)-1, this->tileLength);
+    PipeBarrier<PIPE_V>();
+
+    LocalTensor<int32_t> expertForSourceRowLocalInt32;
+    expertForSourceRowLocalInt32 = sortedExpertForSourceRowLocal.ReinterpretCast<int32_t>();
+    Cast(expertForSourceRowLocalInt32, sortedExpertForSourceRowLocal, RoundMode::CAST_ROUND, this->tileLength);
+    sortDataCopyOutQueue.EnQue<float>(outLocal);
+    sortDataCopyInQueue.FreeTensor(inLocal);
+}
+
+__aicore__ inline void MoeV2SortOneCore::CopyOut()
+{
+    LocalTensor<int32_t> outLocal = sortDataCopyOutQueue.DeQue<int32_t>();
+    DataCopyParams intriParams;
+    intriParams.blockCount = 1;
+    intriParams.blockLen = this->totalLength * sizeof(int32_t);
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    DataCopyCustom<int32_t,true,false>(expandDstToSrcRowGm, outLocal[this->sortNum], intriParams.blockCount, intriParams.blockLen);
+    if (this->needCopy) {
+        DataCopyCustom<int32_t,true,false>(sortedexpertIdxGm, outLocal[0], intriParams.blockCount, intriParams.blockLen);
+    }
+#else
+    DataCopyPad(sortedexpertIdxGm, outLocal[0], intriParams);
+    DataCopyPad(expandDstToSrcRowGm, outLocal[this->sortNum], intriParams);
+#endif
+    sortDataCopyOutQueue.FreeTensor(outLocal);
+}
+
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+__aicore__ inline void MoeV2SortOneCore::ResetIO(GM_ADDR expandedRowIdx, GM_ADDR workspace)
+{
+    sortedexpertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(expandedRowIdx), this->tileLength);
+    expandDstToSrcRowGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(expandedRowIdx), this->tileLength);
+    expertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace) + this->tileLength, this->tileLength);
+    this->expertTokensCountOrCumsumFlag = 0;
+    this->expertTokensBeforeCapacityFlag = 0;
+    this->needCopy = false;
+}
+#endif
+
+template <typename TilingData>
+__aicore__ inline void MoeV2SortOneCore::Init(GM_ADDR expertIdx, GM_ADDR expertTokensCountOrCumsum,
+                                              GM_ADDR expertTokensBeforeCapacity, GM_ADDR workspace,
+                                              const TilingData *tilingData, TPipe *tPipe)
+{
+    this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+    this->tileLength = Align(tilingData->vbsComputeParamsOp.lastCorePerLoopElements, sizeof(int32_t));
+    this->sortNum = Ceil(this->tileLength, ONE_REPEAT_SORT_NUM) * ONE_REPEAT_SORT_NUM;
+    this->totalLength = tilingData->n * tilingData->k;
+    this->coreNum = tilingData->coreNum;
+    this->pipe = tPipe;
+    this->n = tilingData->n;
+    this->k = tilingData->k;
+    this->expertNum = tilingData->expertNum;
+    this->expertTokensCountOrCumsumFlag = tilingData->expertTokensCountOrCumsumFlag;
+    this->expertTokensBeforeCapacityFlag = tilingData->expertTokensBeforeCapacityFlag;
+    this->needCoreNum = tilingData->vbsComputeParamsOp.needCoreNum;
+
+    expertIdxGm.SetGlobalBuffer((__gm__ int32_t *)expertIdx, this->tileLength);
+    sortedexpertIdxGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace), this->tileLength);
+    expandDstToSrcRowGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace) + this->tileLength,
+                                        this->tileLength);
+
+    if (this->blockIdx == this->coreNum - 1) {
+        if (this->expertTokensCountOrCumsumFlag > 0) {
+            expertTokensCountOrCumsumGm.SetGlobalBuffer((__gm__ int32_t *)expertTokensCountOrCumsum,
+                                                        Align(this->expertNum, sizeof(int32_t)));
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+            InitGlobalMemory(expertTokensCountOrCumsumGm, Align(this->expertNum, sizeof(int32_t)), 0);
+#else
+            InitGlobalMemory(expertTokensCountOrCumsumGm, this->expertNum, 0);
+#endif
+        }
+        if (this->expertTokensBeforeCapacityFlag == 1) {
+            expertTokensBeforeCapacityGm.SetGlobalBuffer((__gm__ int32_t *)expertTokensBeforeCapacity,
+                                                         Align(this->expertNum, sizeof(int32_t)));
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+            InitGlobalMemory(expertTokensBeforeCapacityGm, Align(this->expertNum, sizeof(int32_t)), 0);
+#else
+            InitGlobalMemory(expertTokensBeforeCapacityGm, this->expertNum, 0);
+#endif
+        }
+    }
+    // key and value
+    int64_t kvFactor = 2;
+    int64_t buffSize = this->sortNum * sizeof(int32_t) * kvFactor;
+    pipe->InitBuffer(sortDataCopyInQueue, bufferNum, buffSize);
+    pipe->InitBuffer(sortDataCopyOutQueue, bufferNum, buffSize);
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    syncTmpSpaceGm_.SetGlobalBuffer((__gm__ int32_t *)workspace + 2 * this->tileLength, SYNC_LEN);
+    buffSize = GetSortLen<float>(this->sortNum) * sizeof(int32_t);
+    pipe->InitBuffer(tempBuffer, buffSize);
+    pipe->InitBuffer(sortedBuffer, buffSize);
+    LocalTensor<int32_t> syncLocal = tempBuffer.Get<int32_t>();
+    Duplicate<int32_t>(syncLocal, 0, SYNC_LEN);
+    SetWaitFlag<HardEvent::V_MTE3>(HardEvent::V_MTE3);
+    DataCopy(syncTmpSpaceGm_, syncLocal, SYNC_LEN);
+    PipeBarrier<PIPE_ALL>();
+#else
+    pipe->InitBuffer(tempBuffer, buffSize);
+    pipe->InitBuffer(sortedBuffer, buffSize);
+#endif
+}
+
+__aicore__ inline void MoeV2SortOneCore::Process()
+{
+    if (GetBlockIdx() < this->needCoreNum) {
+        CopyIn();
+        SortCompute();
+        CopyOut();
+    }
+#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
+    LocalTensor<int32_t> syncLocal = tempBuffer.Get<int32_t>();
+    AscendC::SyncAll(syncTmpSpaceGm_, syncLocal, GetBlockNum());
+#else
+    this->SyncAll();
+#endif
+}
+} // namespace MoeInitRoutingV2
+#endif // MOE_V2_SORT_ONE_CORE_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_src_to_dst_op.h

@@ -0,0 +1,173 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_src_to_dst_op.h
+ * \brief
+ */
+#ifndef MOE_V2_SRC_TO_DST_H
+#define MOE_V2_SRC_TO_DST_H
+
+#include "moe_v2_common.h"
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+using namespace optiling;
+
+class MoeV2SrcToDstOp {
+public:
+    __aicore__ inline MoeV2SrcToDstOp(){};
+    template <typename TilingData>
+    __aicore__ inline void Init(GM_ADDR expandSrcToDstRow, GM_ADDR workspace, const TilingData *tilingData,
+                                TPipe *tPipe);
+    __aicore__ inline void Process();
+
+private:
+    __aicore__ inline void CopyIn(int64_t progress);
+    __aicore__ inline void Compute(int64_t progress);
+    __aicore__ inline void CopyOut();
+    __aicore__ inline void SyncAll();
+    __aicore__ inline void AssistInit();
+
+private:
+    TPipe *pipe;
+    TQue<QuePosition::VECIN, 1> copyInQueue;
+    TQue<QuePosition::VECOUT, 1> copyOutQueue;
+    TBuf<TPosition::VECCALC> assistBuffer;
+
+    GlobalTensor<int32_t> expandDstToSrcRowGm;
+    GlobalTensor<int32_t> expandSrcToDstRowGm;
+    GlobalTensor<int32_t> assistGm;
+
+    const MoeV2GatherOutComputeTilingData *srcToDstTilingData;
+
+    int64_t coreNum;
+    int64_t blockIdx;
+    int64_t totalLength;
+    int64_t currentLoopRows;
+    int64_t coreRows;
+    int64_t perLoopRows;
+    int64_t lastLoopRows;
+};
+
+__aicore__ inline void MoeV2SrcToDstOp::AssistInit()
+{
+#if defined(ASCENDC_OOM) && ASCENDC_OOM == 1
+    OOMCheckAddrRange(assistGm.GetPhyAddr(), ASSIST_NUM * sizeof(int32_t));
+#endif
+    LocalTensor<int32_t> assistTensor = assistBuffer.Get<int32_t>(ASSIST_NUM);
+    DataCopy(assistTensor, assistGm, ASSIST_NUM);
+    SetWaitFlag<HardEvent::MTE2_V>(HardEvent::MTE2_V);
+    Adds(assistTensor, assistTensor, (int32_t)(this->blockIdx * this->srcToDstTilingData->perCoreRows), ASSIST_NUM);
+}
+
+__aicore__ inline void MoeV2SrcToDstOp::CopyIn(int64_t progress)
+{
+    LocalTensor<int32_t> inLocal = copyInQueue.AllocTensor<int32_t>();
+    DataCopy(inLocal, expandDstToSrcRowGm[progress * perLoopRows], Align(currentLoopRows, sizeof(int32_t)));
+    copyInQueue.EnQue<int32_t>(inLocal);
+}
+
+__aicore__ inline void MoeV2SrcToDstOp::Compute(int64_t progress)
+{
+    LocalTensor<int32_t> outLocal = copyOutQueue.AllocTensor<int32_t>();
+    LocalTensor<int32_t> assistTensor = assistBuffer.Get<int32_t>(ASSIST_NUM);
+
+    PipeBarrier<PIPE_V>();
+    int64_t loops = Ceil(currentLoopRows, ASSIST_INDEX_NUM);
+    for (int64_t i = 0; i < loops; i++) {
+        Adds(outLocal[i * ASSIST_NUM], assistTensor,
+             static_cast<int32_t>(this->perLoopRows * progress + i * ASSIST_INDEX_NUM), ASSIST_NUM);
+    }
+    PipeBarrier<PIPE_V>();
+    copyOutQueue.EnQue<int32_t>(outLocal);
+}
+
+__aicore__ inline void MoeV2SrcToDstOp::CopyOut()
+{
+    LocalTensor<int32_t> inLocal = copyInQueue.DeQue<int32_t>();
+    LocalTensor<int32_t> outLocal = copyOutQueue.DeQue<int32_t>();
+    SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+    DataCopyParams intriParams;
+    intriParams.blockCount = 1;
+    intriParams.blockLen = sizeof(int32_t);
+    uint32_t outOffset;
+    for (int64_t idx = 0; idx < currentLoopRows; idx++) {
+        outOffset = inLocal.GetValue(idx);
+        DataCopyPad(expandSrcToDstRowGm[outOffset], outLocal[idx * INT32_ONE_BLOCK_NUM], intriParams);
+    }
+
+    copyInQueue.FreeTensor(inLocal);
+    copyOutQueue.FreeTensor(outLocal);
+}
+
+__aicore__ inline void MoeV2SrcToDstOp::SyncAll()
+{
+    if (coreNum == 1) {
+        return;
+    }
+#ifndef __CCE_KT_TEST__
+    AscendC::SyncAll();
+#endif
+}
+
+template <typename TilingData>
+__aicore__ inline void MoeV2SrcToDstOp::Init(GM_ADDR expandSrcToDstRow, GM_ADDR workspace, const TilingData *tilingData,
+                                             TPipe *tPipe)
+{
+    int64_t blockNum = GetBlockNum() * 2;
+    this->pipe = tPipe;
+    this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+
+    this->coreNum = tilingData->coreNum;
+    this->totalLength = tilingData->n * tilingData->k;
+    this->srcToDstTilingData = &(tilingData->srcToDstComputeParamsOp);
+
+    if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) {
+        this->coreRows = this->srcToDstTilingData->lastCoreRows;
+        this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows;
+        this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows;
+    } else {
+        this->coreRows = this->srcToDstTilingData->perCoreRows;
+        this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows;
+        this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows;
+    }
+
+    expandSrcToDstRowGm.SetGlobalBuffer((__gm__ int32_t *)expandSrcToDstRow, Align(this->totalLength, sizeof(int32_t)));
+    expandDstToSrcRowGm.SetGlobalBuffer((__gm__ int32_t *)workspace + Align(this->totalLength, sizeof(int32_t)) +
+                                            this->blockIdx * this->srcToDstTilingData->perCoreRows,
+                                        Align(this->coreRows, sizeof(int32_t)));
+    assistGm.SetGlobalBuffer((__gm__ int32_t *)assist, ASSIST_NUM);
+
+    pipe->InitBuffer(copyInQueue, 1, this->perLoopRows * BLOCK_BYTES);
+    pipe->InitBuffer(copyOutQueue, 1, Ceil(this->perLoopRows, ASSIST_NUM) * ASSIST_NUM * BLOCK_BYTES);
+    pipe->InitBuffer(assistBuffer, ASSIST_NUM * sizeof(int32_t));
+}
+
+__aicore__ inline void MoeV2SrcToDstOp::Process()
+{
+    if (this->blockIdx < this->srcToDstTilingData->needCoreNum) {
+        int64_t loops = (coreRows + perLoopRows - 1) / perLoopRows;
+        currentLoopRows = perLoopRows;
+        AssistInit();
+        for (int64_t loop = 0; loop < loops - 1; loop++) {
+            CopyIn(loop);
+            Compute(loop);
+            CopyOut();
+        }
+        currentLoopRows = lastLoopRows;
+        CopyIn(loops - 1);
+        Compute(loops - 1);
+        CopyOut();
+    }
+    this->SyncAll();
+}
+} // namespace MoeInitRoutingV2
+#endif // MOE_V2_SRC_TO_DST_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_src_to_dst_op_simt.h

@@ -0,0 +1,96 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/* !
+ * \file moe_src_to_dst_op_simt.h
+ * \brief
+ */
+#ifndef MOE_V2_SRC_TO_DST_SIMT_H
+#define MOE_V2_SRC_TO_DST_SIMT_H
+
+#include "moe_v2_common.h"
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+using namespace optiling;
+
+class MoeV2SrcToDstOpSimt {
+public:
+    __aicore__ inline MoeV2SrcToDstOpSimt(){};
+    template <typename TilingData>
+    __aicore__ inline void Init(GM_ADDR expandedRowIdx, GM_ADDR expandDstToSrcRow, const TilingData *tilingData);
+    __aicore__ inline void Process();
+
+private:
+    __aicore__ inline void SyncAll();
+    __aicore__ inline void ComputeSimt() const;
+
+private:
+    __gm__ int32_t *expandDstToSrcRowGm_;
+    __gm__ int32_t *expandedRowIdxGm_;
+    const MoeV2GatherOutComputeTilingData *srcToDstTilingData_;
+
+    int64_t coreNum_;
+    int64_t blockIdx_;
+    int64_t totalLength_;
+    int64_t perCoreRows_;
+    int64_t coreRows_;
+    int64_t startIndex_;
+    int64_t threadNum_;
+};
+
+__aicore__ inline void MoeV2SrcToDstOpSimt::SyncAll()
+{
+    if (coreNum_ == 1) {
+        return;
+    }
+    AscendC::SyncAll();
+}
+
+template <typename TilingData>
+__aicore__ inline void MoeV2SrcToDstOpSimt::Init(GM_ADDR expandedRowIdx, GM_ADDR expandDstToSrcRow,
+                                                 const TilingData *tilingData)
+{
+    this->blockIdx_ = get_block_idx() + get_subblockid() * get_block_num();
+    this->coreNum_ = tilingData->coreNum;
+    this->totalLength_ = tilingData->n * tilingData->k;
+    this->srcToDstTilingData_ = &(tilingData->srcToDstComputeParamsOp);
+    this->perCoreRows_ = this->srcToDstTilingData_->perCoreRows;
+    if (this->blockIdx_ == this->srcToDstTilingData_->needCoreNum - 1) {
+        this->coreRows_ = this->srcToDstTilingData_->lastCoreRows;
+    } else {
+        this->coreRows_ = this->srcToDstTilingData_->perCoreRows;
+    }
+    startIndex_ = this->blockIdx_ * this->perCoreRows_;
+    this->threadNum_ = THREAD_NUM < this->coreRows_ ? THREAD_NUM : this->coreRows_;
+
+    expandedRowIdxGm_ = (__gm__ int32_t *)expandedRowIdx;
+    expandDstToSrcRowGm_ = (__gm__ int32_t *)expandDstToSrcRow + Align(this->totalLength_, sizeof(int32_t));
+}
+
+__aicore__ inline void MoeV2SrcToDstOpSimt::ComputeSimt() const
+{
+    for (int32_t index = static_cast<int32_t>(Simt::GetThreadIdx()); index < static_cast<int32_t>(this->coreRows_);
+         index += static_cast<int32_t>(Simt::GetThreadNum())) {
+        int64_t srcIndex = index + this->startIndex_;
+        int64_t dstIndex = expandDstToSrcRowGm_[srcIndex];
+        expandedRowIdxGm_[dstIndex] = srcIndex;
+    }
+}
+
+__aicore__ inline void MoeV2SrcToDstOpSimt::Process()
+{
+    if (this->blockIdx_ < this->srcToDstTilingData_->needCoreNum) {
+        ParallelEXE(this->threadNum_, ComputeSimt);
+    }
+    this->SyncAll();
+}
+} // namespace MoeInitRoutingV2
+#endif // MOE_SRC_TO_DST_SIMT_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/moe_v2_src_to_dst_with_capacity.h

@@ -0,0 +1,279 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+ * CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_v2_src_to_dst_with_capacity.h
+ * \brief
+ */
+#ifndef MOE_V2_SRC_TO_DST_WITH_CAPACITY_H
+#define MOE_V2_SRC_TO_DST_WITH_CAPACITY_H
+
+#include "moe_v2_common.h"
+
+namespace MoeInitRoutingV2 {
+using namespace AscendC;
+using namespace optiling;
+
+template <typename T, typename TilingData>
+class MoeV2SrcToDstWithCapacity {
+public:
+    __aicore__ inline MoeV2SrcToDstWithCapacity(){};
+    __aicore__ inline void Init(GM_ADDR expandedRowIdx, GM_ADDR expandedX, GM_ADDR workspace,
+                                const TilingData *tilingData, TPipe *tPipe);
+    __aicore__ inline void Process();
+
+private:
+    __aicore__ inline void CopyIn(int64_t progress);
+    __aicore__ inline void CopyOut(int64_t progress);
+    __aicore__ inline void CopyOutRemain();
+    __aicore__ inline void SyncAll();
+    __aicore__ inline void AssistInit();
+
+private:
+    TPipe *pipe;
+    TQue<QuePosition::VECIN, 1> copyInQueue;
+    TQue<QuePosition::VECOUT, 1> copyOutQueue;
+    TQue<QuePosition::VECOUT, 1> copyOutZeroQueue;
+
+    GlobalTensor<int32_t> expandDstToSrcRowGm;
+    GlobalTensor<int32_t> expandedRowIdxGm;
+    GlobalTensor<int32_t> expertIdxValueGm;
+    GlobalTensor<int32_t> expandedExpertIdxGm;
+    GlobalTensor<T> expandedXGm;
+
+    LocalTensor<T> outTmpLocal;
+
+    const MoeV2GatherOutComputeTilingData *srcToDstTilingData;
+
+    int64_t coreNum;
+    int64_t blockIdx;
+    int64_t totalLength;
+    int64_t currentLoopRows;
+    int64_t coreRows;
+    int64_t perLoopRows;
+    int64_t lastLoopRows;
+    int64_t rowLoops;
+    int64_t expertCapacity;
+    int64_t expertNum;
+    int64_t cols;
+    int64_t perLoopCols;
+    int64_t lastLoopCols;
+    int64_t colLoops;
+
+    int64_t tokenCount = 0;
+    int32_t lastExpertId = -1;
+    int32_t lastCoreExpertId = 0;
+    int32_t lastCoreExpertIdNum = 0;
+};
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::AssistInit()
+{
+    if constexpr (IsSameType<T, int8_t>::value) {
+        LocalTensor<int16_t> outLocal = copyOutZeroQueue.AllocTensor<int16_t>();
+        Duplicate<int16_t>(outLocal, static_cast<int16_t>(0), this->perLoopCols);
+        copyOutZeroQueue.EnQue<int16_t>(outLocal);
+    } else {
+        LocalTensor<T> outLocal = copyOutZeroQueue.AllocTensor<T>();
+        Duplicate<T>(outLocal, static_cast<T>(0), this->perLoopCols);
+        copyOutZeroQueue.EnQue<T>(outLocal);
+    }
+
+    if (this->blockIdx != 0) {
+        this->lastCoreExpertId = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2);
+        this->lastCoreExpertIdNum = expertIdxValueGm.GetValue((this->blockIdx - 1) * 2 + 1);
+        for (int64_t i = this->blockIdx - 2; i >= 0; i--) {
+            int32_t lastExpertIdx = expertIdxValueGm.GetValue(i * 2);
+            if (lastExpertIdx < this->lastCoreExpertId) {
+                break;
+            }
+            int32_t lastExpertNum = expertIdxValueGm.GetValue(i * 2 + 1);
+            this->lastCoreExpertIdNum += lastExpertNum;
+        }
+    }
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::CopyIn(int64_t progress)
+{
+    LocalTensor<int32_t> inLocal = copyInQueue.AllocTensor<int32_t>();
+    int64_t length = Align(currentLoopRows, sizeof(int32_t));
+    DataCopy(inLocal, expandDstToSrcRowGm[progress * perLoopRows], length);
+    DataCopy(inLocal[length], expandedExpertIdxGm[progress * perLoopRows], length);
+    copyInQueue.EnQue<int32_t>(inLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::CopyOut(int64_t progress)
+{
+    LocalTensor<int32_t> inLocal = copyInQueue.DeQue<int32_t>();
+    LocalTensor<int32_t> outLocal = copyOutQueue.AllocTensor<int32_t>();
+    int64_t length = Align(currentLoopRows, sizeof(int32_t));
+    DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(int32_t)), 0, 0, 0};
+
+    SetWaitFlag<HardEvent::MTE2_S>(HardEvent::MTE2_S);
+    if (this->lastExpertId == -1) {
+        this->lastExpertId = this->lastCoreExpertId;
+        this->tokenCount = this->lastCoreExpertIdNum;
+    }
+    for (int64_t idx = 0; idx < currentLoopRows; idx++) {
+        int32_t expertIdx = inLocal[length].GetValue(idx);
+        SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+        int32_t index = 0;
+        while (this->lastExpertId < expertIdx) {
+            while (this->tokenCount < this->expertCapacity) {
+                index = this->lastExpertId * this->expertCapacity + this->tokenCount;
+                int64_t col = this->perLoopCols;
+                for (int64_t i = 0; i < this->colLoops; i++) {
+                    if (i == this->colLoops - 1) {
+                        col = this->lastLoopCols;
+                    }
+#ifdef __CCE_KT_TEST__
+                    // CPU孪生调试无法使用多核同步，可能导致index为未初始化的脏数据，因此需要特殊处理
+                    if (index * this->cols + i * this->perLoopCols + col * sizeof(T) > expandedXGm.GetSize()) {
+                        continue;
+                    }
+#endif
+                    DataCopyExtParams copyParams1{static_cast<uint16_t>(1), static_cast<uint32_t>(col * sizeof(T)), 0,
+                                                  0, 0};
+                    DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal,
+                                copyParams1);
+                    SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+                }
+                this->tokenCount++;
+            }
+            this->tokenCount = 0;
+            this->lastExpertId++;
+        }
+
+        if (this->tokenCount < this->expertCapacity) {
+            int32_t outOffset = inLocal.GetValue(idx);
+            index = expertIdx * this->expertCapacity + this->tokenCount;
+            outLocal.SetValue(0, index);
+            SetWaitFlag<HardEvent::S_MTE3>(HardEvent::S_MTE3);
+            DataCopyPad(expandedRowIdxGm[outOffset], outLocal, copyParams);
+            SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+            this->tokenCount++;
+        }
+    }
+    copyInQueue.FreeTensor(inLocal);
+    copyOutQueue.FreeTensor(outLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::CopyOutRemain()
+{
+    if (this->blockIdx != this->srcToDstTilingData->needCoreNum - 1) {
+        copyOutZeroQueue.FreeTensor(this->outTmpLocal);
+        return;
+    }
+    while (this->lastExpertId < this->expertNum) {
+        while (this->tokenCount < this->expertCapacity) {
+            int32_t index = this->lastExpertId * this->expertCapacity + this->tokenCount;
+            int64_t col = this->perLoopCols;
+            for (int64_t i = 0; i < this->colLoops; i++) {
+                if (i == this->colLoops - 1) {
+                    col = this->lastLoopCols;
+                }
+                DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(col * sizeof(T)), 0, 0, 0};
+                DataCopyPad(expandedXGm[index * this->cols + i * this->perLoopCols], this->outTmpLocal, copyParams);
+                SetWaitFlag<HardEvent::MTE3_S>(HardEvent::MTE3_S);
+            }
+            this->tokenCount++;
+        }
+        this->tokenCount = 0;
+        this->lastExpertId++;
+    }
+    copyOutZeroQueue.FreeTensor(this->outTmpLocal);
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::SyncAll()
+{
+    if (coreNum == 1) {
+        return;
+    }
+#ifndef __CCE_KT_TEST__
+    AscendC::SyncAll();
+#endif
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::Init(GM_ADDR expandedRowIdx, GM_ADDR expandedX,
+                                                                      GM_ADDR workspace, const TilingData *tilingData,
+                                                                      TPipe *tPipe)
+{
+    int64_t blockNum = GetBlockNum() * 2;
+    this->pipe = tPipe;
+    this->blockIdx = get_block_idx() + get_subblockid() * get_block_num();
+
+    this->coreNum = tilingData->coreNum;
+    this->totalLength = tilingData->n * tilingData->k;
+    this->srcToDstTilingData = &(tilingData->srcToDstCapacityComputeParamsOp);
+    this->expertNum = tilingData->expertNum;
+    this->expertCapacity = tilingData->expertCapacity;
+    this->cols = tilingData->cols;
+
+    if (this->blockIdx == this->srcToDstTilingData->needCoreNum - 1) {
+        this->coreRows = this->srcToDstTilingData->lastCoreRows;
+        this->perLoopRows = this->srcToDstTilingData->lastCorePerLoopRows;
+        this->lastLoopRows = this->srcToDstTilingData->lastCoreLastLoopRows;
+        this->rowLoops = this->srcToDstTilingData->lastCoreLoops;
+    } else {
+        this->coreRows = this->srcToDstTilingData->perCoreRows;
+        this->perLoopRows = this->srcToDstTilingData->perCorePerLoopRows;
+        this->lastLoopRows = this->srcToDstTilingData->perCoreLastLoopRows;
+        this->rowLoops = this->srcToDstTilingData->perCoreLoops;
+    }
+    this->perLoopCols = this->srcToDstTilingData->perLoopCols;
+    this->lastLoopCols = this->srcToDstTilingData->lastLoopCols;
+    this->colLoops = this->srcToDstTilingData->colLoops;
+
+    int64_t length = Align(this->totalLength, sizeof(int32_t));
+    expandedRowIdxGm.SetGlobalBuffer((__gm__ int32_t *)expandedRowIdx, length);
+    expandedXGm.SetGlobalBuffer((__gm__ T *)expandedX, this->expertNum * this->expertCapacity * this->cols);
+
+    expandedExpertIdxGm.SetGlobalBuffer((__gm__ int32_t *)workspace +
+                                            this->blockIdx * this->srcToDstTilingData->perCoreRows,
+                                        Align(this->coreRows, sizeof(int32_t)));
+    expandDstToSrcRowGm.SetGlobalBuffer((__gm__ int32_t *)workspace + length +
+                                            this->blockIdx * this->srcToDstTilingData->perCoreRows,
+                                        Align(this->coreRows, sizeof(int32_t)));
+    expertIdxValueGm.SetGlobalBuffer((__gm__ int32_t *)workspace + length * 2, this->coreNum * 2);
+
+    pipe->InitBuffer(copyInQueue, 1, AlignBytes(this->perLoopRows, sizeof(int32_t)) * 2);
+    pipe->InitBuffer(copyOutQueue, 1, AlignBytes(INT32_ONE_BLOCK_NUM, sizeof(int32_t)));
+    if constexpr (IsSameType<T, int8_t>::value) {
+        pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(int16_t)));
+    } else {
+        pipe->InitBuffer(copyOutZeroQueue, 1, AlignBytes(this->perLoopCols, sizeof(T)));
+    }
+}
+
+template <typename T, typename TilingData>
+__aicore__ inline void MoeV2SrcToDstWithCapacity<T, TilingData>::Process()
+{
+    if (this->blockIdx < this->srcToDstTilingData->needCoreNum) {
+        AssistInit();
+        this->outTmpLocal = copyOutZeroQueue.DeQue<T>();
+        currentLoopRows = perLoopRows;
+        for (int64_t loop = 0; loop < this->rowLoops; loop++) {
+            if (loop == this->rowLoops - 1) {
+                currentLoopRows = lastLoopRows;
+            }
+            CopyIn(loop);
+            CopyOut(loop);
+        }
+        CopyOutRemain();
+    }
+    this->SyncAll();
+}
+} // namespace MoeInitRoutingV2
+#endif // MOE_V2_SRC_TO_DST_WITH_CAPACITY_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/moe_init_routing_v2/tiling_base.h

@@ -0,0 +1,66 @@
+#pragma once
+namespace optiling {
+struct AiCoreParams {
+    uint64_t ubSize;
+    uint64_t blockDim;
+    uint64_t aicNum;
+    
+    uint64_t l1Size;
+    uint64_t l0aSize;
+    uint64_t l0bSize;
+    uint64_t l0cSize;
+};
+
+class TilingBaseClass {
+public:
+    bool DoTiling(
+            int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, 
+            int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag,
+            bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0, 
+            int64_t aivCoreNum, int64_t ubSizePlatForm) 
+    {
+        bool ret = GetShapeAttrsInfo(m, cols, topK, expertCapacity, expertNum, activeNum, dropPadMode, expertTokensCountOrCumsumFlag,
+            expertTokensBeforeCapacityFlag, inuptXDtypeSize, quantMode, scaleDim0);
+
+        if (!ret){
+            return ret;
+        }
+        ret = GetPlatformInfo(aivCoreNum, ubSizePlatForm);
+        if (!ret){
+            return ret;
+        }
+        ret = DoOpTiling();
+        if (!ret){
+            return ret;
+        }
+        ret = GetWorkspaceSize();
+        if (!ret){
+            return ret;
+        }
+        // ret = PostTiling();
+        // if (!ret){
+        //     return ret;
+        // }
+        tilingKey_ = GetTilingKey();
+        
+        return true;
+    }
+
+//protected:
+    virtual bool GetPlatformInfo(int64_t aivCoreNum, int64_t ubSizePlatForm) = 0;
+    virtual bool GetShapeAttrsInfo(int64_t m, int64_t cols, int64_t topK, int64_t expertCapacity, 
+        int64_t expertNum, int64_t activeNum, int64_t dropPadMode, int64_t expertTokensCountOrCumsumFlag,
+        bool expertTokensBeforeCapacityFlag, int64_t inuptXDtypeSize, int64_t quantMode, int64_t scaleDim0) = 0;
+
+    virtual bool DoOpTiling() = 0;
+    virtual bool GetWorkspaceSize() = 0;
+    // virtual bool PostTiling() = 0;
+    virtual uint64_t GetTilingKey() const = 0;
+//protected:
+    uint32_t blockDim_{0};
+    uint64_t workspaceSize_{0};
+    uint64_t tilingKey_{0};
+    AiCoreParams aicoreParams_{0, 0, 0, 0, 0, 0, 0};
+};
+
+}

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/unpermute/moe_token_unpermute.h

@@ -0,0 +1,376 @@
+/**
+ * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*!
+ * \file moe_token_unpermute.h
+ * \brief
+ */
+
+#ifndef MOE_TOKEN_UNPERMUTE
+#define MOE_TOKEN_UNPERMUTE
+
+#include "kernel_operator.h"
+#include "moe_token_unpermute_tiling.h"
+using namespace AscendC;
+
+
+template <typename T1, typename T2, typename T3, bool PROBS> class KernelMoeTokenUnpermute {
+public:
+    __aicore__ inline KernelMoeTokenUnpermute()
+    {
+    }
+
+    __aicore__ inline void Init(GM_ADDR permuted_tokens, GM_ADDR sorted_indices, GM_ADDR probs,
+                                GM_ADDR unpermuted_tokens, const MoeTokenUnpermuteTilingData *__restrict tiling_data);
+    __aicore__ inline void Process();
+
+protected:
+    __aicore__ inline void CalMultiOutToken(const int64_t out_offset, const int64_t out_tokens_number);
+    __aicore__ inline void CalSingleOutToken(const int64_t start_token, const int64_t out_token_idx);
+    __aicore__ inline void CalPartOutToken(const int64_t start_token, const int64_t h_index, const int64_t h_length,
+                                           const int64_t out_token_index);
+    __aicore__ inline void CopyTokenIn(const T2 in_token_index, const int64_t h_index, const int64_t h_length);
+    __aicore__ inline void CalFirstToken(const float prob_value, const int64_t h_length);
+    __aicore__ inline void CalToken(const float prob_value, const int64_t h_length);
+    __aicore__ inline void CopyOut(const int64_t out_token_index, const int64_t h_index, const int64_t h_length);
+
+    TPipe pipe;
+    TQue<QuePosition::VECIN, 1> tokens_inque, indices_inque, probs_inque;
+    TBuf<TPosition::VECCALC> temp_buffer0, temp_buffer1, temp_buffer2;
+    TQue<QuePosition::VECOUT, 1> outque;
+    GlobalTensor<T1> tokensGM, outGM;
+    GlobalTensor<T2> indicesGM;
+    GlobalTensor<T3> probsGM;
+    LocalTensor<T2> indicesLocal;
+    LocalTensor<float> token_tensor0, token_tensor1, probs_tensor;
+    DataCopyPadExtParams<T1> extParams1{false, 0, 0, 0};
+    DataCopyPadExtParams<T2> extParams2{false, 0, 0, 0};
+    DataCopyPadExtParams<T3> extParams3{false, 0, 0, 0};
+    DataCopyExtParams copyParams{1, 0, 0, 0, 0};
+
+    constexpr static uint32_t BLOCK_SIZE = 32;
+    constexpr static uint32_t ALIGN_512 = 512;
+
+    int64_t hidden_size;
+    int64_t top_k;
+    int64_t num_out_tokens;
+    int64_t hidden_splited_length;
+    int64_t hidden_splited_num;
+    int64_t hidden_splited_remain;
+    int64_t tokens_core_length;
+    int64_t tokens_core_remain;
+    int64_t tokens_splited_length;
+    int64_t tokens_splited_num;
+    int64_t tokens_splited_remain;
+    int32_t blockIdx;
+    int32_t blockNum;
+};
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void
+KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::Init(GM_ADDR permuted_tokens, GM_ADDR sorted_indices, GM_ADDR probs,
+                                                 GM_ADDR unpermuted_tokens,
+                                                 const MoeTokenUnpermuteTilingData *__restrict tiling_data)
+{
+    this->blockIdx = get_block_idx();
+    this->blockNum = get_block_num();
+
+    if (blockIdx >= blockNum) {
+        return;
+    }
+    ASSERT(blockNum != 0 && "block dim can not be zero!");
+    // row_input
+    this->hidden_size = tiling_data->hidden_size;
+    this->top_k = tiling_data->top_k;
+    this->num_out_tokens = tiling_data->num_out_tokens;
+    // hidden_tiling
+    this->hidden_splited_length = tiling_data->hidden_splited_length;
+    this->hidden_splited_num = tiling_data->hidden_splited_num;
+    this->hidden_splited_remain = tiling_data->hidden_splited_remain;
+    // token_tiling
+    this->tokens_core_length = tiling_data->tokens_core_length;
+    this->tokens_core_remain = tiling_data->tokens_core_remain;
+    this->tokens_splited_length = tiling_data->tokens_splited_length;
+    this->tokens_splited_num = tiling_data->tokens_splited_num;
+    this->tokens_splited_remain = tiling_data->tokens_splited_remain;
+
+    // Handle the tail block for token_by_core
+    if (this->tokens_core_remain > 0 && blockIdx < this->tokens_core_remain) {
+        this->tokens_core_length += 1;
+        this->tokens_splited_remain += 1;
+    }
+
+    int64_t hidden_splited_length_align512 = (this->hidden_splited_length + ALIGN_512 - 1) & ~(ALIGN_512 - 1);
+
+    int64_t block_length = this->tokens_core_length * this->top_k;
+    int64_t block_splited_length = this->tokens_splited_length * this->top_k;
+
+    int64_t block_offset;
+    if (this->tokens_core_remain > 0) {
+        if (blockIdx < this->tokens_core_remain) {
+            block_offset = block_length * blockIdx;
+        } else {
+            block_offset = (block_length + this->top_k) * this->tokens_core_remain +
+                           block_length * (blockIdx - this->tokens_core_remain);
+        }
+    } else {
+        block_offset = block_length * blockIdx;
+    }
+
+    this->tokensGM.SetGlobalBuffer((__gm__ T1 *)permuted_tokens);
+    this->indicesGM.SetGlobalBuffer((__gm__ T2 *)sorted_indices + block_offset, block_length);
+
+
+    int64_t out_block_offset;
+    if (this->tokens_core_remain > 0) {
+        if (blockIdx < this->tokens_core_remain) {
+            out_block_offset = this->tokens_core_length * blockIdx * hidden_size;
+        } else {
+            out_block_offset = (this->tokens_core_length + 1) * this->tokens_core_remain +
+                               this->tokens_core_length * (blockIdx - this->tokens_core_remain);
+            out_block_offset *= this->hidden_size;
+        }
+    } else {
+        out_block_offset = this->tokens_core_length * blockIdx * hidden_size;
+    }
+
+    this->outGM.SetGlobalBuffer((__gm__ T1 *)unpermuted_tokens + out_block_offset,
+                                this->tokens_core_length * this->hidden_size);
+
+    this->pipe.InitBuffer(tokens_inque, tiling_data->buffer_num, hidden_splited_length_align512 * sizeof(T1));
+    this->pipe.InitBuffer(indices_inque, 1, block_splited_length * (sizeof(T2)));
+    this->pipe.InitBuffer(outque, 1, hidden_splited_length_align512 * sizeof(T1));
+
+    if constexpr (!IsSameType<T1, float>::value) {
+        this->pipe.InitBuffer(temp_buffer0, hidden_splited_length_align512 * sizeof(float) + 256);
+        this->pipe.InitBuffer(temp_buffer1, hidden_splited_length_align512 * sizeof(float));
+        this->token_tensor0 = this->temp_buffer0.template Get<float>();
+        this->token_tensor1 = this->temp_buffer1.template Get<float>();
+    }
+
+    if constexpr (PROBS) {
+        this->probsGM.SetGlobalBuffer((__gm__ T3 *)probs + block_offset, block_length);
+        this->pipe.InitBuffer(probs_inque, 1, block_splited_length * (sizeof(T3)));
+        if constexpr (!IsSameType<T3, float>::value) {
+            this->pipe.InitBuffer(temp_buffer2, block_splited_length * sizeof(float));
+            this->probs_tensor = this->temp_buffer2.template Get<float>();
+        }
+    }
+};
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::Process()
+{
+
+    if (blockIdx >= blockNum) {
+        return;
+    }
+    for (int64_t i = 0; i < this->tokens_splited_num; ++i) {
+        CalMultiOutToken(i * this->tokens_splited_length, this->tokens_splited_length);
+    }
+    // Handle the tail block when tokens_num is not evenly divisible by core count
+    if (this->tokens_splited_remain > 0) {
+        CalMultiOutToken(this->tokens_splited_num * this->tokens_splited_length, this->tokens_splited_remain);
+    }
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CalMultiOutToken(const int64_t out_offset,
+                                                                                    const int64_t out_tokens_number)
+{
+    this->indicesLocal = this->indices_inque.template AllocTensor<T2>();
+    int64_t in_offset = out_offset * this->top_k;
+    this->copyParams.blockLen = out_tokens_number * this->top_k * sizeof(T2);
+    DataCopyPad(this->indicesLocal, this->indicesGM[in_offset], this->copyParams, this->extParams2);
+    this->indices_inque.template EnQue(this->indicesLocal);
+
+    if constexpr (PROBS) {
+        LocalTensor<T3> temp_probs_tensor = this->probs_inque.template AllocTensor<T3>();
+        this->copyParams.blockLen = out_tokens_number * this->top_k * sizeof(T3);
+        DataCopyPad(temp_probs_tensor, this->probsGM[in_offset], this->copyParams, this->extParams3);
+        this->probs_inque.template EnQue(temp_probs_tensor);
+        temp_probs_tensor = this->probs_inque.template DeQue<T3>();
+        if constexpr (!IsSameType<T3, float>::value) {
+            Cast(this->probs_tensor, temp_probs_tensor, RoundMode::CAST_NONE, out_tokens_number * this->top_k);
+            this->probs_inque.FreeTensor(temp_probs_tensor);
+            PipeBarrier<PIPE_V>();
+        } else {
+            this->probs_tensor = temp_probs_tensor;
+        }
+    }
+    this->indicesLocal = this->indices_inque.template DeQue<T2>();
+
+    
+    for (int64_t out_token_idx = 0; out_token_idx < out_tokens_number; ++out_token_idx) {
+        CalSingleOutToken(out_token_idx * this->top_k, out_offset + out_token_idx);
+    }
+    // Free Tensor
+    this->indices_inque.FreeTensor(this->indicesLocal);
+    if constexpr (PROBS && IsSameType<T3, float>::value) {
+        this->probs_inque.FreeTensor(this->probs_tensor);
+    }
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CalSingleOutToken(const int64_t start_token,
+                                                                                     const int64_t out_token_idx)
+{
+    for (int64_t h_index = 0; h_index < this->hidden_splited_num; ++h_index) {
+        CalPartOutToken(start_token, h_index, this->hidden_splited_length, out_token_idx);
+    }
+    // Handle the tail block when a full hidden_size does not fit in one pass
+    if (this->hidden_splited_remain > 0) {
+        CalPartOutToken(start_token, this->hidden_splited_num, this->hidden_splited_remain, out_token_idx);
+    }
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void
+KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CalPartOutToken(const int64_t start_token, const int64_t h_index,
+                                                            const int64_t h_length, const int64_t out_token_index)
+{
+    if constexpr (IsSameType<T1, float>::value) {
+        this->token_tensor0 = this->outque.template AllocTensor<T1>();
+    }
+    int64_t end_token = start_token + this->top_k;
+    T2 cal_token_idx = this->indicesLocal.GetValue(start_token);
+
+    // Handle the first token
+    if (cal_token_idx < this->num_out_tokens) {
+        float probsValue = 0;
+        if constexpr (PROBS) {
+            probsValue = this->probs_tensor.GetValue(start_token);
+        }
+
+        CopyTokenIn(cal_token_idx, h_index, h_length);
+        PipeBarrier<PIPE_V>();
+        CalFirstToken(probsValue, h_length);
+    } else {
+        PipeBarrier<PIPE_V>();
+        Duplicate(this->token_tensor0, static_cast<float>(0), h_length);
+    }
+
+    // Handle the remaining tokens
+    for (int64_t token_index = start_token + 1; token_index < end_token; ++token_index) {
+        cal_token_idx = this->indicesLocal.GetValue(token_index);
+        if (cal_token_idx < this->num_out_tokens) {
+            float probsValue = 0;
+            if constexpr (PROBS) {
+                probsValue = this->probs_tensor.GetValue(token_index);
+            }
+        
+            CopyTokenIn(cal_token_idx, h_index, h_length);
+            PipeBarrier<PIPE_V>();
+            CalToken(probsValue, h_length);
+        }
+    }
+
+    // Write out the computed result
+    CopyOut(out_token_index, h_index, h_length);
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CopyTokenIn(const T2 in_token_index,
+                                                                               const int64_t h_index,
+                                                                               const int64_t h_length)
+{
+    LocalTensor<T1> tokensLocal = this->tokens_inque.template AllocTensor<T1>();
+    int64_t offset = in_token_index * this->hidden_size + h_index * this->hidden_splited_length;
+
+    if (likely((h_length * sizeof(T1)) % BLOCK_SIZE == 0)) {
+        DataCopy(tokensLocal, this->tokensGM[offset], h_length);
+    } else {
+        this->copyParams.blockLen = h_length * sizeof(T1);
+        DataCopyPad(tokensLocal, this->tokensGM[offset], this->copyParams, this->extParams1);
+    }
+
+    this->tokens_inque.template EnQue(tokensLocal);
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CalFirstToken(const float prob_value,
+                                                                                 const int64_t h_length)
+{
+    LocalTensor<T1> tokensLocal = this->tokens_inque.template DeQue<T1>();
+
+    if constexpr (!IsSameType<T1, float>::value) {
+        Cast(this->token_tensor0, tokensLocal, RoundMode::CAST_NONE, h_length);
+    } else {
+        uint64_t byteAlign32 = (h_length * sizeof(float) + BLOCK_SIZE - 1) & ~(BLOCK_SIZE - 1);
+        DataCopy(this->token_tensor0, tokensLocal, byteAlign32 / sizeof(float));
+    }
+
+    this->tokens_inque.FreeTensor(tokensLocal);
+
+    if constexpr (PROBS) {
+        PipeBarrier<PIPE_V>();
+        Muls(this->token_tensor0, this->token_tensor0, prob_value, h_length);
+    }
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CalToken(const float prob_value,
+                                                                            const int64_t h_length)
+{
+    LocalTensor<T1> tokensLocal = this->tokens_inque.template DeQue<T1>();
+
+    if constexpr (!IsSameType<T1, float>::value) {
+        Cast(this->token_tensor1, tokensLocal, RoundMode::CAST_NONE, h_length);
+        this->tokens_inque.FreeTensor(tokensLocal);
+        if constexpr (PROBS) {
+            PipeBarrier<PIPE_V>();
+            Muls(this->token_tensor1, this->token_tensor1, prob_value, h_length);
+        }
+        PipeBarrier<PIPE_V>();
+        Add(this->token_tensor0, this->token_tensor0, this->token_tensor1, h_length);
+    } else {
+        if constexpr (PROBS) {
+            Muls(tokensLocal, tokensLocal, prob_value, h_length);
+            PipeBarrier<PIPE_V>();
+        }
+        Add(this->token_tensor0, this->token_tensor0, tokensLocal, h_length);
+        this->tokens_inque.FreeTensor(tokensLocal);
+    }
+}
+
+template <typename T1, typename T2, typename T3, bool PROBS>
+__aicore__ inline void KernelMoeTokenUnpermute<T1, T2, T3, PROBS>::CopyOut(const int64_t out_token_index,
+                                                                           const int64_t h_index,
+                                                                           const int64_t h_length)
+{
+    LocalTensor<T1> temp_out_tensors;
+    if constexpr (!IsSameType<T1, float>::value) {
+        temp_out_tensors = this->outque.template AllocTensor<T1>();
+        PipeBarrier<PIPE_V>();
+        Cast(temp_out_tensors, this->token_tensor0, RoundMode::CAST_RINT, h_length);
+    } else {
+        temp_out_tensors = this->token_tensor0;
+    }
+
+    this->outque.template EnQue<T1>(temp_out_tensors);
+    temp_out_tensors = this->outque.template DeQue<T1>();
+
+    int64_t offset = out_token_index * this->hidden_size + h_index * this->hidden_splited_length;
+    if (likely((h_length * sizeof(T1)) % BLOCK_SIZE == 0)) {
+        DataCopy(this->outGM[offset], temp_out_tensors, h_length);
+    } else {
+        this->copyParams.blockLen = h_length * sizeof(T1);
+        DataCopyPad(this->outGM[offset], temp_out_tensors, this->copyParams);
+    }
+
+    this->outque.FreeTensor(temp_out_tensors);
+}
+#endif // MOE_TOKEN_UNPERMUTE

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/unpermute/moe_token_unpermute_tiling.h

@@ -0,0 +1,38 @@
+#ifndef MOE_TOKEN_UNPERMUTE_TILING
+#define MOE_TOKEN_UNPERMUTE_TILING
+
+struct MoeTokenUnpermuteTilingData {
+    int64_t hidden_size;
+    int64_t top_k;
+    int64_t num_out_tokens;
+    int64_t hidden_splited_length;
+    int64_t hidden_splited_num;
+    int64_t hidden_splited_remain;
+    int64_t tokens_core_length;
+    int64_t tokens_core_remain;
+    int64_t tokens_splited_length;
+    int64_t tokens_splited_num;
+    int64_t tokens_splited_remain;
+    int64_t buffer_num;
+};
+
+__forceinline__ [host, aicore] void
+MoeTokenUnpermuteTiling(int32_t m, int32_t n, int32_t topK, MoeTokenUnpermuteTilingData &tilingData, uint32_t coreNum)
+{
+    #define I64(x) static_cast<int64_t>(x)
+    tilingData.hidden_size = I64(n);
+    tilingData.top_k = I64(topK);
+    tilingData.num_out_tokens = I64(m);
+    tilingData.hidden_splited_length = tilingData.hidden_size;
+    tilingData.hidden_splited_num = 1;
+    tilingData.hidden_splited_remain = 0;
+    uint32_t outTokens = m / topK;
+    tilingData.tokens_core_length = I64(outTokens / coreNum);
+    tilingData.tokens_core_remain = I64(outTokens % coreNum);
+    tilingData.tokens_splited_length = I64(min(tilingData.tokens_core_length, 600));
+    tilingData.tokens_splited_num = I64(tilingData.tokens_core_length / tilingData.tokens_splited_length);
+    tilingData.tokens_splited_remain = I64(tilingData.tokens_core_length % tilingData.tokens_splited_length);
+    tilingData.buffer_num = 4;
+}
+
+#endif

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/block_epilogue_pertoken_row.hpp

@@ -0,0 +1,208 @@
+/*
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_ROW_HPP
+#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_ROW_HPP
+
+#include "catlass/catlass.hpp"
+#include "catlass/arch/resource.hpp"
+#include "catlass/epilogue/dispatch_policy.hpp"
+#include "catlass/gemm_coord.hpp"
+#include "catlass/matrix_coord.hpp"
+#include "catlass/layout/layout.hpp"
+#include "catlass/detail/callback.hpp"
+#include "catlass/epilogue/block/block_epilogue.hpp"
+
+namespace Catlass::Epilogue::Block {
+
+// float scale, dequant per expert
+template <
+    uint32_t UB_STAGES_,
+    class CType_,
+    class LayoutPerTokenScale_,
+    class DType_,
+    class TileCopy_
+>
+class BlockEpilogue <
+    EpilogueAtlasA2PerTokenDequant<UB_STAGES_>,
+    CType_,
+    Gemm::GemmType<float, LayoutPerTokenScale_>,
+    DType_,
+    TileCopy_
+> {
+public:
+    using DispatchPolicy = EpilogueAtlasA2PerTokenDequant<UB_STAGES_>;
+    using ArchTag = typename DispatchPolicy::ArchTag;
+    static constexpr uint32_t UB_STAGES = UB_STAGES_;
+
+    // Data infos
+    using ElementC = typename CType_::Element;
+    using LayoutC = typename CType_::Layout;
+    using ElementPerTokenScale = float;
+    using LayoutPerTokenScale = LayoutPerTokenScale_;
+    using ElementD = typename DType_::Element;
+    using LayoutD = typename DType_::Layout;
+
+    // Check data infos
+    static_assert(
+        (std::is_same_v<ElementC, half> || std::is_same_v<ElementC, bfloat16_t>) && 
+        (std::is_same_v<ElementD, half> || std::is_same_v<ElementD, bfloat16_t>),
+        "The element type template parameters of BlockEpilogue are wrong"
+    );
+    static_assert(
+        std::is_same_v<LayoutC, layout::RowMajor> && 
+            std::is_same_v<LayoutPerTokenScale, layout::VectorLayout> && std::is_same_v<LayoutD, layout::RowMajor>,
+        "The layout template parameters of BlockEpilogue are wrong"
+    );
+
+
+    // Tile copy
+    using CopyGmToUbC = typename TileCopy_::CopyGmToUbC;
+    using CopyUbToGmD = typename TileCopy_::CopyUbToGmD;
+
+    struct Params {
+        __gm__ int32_t *ptrTokenPerExpert{nullptr};
+        int32_t EP;
+        int32_t expertPerRank;
+
+        CATLASS_DEVICE
+        Params() {};
+
+        CATLASS_DEVICE
+        Params(int32_t EP_, int32_t expertPerRank_, __gm__ int32_t *ptrTokenPerExpert_) : ptrTokenPerExpert(ptrTokenPerExpert_), EP(EP_), expertPerRank(expertPerRank_) {}
+    };
+
+    CATLASS_DEVICE
+    BlockEpilogue(Arch::Resource<ArchTag> const &resource, Params const &params = Params{}) : params(params)
+    {
+        size_t ubOffset = 4096;
+        int32_t eventVMTE2 = 0;
+        int32_t eventMTE2V = 0;
+        int32_t eventMTE3V = 0;
+        int32_t eventVMTE3 = 0;
+        constexpr int32_t blockN = 12000;
+        for (uint32_t i = 0; i < UB_STAGES; ++i) {
+            ubCList[i] = resource.ubBuf.template GetBufferByByte<ElementC>(ubOffset);
+            ubOffset += blockN * sizeof(ElementC);
+            ubDList[i] = resource.ubBuf.template GetBufferByByte<ElementD>(ubOffset);
+            ubOffset += blockN * sizeof(ElementD);
+
+            eventUbCVMTE2List[i] = eventVMTE2++;
+            eventUbCMTE2VList[i] = eventMTE2V++;
+            eventUbDMTE3VList[i] = eventMTE3V++;
+            eventUbDVMTE3List[i] = eventVMTE3++;
+
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[i]);
+            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[i]);
+            ubCFp32List[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += blockN * sizeof(float);
+        }
+    }
+    CATLASS_DEVICE
+    void Finalize()
+    {
+        for (uint32_t i = 0; i < UB_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[i]);
+            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[i]);
+        }
+    }
+    CATLASS_DEVICE
+    ~BlockEpilogue()
+    {
+        
+    }
+
+    CATLASS_DEVICE
+    void UpdateParams(Params const &params_)
+    {
+        params = params_;
+    }
+
+    CATLASS_DEVICE
+    void operator() (
+        AscendC::GlobalTensor<ElementC> const &gmC,
+        MatrixCoord const &shapeC,
+        AscendC::GlobalTensor<ElementPerTokenScale> const &gmPerTokenScale,
+        AscendC::GlobalTensor<ElementD> const &gmD
+    )
+    {
+        uint32_t blockM = shapeC.row();
+        uint32_t blockN = shapeC.column();
+
+        uint32_t tileLoops = blockM;
+
+        for (uint32_t loopIdx = 0; loopIdx < tileLoops; loopIdx ++) {
+            auto gmTileC = gmC[loopIdx * blockN];
+            auto &ubC = ubCList[ubListId];
+            auto &ubCFp32 = ubCFp32List[ubListId];
+            auto &ubMul = ubMulList[ubListId];
+            auto &ubD = ubDList[ubListId];
+            auto gmTileD = gmD[loopIdx * blockN];
+            LayoutC layoutUbC{1, blockN};
+
+            // Move C from GM workspace to UB
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[ubListId]);
+            copyGmToUbC(ubC, gmTileC, layoutUbC, layoutUbC);
+            AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(eventUbCMTE2VList[ubListId]);
+
+            // Cast C to FP32 in UB
+            AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(eventUbCMTE2VList[ubListId]);
+            AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_NONE, blockN);
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[ubListId]);
+
+            // Get per-token scale from row loopIdx of gmPerTokenScale
+            ElementPerTokenScale perTokenScale = gmPerTokenScale(loopIdx);
+
+            AscendC::SetFlag<AscendC::HardEvent::S_V>(0);
+            AscendC::WaitFlag<AscendC::HardEvent::S_V>(0);
+            // Multiply FP32 C by the per-token scale
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Muls(ubCFp32, ubCFp32, perTokenScale, blockN);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            // Cast the muls result back to fp16/bf16
+            LayoutD layoutUbD{1, blockN};
+            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[ubListId]);
+
+            AscendC::Cast(ubD, ubCFp32, AscendC::RoundMode::CAST_RINT, blockN);
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(eventUbDVMTE3List[ubListId]);
+            
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(eventUbDVMTE3List[ubListId]);
+            copyUbToGmD(gmTileD, ubD, layoutUbD, layoutUbD);
+            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[ubListId]);
+
+            ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0;
+        }
+    }
+
+private:
+    Params params;
+
+    AscendC::LocalTensor<ElementC> ubCList[UB_STAGES];
+    AscendC::LocalTensor<ElementD> ubDList[UB_STAGES];
+
+    int32_t eventUbCVMTE2List[UB_STAGES];
+    int32_t eventUbCMTE2VList[UB_STAGES];
+    int32_t eventUbDMTE3VList[UB_STAGES];
+    int32_t eventUbDVMTE3List[UB_STAGES];
+
+    uint32_t ubListId{0};
+
+    AscendC::LocalTensor<float> ubCFp32List[UB_STAGES];
+    AscendC::LocalTensor<float> ubMulList[UB_STAGES];
+
+
+    CopyGmToUbC copyGmToUbC;
+    CopyUbToGmD copyUbToGmD;
+};
+
+}  // namespace Catlass::Epilogue::Block
+
+#endif  // CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_ROW_HPP

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/block_epilogue_pertoken_swiglu.hpp

@@ -0,0 +1,402 @@
+/*
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_SWIGLU_HPP
+#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_SWIGLU_HPP
+
+#include "catlass/catlass.hpp"
+#include "catlass/arch/resource.hpp"
+#include "catlass/epilogue/dispatch_policy.hpp"
+#include "catlass/gemm_coord.hpp"
+#include "catlass/matrix_coord.hpp"
+#include "catlass/layout/layout.hpp"
+#include "catlass/detail/callback.hpp"
+
+namespace Catlass::Epilogue::Block {
+
+// float scale, dequant per expert
+template <
+    uint32_t UB_STAGES_,
+    class CType_,
+    class LayoutPerTokenScale_,
+    class DType_,
+    class TileElemWiseMuls_,
+    class TileCopy_
+>
+class BlockEpilogue <
+    EpilogueAtlasA2PerTokenDequantSwigluQuant<UB_STAGES_>,
+    CType_,
+    Gemm::GemmType<float, LayoutPerTokenScale_>,
+    DType_,
+    TileElemWiseMuls_,
+    TileCopy_
+> {
+public:
+    using DispatchPolicy = EpilogueAtlasA2PerTokenDequantSwigluQuant<UB_STAGES_>;
+    using ArchTag = typename DispatchPolicy::ArchTag;
+    static constexpr uint32_t UB_STAGES = UB_STAGES_;
+
+    // Data infos
+    using ElementC = typename CType_::Element;
+    using LayoutC = typename CType_::Layout;
+    using ElementPerTokenScale = float;
+    using LayoutPerTokenScale = LayoutPerTokenScale_;
+    using ElementD = typename DType_::Element;
+    using LayoutD = typename DType_::Layout;
+
+    // Check data infos
+    static_assert(
+        (std::is_same_v<ElementC, half> || std::is_same_v<ElementC, bfloat16_t>) &&
+        (std::is_same_v<ElementD, float> || std::is_same_v<ElementD, int8_t> || std::is_same_v<ElementD, half> || std::is_same_v<ElementD, bfloat16_t>),
+        "The element type template parameters of BlockEpilogue are wrong"
+    );
+    static_assert(
+        std::is_same_v<LayoutC, layout::RowMajor> &&
+            std::is_same_v<LayoutPerTokenScale, layout::VectorLayout> && std::is_same_v<LayoutD, layout::RowMajor>,
+        "The layout template parameters of BlockEpilogue are wrong"
+    );
+
+    // Tile copy
+    using CopyGmToUbC = typename TileCopy_::CopyGmToUbC;
+    using CopyUbToGmD = typename TileCopy_::CopyUbToGmD;
+    using CopyUbToGmDequantScale = Epilogue::Tile::CopyUb2Gm<ArchTag, Gemm::GemmType<ElementPerTokenScale, LayoutPerTokenScale>>;
+
+    struct Params {
+        __gm__ ElementPerTokenScale *ptrPerTokenScale{nullptr};
+        LayoutPerTokenScale layoutPerTokenScale{};
+        __gm__ ElementD *ptrD{nullptr};
+        LayoutD layoutD{};
+
+        CATLASS_DEVICE
+        Params() {};
+
+        CATLASS_DEVICE
+        Params(__gm__ ElementPerTokenScale *ptrPerTokenScale_, LayoutPerTokenScale const &layoutPerTokenScale_,
+            __gm__ ElementD *ptrD_, LayoutD const &layoutD_
+        ) : ptrPerTokenScale(ptrPerTokenScale_), layoutPerTokenScale(layoutPerTokenScale_),
+            ptrD(ptrD_), layoutD(layoutD_) {}
+    };
+
+    CATLASS_DEVICE
+    BlockEpilogue(Arch::Resource<ArchTag> const &resource, int32_t n, Params const &params = Params{}) : params(params)
+    {
+        size_t ubOffset = 0;
+        int32_t eventVMTE2 = 0;
+        int32_t eventMTE2V = 0;
+        int32_t eventMTE3V = 0;
+        int32_t eventVMTE3 = 0;
+        uint32_t blockN = n;
+        uint32_t ChunkTileLen = blockN / 2;
+        uint32_t HalfChunkTileLen = ChunkTileLen / 2;
+
+        for (uint32_t i = 0; i < UB_STAGES; ++i) {
+            ubCList[i] = resource.ubBuf.template GetBufferByByte<ElementC>(ubOffset);
+            ubOffset += blockN * sizeof(ElementC);
+            ubDList[i] = resource.ubBuf.template GetBufferByByte<ElementD>(ubOffset);
+            ubOffset += blockN * sizeof(ElementD);
+            ubCFp32List[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += blockN * sizeof(float);
+            ubCFp32ChunkNList[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += ChunkTileLen * sizeof(float);
+            ubCFp32ChunkNAbsList[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += ChunkTileLen * sizeof(float);
+            ubCFp32ChunkNMaxList[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += HalfChunkTileLen * sizeof(float);
+            ubQuantS32List[i] = ubCFp32ChunkNAbsList[i].template ReinterpretCast<int32_t>();
+            ubQuantF16List[i] = ubCFp32ChunkNAbsList[i].template ReinterpretCast<half>();
+
+            eventUbCVMTE2List[i] = eventVMTE2++;
+            eventUbCMTE2VList[i] = eventMTE2V++;
+            eventUbDMTE3VList[i] = eventMTE3V++;
+            eventUbDVMTE3List[i] = eventVMTE3++;
+
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[i]);
+            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[i]);
+        }
+
+        ubPerTokenScaleOutput = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+    }
+    CATLASS_DEVICE
+    void Finalize()
+    {
+        for (uint32_t i = 0; i < UB_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[i]);
+            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[i]);
+        }
+    }
+    CATLASS_DEVICE
+    ~BlockEpilogue()
+    {
+    }
+
+    CATLASS_DEVICE
+    void UpdateParams(Params const &params_)
+    {
+        params = params_;
+    }
+    // 每个tile就是1*7168，每个block是一个expert的所有token=[group[i], 7168]
+    CATLASS_DEVICE
+    void operator() (
+        AscendC::GlobalTensor<ElementC> const &gmC,
+        MatrixCoord const &shapeC,
+        AscendC::GlobalTensor<ElementPerTokenScale> const &gmPerTokenScale1,
+        AscendC::GlobalTensor<ElementD> const &gmD,
+        AscendC::GlobalTensor<ElementPerTokenScale> const &gmPerTokenScale2,
+
+        uint32_t epilogueCoreNum = 40,
+        Callback &&callback = Callback{}
+    )
+    {
+        callback();
+        uint32_t blockM = shapeC.row();
+        uint32_t blockN = shapeC.column();
+
+        uint32_t tileLoops = blockM;
+        uint32_t subblockIdx = get_block_idx() + get_subblockid() * get_block_num();
+
+        uint32_t subblockNum = get_block_num() * 2;
+        uint32_t moveDataCoreNum = subblockNum - epilogueCoreNum;
+
+        if (subblockIdx < moveDataCoreNum) {
+            return;
+        }
+        uint32_t epilogueCoreIdx = subblockIdx - moveDataCoreNum;
+
+        uint32_t perCoreData =  blockM / epilogueCoreNum;
+        uint32_t remainderData = blockM % epilogueCoreNum;
+
+        uint32_t tasksForIdx  = epilogueCoreIdx < remainderData ? perCoreData + 1 : perCoreData;
+        uint32_t loopStartIdx = epilogueCoreIdx * perCoreData + (epilogueCoreIdx < remainderData? epilogueCoreIdx : remainderData);
+
+        uint32_t alignedPerCoreData = RoundUp<BYTE_PER_BLK / sizeof(ElementPerTokenScale)>(perCoreData + 1);
+
+        uint32_t ChunkTileLen = blockN / 2;
+        uint32_t HalfChunkTileLen = ChunkTileLen / 2;
+
+
+        for (uint32_t loopIdx = loopStartIdx; loopIdx < loopStartIdx + tasksForIdx; ++loopIdx) {
+
+            auto gmTileC = gmC[loopIdx * blockN];
+
+            auto &ubC = ubCList[ubListId];
+            auto &ubD = ubDList[ubListId];
+
+            auto &ubCFp32 = ubCFp32List[ubListId];
+            auto &ubCFp32ChunkN = ubCFp32ChunkNList[ubListId];
+            auto &ubAbs = ubCFp32ChunkNAbsList[ubListId];
+            // auto &ubMax = ubCFp32ChunkNMaxList[ubListId];
+            auto &ubReduceMax = ubCFp32ChunkNMaxList[ubListId];
+            auto &ubOutputTmp = ubAbs;
+            auto &sharedUbTmpBuffer = ubReduceMax;
+            auto &ubQuantS32 = ubQuantS32List[ubListId];
+            auto &ubQuantF16 = ubQuantF16List[ubListId];
+
+            auto gmTileD = gmD[loopIdx * ChunkTileLen];
+            LayoutC layoutUbC{1, blockN};
+
+            // 把C从GM workspace搬到UB
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[ubListId]);
+            copyGmToUbC(ubC, gmTileC, layoutUbC, layoutUbC);
+            AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(eventUbCMTE2VList[ubListId]);
+
+            // 在UB上做把C cast成FP32
+            AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(eventUbCMTE2VList[ubListId]);
+            AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_NONE, blockN);
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[ubListId]);
+
+            // 获取pertoken scale值，gmPerTokenScale的第loopIdx行
+            ElementPerTokenScale perTokenScale = gmPerTokenScale1(loopIdx);
+
+            AscendC::SetFlag<AscendC::HardEvent::S_V>(0);
+            AscendC::WaitFlag<AscendC::HardEvent::S_V>(0);
+            // pertoken scale值与FP32的C做Muls乘法
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Muls(ubCFp32, ubCFp32, perTokenScale, blockN);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            //swiglue计算过程
+            AscendC::Muls(ubCFp32ChunkN, ubCFp32, -1.0f, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Exp(ubCFp32ChunkN, ubCFp32ChunkN, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Adds(ubCFp32ChunkN, ubCFp32ChunkN, 1.0f, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            //TODO除的时候是否会对之后的数据有影响；
+            AscendC::Div(ubCFp32ChunkN, ubCFp32, ubCFp32ChunkN, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Mul(ubCFp32ChunkN, ubCFp32ChunkN, ubCFp32[ChunkTileLen], ChunkTileLen);
+
+            //quant过程，两种方式区别；
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Abs(ubAbs, ubCFp32ChunkN, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            AscendC::ReduceMax<float>(ubReduceMax, ubAbs, sharedUbTmpBuffer, ChunkTileLen, false);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            AscendC::SetFlag<AscendC::HardEvent::V_S>(0);
+            AscendC::WaitFlag<AscendC::HardEvent::V_S>(0);
+
+            //TODO两种计算方法的效率比较
+            ElementPerTokenScale GMubDequantScale = ubReduceMax.GetValue(0);
+            AscendC::SetFlag<AscendC::HardEvent::S_V>(0);
+
+            auto ubPerTokenScaleOutputOffset = loopIdx - loopStartIdx;
+            ubPerTokenScaleOutput.SetValue(ubPerTokenScaleOutputOffset, GMubDequantScale / 127.f);
+
+            AscendC::WaitFlag<AscendC::HardEvent::S_V>(0);
+            AscendC::Muls(ubOutputTmp, ubCFp32ChunkN, 127.f / GMubDequantScale, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            AscendC::Cast(ubQuantS32, ubOutputTmp, AscendC::RoundMode::CAST_RINT, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::SetDeqScale(static_cast<half>(1.0));
+            AscendC::Cast(ubQuantF16, ubQuantS32, AscendC::RoundMode::CAST_RINT, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventUbDVMTE3List[ubListId]);
+            AscendC::Cast(ubD, ubQuantF16, AscendC::RoundMode::CAST_RINT, ChunkTileLen);
+            // AscendC::Muls(ubD, ubCFp32ChunkN, 127.f / GMubDequantScale, ChunkTileLen);
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(eventUbDMTE3VList[ubListId]);
+
+            LayoutD layoutUbD{1, ChunkTileLen};
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(eventUbDVMTE3List[ubListId]);
+            copyUbToGmD(gmTileD, ubD, layoutUbD, layoutUbD);
+            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[ubListId]);
+            ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0;
+        }
+
+        if(tasksForIdx > 0){
+            LayoutPerTokenScale layoutGmPerTokenScale2{tasksForIdx};
+
+            AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
+            AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
+
+            copyUbToGmDequantScale(gmPerTokenScale2[loopStartIdx], ubPerTokenScaleOutput[0], layoutGmPerTokenScale2, layoutGmPerTokenScale2);
+        }
+
+
+    }
+
+    CATLASS_DEVICE
+    void operator() (
+        AscendC::GlobalTensor<ElementC> const &gmC,
+        MatrixCoord const &shapeC,
+        AscendC::GlobalTensor<ElementD> const &gmD,
+        uint32_t epilogueCoreNum = 40,
+        Callback &&callback = Callback{}
+    )
+    {
+        callback();
+        uint32_t blockM = shapeC.row();
+        uint32_t blockN = shapeC.column();
+
+        uint32_t tileLoops = blockM;
+        uint32_t subblockIdx = get_block_idx() + get_subblockid() * get_block_num();
+        //uint32_t subblockIdx = get_block_idx() * 2 + get_subblockid();
+
+        uint32_t subblockNum = get_block_num() * 2;
+        uint32_t moveDataCoreNum = subblockNum - epilogueCoreNum;
+
+        if (subblockIdx < moveDataCoreNum) {
+            return;
+        }
+        uint32_t epilogueCoreIdx = subblockIdx - moveDataCoreNum;
+
+
+        uint32_t perCoreData =  blockM / epilogueCoreNum;
+        uint32_t remainderData = blockM % epilogueCoreNum;
+
+        uint32_t tasksForIdx  = epilogueCoreIdx < remainderData ? perCoreData + 1 : perCoreData;
+        uint32_t loopStartIdx = epilogueCoreIdx * perCoreData + (epilogueCoreIdx < remainderData? epilogueCoreIdx : remainderData);
+
+        uint32_t alignedPerCoreData = RoundUp<BYTE_PER_BLK / sizeof(ElementPerTokenScale)>(perCoreData + 1);
+
+        uint32_t ChunkTileLen = blockN / 2;
+        uint32_t HalfChunkTileLen = ChunkTileLen / 2;
+
+
+        for (uint32_t loopIdx = loopStartIdx; loopIdx < loopStartIdx + tasksForIdx; ++loopIdx) {
+
+            auto gmTileC = gmC[loopIdx * blockN];
+
+            auto &ubC = ubCList[ubListId];
+            auto &ubD = ubDList[ubListId];
+
+            auto &ubCFp32 = ubCFp32List[ubListId];
+            auto &ubCFp32ChunkN = ubCFp32ChunkNList[ubListId];
+
+            auto gmTileD = gmD[loopIdx * ChunkTileLen];
+            LayoutC layoutUbC{1, blockN};
+
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[ubListId]);
+            copyGmToUbC(ubC, gmTileC, layoutUbC, layoutUbC);
+            AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(eventUbCMTE2VList[ubListId]);
+
+            AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(eventUbCMTE2VList[ubListId]);
+            AscendC::Cast(ubCFp32, ubC, AscendC::RoundMode::CAST_NONE, blockN);
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventUbCVMTE2List[ubListId]);
+
+            AscendC::Muls(ubCFp32ChunkN, ubCFp32, -1.0f, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Exp(ubCFp32ChunkN, ubCFp32ChunkN, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Adds(ubCFp32ChunkN, ubCFp32ChunkN, 1.0f, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Div(ubCFp32ChunkN, ubCFp32, ubCFp32ChunkN, ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+            AscendC::Mul(ubCFp32ChunkN, ubCFp32ChunkN, ubCFp32[ChunkTileLen], ChunkTileLen);
+            AscendC::PipeBarrier<PIPE_V>();
+
+            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventUbDVMTE3List[ubListId]);
+            AscendC::Cast(ubD, ubCFp32ChunkN, AscendC::RoundMode::CAST_ROUND, ChunkTileLen);
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(eventUbDMTE3VList[ubListId]);
+
+            LayoutD layoutUbD{1, ChunkTileLen};
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(eventUbDVMTE3List[ubListId]);
+            // copyUbToGmD(gmTileD, ubCFp32ChunkN, layoutUbD, layoutUbD);
+            copyUbToGmD(gmTileD, ubD, layoutUbD, layoutUbD);
+            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventUbDMTE3VList[ubListId]);
+            ubListId = (ubListId + 1 < UB_STAGES) ? (ubListId + 1) : 0;
+        }
+
+    }
+
+
+private:
+    Params params;
+
+    AscendC::LocalTensor<ElementC> ubCList[UB_STAGES];
+    AscendC::LocalTensor<ElementD> ubDList[UB_STAGES];
+
+    int32_t eventUbCVMTE2List[UB_STAGES];
+    int32_t eventUbCMTE2VList[UB_STAGES];
+    int32_t eventUbDMTE3VList[UB_STAGES];
+    int32_t eventUbDVMTE3List[UB_STAGES];
+
+    uint32_t ubListId{0};
+
+    AscendC::LocalTensor<float> ubCFp32List[UB_STAGES];
+    AscendC::LocalTensor<float> ubCFp32ChunkNList[UB_STAGES];
+    AscendC::LocalTensor<float> ubCFp32ChunkNAbsList[UB_STAGES];
+    AscendC::LocalTensor<float> ubCFp32ChunkNMaxList[UB_STAGES];
+    AscendC::LocalTensor<int32_t> ubQuantS32List[UB_STAGES];
+    AscendC::LocalTensor<half> ubQuantF16List[UB_STAGES];
+    AscendC::LocalTensor<float> ubPerTokenScaleOutput;
+
+
+    CopyGmToUbC copyGmToUbC;
+    CopyUbToGmD copyUbToGmD;
+    CopyUbToGmDequantScale copyUbToGmDequantScale;
+};
+
+}  // namespace Catlass::Epilogue::Block
+
+#endif  // CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_SWIGLU_HPP

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/block_epilogue_pertoken_v2.hpp

@@ -0,0 +1,201 @@
+#ifndef CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_V2_ONLY_HPP
+#define CATLASS_EPILOGUE_BLOCK_EPILOGUE_PER_TOKEN_V2_ONLY_HPP
+
+#include "catlass/catlass.hpp"
+#include "catlass/arch/resource.hpp"
+#include "catlass/epilogue/dispatch_policy.hpp"
+#include "catlass/gemm_coord.hpp"
+#include "catlass/matrix_coord.hpp"
+#include "catlass/layout/layout.hpp"
+#include "catlass/detail/callback.hpp"
+
+#include "hccl_shmem.hpp"
+#include "layout3d.hpp"
+
+namespace Catlass::Epilogue::Block {
+template <
+    uint32_t UB_STAGES_,
+    class CType_,
+    class LayoutPerTokenScale_,
+    class DType_,
+    class TileCopy_
+>
+class BlockEpilogue <
+    EpilogueAtlasA2PerTokenDequantV2<UB_STAGES_>,
+    CType_,
+    Gemm::GemmType<float, LayoutPerTokenScale_>,
+    DType_,
+    TileCopy_
+> {
+public:
+    using DispatchPolicy = EpilogueAtlasA2PerTokenDequantV2<UB_STAGES_>;
+    using ArchTag = typename DispatchPolicy::ArchTag;
+    static constexpr uint32_t UB_STAGES = UB_STAGES_;
+
+    // Data infos
+    using ElementC = typename CType_::Element;
+    using LayoutC = typename CType_::Layout;
+    using ElementPerTokenScale = float;
+    using LayoutPerTokenScale = LayoutPerTokenScale_;
+    using ElementD = typename DType_::Element;
+    using LayoutD = typename DType_::Layout;
+
+    //using CopyScaleGmToUb = Epilogue::Tile::CopyGm2Ub<ArchTag, Gemm::GemmType<float, layout::RowMajor>>;
+    using CopyScaleGmToUb = Epilogue::Tile::CopyGm2Ub<ArchTag, Gemm::GemmType<float, layout::VectorLayout>>;
+        // Tile copy
+    using CopyGmToUbC = typename TileCopy_::CopyGmToUbC;
+    using CopyUbToGmD = typename TileCopy_::CopyUbToGmD;
+
+    struct Params {
+        __gm__ int32_t *ptrTokenPerExpert{nullptr};
+        int32_t EP;
+        int32_t expertPerRank;
+        int32_t n2;
+        LayoutC layoutC;
+        int32_t n0;
+        int32_t rank;
+        HcclShmem shmem;
+        int32_t offsetD;
+
+        CATLASS_DEVICE
+        Params() {};
+        CATLASS_DEVICE
+        Params(int32_t EP_, int32_t expertPerRank_, int32_t rank_, __gm__ int32_t *ptrTokenPerExpert_,
+        LayoutC layoutC_, int32_t n2_, int32_t n0_, HcclShmem& shmem_, int32_t offsetD_) :
+        ptrTokenPerExpert(ptrTokenPerExpert_), EP(EP_),
+        expertPerRank(expertPerRank_),rank(rank_), layoutC(layoutC_), n2(n2_), n0(n0_),
+        shmem(shmem_), offsetD(offsetD_)
+         {}
+    };
+
+
+    CATLASS_DEVICE
+    BlockEpilogue(Arch::Resource<ArchTag> const &resource, Params const &params = Params{}) : params(params)
+    {
+        //ub:192KB
+        n0 = params.n0;
+        size_t ubOffset = 0;
+        for(int32_t i = 0; i < 2; i++) {
+            ubCList[i] = resource.ubBuf.template GetBufferByByte<ElementC>(ubOffset);
+            ubOffset += max_len * sizeof(ElementC);
+            ubDList[i] = resource.ubBuf.template GetBufferByByte<ElementD>(ubOffset);
+            ubOffset += max_len * sizeof(ElementD);
+            ubFp32List[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += max_len * sizeof(float);
+            scaleUbList[i] = resource.ubBuf.template GetBufferByByte<float>(ubOffset);
+            ubOffset += (max_len / n0) * sizeof(float);
+            source_scale_offset[i] = -1;
+        }
+        tokenPerExpert.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(params.ptrTokenPerExpert));
+        tokenPerExpertLayout = Layout3D(AlignUp(params.EP * params.expertPerRank, ALIGN_128), params.expertPerRank);
+        is_ping = true;
+    }
+
+    CATLASS_DEVICE
+    void Finalize()
+    {
+        AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
+        AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID1);
+    }
+    CATLASS_DEVICE
+    ~BlockEpilogue()
+    {
+
+    }
+
+    CATLASS_DEVICE
+    void operator() (
+        AscendC::GlobalTensor<ElementC> const &gmC,
+        GemmCoord& blockCoord,
+        GemmCoord& actualBlockShape,
+        int32_t groupIdx,
+        int32_t preSrcExpertSum,
+        AscendC::GlobalTensor<int32_t> preSumBeforeRank
+    ){
+        is_ping = !is_ping;
+        auto event_id = is_ping ? EVENT_ID0 : EVENT_ID1;
+
+        auto &ubC = ubCList[is_ping];
+        int32_t gmCOffset = preSrcExpertSum * params.n2 + blockCoord.m() * params.n2 + blockCoord.n();
+        auto gmTileC = gmC[gmCOffset];
+
+        LayoutC layoutGM{actualBlockShape.m(), actualBlockShape.n(), params.n2};
+        LayoutC layoutUB{actualBlockShape.m(), actualBlockShape.n(), n0};
+
+
+        AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(event_id); //for debug
+        copyGmToUbC(ubC, gmTileC, layoutUB, layoutGM);
+        AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE3>(event_id); //for debug
+
+        int32_t lenTile = actualBlockShape.m();
+        int32_t stTile = blockCoord.m();
+        int32_t edTile = stTile + lenTile;
+        int32_t preSumRankInExpert = 0;
+        int32_t tileOffset = 0;
+
+        AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE3>(event_id);
+        for (int32_t dstEpIdx = 0; dstEpIdx < params.EP; dstEpIdx ++) {
+            int32_t lenRankInExpert = tokenPerExpert(tokenPerExpertLayout(dstEpIdx, params.rank, groupIdx));
+            int32_t dstExpertOffset = preSumBeforeRank(dstEpIdx * params.expertPerRank + groupIdx);
+            int32_t stRankInExpert = preSumRankInExpert;
+            int32_t edRankInExpert = stRankInExpert + lenRankInExpert;
+            preSumRankInExpert += lenRankInExpert;
+            if (stRankInExpert >= edTile) {
+                break;
+            }
+            else if (edRankInExpert <= stTile) {
+                continue;
+            }
+            int32_t stData = max(stRankInExpert, stTile);
+            int32_t edData = min(edRankInExpert, edTile);
+            uint32_t lenData = edData - stData;
+            if (lenData <= 0){
+                continue;
+            }
+            
+            uint32_t dstOffsetInExpert = 0;
+            if (stTile > stRankInExpert) {
+                dstOffsetInExpert = stTile - stRankInExpert;
+            }
+            AscendC::GlobalTensor<ElementD> gmRemotePeer;
+            __gm__ void* dstPeermemPtr = params.shmem(params.offsetD, dstEpIdx);
+            gmRemotePeer.SetGlobalBuffer(reinterpret_cast<__gm__ ElementD*>(dstPeermemPtr));
+            MatrixCoord dstOffset{dstOffsetInExpert + dstExpertOffset, blockCoord.n()};
+            int64_t gmDstOffset = params.layoutC.GetOffset(dstOffset);
+            auto gmTileD = gmRemotePeer[gmDstOffset];
+            LayoutC layoutGM2{lenData, actualBlockShape.n(), params.n2};
+            LayoutC layoutUB2{lenData, actualBlockShape.n(), n0};
+            copyUbToGmD(gmTileD, ubC[tileOffset *  n0], layoutGM2, layoutUB2);
+            tileOffset += lenData;
+        }
+
+        AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(event_id);
+
+    }
+
+private:
+
+    Params params;
+    AscendC::LocalTensor<ElementC> ubCList[UB_STAGES];
+    AscendC::LocalTensor<ElementD> ubDList[UB_STAGES];
+    AscendC::LocalTensor<float> ubFp32List[UB_STAGES];
+    AscendC::LocalTensor<float> scaleUbList[UB_STAGES];
+    int32_t source_scale_offset[UB_STAGES];
+
+    int32_t max_len = 8 * 32 / 4 * 128;
+    int32_t n0;
+    bool is_ping = false;
+
+
+    int32_t repeat = 128;
+
+
+    CopyGmToUbC copyGmToUbC;
+    CopyUbToGmD copyUbToGmD;
+
+    CopyScaleGmToUb copyScaleGmToUb;
+    AscendC::GlobalTensor<int32_t> tokenPerExpert;
+    Layout3D tokenPerExpertLayout;
+};
+}
+#endif

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/block_mmad_preload_async_fixpipe_quant.hpp

@@ -0,0 +1,556 @@
+/*
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_FIXPIPE_QUANT_HPP
+#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_FIXPIPE_QUANT_HPP
+
+#include "catlass/catlass.hpp"
+#include "catlass/arch/resource.hpp"
+#include "catlass/coord.hpp"
+#include "catlass/gemm_coord.hpp"
+#include "catlass/gemm/dispatch_policy.hpp"
+#include "catlass/gemm/helper.hpp"
+#include "dispatch_policy_custom.hpp"
+
+
+namespace Catlass::Gemm::Block {
+
+template<AscendC::HardEvent event>
+__aicore__ inline void SyncFlagFunc(int32_t eventID)
+{
+    AscendC::SetFlag<event>(eventID);
+    AscendC::WaitFlag<event>(eventID);
+}
+
+template <
+    uint32_t PRELOAD_STAGES_,
+    uint32_t L1_STAGES_,
+    uint32_t L0A_STAGES_,
+    uint32_t L0B_STAGES_,
+    uint32_t L0C_STAGES_,
+    bool ENABLE_UNIT_FLAG_,
+    bool ENABLE_SHUFFLE_K_,
+    class L1TileShape_,
+    class L0TileShape_,
+    class AType_,
+    class BType_,
+    class CType_,
+    class BiasType_,
+    class TileCopy_,
+    class TileMmad_
+>
+struct BlockMmad <
+    MmadAtlasA2PreloadAsyncFixpipe<
+        PRELOAD_STAGES_,
+        L1_STAGES_,
+        L0A_STAGES_,
+        L0B_STAGES_,
+        L0C_STAGES_,
+        ENABLE_UNIT_FLAG_,
+        ENABLE_SHUFFLE_K_
+    >,
+    L1TileShape_,
+    L0TileShape_,
+    AType_,
+    BType_,
+    CType_,
+    BiasType_,
+    TileCopy_,
+    TileMmad_
+> {
+public:
+    // Type Aliases
+    using DispatchPolicy = MmadAtlasA2PreloadAsyncFixpipe<
+        PRELOAD_STAGES_,
+        L1_STAGES_,
+        L0A_STAGES_,
+        L0B_STAGES_,
+        L0C_STAGES_,
+        ENABLE_UNIT_FLAG_,
+        ENABLE_SHUFFLE_K_
+    >;
+    using ArchTag = typename DispatchPolicy::ArchTag;
+    using L1TileShape = L1TileShape_;
+    using L0TileShape = L0TileShape_;
+    using ElementA = typename AType_::Element;
+    using LayoutA = typename AType_::Layout;
+    using ElementB = typename BType_::Element;
+    using LayoutB = typename BType_::Layout;
+    using ElementC = typename CType_::Element;
+    using LayoutC = typename CType_::Layout;
+    using TileMmad = TileMmad_;
+    using CopyGmToL1A = typename TileCopy_::CopyGmToL1A;
+    using CopyGmToL1B = typename TileCopy_::CopyGmToL1B;
+    using CopyGmToL1S = Gemm::Tile::CopyGmToL1<ArchTag, Gemm::GemmType<uint64_t, layout::VectorLayout>>;
+    using CopyL1ToFP = typename Gemm::Tile::QuantTileCopy<ArchTag, AType_, BType_, CType_, void, Catlass::Gemm::Tile::ScaleGranularity::PER_CHANNEL>::CopyL1ToFP;
+    using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A;
+    using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B;
+    
+    using ElementAccumulator =
+        typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
+    using CopyL0CToGm = typename std::conditional<
+        std::is_same_v<ElementA, int8_t>,
+        Gemm::Tile::CopyL0CToGm<ArchTag, ElementAccumulator, CType_, Gemm::Tile::ScaleGranularity::PER_CHANNEL>,
+        typename TileCopy_::CopyL0CToGm
+    >::type;
+    using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc;
+    using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc;
+    using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst;
+    using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst;
+    using LayoutCInL0 = layout::zN;
+
+    using L1AAlignHelper = Gemm::helper::L1AlignHelper<ElementA, LayoutA>;
+    using L1BAlignHelper = Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
+
+    static constexpr uint32_t PRELOAD_STAGES = DispatchPolicy::PRELOAD_STAGES;
+    static constexpr uint32_t L1_STAGES = DispatchPolicy::L1_STAGES;
+    static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES;
+    static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES;
+    static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES;
+
+    static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG;
+    static constexpr bool ENABLE_SHUFFLE_K = DispatchPolicy::ENABLE_SHUFFLE_K;
+
+    // L1 tile size
+    static constexpr uint32_t L1A_TILE_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA);
+    static constexpr uint32_t L1B_TILE_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB);
+    static constexpr uint32_t L1S_TILE_SIZE = L1TileShape::N * sizeof(int64_t);
+    // L0 tile size
+    static constexpr uint32_t L0A_TILE_SIZE = L0TileShape::M * L0TileShape::K * sizeof(ElementA);
+    static constexpr uint32_t L0B_TILE_SIZE = L0TileShape::K * L0TileShape::N * sizeof(ElementB);
+    static constexpr uint32_t L0C_TILE_SIZE = L1TileShape::M * L1TileShape::N * sizeof(ElementAccumulator);
+
+    // Check LayoutC
+    static_assert(std::is_same_v<LayoutC, layout::RowMajor>, "LayoutC only support RowMajor yet!");
+
+    // Check L1TileShape
+    static_assert(
+        (std::is_same_v<ElementA, int8_t> 
+            ? (L1A_TILE_SIZE + L1B_TILE_SIZE + L1S_TILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE
+            : (L1A_TILE_SIZE + L1B_TILE_SIZE) * L1_STAGES <= ArchTag::L1_SIZE),
+        "L1TileShape exceeding the L1 space for the given data type"
+    );
+
+    // Check L0TileShape
+    static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!");
+    static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!");
+    static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!");
+
+    static_assert(L1TileShape::M == L0TileShape::M && L1TileShape::N == L0TileShape::N,
+        "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet");
+
+    static constexpr auto L1A_LAYOUT = LayoutAInL1::template MakeLayout<ElementA>(
+        L1TileShape::M, L1TileShape::K);
+    static constexpr auto L1B_LAYOUT = LayoutBInL1::template MakeLayout<ElementB>(
+        L1TileShape::K, L1TileShape::N);
+
+    CATLASS_DEVICE
+    BlockMmad(Arch::Resource<ArchTag> &resource, __gm__ int32_t* flagPtr = nullptr, int32_t expertPerRank = 0, uint32_t l1BufAddrStart = 0, uint32_t FpAddrStart = 0)
+    {
+        syncGroupIdx = 0;
+        ptrSoftFlagBase_ = flagPtr;
+        expertPerRank_ = expertPerRank;
+        InitL1(resource, l1BufAddrStart);
+        InitFpBuf(resource, FpAddrStart);
+        InitL0A(resource);
+        InitL0B(resource);
+        InitL0C(resource);
+    }
+
+    CATLASS_DEVICE
+    ~BlockMmad()
+    {
+        SynchronizeBlock();
+        for (uint32_t i = 0; i < L1_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[i]);
+            AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[i]);
+        }
+        for (uint32_t i = 0; i < L0A_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[i]);
+        }
+        for (uint32_t i = 0; i < L0B_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[i]);
+        }
+        for (uint32_t i = 0; i < L0C_STAGES; ++i) {
+            AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList[i]);
+        }
+        if constexpr (std::is_same_v<ElementA, int8_t>) {
+            AscendC::WaitFlag<AscendC::HardEvent::FIX_MTE2>(0);
+        }
+    }
+
+    CATLASS_DEVICE
+    void operator()(
+        AscendC::GlobalTensor<ElementA> const &gmBlockA, LayoutA const &layoutA,
+        AscendC::GlobalTensor<ElementB> const &gmBlockB, LayoutB const &layoutB,
+        AscendC::GlobalTensor<ElementC> const &gmBlockC, LayoutC const &layoutC,
+        AscendC::GlobalTensor<uint64_t> const &gmBlockS, layout::VectorLayout const &layoutScale,
+        GemmCoord const &actualShape, int32_t syncLoopIdx = -1, int32_t flag = 0
+    )
+    {
+        uint32_t kTileCount = CeilDiv<L1TileShape::K>(actualShape.k());
+
+        uint32_t mRound = RoundUp<L1AAlignHelper::M_ALIGNED>(actualShape.m());
+        uint32_t nRound = RoundUp<L1BAlignHelper::N_ALIGNED>(actualShape.n());
+
+        uint32_t startTileIdx = 0;
+        if constexpr (ENABLE_SHUFFLE_K) {
+            startTileIdx = AscendC::GetBlockIdx() % kTileCount;
+        }
+
+        for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) {
+            uint32_t kTileIdx = (startTileIdx + kLoopIdx < kTileCount) ?
+                (startTileIdx + kLoopIdx) : (startTileIdx + kLoopIdx - kTileCount);
+
+            uint32_t kActual = (kTileIdx < kTileCount - 1) ?
+                L1TileShape::K : (actualShape.k() - kTileIdx * L1TileShape::K);
+
+            // Emission load instruction from GM to L1
+            MatrixCoord gmTileAOffset{0, kTileIdx * L1TileShape::K};
+            MatrixCoord gmTileBOffset{kTileIdx * L1TileShape::K, 0};
+            auto gmTileA = gmBlockA[layoutA.GetOffset(gmTileAOffset)];
+            auto gmTileB = gmBlockB[layoutB.GetOffset(gmTileBOffset)];
+            // Load first matrix A tile from GM to L1
+            AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[l1ListId]);
+            auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), kActual));
+            copyGmToL1A(l1ATensorList[l1ListId], gmTileA, L1A_LAYOUT, layoutTileA);
+            AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList[l1ListId]);
+            // Load first matrix B tile from GM to L1
+            AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[l1ListId]);
+            auto layoutTileB = layoutB.GetTileLayout(MakeCoord(kActual, actualShape.n()));
+            copyGmToL1B(l1BTensorList[l1ListId], gmTileB, L1B_LAYOUT, layoutTileB);
+            AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList[l1ListId]);
+
+            // If the number of preload instructions reaches the upper limit, perform an mmad calculation on L1 tile
+            if (preloadCount == PRELOAD_STAGES) {
+                L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]);
+            }
+
+            // Store the current load status
+            uint32_t preloadL1TileMmadParamsId = (l1TileMmadParamsId + preloadCount < PRELOAD_STAGES) ?
+                (l1TileMmadParamsId + preloadCount) : (l1TileMmadParamsId + preloadCount - PRELOAD_STAGES);
+            auto &l1TileMmadParams = l1TileMmadParamsList[preloadL1TileMmadParamsId];
+            l1TileMmadParams.l1ListId = l1ListId;
+            l1TileMmadParams.mRound = mRound;
+            l1TileMmadParams.nRound = nRound;
+            l1TileMmadParams.kActual = kActual;
+            l1TileMmadParams.isKLoopFirst = (kLoopIdx == 0);
+            l1TileMmadParams.isKLoopLast = (kLoopIdx == kTileCount - 1);
+            l1TileMmadParams.flag = flag;
+            if (kLoopIdx == kTileCount - 1) {
+                l1TileMmadParams.gmBlockC = gmBlockC;
+                l1TileMmadParams.gmBlockS = gmBlockS;
+                l1TileMmadParams.layoutCInGm = layoutC.GetTileLayout(actualShape.GetCoordMN());
+                l1TileMmadParams.layoutScale = layoutScale;
+                l1TileMmadParams.syncLoopIdx = syncLoopIdx;
+            }
+
+            if (preloadCount < PRELOAD_STAGES) {
+                ++preloadCount;
+            } else {
+                l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0;
+            }
+            l1ListId = (l1ListId + 1 < L1_STAGES) ? (l1ListId + 1) : 0;
+        }
+    }
+
+    CATLASS_DEVICE
+    void SynchronizeBlock()
+    {
+        while (preloadCount > 0) {
+            L1TileMmad(l1TileMmadParamsList[l1TileMmadParamsId]);
+            l1TileMmadParamsId = (l1TileMmadParamsId + 1 < PRELOAD_STAGES) ? (l1TileMmadParamsId + 1) : 0;
+            --preloadCount;
+        }
+    }
+
+    CATLASS_DEVICE
+    void Finalize(int32_t target, int32_t flag = 0)
+    {
+        if (ptrSoftFlagBase_ != nullptr) {
+            if (target < 0) {
+                return;
+            }
+            AscendC::SetFlag<AscendC::HardEvent::FIX_MTE3>(EVENT_ID0);
+            AscendC::WaitFlag<AscendC::HardEvent::FIX_MTE3>(EVENT_ID0);
+            AscendC::GlobalTensor<int32_t> flagGlobal;
+            flagGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(ptrSoftFlagBase_) + (expertPerRank_ + AscendC::GetBlockIdx()) * FLAGSTRIDE);
+            AscendC::DataCopy(flagGlobal, l1FTensor[target * 16], FLAGSTRIDE);
+        }
+        else {
+            for(;syncGroupIdx <= target; syncGroupIdx++) {
+                int32_t flagId = syncGroupIdx / 15 + flag;
+                AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(flagId);
+            }
+        }
+    }
+private:
+    struct L1TileMmadParams {
+        uint32_t l1ListId;
+        uint32_t mRound;
+        uint32_t nRound;
+        uint32_t kActual;
+        bool isKLoopFirst;
+        bool isKLoopLast;
+        AscendC::GlobalTensor<ElementC> gmBlockC;
+        AscendC::GlobalTensor<uint64_t> gmBlockS;
+        LayoutC layoutCInGm;
+        layout::VectorLayout layoutScale;
+        int32_t syncLoopIdx;
+        int32_t flag;
+        CATLASS_DEVICE
+        L1TileMmadParams() = default;
+    };
+
+    CATLASS_DEVICE
+    void InitL1(Arch::Resource<ArchTag> &resource, uint32_t l1BufAddrStart)
+    {
+        uint32_t l1AOffset = l1BufAddrStart;
+        uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1_STAGES;
+
+        for (uint32_t i = 0; i < L1_STAGES; ++i) {
+            l1ATensorList[i] = resource.l1Buf.template GetBufferByByte<ElementA>(l1AOffset + L1A_TILE_SIZE * i);
+            l1BTensorList[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1BOffset + L1B_TILE_SIZE * i);
+            l1AEventList[i] = i;
+            l1BEventList[i] = i + L1_STAGES;
+            AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[i]);
+            AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[i]);
+        }
+        uint32_t l1SOffset = l1BOffset + L1B_TILE_SIZE * L1_STAGES;
+        if constexpr (std::is_same_v<ElementA, int8_t>) {
+            l1STensor = resource.l1Buf.template GetBufferByByte<uint64_t>(l1SOffset);
+            AscendC::SetFlag<AscendC::HardEvent::FIX_MTE2>(0);
+        }
+        if (ptrSoftFlagBase_ != nullptr) {
+            // Initialize the flag matrix (structure as below):
+            // 1 0 0 0 0 0 0 0
+            // 2 0 0 0 0 0 0 0
+            // ...
+            // 16 0 0 0 0 0 0 0
+            // Then move it to L1
+            uint32_t l1FOffset = l1SOffset + L1S_TILE_SIZE;
+            l1FTensor = resource.l1Buf.template GetBufferByByte<int32_t>(l1FOffset);
+            AscendC::GlobalTensor<int32_t> flagBase;
+            flagBase.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(ptrSoftFlagBase_));
+            AscendC::DataCopy(l1FTensor, flagBase, expertPerRank_ * FLAGSTRIDE);
+        }
+    }
+
+    CATLASS_DEVICE
+    void InitFpBuf(Arch::Resource<ArchTag> &resource, uint32_t FpAddrStart)
+    {
+        uint32_t FpOffset = FpAddrStart;
+        fixpipeBuf = resource.fpBuf.template GetBufferByByte<uint64_t>(FpOffset);
+    }
+
+    CATLASS_DEVICE
+    void InitL0A(Arch::Resource<ArchTag> &resource)
+    {
+        for (uint32_t i = 0; i < L0A_STAGES; ++i) {
+            l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_TILE_SIZE * i);
+            l0AEventList[i] = i;
+            AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[i]);
+        }
+    }
+
+    CATLASS_DEVICE
+    void InitL0B(Arch::Resource<ArchTag> &resource)
+    {
+        for (uint32_t i = 0; i < L0B_STAGES; ++i) {
+            l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_TILE_SIZE * i);
+            l0BEventList[i] = i + L0A_STAGES;
+            AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[i]);
+        }
+    }
+
+    CATLASS_DEVICE
+    void InitL0C(Arch::Resource<ArchTag> &resource)
+    {
+        for (uint32_t i = 0; i < L0C_STAGES; ++i) {
+            l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(L0C_TILE_SIZE * i);
+            l0CEventList[i] = i;
+            AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList[i]);
+        }
+    }
+
+    CATLASS_DEVICE
+    void L1TileMmad(L1TileMmadParams const &params)
+    {
+        uint32_t mPartLoop = CeilDiv<L0TileShape::M>(params.mRound);
+        uint32_t nPartLoop = CeilDiv<L0TileShape::N>(params.nRound);
+        uint32_t kPartLoop = CeilDiv<L0TileShape::K>(params.kActual);
+        auto &l1ATensor = l1ATensorList[params.l1ListId];
+        auto &l1BTensor = l1BTensorList[params.l1ListId];
+
+        auto &l0CTensor = l0CTensorList[l0CListId];
+        LayoutCInL0 layoutCInL0 = LayoutCInL0::MakeLayoutInL0C(MakeCoord(params.mRound, params.nRound));
+
+        if constexpr (!ENABLE_UNIT_FLAG) {
+            if (params.isKLoopFirst) {
+                AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList[l0CListId]);
+            }
+        }
+
+        for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) {
+            uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ?
+                L0TileShape::M : (params.mRound - mPartIdx * L0TileShape::M);
+
+            for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) {
+                uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ?
+                    L0TileShape::K : (params.kActual - kPartIdx * L0TileShape::K);
+
+                auto &l0ATile = l0ATensorList[l0AListId];
+                auto layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mPartActual, kPartActual);
+                auto l1AOffset = MakeCoord(mPartIdx, kPartIdx) * L0TileShape::ToCoordMK();
+                auto l1ATile = l1ATensor[L1A_LAYOUT.GetOffset(l1AOffset)];
+
+                AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]);
+                if ((mPartIdx == 0) && (kPartIdx == 0)) {
+                    AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList[params.l1ListId]);
+                }
+                copyL1ToL0A(l0ATile, l1ATile, layoutAInL0, L1A_LAYOUT);
+                if ((mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1)) {
+                    AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList[params.l1ListId]);
+                }
+
+                for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) {
+                    uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ?
+                        L0TileShape::N : (params.nRound - nPartIdx * L0TileShape::N);
+
+                    auto &l0BTile = l0BTensorList[l0BListId];
+                    auto layoutBInL0 = LayoutBInL0::template MakeLayout<ElementB>(kPartActual, nPartActual);
+                    auto l1BOffset = MakeCoord(kPartIdx, nPartIdx) * L0TileShape::ToCoordKN();
+                    auto l1BTile = l1BTensor[L1B_LAYOUT.GetOffset(l1BOffset)];
+
+                    AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[l0BListId]);
+                    if ((kPartIdx == 0) && (nPartIdx == 0)) {
+                        AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList[params.l1ListId]);
+                    }
+                    copyL1ToL0B(l0BTile, l1BTile, layoutBInL0, L1B_LAYOUT);
+                    if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) {
+                        AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList[params.l1ListId]);
+                    }
+
+                    AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
+
+                    auto l0COffset = MakeCoord(mPartIdx, nPartIdx) * L0TileShape::ToCoordMN();
+                    auto l0CTile = l0CTensor[layoutCInL0.GetOffset(l0COffset)];
+
+                    AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
+                    // If the current tile is the first tile on the k axis, the accumulator needs to be reset to 0
+                    bool initC = (params.isKLoopFirst && (kPartIdx == 0));
+                    // If the unit flag is enabled, the unit flag is set according to the calculation progress
+                    uint8_t unitFlag = 0b00;
+                    if constexpr (ENABLE_UNIT_FLAG) {
+                        if (params.isKLoopLast &&
+                            (mPartIdx == mPartLoop - 1) && (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) {
+                            unitFlag = 0b11;
+                        } else {
+                            unitFlag = 0b10;
+                        }
+                    }
+                    tileMmad(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag);
+
+                    AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList[l0BListId]);
+                    l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0;
+                }
+                AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList[l0AListId]);
+                l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0;
+            }
+        }
+
+        if (params.isKLoopLast) {
+            auto layoutCInGm = params.layoutCInGm;
+            if constexpr (std::is_same_v<ElementA, int8_t>) {
+                auto layoutScale = params.layoutScale;
+                auto layoutTileS = layoutScale.GetTileLayout(MakeCoord(layoutCInGm.shape(1)));
+                AscendC::WaitFlag<AscendC::HardEvent::FIX_MTE2>(0);
+                copyGmToL1S(l1STensor, params.gmBlockS, layoutTileS, layoutTileS);
+                AscendC::SetFlag<AscendC::HardEvent::MTE2_FIX>(0);
+                AscendC::WaitFlag<AscendC::HardEvent::MTE2_FIX>(0);
+
+                copyL1ToFP(fixpipeBuf, l1STensor, layoutTileS, layoutTileS);
+                AscendC::SetFlag<AscendC::HardEvent::FIX_MTE2>(0);
+                AscendC::PipeBarrier<PIPE_FIX>();
+            }
+            if constexpr (!ENABLE_UNIT_FLAG) {
+                AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventList[l0CListId]);
+                AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventList[l0CListId]);
+                if constexpr (std::is_same_v<ElementA, int8_t>) {
+                    copyL0CToGm(params.gmBlockC, l0CTensor, l1STensor, layoutCInGm, layoutCInL0);
+                } else {
+                    copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0);
+                }
+                AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList[l0CListId]);
+            } else {
+                if constexpr (std::is_same_v<ElementA, int8_t>) {
+                    copyL0CToGm(params.gmBlockC, l0CTensor, l1STensor, layoutCInGm, layoutCInL0, 0b11);
+                } else {
+                    copyL0CToGm(params.gmBlockC, l0CTensor, layoutCInGm, layoutCInL0, 0b11);
+                }
+            }
+            l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0;
+            if constexpr (std::is_same_v<ElementA, int8_t>) {
+                AscendC::SetFlag<AscendC::HardEvent::FIX_MTE2>(0);
+            }
+            #ifdef __TILE_SYNC__
+            if (params.flag > 0) {
+                int32_t flagId = params.flag + params.syncLoopIdx / 8;
+                AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(flagId);
+            }
+            #else
+            Finalize(params.syncLoopIdx, params.flag);
+            #endif
+        }
+    }
+
+    AscendC::LocalTensor<uint64_t> fixpipeBuf;
+
+    AscendC::LocalTensor<ElementA> l1ATensorList[L1_STAGES];
+    AscendC::LocalTensor<ElementB> l1BTensorList[L1_STAGES];
+    AscendC::LocalTensor<uint64_t> l1STensor;
+    AscendC::LocalTensor<int32_t> l1FTensor;
+    int32_t syncGroupIdx;
+    int32_t l1AEventList[L1_STAGES];
+    int32_t l1BEventList[L1_STAGES];
+    uint32_t l1ListId{0};
+
+    AscendC::LocalTensor<ElementA> l0ATensorList[L0A_STAGES];
+    int32_t l0AEventList[L0A_STAGES];
+    uint32_t l0AListId{0};
+
+    AscendC::LocalTensor<ElementB> l0BTensorList[L0B_STAGES];
+    int32_t l0BEventList[L0B_STAGES];
+    uint32_t l0BListId{0};
+
+    AscendC::LocalTensor<ElementAccumulator> l0CTensorList[L0C_STAGES_];
+    int32_t l0CEventList[L0C_STAGES_];
+    uint32_t l0CListId{0};
+
+    L1TileMmadParams l1TileMmadParamsList[PRELOAD_STAGES];
+    uint32_t l1TileMmadParamsId{0};
+    uint32_t preloadCount{0};
+
+    TileMmad tileMmad;
+    CopyGmToL1A copyGmToL1A;
+    CopyGmToL1B copyGmToL1B;
+    CopyGmToL1S copyGmToL1S;
+    CopyL1ToL0A copyL1ToL0A;
+    CopyL1ToL0B copyL1ToL0B;
+    CopyL0CToGm copyL0CToGm;
+
+    __gm__ int32_t* ptrSoftFlagBase_ = nullptr;
+    int32_t expertPerRank_;
+    CopyL1ToFP copyL1ToFP;
+};
+
+}  // namespace Catlass::Gemm::Block
+
+#endif  // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PRELOAD_FIXPIPE_QUANT_HPP

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/const_args.hpp

@@ -0,0 +1,13 @@
+
+#ifndef CONST_ARGS_HPP
+#define CONST_ARGS_HPP
+constexpr static uint64_t MB_SIZE = 1024 * 1024UL;
+constexpr static int32_t NUMS_PER_FLAG = 16;
+constexpr static int32_t CACHE_LINE = 512;
+constexpr static int32_t FLAGSTRIDE = 16;
+constexpr static int32_t RESET_VAL = 0xffff;
+constexpr static int32_t ALIGN_128 = 128;
+constexpr uint32_t MAX_EXPERTS_PER_RANK = 32;
+constexpr static int32_t UB_ALIGN = 32;
+constexpr uint16_t CROSS_CORE_FLAG_MAX_SET_COUNT = 15;
+#endif

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/copy_gm_to_l1_custom.hpp

@@ -0,0 +1,40 @@
+#ifndef COPY_GM_TO_L1_CUSTOM_HPP
+#define COPY_GM_TO_L1_CUSTOM_HPP
+
+namespace Catlass::Gemm::Tile {
+    /// Partial specialization for nZ in and nZ out.
+    template <
+        class ArchTag,
+        class Element
+    >
+    struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::VectorLayout>> {
+        using LayoutDst = layout::VectorLayout;
+        using LayoutSrc = layout::VectorLayout;
+
+        static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);   // int64, 32/8=4
+
+        // Methods  
+
+        CATLASS_DEVICE
+        CopyGmToL1() {};
+
+        CATLASS_DEVICE
+        void operator()(
+            AscendC::LocalTensor<Element> const &dstTensor,
+            AscendC::GlobalTensor<Element> const &srcTensor,
+            LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
+        {
+            uint32_t blockCount = 1;
+            uint32_t blockLen = CeilDiv<ELE_NUM_PER_C0>(layoutSrc.shape(0));
+
+            AscendC::DataCopyParams repeatParams;
+
+            repeatParams.blockCount = blockCount;
+            repeatParams.blockLen = blockLen;
+            repeatParams.srcStride = 0;
+            repeatParams.dstStride = 0;
+            AscendC::DataCopy(dstTensor, srcTensor, repeatParams);
+        }
+    };
+}
+#endif // COPY_GM_TO_L1_CUSTOM_HPP

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/copy_l0c_to_gm_custom.hpp

@@ -0,0 +1,47 @@
+#ifndef COPY_L0C_TO_GM_CUSTOM_HPP
+#define COPY_L0C_TO_GM_CUSTOM_HPP
+
+namespace Catlass::Gemm::Tile {
+    template <
+        class ElementAccumulator_,
+        class ElementDst_,
+        bool ReluEnable_
+    >
+    struct CopyL0CToGm<Catlass::Arch::AtlasA2,
+                    ElementAccumulator_,
+                    Gemm::GemmType<ElementDst_, layout::RowMajor>,
+                    ScaleGranularity::PER_CHANNEL,
+                    ReluEnable_>
+    {
+        using ArchTag = Catlass::Arch::AtlasA2;
+        using ElementDst = ElementDst_;
+        using ElementSrc = ElementAccumulator_;
+        using LayoutSrc = Catlass::layout::zN;
+        using LayoutDst = Catlass::layout::RowMajor;
+        static constexpr auto quantPre = CopyL0CToGmQuantMode<ArchTag, ElementSrc, ElementDst,
+            ScaleGranularity::PER_CHANNEL>::VALUE;
+        static constexpr auto reluEn = ReluEnable_;
+
+        CATLASS_DEVICE
+        void operator()(AscendC::GlobalTensor<ElementDst> const &dst, AscendC::LocalTensor<ElementSrc> const &src, AscendC::LocalTensor<uint64_t> cbufWorkspace,
+            LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0)
+        {
+            AscendC::FixpipeParamsV220 intriParams;
+
+            // Fixpipe layout information
+            intriParams.nSize = dstLayout.shape(1);
+            intriParams.mSize = dstLayout.shape(0);
+            intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0);
+            intriParams.dstStride = dstLayout.stride(0);
+
+            // Fixpipe auxiliary arguments
+            intriParams.quantPre = quantPre;
+            intriParams.reluEn = reluEn;
+            intriParams.unitFlag = unitFlag;
+
+            // Call AscendC Fixpipe
+            AscendC::Fixpipe<ElementDst, ElementSrc, AscendC::CFG_ROW_MAJOR>(dst, src, cbufWorkspace, intriParams);
+        }
+    };
+}
+#endif // COPY_L0C_TO_GM_CUSTOM_HPP

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/dispatch_policy_custom.hpp

@@ -0,0 +1,53 @@
+#ifndef DISPATH_POLICY_CUSTOM_HPP
+#define DISPATH_POLICY_CUSTOM_HPP
+
+namespace Catlass::Gemm {
+    template <bool ENABLE_UNIT_FLAG_ = false, bool ENABLE_SHUFFLE_K_ = false>
+    struct MmadAtlasA2PreloadFixpipeQuant : public MmadAtlasA2 {
+        static constexpr uint32_t STAGES = 2;
+        static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
+        static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_;
+    };
+
+    template <uint32_t PRELOAD_STAGES_, uint32_t L1_STAGES_, uint32_t L0A_STAGES_, uint32_t L0B_STAGES_,
+        uint32_t L0C_STAGES_, bool ENABLE_UNIT_FLAG_, bool ENABLE_SHUFFLE_K_>
+    struct MmadAtlasA2PreloadAsyncFixpipe :
+        public MmadAtlasA2PreloadAsync<
+            PRELOAD_STAGES_,
+            L1_STAGES_, 
+            L0A_STAGES_,
+            L0B_STAGES_,
+            L0C_STAGES_,
+            ENABLE_UNIT_FLAG_,
+            ENABLE_SHUFFLE_K_
+        > {
+    };
+}
+
+namespace Catlass::Epilogue {
+        
+    template <uint32_t UB_STAGES_>
+    struct EpilogueAtlasA2UnQuant {
+        using ArchTag = Arch::AtlasA2;
+        static constexpr uint32_t UB_STAGES = UB_STAGES_;
+    };
+
+    template <uint32_t UB_STAGES_>
+    struct EpilogueAtlasA2PerTokenDequantQuant {
+        using ArchTag = Arch::AtlasA2;
+        static constexpr uint32_t UB_STAGES = UB_STAGES_;
+    };
+    
+    template <uint32_t UB_STAGES_>
+    struct EpilogueAtlasA2PerTokenDequantSwigluQuant {
+        using ArchTag = Arch::AtlasA2;
+        static constexpr uint32_t UB_STAGES = UB_STAGES_;
+    };
+
+    template <uint32_t UB_STAGES_>
+    struct EpilogueAtlasA2PerTokenDequantV2 {
+        using ArchTag = Arch::AtlasA2;
+        static constexpr uint32_t UB_STAGES = UB_STAGES_;
+    };
+}
+#endif // DISPATH_POLICY_CUSTOM_HPP

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/get_tensor_addr.hpp

@@ -0,0 +1,16 @@
+#ifndef GET_TENSOR_ADDR_HPP
+#define GET_TENSOR_ADDR_HPP
+#include "kernel_operator.h"
+
+#define FORCE_INLINE_AICORE inline __attribute__((always_inline)) __aicore__
+
+template <typename T>
+FORCE_INLINE_AICORE __gm__ T* GetTensorAddr(uint32_t index, GM_ADDR tensorPtr) {
+    __gm__ uint64_t* dataAddr = reinterpret_cast<__gm__ uint64_t*>(tensorPtr);
+    uint64_t tensorPtrOffset = *dataAddr;  // The offset of the data address from the first address.
+    // Moving 3 bits to the right means dividing by sizeof(uint64 t).
+    __gm__ uint64_t* retPtr = dataAddr + (tensorPtrOffset >> 3);
+    return reinterpret_cast<__gm__ T*>(*(retPtr + index));
+}
+
+#endif // GET_TENSOR_ADDR_HPP

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/hccl_shmem.hpp

@@ -0,0 +1,327 @@
+#ifndef SYNC_UTIL_HPP
+#define SYNC_UTIL_HPP
+
+
+#include "kernel_operator.h"
+#include "const_args.hpp"
+
+#ifdef HCCL_COMM
+#include "moe_distribute_base.h"
+using namespace AscendC::HcclContextDef;
+
+#else
+#include "shmem_api.h"
+#endif
+
+#define FORCE_INLINE_AICORE inline __attribute__((always_inline)) __aicore__
+constexpr int32_t MAX_RANK_SIZE = 32;
+constexpr int32_t SHMEM_MEM = 700 * MB_SIZE;
+
+constexpr uint16_t SEND_SYNC_EVENT_ID = 9;
+constexpr uint16_t RECV_SYNC_EVENT_ID = 10;
+
+constexpr uint32_t SELF_STATE_OFFSET = 256 * 1024;
+constexpr uint32_t STATE_OFFSET = 512;
+
+FORCE_INLINE_AICORE void AicSyncAll() {
+    AscendC::CrossCoreSetFlag<0x0, PIPE_FIX>(8);
+    AscendC::CrossCoreWaitFlag<0x0>(8);
+}
+
+template<typename T>
+FORCE_INLINE_AICORE void gm_store(__gm__ T *addr, T val) {
+    *((__gm__ T *)addr) = val;
+}
+
+template<typename T>
+FORCE_INLINE_AICORE T gm_load(__gm__ T *cache) {
+    return *((__gm__ T *)cache);
+}
+
+template<typename T>
+FORCE_INLINE_AICORE void gm_dcci(__gm__ T * addr) {
+    using namespace AscendC;
+    GlobalTensor<uint8_t> global;
+    global.SetGlobalBuffer(reinterpret_cast<GM_ADDR>(addr));
+
+    // Important: add hint to avoid dcci being optimized by compiler
+    __asm__ __volatile__("");
+    DataCacheCleanAndInvalid<uint8_t, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(global);
+    __asm__ __volatile__("");
+}
+
+FORCE_INLINE_AICORE int32_t gm_signal_wait_until_eq_for_barrier(__gm__ int32_t *sig_addr, int32_t cmp_val) {
+    do {
+        gm_dcci((__gm__ uint8_t *)sig_addr);
+        if (*sig_addr == cmp_val) {
+            return *sig_addr;
+        }
+        if (*sig_addr == cmp_val + 1) {
+            return *sig_addr;
+        }
+    } while (true);
+    return -1;
+}
+
+FORCE_INLINE_AICORE void gm_signal_wait_until_ne(__gm__ int32_t *sig_addr, int32_t cmp_val) {
+    do {
+        AscendC::LocalTensor<int32_t> ub;
+        ub.address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::VECIN);
+        ub.address_.bufferAddr = 0;
+        AscendC::GlobalTensor<int32_t> sig;
+        sig.SetGlobalBuffer(sig_addr);
+        AscendC::DataCopy(ub, sig, 8);
+        AscendC::SetFlag<AscendC::HardEvent::MTE2_S>(EVENT_ID0);
+        AscendC::WaitFlag<AscendC::HardEvent::MTE2_S>(EVENT_ID0);
+        if (ub(0) != cmp_val) {
+            return;
+        }
+    } while (true);
+    return;
+}
+
+
+class HcclShmem {
+public:
+    #ifdef HCCL_COMM    // HCCL needs to initialize the HCCL context
+        __gm__ HcclOpResParamCustom *WinContext_{nullptr};
+        Hccl<HCCL_SERVER_TYPE_AICPU> hccl_;
+        AscendC::LocalTensor<int32_t> ub;
+        FORCE_INLINE_AICORE
+        HcclShmem(){
+            auto contextGM0 = AscendC::GetHcclContext<HCCL_GROUP_ID_0>();
+            WinContext_ = (__gm__ HcclOpResParamCustom *)contextGM0;
+
+            m_rank = WinContext_->localUsrRankId;
+            m_rankSize = WinContext_->rankSize;
+            m_segmentSize = WinContext_->winSize;
+        }
+    #else
+        FORCE_INLINE_AICORE
+        HcclShmem(){
+            m_segmentSize = SHMEM_MEM;
+        }
+        FORCE_INLINE_AICORE 
+        void initShmem(GM_ADDR symmetricPtr_, size_t rank, size_t rankSize) {
+            symmetricPtr = symmetricPtr_;
+            m_rank = rank;
+            m_rankSize = rankSize;
+        }
+    #endif
+
+    FORCE_INLINE_AICORE
+    GM_ADDR operator() () const {   // No parameters: return pointer to local peermem
+        #ifdef HCCL_COMM
+            return (GM_ADDR)(WinContext_->localWindowsIn);
+        #else
+            return reinterpret_cast<GM_ADDR>(shmem_ptr(symmetricPtr, m_rank));
+        #endif
+    }
+
+    FORCE_INLINE_AICORE
+    GM_ADDR operator() (int32_t index) const {  // With index parameter: return pointer to the base address of remote peermem
+        #ifdef HCCL_COMM
+            return (GM_ADDR)((index == m_rank) ? WinContext_->localWindowsIn :
+                                    ((HcclRankRelationResV2Custom *)(WinContext_->remoteRes[index].nextDevicePtr))->windowsIn);
+        #else
+            return reinterpret_cast<GM_ADDR>(shmem_ptr(symmetricPtr, index));
+        #endif
+    }
+
+    FORCE_INLINE_AICORE
+    GM_ADDR operator () (int64_t offset, int32_t rankId) const  {  
+        #ifdef HCCL_COMM
+            if (offset < 0 || offset >= m_segmentSize) {
+                return nullptr;
+            }
+            if (rankId < 0 || rankId >= m_rankSize) {
+                return nullptr;
+            }
+            return (GM_ADDR)((rankId == m_rank) ? WinContext_->localWindowsIn :
+                                    ((HcclRankRelationResV2Custom *)(WinContext_->remoteRes[rankId].nextDevicePtr))->windowsIn) + offset;
+        #else
+            return reinterpret_cast<GM_ADDR>(shmem_ptr((symmetricPtr + offset), rankId));
+        #endif
+    }
+
+
+
+    FORCE_INLINE_AICORE
+    size_t SegmentSize() const {
+        return m_segmentSize;
+    }
+
+    FORCE_INLINE_AICORE
+    int32_t RankSize() const {
+        return m_rankSize;
+    }
+
+
+    FORCE_INLINE_AICORE
+    ~HcclShmem() {
+    }
+
+
+    FORCE_INLINE_AICORE
+    void CrossRankSync() {
+        uint64_t flag_offset = (m_segmentSize - MB_SIZE) / sizeof(int32_t);
+        __gm__ int32_t* sync_counter = (__gm__ int32_t*)(*this)() + flag_offset;
+        __gm__ int32_t* sync_base = (__gm__ int32_t*)(*this)() + flag_offset + 2048;
+        int count = gm_load(sync_base) + 1;
+        int vec_id = AscendC::GetBlockIdx();
+        int vec_size = AscendC::GetBlockNum() * AscendC::GetTaskRation();
+        for(int i = vec_id; i < m_rankSize; i += vec_size) {
+            __gm__ int32_t* sync_remote = (__gm__ int32_t*)((*this)(i)) + flag_offset + m_rank * 16;
+            gm_store(sync_remote, count);
+            gm_dcci((__gm__ uint8_t*)sync_remote);
+            auto sync_check = sync_counter + i * 16;
+            gm_signal_wait_until_eq_for_barrier(sync_check, count);
+        }
+
+        AscendC::SyncAll<true>();
+        gm_store(sync_base, count);
+    }
+
+
+    FORCE_INLINE_AICORE
+    void InitStatusTargetSum()
+    {
+        using namespace AscendC;
+        uint64_t flag_offset = (m_segmentSize - MB_SIZE) + SELF_STATE_OFFSET;
+        //uint64_t self_state_offset = (m_segmentSize - 2 * MB_SIZE);
+        // ep state
+        //uint32_t coreIdx = get_block_idx();;
+        uint32_t coreIdx = GetBlockIdx();
+        GlobalTensor<int32_t> selfStatusTensor;
+        selfStatusTensor.SetGlobalBuffer((__gm__ int32_t *)((*this)() + flag_offset));
+        __asm__ __volatile__("");
+        DataCacheCleanAndInvalid<int32_t, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(selfStatusTensor[coreIdx * UB_ALIGN]);
+        __asm__ __volatile__("");
+        int32_t state = selfStatusTensor(coreIdx * UB_ALIGN);
+        if (state == 0) {
+            sumTarget_ = static_cast<float>(1.0);
+            selfStatusTensor(coreIdx * UB_ALIGN) = 0x3F800000;  // 1.0f
+            epStateValue_ = 0x3F800000;                          // 1.0f
+        } else {
+            sumTarget_ = static_cast<float>(0.0);
+            selfStatusTensor(coreIdx * UB_ALIGN) = 0;
+            epStateValue_ = 0;
+        }
+        __asm__ __volatile__("");
+        DataCacheCleanAndInvalid<int32_t, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(selfStatusTensor[coreIdx * UB_ALIGN]);
+        __asm__ __volatile__("");
+    }
+
+    FORCE_INLINE_AICORE
+    void CrossRankSyncV2Set(AscendC::LocalTensor<int32_t> ctrBuffer) {
+        //subblockid = 0
+        uint32_t stateOffset_ =  STATE_OFFSET;
+        // uint32_t epStateOffsetOnWin_ = m_rank * stateOffset_;
+        
+        uint64_t flag_offset = (m_segmentSize - MB_SIZE) + m_rank * stateOffset_;
+        //uint64_t flag_offset = (m_segmentSize - MB_SIZE);
+        int vec_size = get_block_num();
+        int vec_id = get_block_idx();
+ 
+        AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(RECV_SYNC_EVENT_ID);
+        AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(SEND_SYNC_EVENT_ID);
+        AscendC::CrossCoreWaitFlag(SEND_SYNC_EVENT_ID);
+        pipe_barrier(PIPE_ALL);
+ 
+        ctrBuffer.SetValue(0, epStateValue_);
+        AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
+        AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
+        for (uint32_t dstEpIdx = vec_id; dstEpIdx < m_rankSize; dstEpIdx += vec_size) {
+            AscendC::GlobalTensor<int32_t> gmDstStates;
+            gmDstStates.SetGlobalBuffer((__gm__ int32_t*)((*this)(flag_offset, dstEpIdx)));
+            DataCopy(gmDstStates, ctrBuffer, 8);
+        }
+        AscendC::CrossCoreWaitFlag(RECV_SYNC_EVENT_ID);
+    }
+
+    FORCE_INLINE_AICORE
+    void CrossRankSyncV2Wait(AscendC::LocalTensor<float> statusTensor, AscendC::LocalTensor<float> gatherMaskOutTensor,
+        AscendC::LocalTensor<uint32_t> gatherTmpTensor, AscendC::LocalTensor<float> statusSumOutTensor) {
+
+        uint64_t flag_offset = (m_segmentSize - MB_SIZE);
+        int vec_size = get_block_num();
+        int vec_id = get_block_idx();
+        uint32_t stateOffset_ =  STATE_OFFSET;
+
+        uint32_t sendRankNum_ = m_rankSize / vec_size;
+        uint32_t remainderRankNum = m_rankSize % vec_size;
+        uint32_t startRankId_ = sendRankNum_ * vec_id;
+        if (vec_id < remainderRankNum) {
+            sendRankNum_++;
+            startRankId_ += vec_id;
+        } else {
+            startRankId_ += remainderRankNum;
+        }
+        uint32_t endRankId_ = startRankId_ + sendRankNum_;
+        AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(SEND_SYNC_EVENT_ID);
+
+        AscendC::GlobalTensor<float> epStatusSpaceGlobalTensor_;
+        epStatusSpaceGlobalTensor_.SetGlobalBuffer((__gm__ float *)((*this)() + flag_offset));
+
+        if (startRankId_ < m_rankSize) {
+            AscendC::PipeBarrier<PIPE_ALL>();
+            gatherTmpTensor.SetValue(0, 1);
+            uint32_t mask = 1;  // gatherMask + sum
+            uint64_t rsvdCnt = 0;
+            // DataCopyParams intriParams{static_cast<uint16_t>(sendRankNum_), 1,
+            //                         static_cast<uint16_t>((moeSendNum_ > 512) ? 7 : 15), 0}; 
+            AscendC::DataCopyParams intriParams{static_cast<uint16_t>(sendRankNum_), 1,
+                                    static_cast<uint16_t>(15), 0}; 
+
+            float sumOfFlag = static_cast<float>(-1.0);
+            float minTarget = (sumTarget_ * sendRankNum_) - (float)0.5;
+            float maxTarget = (sumTarget_ * sendRankNum_) + (float)0.5;
+            AscendC::SumParams sumParams{1, sendRankNum_, sendRankNum_};
+
+            AscendC::SetFlag<AscendC::HardEvent::S_V>(EVENT_ID0);
+            AscendC::WaitFlag<AscendC::HardEvent::S_V>(EVENT_ID0);
+
+            while ((sumOfFlag < minTarget) || (sumOfFlag > maxTarget)) {
+                AscendC::DataCopy<float>(statusTensor, epStatusSpaceGlobalTensor_[startRankId_ * stateOffset_ / sizeof(float)],
+                                intriParams);
+                AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
+                AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
+
+                GatherMask(gatherMaskOutTensor, statusTensor, gatherTmpTensor, true, mask,
+                        {1, (uint16_t)sendRankNum_, 1, 0}, rsvdCnt);
+
+                AscendC::PipeBarrier<PIPE_V>();
+                AscendC::Sum(statusSumOutTensor, gatherMaskOutTensor, sumParams);
+                AscendC::SetFlag<AscendC::HardEvent::V_S>(EVENT_ID0);
+                AscendC::WaitFlag<AscendC::HardEvent::V_S>(EVENT_ID0);
+                sumOfFlag = statusSumOutTensor.GetValue(0);
+            }
+        }
+
+        AscendC::CrossCoreSetFlag<0x0, PIPE_MTE3>(RECV_SYNC_EVENT_ID);
+        AscendC::CrossCoreWaitFlag(RECV_SYNC_EVENT_ID);
+
+        //unpermute
+        AscendC::CrossCoreWaitFlag(SEND_SYNC_EVENT_ID);
+    }
+
+
+    FORCE_INLINE_AICORE
+    __gm__ int32_t* SyncBaseAddr() {
+        uint64_t flag_offset = (m_segmentSize - MB_SIZE) / sizeof(int32_t);
+        return (__gm__ int32_t*)(*this)() + flag_offset + 2048;
+    }
+
+private:
+    GM_ADDR symmetricPtr;
+    int32_t m_rank;
+    int32_t m_rankSize;
+    size_t m_segmentSize;
+    float sumTarget_{0.0};
+    int32_t epStateValue_;
+};
+
+
+
+
+#endif

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/layout3d.hpp

@@ -0,0 +1,20 @@
+#ifndef LAYOUT_3D_HPP
+#define LAYOUT_3D_HPP
+#include "kernel_operator.h"
+#include "catlass/catlass.hpp"
+class Layout3D {
+    int64_t strides[2];
+    public:
+    CATLASS_DEVICE
+    Layout3D() {}
+    CATLASS_DEVICE
+    Layout3D(int64_t stride0, int64_t stride1) {
+        strides[0] = stride0;
+        strides[1] = stride1;
+    }
+    CATLASS_DEVICE
+    int64_t operator() (int64_t dim0, int64_t dim1, int64_t dim2) {
+        return dim0 * strides[0] + dim1 * strides[1] + dim2;
+    }
+};
+#endif // LAYOUT_3D_HPP

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/moe_distribute_base.h

@@ -0,0 +1,365 @@
+/**
+ * This program is free software, you can redistribute it and/or modify.
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file moe_distribute_base.h
+ * \brief
+ */
+
+#ifndef MOE_DISTRIBUTE_BASE_H
+#define MOE_DISTRIBUTE_BASE_H
+
+#include "kernel_operator.h"
+
+constexpr uint32_t LOCAL_NOTIFY_MAX_NUM = 64;
+constexpr uint32_t LOCAL_STREAM_MAX_NUM = 19U;
+constexpr uint32_t AICPU_OP_NOTIFY_MAX_NUM = 2;
+constexpr uint32_t AICPU_MAX_RANK_NUM = 128 * 1024;
+constexpr uint32_t TIME_CYCLE = 50; // 系统cycle数转换成时间的基准单位，固定为50
+
+struct HcclSignalInfo {
+    uint64_t resId; // 在代表event时为eventid，notify时为notifyid
+    uint64_t addr;
+    uint32_t devId;
+    uint32_t tsId;
+    uint32_t rankId;
+    uint32_t flag;
+};
+
+struct ListCommon {
+    uint64_t nextHost;
+    uint64_t preHost;
+    uint64_t nextDevice;
+    uint64_t preDevice;
+};
+
+struct HcclStreamInfo {
+    int32_t streamIds;
+    uint32_t sqIds;
+    uint32_t cqIds;      // 记录物理cqId
+    uint32_t logicCqids; // 记录逻辑cqId
+};
+
+struct LocalResInfoV2 {
+    uint32_t streamNum;
+    uint32_t signalNum;
+    HcclSignalInfo localSignals[LOCAL_NOTIFY_MAX_NUM];
+    HcclStreamInfo streamInfo[LOCAL_STREAM_MAX_NUM];
+    HcclStreamInfo mainStreamInfo;
+    HcclSignalInfo aicpuOpNotify[AICPU_OP_NOTIFY_MAX_NUM];  // 集合通信AICPU展开资源
+    ListCommon nextTagRes;                                  // HccltagLocalResV2
+};
+
+enum class rtFloatOverflowMode_t {
+    RT_OVERFLOW_MODE_SATURATION = 0,
+    RT_OVERFLOW_MODE_INFNAN,
+    RT_OVERFLOW_MODE_UNDEF,
+};
+
+struct AlgoTopoInfo {
+    uint32_t userRank;     // 通信域 RankID
+    uint32_t userRankSize; // 通信域的Rank数量
+    int32_t deviceLogicId;
+    bool isSingleMeshAggregation;
+    uint32_t deviceNumPerAggregation;  // 每个Module中的Device数量
+    uint32_t superPodNum;              // 集群中总的超节点数
+    uint32_t devicePhyId;
+    uint32_t topoType; // TopoType
+    uint32_t deviceType;
+    uint32_t serverNum;
+    uint32_t meshAggregationRankSize;
+    uint32_t multiModuleDiffDeviceNumMode;
+    uint32_t multiSuperPodDiffServerNumMode;
+    uint32_t realUserRank;
+    bool isDiffDeviceModule;
+    bool isDiffDeviceType;
+    uint32_t gcdDeviceNumPerAggregation;
+    uint32_t moduleNum;
+    uint32_t isUsedRdmaRankPairNum;
+    uint64_t isUsedRdmaRankPair;
+    uint32_t pairLinkCounterNum;
+    uint64_t pairLinkCounter;
+    uint32_t nicNum;
+    uint64_t nicList;           // niclist数组指针
+    uint64_t complanRankLength; // complanRank占用的字节数
+    uint64_t complanRank;       // 指针
+    uint64_t bridgeRankNum;     // bridgeRank占用的个数
+    uint64_t bridgeRank;        // 指针
+    uint64_t serverAndsuperPodRankLength; // serverAndsuperPodRank占用的字节数
+    uint64_t serverAndsuperPodRank; // 指针
+};
+
+struct HcclOpConfig {
+    uint8_t deterministic; //确定性计算开关
+    uint8_t retryEnable;   // 是否重执行
+    uint8_t highPerfEnable;
+    uint8_t padding[5];    // 大小需要64By对齐，未来添加参数时减小padding
+    uint8_t linkTimeOut[8]; // 发送超时时长
+    uint64_t notifyWaitTime; // 超时时长，同HCCL_EXEC_TIMEOUT
+    uint32_t retryHoldTime;
+    uint32_t retryIntervalTime;
+    bool interHccsDisable = false; //使能rdma开关
+    rtFloatOverflowMode_t floatOverflowMode = rtFloatOverflowMode_t::RT_OVERFLOW_MODE_UNDEF;
+    uint32_t multiQpThreshold = 512;  // 多QP每个QP分担数据量最小阈值
+};
+
+struct HcclMC2WorkSpace {
+    uint64_t workSpace;
+    uint64_t workSpaceSize;
+};
+
+struct RemoteResPtr {
+    uint64_t nextHostPtr;
+    uint64_t nextDevicePtr;
+};
+
+struct HDCommunicateParams {
+    uint64_t hostAddr { 0 };
+    uint64_t deviceAddr { 0 };
+    uint64_t readCacheAddr { 0 };
+    uint32_t devMemSize{ 0 };
+    uint32_t buffLen{ 0 };
+    uint32_t flag{ 0 };
+};
+
+struct HcclRankRelationResV2Custom {
+    uint32_t remoteUsrRankId;
+    uint32_t remoteWorldRank;
+    uint64_t windowsIn;
+    uint64_t windowsOut;
+    uint64_t windowsExp;
+    ListCommon nextTagRes;
+};
+
+struct HcclOpResParamCustom {
+    // 本地资源
+    HcclMC2WorkSpace mc2WorkSpace;
+    uint32_t localUsrRankId; // usrrankid
+    uint32_t rankSize;       // 通信域内total rank个数
+    uint64_t winSize; // 每个win大小，静态图时，可能是0，如果通信域内也有动态图，则可能为非0
+    uint64_t localWindowsIn; // 全F为无效值
+    uint64_t localWindowsOut; // 全F为无效值
+    char hcomId[128];
+    // aicore识别remote window
+    uint64_t winExpSize;
+    uint64_t localWindowsExp;
+    uint32_t rWinStart; // 为HcclRankRelationRes起始位置
+    uint32_t rWinOffset; // 为HcclRemoteRes的大小
+    uint64_t version;
+    LocalResInfoV2 localRes;
+    AlgoTopoInfo topoInfo;
+
+    // 外部配置参数
+    HcclOpConfig config;
+    uint64_t hostStateInfo;
+    uint64_t aicpuStateInfo;
+    uint64_t lockAddr;
+    uint32_t rsv[16];
+    uint32_t notifysize;                         // RDMA场景使用，910B/910_93为4B，其余芯片为8B
+    uint32_t remoteResNum;                       // 有效的remoteResNum
+    RemoteResPtr remoteRes[AICPU_MAX_RANK_NUM];  //数组指针，指向HcclRankRelationResV2，下标为remoteUserRankId
+
+    // communicate retry
+    HDCommunicateParams kfcControlTransferH2DParams;
+    HDCommunicateParams kfcStatusTransferD2HParams;
+    uint64_t tinyMem;  // for all2all
+    uint64_t tinyMemSize;
+    // 零拷贝场景使用
+    uint64_t zeroCopyHeadPtr;
+    uint64_t zeroCopyTailPtr;
+    uint64_t zeroCopyRingBuffer;
+    uint64_t zeroCopyIpcPtrs[16];                // 保存集合通信时每个对端的输入输出内存地址
+    uint32_t zeroCopyDevicePhyId[16];            // 保存每个rank对应的物理卡Id
+
+    bool utraceStatusFlag;
+};
+
+// Transport 内存类型
+enum class HcclAiRMAMemType : uint32_t {
+    LOCAL_INPUT = 0,
+    REMOTE_INPUT,
+
+    LOCAL_OUTPUT,
+    REMOTE_OUTPUT,
+
+    // 可透传更多的内存，可在MAX_NUM之前追加，例如：
+    // LOCAL_EXP,
+    // REMOTE_EXP,
+    MAX_NUM
+};
+
+// Transport 内存信息
+struct HcclAiRMAMemInfo {
+    uint32_t memMaxNum{0};  // 最大内存数量，等于 HcclAiRMAMemType::MAX_NUM
+    uint32_t sizeOfMemDetails{0};  // sizeof(MemDetails)，用于内存校验和偏移计算
+    uint64_t memDetailPtr{0};  // MemDetails数组首地址, 个数: HcclAiRMAMemType::MAX_NUM
+    // 可往后追加字段
+};
+
+// 全部 Transport QP/Mem 信息
+struct HcclAiRMAInfo {
+    uint32_t curRankId{0};  // 当前rankId
+    uint32_t rankNum{0};  // rank数量
+    uint32_t qpNum{0};  // 单个Transport的QP数量
+
+    uint32_t sizeOfAiRMAWQ{0};  // sizeof(HcclAiRMAWQ)
+    uint32_t sizeOfAiRMACQ{0};  // sizeof(HcclAiRMACQ)
+    uint32_t sizeOfAiRMAMem{0};  // sizeof(HcclAiRMAMemInfo)
+
+    // HcclAiRMAWQ二维数组首地址
+    // QP个数: rankNum * qpNum
+    // 计算偏移获取SQ指针：sqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMAWQ
+    // 0 <= qpIndex < qpNum
+    uint64_t sqPtr{0};
+
+    // HcclAiRMACQ二维数组首地址
+    // QP个数: rankNum * qpNum
+    // 计算偏移获取SCQ指针：scqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMACQ
+    // 0 <= qpIndex < qpNum
+    uint64_t scqPtr{0};
+
+    // HcclAiRMAWQ二维数组首地址
+    // QP个数: rankNum * qpNum
+    // 计算偏移获取RQ指针：rqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMAWQ
+    // 0 <= qpIndex < qpNum
+    uint64_t rqPtr{0};
+
+    // HcclAiRMACQ二维数组首地址
+    // QP个数: rankNum * qpNum
+    // 计算偏移获取RCQ指针: rcqPtr + (dstRankId * qpNum + qpIndex) * sizeOfAiRMACQ
+    // 0 <= qpIndex < qpNum
+    uint64_t rcqPtr{0};
+
+    // HcclAivMemInfo一维数组
+    // 内存信息个数: rankNum
+    // 计算偏移获取内存信息指针: memPtr + rankId * sizeOfAiRMAMem
+    // srcRankId 获取自身内存信息，dstRankId 获取 Transport 内存信息
+    uint64_t memPtr{0};
+    // 可往后追加字段
+};
+struct CombinedCapability {
+    uint64_t dataplaneModeBitmap;
+};
+
+struct HcclA2CombineOpParam {
+    uint64_t workSpace;                         // Address for communication between client and server,
+                                                // hccl requests and clears
+    uint64_t workSpaceSize;                     // Space for communication between client and server
+    uint32_t rankId;                            // id of this rank
+    uint32_t rankNum;                           // num of ranks in this comm group
+    uint64_t winSize;                           // size of each windows memory
+    uint64_t windowsIn[AscendC::HCCL_MAX_RANK_NUM];      // windows address for input, windowsIn[rankId] corresponds
+                                                // to the local card address,
+                                                // and others are cross-card mapping addresses.
+    uint64_t windowsOut[AscendC::HCCL_MAX_RANK_NUM];     // windows address for output, windowsOut[rankId] corresponds
+                                                // to the local card address,
+                                                // and others are cross-card mapping addresses.
+    uint8_t res[8328];
+    uint8_t multiFlag;
+    __gm__ AscendC::IbVerbsData *data;
+    uint64_t dataSize;
+    // 追加字段
+    uint64_t sizeOfAiRMAInfo; // sizeof(HcclAiRMAInfo)
+    uint64_t aiRMAInfo; // HcclAiRMAInfo* 单个结构体指针
+
+    CombinedCapability* capability;             // address of the communication capability information structure on the Device
+    uint64_t capabilitySize;                    // size of the communication capability information structure
+};
+enum class DataplaneMode : uint32_t {
+    HOST = 0,
+    AICPU = 1,
+    AIV = 2,
+};
+
+enum class DBMode : int32_t {
+    INVALID_DB = -1,
+    HW_DB = 0,
+    SW_DB
+};
+
+struct HcclAiRMAWQ {
+    uint32_t wqn{0};
+    uint64_t bufAddr{0};
+    uint32_t wqeSize{0};
+    uint32_t depth{0};
+    uint64_t headAddr{0};
+    uint64_t tailAddr{0};
+    DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw
+    uint64_t dbAddr{0};
+    uint32_t sl{0};
+};
+
+struct HcclAiRMACQ {
+    uint32_t cqn{0};
+    uint64_t bufAddr{0};
+    uint32_t cqeSize{0};
+    uint32_t depth{0};
+    uint64_t headAddr{0};
+    uint64_t tailAddr{0};
+    DBMode dbMode{DBMode::INVALID_DB}; // 0-hw/1-sw
+    uint64_t dbAddr{0};
+};
+
+struct hns_roce_rc_sq_wqe {
+    uint32_t byte_4;
+    uint32_t msg_len;
+    uint32_t immtdata;
+    uint32_t byte_16;
+    uint32_t byte_20;
+    uint32_t rkey;
+    uint64_t remoteVA;
+};
+
+
+struct hns_roce_lite_wqe_data_seg {
+    uint32_t len;
+    uint32_t lkey;
+    uint64_t localVA;
+};
+
+__aicore__ inline void cacheWriteThrough(__gm__ uint8_t* sourceAddr, uint64_t length) {
+    __gm__ uint8_t* start =
+        (__gm__ uint8_t*)((uint64_t)sourceAddr / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE);
+    __gm__ uint8_t* end =
+        (__gm__ uint8_t*)(((uint64_t)sourceAddr + length) / AscendC::CACHE_LINE_SIZE * AscendC::CACHE_LINE_SIZE);
+    AscendC::GlobalTensor<uint8_t> global;
+    global.SetGlobalBuffer(start);
+    for (uint32_t i = 0; i <= end - start; i += AscendC::CACHE_LINE_SIZE) {
+        AscendC::DataCacheCleanAndInvalid<uint8_t, AscendC::CacheLine::SINGLE_CACHE_LINE,
+            AscendC::DcciDst::CACHELINE_OUT>(global[i]);
+    }
+}
+__aicore__ inline DataplaneMode GetDataplaneMode(GM_ADDR contextGM0) {
+    __gm__ HcclA2CombineOpParam *winContext_ = (__gm__ HcclA2CombineOpParam *)contextGM0;
+    CombinedCapability* capability = winContext_->capability;
+    uint64_t capabilitySize = winContext_->capabilitySize;
+    DataplaneMode dataplaneMode = DataplaneMode::AICPU;
+    if (capability == 0) {
+        return dataplaneMode;
+    }
+    uint64_t dataplaneModeBitmap = capability->dataplaneModeBitmap;
+    if ((dataplaneModeBitmap & 0x04) == 0x04) {
+        dataplaneMode = DataplaneMode::AIV;
+    }
+    return dataplaneMode;
+}
+
+__aicore__ inline int64_t GetCurrentTimestampUs()
+{
+    return AscendC::GetSystemCycle() / TIME_CYCLE;
+}
+
+__aicore__ inline void RecordRankCommDuration(AscendC::LocalTensor<int32_t> performanceInfoU32Tensor, uint32_t rankId, int64_t startTime)
+{
+    int64_t endTime = GetCurrentTimestampUs();
+    int32_t duration = static_cast<int32_t>(endTime - startTime); // int32_t可以表示2^31(us)，约35min在实际场景下满足需要
+    performanceInfoU32Tensor.SetValue(rankId * sizeof(int64_t) / sizeof(int32_t), duration); // 使用int32_t是因为atomicAdd不支持int64_t类型，这里只赋值到int64_t的低32位。
+}
+#endif // MOE_DISTRIBUTE_BASE_H

csrc/mc2/dispatch_ffn_combine_bf16/op_kernel/utils/select_helper.hpp

@@ -0,0 +1,25 @@
+#ifndef SELECT_HELPER_HPP
+#define SELECT_HELPER_HPP
+
+#include "catlass/layout/layout.hpp"
+using namespace AscendC;
+using namespace Catlass;
+
+template <typename Layout, typename ElementType, typename = void>
+struct LayoutBInitializer {
+    CATLASS_DEVICE
+    static Layout create(uint32_t k, uint32_t n) {
+        return Layout{k, n};
+    }
+};
+
+template <typename Layout, typename ElementType>
+struct LayoutBInitializer<Layout, ElementType,
+    std::enable_if_t<std::is_same_v<Layout, layout::zN>>
+> {
+    CATLASS_DEVICE
+    static Layout create(uint32_t k, uint32_t n) {
+        return Layout::template MakeLayout<ElementType>(k, n);
+    }
+};
+#endif // SELECT_HELPER_HPP

csrc/mc2/dispatch_ffn_combine_w4_a8/CMakeLists.txt

@@ -0,0 +1,19 @@
+# -----------------------------------------------------------------------------------------------------------
+# Copyright (c) 2025 Huawei Technologies Co., Ltd.
+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+# CANN Open Software License Agreement Version 2.0 (the "License").
+# Please refer to the License for details. You may not use this file except in compliance with the License.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+# See LICENSE in the root of the software repository for the full text of the License.
+# -----------------------------------------------------------------------------------------------------------
+
+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
+if(NOT ENABLE_TEST AND NOT BENCHMARK)
+    list(REMOVE_ITEM CURRENT_DIRS tests)
+endif()
+foreach(SUB_DIR ${CURRENT_DIRS})
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
+        add_subdirectory(${SUB_DIR})
+    endif()
+endforeach()

csrc/mc2/dispatch_ffn_combine_w4_a8/op_host/CMakeLists.txt

@@ -0,0 +1,38 @@
+set(_DISPATCH_FFN_INC_OPTS)
+if (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/aarch64-linux/ascendc/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/aarch64-linux/ascendc/include)
+elseif (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/arm64-linux/ascendc/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/arm64-linux/ascendc/include)
+elseif (EXISTS ${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/ascendc/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${ASCEND_CANN_PACKAGE_PATH}/${CMAKE_SYSTEM_PROCESSOR}-linux/ascendc/include)
+endif()
+if (EXISTS ${CMAKE_SOURCE_DIR}/third_party/catlass/include)
+    list(APPEND _DISPATCH_FFN_INC_OPTS -I${CMAKE_SOURCE_DIR}/third_party/catlass/include)
+endif()
+
+add_op_to_compiled_list()
+
+if (BUILD_OPEN_PROJECT)
+    target_sources(op_host_aclnnInner PRIVATE
+        dispatch_ffn_combine_w4_a8_def.cpp
+    )
+endif()
+
+add_ops_compile_options(
+    OP_NAME DispatchFFNCombineW4A8
+    OPTIONS
+        --cce-auto-sync=on
+        -Wno-deprecated-declarations
+        -DHCCL_COMM
+        -DCATLASS_ARCH=2201
+        ${_DISPATCH_FFN_INC_OPTS}
+)
+
+if (NOT BUILD_OPS_RTY_KERNEL)
+    add_modules_sources(OPTYPE dispatch_ffn_combine_w4_a8 ACLNNTYPE aclnn_inner)
+    target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
+        ${CMAKE_CURRENT_SOURCE_DIR}
+        ${CMAKE_CURRENT_SOURCE_DIR}/../op_kernel
+    )
+endif()
+

csrc/mc2/dispatch_ffn_combine_w4_a8/op_host/dispatch_ffn_combine_w4_a8_def.cpp

@@ -0,0 +1,108 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file dispatch_ffn_combine_w4_a8_def.cpp
+ * \brief
+ */
+#include "register/op_def_registry.h"
+
+namespace ops {
+class DispatchFFNCombineW4A8 : public OpDef {
+ public:
+  explicit DispatchFFNCombineW4A8(const char *name) : OpDef(name) {
+    this->Input("a")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("w1")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_INT4, ge::DT_INT4, ge::DT_INT4, ge::DT_INT4, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ})
+        .IgnoreContiguous();
+    this->Input("w2")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_INT4, ge::DT_INT4, ge::DT_INT4, ge::DT_INT4, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_FRACTAL_NZ, ge::FORMAT_FRACTAL_NZ})
+        .IgnoreContiguous();
+    this->Input("expertIdx")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("scale1")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("scale2")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("bias1")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("bias2")
+        .ParamType(DYNAMIC)
+        .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("probs")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Input("xActiveMaskOptional")
+        .ParamType(OPTIONAL)
+        .DataType({ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+
+    // 输出
+    this->Output("out")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+    this->Output("expert_token_nums")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
+        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
+        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
+
+    this->Attr("group").AttrType(REQUIRED).String();
+    this->Attr("M").AttrType(OPTIONAL).Int();
+    this->Attr("transB").AttrType(OPTIONAL).Bool(false);
+    this->Attr("weightNz").AttrType(OPTIONAL).Bool(false);
+    this->Attr("swigluLimit").AttrType(OPTIONAL).Float(0.0f);
+
+    OpAICoreConfig aicore_config;
+    aicore_config.DynamicCompileStaticFlag(true)
+        .DynamicFormatFlag(true)
+        .DynamicRankSupportFlag(true)
+        .DynamicShapeSupportFlag(true)
+        .NeedCheckSupportFlag(false)
+        .PrecisionReduceFlag(true)
+        .ExtendCfgInfo("aclnnSupport.value", "support_aclnn")
+        .ExtendCfgInfo("jitCompile.flag", "static_false")
+        .ExtendCfgInfo("multiKernelSupportDynamicGraph.value", "multi_kernel");
+    this->AICore().AddConfig("ascend910_93", aicore_config);
+    this->MC2().HcclGroup("group");
+  }
+};
+
+OP_ADD(DispatchFFNCombineW4A8);
+}  // namespace ops

csrc/mc2/dispatch_ffn_combine_w4_a8/op_host/dispatch_ffn_combine_w4_a8_proto.cpp

@@ -0,0 +1,37 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+
+/*!
+ * \file dispatch_ffn_w4_a8_proto.cpp
+ * \brief
+ */
+#include <graph/utils/type_utils.h>
+#include <register/op_impl_registry.h>
+
+using namespace ge;
+namespace ops {
+const size_t ATTR_GROUP = 0;
+const size_t ATTR_RANK_SIZE = 1;
+const size_t SUPPORT_DIM_SIZE = 2;
+
+static ge::graphStatus InferShapeDispatchFFNCombineW4A8(gert::InferShapeContext* context) {
+  (void) context;
+  return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus InferDataTypeDispatchFFNCombineW4A8(gert::InferDataTypeContext* context) {
+  (void) context;
+  return ge::GRAPH_SUCCESS;
+}
+
+IMPL_OP_INFERSHAPE(DispatchFFNCombineW4A8)
+  .InferShape(InferShapeDispatchFFNCombineW4A8)
+  .InferDataType(InferDataTypeDispatchFFNCombineW4A8);
+}  // namespace ops

csrc/mc2/dispatch_ffn_combine_w4_a8/op_host/dispatch_ffn_combine_w4_a8_tiling.cpp

@@ -0,0 +1,294 @@
+/**
+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ */
+/*!
+ * \file dispatch_ffn_combine_tiling_w4_a8_tiling.cpp
+ * \brief
+ */
+#include "vector"
+#include "register/tilingdata_base.h"
+#include "tiling/tiling_api.h"
+#include "tiling_base/error_log.h"
+#include "hcom_topo_info.h"
+#include "register/op_def_registry.h"
+#include "dispatch_ffn_combine_w4_a8_tiling.h"
+#include <vector>
+#include <map>
+#include <algorithm>
+#include "moe_init_routing_quant_v2/moe_init_routing_quant_v2_tiling.h"
+
+using namespace AscendC;
+using namespace ge;
+
+namespace {
+    const char *K_INNER_DEBUG = "DispatchFFNCombineW4A8 Tiling Debug";
+    constexpr uint32_t ATTR_GROUP_INDEX = 0;
+    constexpr uint32_t ATTR_MAX_OUTPUT_SIZE_INDEX = 1;
+    constexpr uint32_t ATTR_IS_TRANS_B = 2;
+    constexpr uint32_t ATTR_WEIGHT_NZ = 3;
+    constexpr uint32_t ATTR_SWIGLU_LIMIT = 4;
+    constexpr uint64_t INIT_TILINGKEY = 1000000;
+    constexpr uint64_t TILINGKEY_TRANS_B = 1U;
+    constexpr uint64_t TILINGKEY_WEIGHT_NZ = 10;
+    constexpr uint32_t X_INDEX = 0;
+    constexpr uint32_t WEIGHT_INDEX = 1;
+    constexpr uint32_t WEIGHT2_INDEX = 2;
+    constexpr uint32_t EXPERTID_INDEX = 3;
+    constexpr uint32_t X_ACTIVE_MASK_INDEX = 9;
+    constexpr uint32_t BLOCK_NUM = 20;
+    constexpr uint32_t SYSTEM_NEED_WORKSPACE = 16 * 1024 * 1024;
+}
+
+namespace optiling {
+
+static ge::graphStatus DispatchFFNCombineW4A8CheckAttrAndSetTiling(gert::TilingContext *context, DispatchFFNCombineW4A8Info& info)
+{
+    auto attrs = context->GetAttrs();
+    OP_TILING_CHECK(attrs == nullptr, OP_LOGE(K_INNER_DEBUG, "attrs is null."), return ge::GRAPH_FAILED);
+
+    auto groupPtr = attrs->GetAttrPointer<char>(static_cast<int>(ATTR_GROUP_INDEX));
+    auto maxOutputSizePtr = attrs->GetAttrPointer<int>(ATTR_MAX_OUTPUT_SIZE_INDEX);
+    auto is_trans_b = attrs->GetAttrPointer<bool>(ATTR_IS_TRANS_B);
+    auto weight_nz = attrs->GetAttrPointer<bool>(ATTR_WEIGHT_NZ);
+    auto swiglu_limit = attrs->GetAttrPointer<float>(ATTR_SWIGLU_LIMIT);
+    OP_TILING_CHECK(groupPtr == nullptr || strlen(groupPtr) == 0,
+    OP_LOGE(K_INNER_DEBUG, "group is invalid."), return GRAPH_FAILED);
+
+    OP_TILING_CHECK(is_trans_b == nullptr,
+        OP_LOGE(K_INNER_DEBUG, "is_trans_b is invalid."), return GRAPH_FAILED);
+    OP_TILING_CHECK(weight_nz == nullptr,
+        OP_LOGE(K_INNER_DEBUG, "weight_nz is invalid."), return GRAPH_FAILED);
+
+    info.maxOutputSize = *maxOutputSizePtr;
+    info.isTransposeB = *is_trans_b;
+    info.isWeightNz = *weight_nz;
+    info.swigluLimit = swiglu_limit != nullptr ? *swiglu_limit : 0.0f;
+
+    int64_t rankSize;
+    (void)ge::HcomTopoInfo::Instance().GetGroupRankSize(groupPtr, rankSize);
+    info.worldSize = rankSize;
+
+    OP_LOGD(K_INNER_DEBUG, "maxOutputSize=%d ", info.maxOutputSize);
+    OP_LOGD(K_INNER_DEBUG, "rankSize=%d ", info.worldSize);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus DispatchFFNCombineW4A8CheckShapeAndSetTiling(gert::TilingContext *context, DispatchFFNCombineW4A8Info &info)
+{
+    const gert::StorageShape *aStorageShape = context->GetInputShape(X_INDEX);
+    auto expertIdxTensor = context->GetDynamicInputTensor(EXPERTID_INDEX, 0);
+    uint32_t M = aStorageShape->GetStorageShape().GetDim(0);
+    uint32_t K = aStorageShape->GetStorageShape().GetDim(1);
+
+    auto wTensor = context->GetDynamicInputTensor(WEIGHT_INDEX, 0);
+    auto wDataType = wTensor->GetDataType();
+    uint32_t wTensorDims = wTensor->GetOriginShape().GetDimNum();
+    uint32_t N = wTensor->GetStorageShape().GetDim(wTensorDims - 1);
+    if(ge::DT_INT32 == wDataType)
+    {
+        N = N * 8;
+    }
+
+    uint32_t topK = expertIdxTensor->GetStorageShape().GetDim(1);
+    uint32_t listLen = 0;
+    while (true) {
+        auto wTensorT = context->GetDynamicInputTensor(WEIGHT_INDEX, ++listLen);
+        if (wTensorT == nullptr) {break;}
+    }
+
+    uint32_t expertPerRank;
+    if (listLen == 1) {
+        expertPerRank = wTensor->GetStorageShape().GetDim(0);
+    } else {
+        expertPerRank = listLen;
+    }
+
+    info.M = M;
+    info.N = N;
+    info.K = K;
+    info.expertPerRank = expertPerRank;
+    info.topK = topK;
+    info.listLen = listLen;
+    OP_LOGD(K_INNER_DEBUG, "M=%d ", info.M);
+    OP_LOGD(K_INNER_DEBUG, "K=%d ", info.K);
+    OP_LOGD(K_INNER_DEBUG, "N=%d ", info.N);
+    OP_LOGD(K_INNER_DEBUG, "expertPerRank=%d ", info.expertPerRank);
+    OP_LOGD(K_INNER_DEBUG, "topK=%d ", info.topK);
+    OP_LOGD(K_INNER_DEBUG, "listLen=%d ", info.listLen);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus DispatchFFNCombineW4A8GetPlatformInfoAndSetTiling(gert::TilingContext *context, DispatchFFNCombineW4A8Info& info)
+{
+    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
+    uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
+    uint64_t ubSize = 0U;
+    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
+    info.aivNum = aivNum;
+    info.totalUbSize = ubSize;
+
+    OP_LOGD(K_INNER_DEBUG, "aivNum=%d", info.aivNum);
+    OP_LOGD(K_INNER_DEBUG, "ubSize=%lu", info.totalUbSize);
+
+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus CheckXActiveMaskShape(gert::TilingContext *context, const char *nodeName, DispatchFFNCombineW4A8Info &info)
+{
+    const gert::StorageShape* xActiveMaskStorageShape = context->GetOptionalInputShape(X_ACTIVE_MASK_INDEX);
+    if (xActiveMaskStorageShape != nullptr) {
+        OP_TILING_CHECK(xActiveMaskStorageShape->GetStorageShape().GetDimNum() != 1,
+            OP_LOGE(nodeName, "xActiveMask shape dims must be 1, but current dim num is %lu.",
+                xActiveMaskStorageShape->GetStorageShape().GetDimNum()), return ge::GRAPH_FAILED);
+        const int64_t xActiveMaskDim0 = xActiveMaskStorageShape->GetStorageShape().GetDim(0);
+        OP_TILING_CHECK(xActiveMaskDim0 != static_cast<int64_t>(info.M),
+            OP_LOGE(nodeName, "xActiveMask Dim0 must be M(%u), but current dim is %lu.", info.M, xActiveMaskDim0),
+            return ge::GRAPH_FAILED);
+    }
+    return ge::GRAPH_SUCCESS;
+}
+
+void SetTilingData(CoCTiling &cocTilingData, DispatchFFNCombineW4A8Info &info)
+{
+    cocTilingData.m0 = 128;
+    cocTilingData.k0 = 256;
+    cocTilingData.n0 = 256;
+    cocTilingData.swizzleDirect = 1;
+    cocTilingData.swizzleOffset = 7;
+    cocTilingData.ubMoveNum = 16 * 1024;
+    cocTilingData.pValue = 1;
+    cocTilingData.commNpuSplit = info.worldSize;
+    cocTilingData.commDataSplit = 1;
+    cocTilingData.lenPerLoop = cocTilingData.m0 * cocTilingData.n0 / 2;
+}
+
+static ge::graphStatus DispatchFFNCombineW4A8TilingFuncImpl(gert::TilingContext *context)
+{
+    const char *nodeName = context->GetNodeName();
+    OP_LOGI(nodeName, "Enter DispatchFFNCombineW4A8 tiling func.");
+
+    // 1. tilingData
+    DispatchFFNCombineW4A8TilingData *tilingData = context->GetTilingData<DispatchFFNCombineW4A8TilingData>();
+    OP_TILING_CHECK(tilingData == nullptr, OP_LOGE(nodeName, "tilingData is nullptr."),
+        return ge::GRAPH_FAILED);
+    OP_LOGI(nodeName, "DispatchFFNCombineW4A8 get tilingData.");
+    DispatchFFNCombineW4A8Info& info = tilingData->dispatchFFNCombineW4A8Info;
+    OP_LOGI(nodeName, "DispatchFFNCombineW4A8 get tilingData info.");
+
+    OP_TILING_CHECK(DispatchFFNCombineW4A8CheckAttrAndSetTiling(context, info) != ge::GRAPH_SUCCESS,
+        OP_LOGE(context->GetNodeName(), "DispatchFFNCombineW4A8 CheckAttrAndSetTiling Failed"),
+        return ge::GRAPH_FAILED);
+    OP_TILING_CHECK(DispatchFFNCombineW4A8CheckShapeAndSetTiling(context, info) != ge::GRAPH_SUCCESS,
+        OP_LOGE(context->GetNodeName(), "DispatchFFNCombineW4A8 CheckShapeAndSetTiling Failed"),
+        return ge::GRAPH_FAILED);
+    OP_TILING_CHECK(DispatchFFNCombineW4A8GetPlatformInfoAndSetTiling(context, info) != ge::GRAPH_SUCCESS,
+        OP_LOGE(context->GetNodeName(), "DispatchFFNCombineW4A8 GetPlatformInfoAndSetTiling Failed"),
+        return ge::GRAPH_FAILED);
+    OP_TILING_CHECK(CheckXActiveMaskShape(context, nodeName, info) != ge::GRAPH_SUCCESS,
+        OP_LOGE(context->GetNodeName(), "DispatchFFNCombineW4A8 CheckXActiveMaskShape Failed"),
+        return ge::GRAPH_FAILED);
+
+    SetTilingData(tilingData->cocTiling, info);
+
+    // 2. set blockDim
+    uint32_t blockDim = 1U;
+    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
+    auto aicNum = ascendcPlatform.GetCoreNumAic();
+    auto aivNum = ascendcPlatform.GetCoreNumAiv();
+    blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum);
+    context->SetBlockDim(blockDim);
+
+    // 3. set tilingKey
+    uint64_t tilingKey = INIT_TILINGKEY;
+    tilingKey += info.isTransposeB ? TILINGKEY_TRANS_B : 0;
+    tilingKey += info.isWeightNz ? TILINGKEY_WEIGHT_NZ : 0;
+    context->SetTilingKey(tilingKey);
+
+    OP_LOGD(K_INNER_DEBUG, "tilingKey=%d", tilingKey);
+
+    optiling::MoeInitRoutingQuantV2TilingBase moeInitRoutingQuantV2TilingBase;
+    int64_t inuptXDtypeSize = sizeof(int16_t);
+    int64_t scaleDim0 = 0;
+    int64_t ubSize = 196352;
+    int64_t expertCapacity = 0;
+    int64_t expertNum = info.expertPerRank * info.worldSize + 1;
+    int64_t activeNum = info.M * info.topK;
+    int64_t dropPadMode = 0;
+    int64_t expertTokensCountOrCumsumFlag = 2;
+    bool expertTokensBeforeCapacityFlag = false;
+    int64_t quantMode = 1;
+    uint32_t aivNumInitRouting = 2 * BLOCK_NUM;
+    moeInitRoutingQuantV2TilingBase.DoTiling(info.M, info.K, info.topK, expertCapacity, expertNum, activeNum, dropPadMode,
+        expertTokensCountOrCumsumFlag, expertTokensBeforeCapacityFlag, inuptXDtypeSize, quantMode, scaleDim0, aivNumInitRouting, ubSize);
+    uint64_t initRoutingQuantTilingKey = moeInitRoutingQuantV2TilingBase.tilingKey_;
+    size_t initRoutingWorkspace = moeInitRoutingQuantV2TilingBase.workspaceSize_;
+
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData = moeInitRoutingQuantV2TilingBase.quantTilingData;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.vbsComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.vbsComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.vmsMiddleComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.vmsMiddleComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.sortOutComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.sortOutComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.srcToDstComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.srcToDstComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.srcToDstCapacityComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.srcToDstCapacityComputeParamsOp;
+    tilingData->cocTiling.moeInitRoutingQuantV2TilingData.gatherOutComputeParamsOp = moeInitRoutingQuantV2TilingBase.quantTilingData.gatherOutComputeParamsOp;
+    tilingData->cocTiling.initRoutingQuantTilingKey = initRoutingQuantTilingKey;
+    // OP_LOGE(initRoutingTilingKey, " initRoutingTilingKey.");
+    OP_LOGD(K_INNER_DEBUG, "tilingKey=%ld", initRoutingQuantTilingKey);
+
+    // 4. workspace
+    size_t *workSpaces = context->GetWorkspaceSizes(1);
+    OP_TILING_CHECK(workSpaces == nullptr, OP_LOGE(nodeName, "workSpaces is nullptr."),
+        return ge::GRAPH_FAILED);
+
+    uint32_t n2 = info.K;
+    uint32_t k2 = info.N / 2;
+
+    // ptrC/ptrC2 alias reuse: allocate max(N, n2), not N + n2
+    // ptrA1Int4/ptrA2Int4 alias reuse: allocate max(K, k2), not K + k2
+    // (lifetimes separated by SyncAll barrier after GMM1 loop, same pattern as W8A8 ptrA/ptrPermutedToken)
+    uint64_t cocWorkspace = (info.M + 256 - 1) / 256 * 256 * info.topK * sizeof(int32_t) +
+                            info.worldSize * info.worldSize * info.expertPerRank * sizeof(int32_t) * 2 +
+                            info.maxOutputSize * sizeof(float) * 2 +
+                            info.maxOutputSize * std::max(info.N, n2) * sizeof(int16_t) * 2 +
+                            info.maxOutputSize * std::max(info.K, k2) +
+                            info.worldSize * sizeof(int32_t) * 16;
+
+    workSpaces[0] = SYSTEM_NEED_WORKSPACE + std::max(cocWorkspace, initRoutingWorkspace);
+
+
+    // 5. communication
+    auto attrs = context->GetAttrs();
+    auto group = attrs->GetAttrPointer<char>(static_cast<int>(ATTR_GROUP_INDEX));
+    uint32_t opType = 8U;
+    std::string algConfig = "AlltoAll=level0:fullmesh;level1:pairwise";
+    AscendC::Mc2CcTilingConfig mc2CcTilingConfig(group, opType, algConfig);
+    mc2CcTilingConfig.GetTiling(tilingData->mc2InitTiling);
+    mc2CcTilingConfig.GetTiling(tilingData->mc2CcTiling);
+
+    OP_LOGI(nodeName, "Leave DispatchFFNCombineW4A8 tiling func.");
+    return ge::GRAPH_SUCCESS;
+}
+
+static ge::graphStatus DispatchFFNCombineW4A8TilingFunc(gert::TilingContext* context)
+{
+    return DispatchFFNCombineW4A8TilingFuncImpl(context);
+}
+
+struct DispatchFFNCombineW4A8CompileInfo {};
+ge::graphStatus TilingParseForDispatchFFNCombineW4A8(gert::TilingParseContext *context)
+{
+    (void)context;
+    return ge::GRAPH_SUCCESS;
+}
+
+IMPL_OP_OPTILING(DispatchFFNCombineW4A8)
+    .Tiling(DispatchFFNCombineW4A8TilingFunc)
+    .TilingParse<DispatchFFNCombineW4A8CompileInfo>(TilingParseForDispatchFFNCombineW4A8);
+} // namespace optiling

csrc/mc2/dispatch_ffn_combine_w4_a8/op_host/hcom_topo_info.h

@@ -0,0 +1,72 @@
+/* Copyright (c) 2025 Huawei Technologies Co., Ltd.
+ * This file is a part of the CANN Open Software.
+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
+ * Please refer to the License for details. You may not use this file except in compliance with the License.
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+ * See LICENSE in the root of the software repository for the full text of the License.
+ * ===================================================================================================================*/
+
+#ifndef METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_
+#define METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_
+
+#include <unordered_map>
+#include <mutex>
+
+using Status = int32_t;
+
+namespace ge {
+static constexpr uint32_t COMM_MESH = 0b1U;
+static constexpr uint32_t COMM_SWITCH = (COMM_MESH << 1U);
+static constexpr uint32_t COMM_RING = (COMM_MESH << 2U);
+static constexpr uint32_t COMM_PAIRWISE = (COMM_MESH << 3U);
+class HcomTopoInfo {
+ public:
+  enum class TopoLevel {
+    L0 = 0,
+    L1,
+    MAX,
+  };
+  struct TopoLevelDesc {
+    uint32_t comm_sets;
+    uint32_t rank_size;
+  };
+  using TopoDescs = TopoLevelDesc[static_cast<int32_t>(TopoLevel::MAX)];
+  struct TopoInfo {
+    int64_t rank_size;
+    void *notify_handle;
+    TopoDescs topo_level_descs;
+  };
+  static HcomTopoInfo &Instance();
+  bool TopoInfoHasBeenSet(const char_t *group);
+  bool TryGetGroupTopoInfo(const char_t *group, TopoInfo &info);
+  Status SetGroupTopoInfo(const char_t *group, const TopoInfo &info);
+  Status GetGroupRankSize(const char_t *group, int64_t &rank_size);
+  TopoDescs *GetGroupTopoDesc(const char_t *group);
+  Status GetGroupNotifyHandle(const char_t *group, void *&notify_handle);
+  void UnsetGroupTopoInfo(const char_t *group) {
+    const std::lock_guard<std::mutex> lock(mutex_);
+    (void) rank_info_.erase(group);
+  }
+
+  Status SetGroupOrderedStream(const char_t *group, void *stream);
+  Status GetGroupOrderedStream(const char_t *group, void *&stream);
+  void UnsetGroupOrderedStream(const char_t *group) {
+    const std::lock_guard<std::mutex> lock(mutex_);
+    (void) group_to_ordered_stream_.erase(group);
+  };
+
+  Status SetGroupOrderedStream(const int32_t device_id, const char_t *group, void *stream);
+  Status GetGroupOrderedStream(const int32_t device_id, const char_t *group, void *&stream);
+  void UnsetGroupOrderedStream(const int32_t device_id, const char_t *group);
+ private:
+  HcomTopoInfo() = default;
+  ~HcomTopoInfo() = default;
+  std::unordered_map<std::string, TopoInfo> rank_info_;
+  std::mutex mutex_;
+  std::unordered_map<std::string, void*> group_to_ordered_stream_; // Ordered stream for the communication domain
+  std::unordered_map<int32_t, std::unordered_map<std::string, void*>> device_id_to_group_to_ordered_stream_; // Ordered stream for the communication domain
+};
+}
+
+#endif // METADEF_CXX_INC_EXTERNAL_HCOM_HCOM_TOPO_INFO_H_