[INFO] Initializing environment for https://gitcode.com/pre-commit/pre-commit-hooks. [WARNING] repo `https://gitcode.com/pre-commit/pre-commit-hooks` uses deprecated stage names (commit, push) which will be removed in a future version. Hint: often `pre-commit autoupdate --repo https://gitcode.com/pre-commit/pre-commit-hooks` will fix this. if it does not -- consider reporting an issue to that repo. [INFO] Initializing environment for https://gitcode.com/pre-commit-clang/mirrors-clang-format. [INFO] Initializing environment for https://gitcode.com/gh_mirrors/ru/ruff-pre-commit. [INFO] Initializing environment for https://gitcode.com/gh_mirrors/co/codespell. [INFO] Installing environment for https://gitcode.com/pre-commit/pre-commit-hooks. [INFO] Once installed this environment will be reused. [INFO] This may take a few minutes... [INFO] Installing environment for https://gitcode.com/pre-commit-clang/mirrors-clang-format. [INFO] Once installed this environment will be reused. [INFO] This may take a few minutes... [INFO] Installing environment for https://gitcode.com/gh_mirrors/ru/ruff-pre-commit. [INFO] Once installed this environment will be reused. [INFO] This may take a few minutes... [INFO] Installing environment for https://gitcode.com/gh_mirrors/co/codespell. [INFO] Once installed this environment will be reused. [INFO] This may take a few minutes... trim trailing whitespace.................................................Failed - hook id: trailing-whitespace - exit code: 1 - files were modified by this hook Fixing mc2/mega_moe/op_host/op_tiling/arch22/mega_moe_tiling_arch22.cpp Fixing torch_extension/cann_ops_transformer/ops/csrc/mega_moe.cpp fix end of files.........................................................Passed check yaml...........................................(no files to check)Skipped check for added large files..............................................Passed check for merge conflicts................................................Passed detect private key.......................................................Passed check json...........................................(no files to check)Skipped clang-format.............................................................Failed - hook id: clang-format - files were modified by this hook Formatting [1/4] mc2/mega_moe/op_kernel/arch22/mega_moe_kernel_a3.hpp Formatting [2/4] mc2/mega_moe/op_host/op_tiling/arch22/mega_moe_tiling_arch22.cpp Formatting [3/4] mc2/mega_moe/op_kernel/arch22/utils/hccl_shmem.hpp Formatting [4/4] torch_extension/cann_ops_transformer/ops/csrc/mega_moe.cpp ruff check...........................................(no files to check)Skipped ruff format..........................................(no files to check)Skipped codespell................................................................Passed All changes made by hooks: diff --git a/mc2/mega_moe/op_host/op_tiling/arch22/mega_moe_tiling_arch22.cpp b/mc2/mega_moe/op_host/op_tiling/arch22/mega_moe_tiling_arch22.cpp index db7bd0d..4e42ad9 100644 --- a/mc2/mega_moe/op_host/op_tiling/arch22/mega_moe_tiling_arch22.cpp +++ b/mc2/mega_moe/op_host/op_tiling/arch22/mega_moe_tiling_arch22.cpp @@ -35,99 +35,96 @@ using namespace ge; using namespace mc2tiling; namespace MegaMoeA2A3Tiling { - const char *K_INNER_DEBUG = "MegaMoeA2A3 Tiling Debug"; - - // 算子属性索引 - constexpr uint32_t ATTR_MOE_EXPERT_NUM_INDEX = 0; // moe 专家数 - constexpr uint32_t ATTR_EP_WORLD_SIZE_INDEX = 1; // EP 并行 world size - constexpr uint32_t ATTR_CCL_BUFFER_SIZE_INDEX = 2; // HCCL 通信缓冲区大小 - constexpr uint32_t ATTR_MAX_RECV_TOKEN_NUM_INDEX = 3; // 最大接收 token 数（用于预分配 workspace） - constexpr uint32_t ATTR_DISPATCH_QUANT_MODE_INDEX = 4; // 分发阶段量化模式 - constexpr uint32_t ATTR_DISPATCH_QUANT_OUT_TYPE_INDEX = 5; // 分发阶段量化输出数据类型 - constexpr uint32_t ATTR_COMBINE_QUANT_MODE_INDEX = 6; // 合并阶段量化模式 - constexpr uint32_t ATTR_COMM_ALG_INDEX = 7; // 通信算法配置 - constexpr uint32_t ATTR_NUM_MAX_TOKENS_PER_RANK_INDEX = 8; // 每个 rank 的最大 token 数(bs数量) - constexpr uint32_t ATTR_ACTIVATION_INDEX = 9; // 激活函数类型（如 "swiglu"） - constexpr uint32_t ATTR_ACTIVATION_CLAMP_INDEX = 10; // 激活函数 clamp 值 - constexpr uint32_t ATTR_ACTIVATION_OUT_DTYPE_INDEX = 11; // 激活函数输出数据类型 - constexpr uint32_t ATTR_TRANSPOSE_WEIGHT1_INDEX = 12; // weight1 是否转置 - constexpr uint32_t ATTR_TRANSPOSE_WEIGHT2_INDEX = 13; // weight2 是否转置 - constexpr uint32_t ATTR_WEIGHT1_INTERLEAVE_INDEX = 14; // weight1 交错模式 - - // 输入 tensor 索引 - constexpr uint32_t CONTEXT_INDEX = 0; // context，shape = [?] - constexpr uint32_t X_INDEX = 1; // 输入 token x，shape = [M, K] - constexpr uint32_t TOPK_IDS_INDEX = 2; // topk 专家索引，shape = [bs, topK] - constexpr uint32_t TOPK_WEIGHTS_INDEX = 3; // topk 权重，shape = [bs, topK] - constexpr uint32_t WEIGHT1_INDEX = 4; // FFN 第一层权重（动态输入，每个专家一组） - constexpr uint32_t WEIGHT2_INDEX = 5; // FFN 第二层权重（动态输入，每个专家一组） - constexpr uint32_t WEIGHT_SCALES1_INDEX = 6; // FFN 第一层权重量化 scale - constexpr uint32_t WEIGHT_SCALES2_INDEX = 7; // FFN 第二层权重量化 scale - constexpr uint32_t BIAS1_INDEX = 8; // GMM1 的 bias（OPTIONAL） - constexpr uint32_t BIAS2_INDEX = 9; // GMM2 的 bias（OPTIONAL） - constexpr uint32_t X_ACTIVE_MASK_INDEX = 10; // token 活跃掩码（OPTIONAL） - constexpr uint32_t SCALES_INDEX = 11; // 额外 scale（OPTIONAL） - - // 输出 tensor 索引 - constexpr uint32_t OUTPUT_Y_INDEX = 0; // 输出 y，shape = [M, K] - constexpr uint32_t OUTPUT_EXPERT_TOKEN_NUMS_INDEX = 1; // 专家 token 数量，shape = [expert_num] - - constexpr uint32_t SYSTEM_NEED_WORKSPACE = 16 * 1024 * 1024; - constexpr uint64_t MB_SIZE = 1024UL * 1024UL; - constexpr uint64_t RESERVED_SPACE_SIZE = 10 * 1024 * 1024; - - // CCL Buffer 相关常量 - constexpr int64_t PEERMEM_DATA_OFFSET = 1024 * 60LL; // （预留）60KB 固定偏移 - constexpr int64_t ALIGN_128 = 128LL; - constexpr int64_t ALIGN_512 = 512LL; - // 与 kernel 侧 SYNC_STATE_RESERVED_SIZE 保持一致： - // 前 256KB 内放 per-rank 状态，后 256KB 放 per-block self state。 - constexpr int64_t SYNC_STATE_RESERVED_SIZE = 512 * 1024LL; - - // 维度范围限制 - constexpr int64_t MIN_BS = 1; - constexpr int64_t MAX_BS = 4096; - constexpr int64_t MIN_HIDDEN_SIZE = 1024; - constexpr int64_t MAX_HIDDEN_SIZE = 8192; - constexpr int64_t MIN_INTERMEDIATE_HIDDEN = 1024; - constexpr int64_t MAX_INTERMEDIATE_HIDDEN = 3072; - constexpr int64_t MIN_TOPK = 1; - constexpr int64_t MAX_TOPK = 16; - constexpr int64_t MIN_EXPERT_PER_RANK = 1; - constexpr int64_t MAX_EXPERT_PER_RANK = 128; - constexpr int64_t HIDDEN_SIZE_ALIGN = 512; - - // 属性范围限制 - constexpr int64_t MIN_MOE_EXPERT_NUM = 1; - constexpr int64_t MAX_MOE_EXPERT_NUM = 1024; - constexpr int64_t VALID_EP_WORLD_SIZE[] = {2, 4, 8, 16, 32, 64}; - - constexpr uint32_t TWO_DIMS = 2U; - constexpr uint32_t ONE_DIM = 1U; - constexpr uint32_t THREE_DIMS = 3U; - - // - constexpr int64_t DISPATCH_QUANT_MODE_NO_QUANT = 0; - constexpr int64_t DISPATCH_QUANT_MODE_PER_TENSOR = 2; - -static int64_t CalcLeastCclBufferSize(int64_t maxRecvTokenNum, int64_t h, - int64_t epWorldSize, int64_t expertPerRank, - bool isQuantRouting, bool isW4A8, bool isA3, - int64_t bs, int64_t topK) +const char *K_INNER_DEBUG = "MegaMoeA2A3 Tiling Debug"; + +// 算子属性索引 +constexpr uint32_t ATTR_MOE_EXPERT_NUM_INDEX = 0; // moe 专家数 +constexpr uint32_t ATTR_EP_WORLD_SIZE_INDEX = 1; // EP 并行 world size +constexpr uint32_t ATTR_CCL_BUFFER_SIZE_INDEX = 2; // HCCL 通信缓冲区大小 +constexpr uint32_t ATTR_MAX_RECV_TOKEN_NUM_INDEX = 3; // 最大接收 token 数（用于预分配 workspace） +constexpr uint32_t ATTR_DISPATCH_QUANT_MODE_INDEX = 4; // 分发阶段量化模式 +constexpr uint32_t ATTR_DISPATCH_QUANT_OUT_TYPE_INDEX = 5; // 分发阶段量化输出数据类型 +constexpr uint32_t ATTR_COMBINE_QUANT_MODE_INDEX = 6; // 合并阶段量化模式 +constexpr uint32_t ATTR_COMM_ALG_INDEX = 7; // 通信算法配置 +constexpr uint32_t ATTR_NUM_MAX_TOKENS_PER_RANK_INDEX = 8; // 每个 rank 的最大 token 数(bs数量) +constexpr uint32_t ATTR_ACTIVATION_INDEX = 9; // 激活函数类型（如 "swiglu"） +constexpr uint32_t ATTR_ACTIVATION_CLAMP_INDEX = 10; // 激活函数 clamp 值 +constexpr uint32_t ATTR_ACTIVATION_OUT_DTYPE_INDEX = 11; // 激活函数输出数据类型 +constexpr uint32_t ATTR_TRANSPOSE_WEIGHT1_INDEX = 12; // weight1 是否转置 +constexpr uint32_t ATTR_TRANSPOSE_WEIGHT2_INDEX = 13; // weight2 是否转置 +constexpr uint32_t ATTR_WEIGHT1_INTERLEAVE_INDEX = 14; // weight1 交错模式 + +// 输入 tensor 索引 +constexpr uint32_t CONTEXT_INDEX = 0; // context，shape = [?] +constexpr uint32_t X_INDEX = 1; // 输入 token x，shape = [M, K] +constexpr uint32_t TOPK_IDS_INDEX = 2; // topk 专家索引，shape = [bs, topK] +constexpr uint32_t TOPK_WEIGHTS_INDEX = 3; // topk 权重，shape = [bs, topK] +constexpr uint32_t WEIGHT1_INDEX = 4; // FFN 第一层权重（动态输入，每个专家一组） +constexpr uint32_t WEIGHT2_INDEX = 5; // FFN 第二层权重（动态输入，每个专家一组） +constexpr uint32_t WEIGHT_SCALES1_INDEX = 6; // FFN 第一层权重量化 scale +constexpr uint32_t WEIGHT_SCALES2_INDEX = 7; // FFN 第二层权重量化 scale +constexpr uint32_t BIAS1_INDEX = 8; // GMM1 的 bias（OPTIONAL） +constexpr uint32_t BIAS2_INDEX = 9; // GMM2 的 bias（OPTIONAL） +constexpr uint32_t X_ACTIVE_MASK_INDEX = 10; // token 活跃掩码（OPTIONAL） +constexpr uint32_t SCALES_INDEX = 11; // 额外 scale（OPTIONAL） + +// 输出 tensor 索引 +constexpr uint32_t OUTPUT_Y_INDEX = 0; // 输出 y，shape = [M, K] +constexpr uint32_t OUTPUT_EXPERT_TOKEN_NUMS_INDEX = 1; // 专家 token 数量，shape = [expert_num] + +constexpr uint32_t SYSTEM_NEED_WORKSPACE = 16 * 1024 * 1024; +constexpr uint64_t MB_SIZE = 1024UL * 1024UL; +constexpr uint64_t RESERVED_SPACE_SIZE = 10 * 1024 * 1024; + +// CCL Buffer 相关常量 +constexpr int64_t PEERMEM_DATA_OFFSET = 1024 * 60LL; // （预留）60KB 固定偏移 +constexpr int64_t ALIGN_128 = 128LL; +constexpr int64_t ALIGN_512 = 512LL; +// 与 kernel 侧 SYNC_STATE_RESERVED_SIZE 保持一致： +// 前 256KB 内放 per-rank 状态，后 256KB 放 per-block self state。 +constexpr int64_t SYNC_STATE_RESERVED_SIZE = 512 * 1024LL; + +// 维度范围限制 +constexpr int64_t MIN_BS = 1; +constexpr int64_t MAX_BS = 4096; +constexpr int64_t MIN_HIDDEN_SIZE = 1024; +constexpr int64_t MAX_HIDDEN_SIZE = 8192; +constexpr int64_t MIN_INTERMEDIATE_HIDDEN = 1024; +constexpr int64_t MAX_INTERMEDIATE_HIDDEN = 3072; +constexpr int64_t MIN_TOPK = 1; +constexpr int64_t MAX_TOPK = 16; +constexpr int64_t MIN_EXPERT_PER_RANK = 1; +constexpr int64_t MAX_EXPERT_PER_RANK = 128; +constexpr int64_t HIDDEN_SIZE_ALIGN = 512; + +// 属性范围限制 +constexpr int64_t MIN_MOE_EXPERT_NUM = 1; +constexpr int64_t MAX_MOE_EXPERT_NUM = 1024; +constexpr int64_t VALID_EP_WORLD_SIZE[] = {2, 4, 8, 16, 32, 64}; + +constexpr uint32_t TWO_DIMS = 2U; +constexpr uint32_t ONE_DIM = 1U; +constexpr uint32_t THREE_DIMS = 3U; + +// +constexpr int64_t DISPATCH_QUANT_MODE_NO_QUANT = 0; +constexpr int64_t DISPATCH_QUANT_MODE_PER_TENSOR = 2; + +static int64_t CalcLeastCclBufferSize(int64_t maxRecvTokenNum, int64_t h, int64_t epWorldSize, int64_t expertPerRank, + bool isQuantRouting, bool isW4A8, bool isA3, int64_t bs, int64_t topK) { // ccl buff的承载的数据块 1（winIn）： // TPE = epWorldSize × CeilAlign(epWorldSize × MAX_EXPERT_PER_RANK + 1, 128) × 4B - int64_t offsetTokenPerExpert = epWorldSize * - ops::CeilAlign(epWorldSize * MAX_EXPERT_PER_RANK + 1, ALIGN_128) * + int64_t offsetTokenPerExpert = epWorldSize * ops::CeilAlign(epWorldSize * MAX_EXPERT_PER_RANK + 1, ALIGN_128) * static_cast(sizeof(int32_t)); // ccl buff的承载的数据块 2： // ============================== winIn ============================== // FFN的左矩阵，非量化则token为 h×2 字节 (bf16)，量化则为 h+512 字节 (int8) int64_t offsetAAfterDispatch = 0; - if (isA3) { // A3打包发送，+32与+512中取大值 + if (isA3) { // A3打包发送，+32与+512中取大值 offsetAAfterDispatch = bs * topK * (isQuantRouting ? (h + ALIGN_512) : h * sizeof(int16_t)); - } else { // A2分开发送 + } else { // A2分开发送 offsetAAfterDispatch = maxRecvTokenNum * (isQuantRouting ? (h + ALIGN_512) : h * sizeof(int16_t)); } // FFN的输出在分发后，接收的空间大小 @@ -143,12 +140,12 @@ static int64_t CalcLeastCclBufferSize(int64_t maxRecvTokenNum, int64_t h, int64_t offsetTensor = std::max(winInTensorSize, winOutTensorSize); // pertokenscale的额外空间 if (isQuantRouting) { - offsetTensor += (isA3 ? bs * topK : maxRecvTokenNum) * sizeof(float); // pertokenScale + offsetTensor += (isA3 ? bs * topK : maxRecvTokenNum) * sizeof(float); // pertokenScale } // ccl buff的承载的数据块 3（winIn）： // 同步flag - int64_t offsetFlag = epWorldSize * ALIGN_512; // CrossRankSync所用空间 + int64_t offsetFlag = epWorldSize * ALIGN_512; // CrossRankSync所用空间 if (!isA3) { // A2: dispatch flag(EP×E×64B) + allgather flag(EP×64B) int64_t dispatchFlag = epWorldSize * MAX_EXPERT_PER_RANK * 64; @@ -165,45 +162,41 @@ static int64_t CalcLeastCclBufferSize(int64_t maxRecvTokenNum, int64_t h, static ge::graphStatus CheckMoeExpertNumAttr(const int64_t *ptr) { - OP_TILING_CHECK(ptr == nullptr, - OP_LOGE(K_INNER_DEBUG, "moeExpertNum is nullptr."), return GRAPH_FAILED); + OP_TILING_CHECK(ptr == nullptr, OP_LOGE(K_INNER_DEBUG, "moeExpertNum is nullptr."), return GRAPH_FAILED); OP_TILING_CHECK(*ptr < MIN_MOE_EXPERT_NUM || *ptr > MAX_MOE_EXPERT_NUM, - OP_LOGE(K_INNER_DEBUG, "moeExpertNum should be in [%ld, %ld], but got %ld.", - MIN_MOE_EXPERT_NUM, MAX_MOE_EXPERT_NUM, *ptr), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "moeExpertNum should be in [%ld, %ld], but got %ld.", MIN_MOE_EXPERT_NUM, + MAX_MOE_EXPERT_NUM, *ptr), + return GRAPH_FAILED); return ge::GRAPH_SUCCESS; } static ge::graphStatus CheckEpWorldSizeAttr(const int64_t *ptr) { - OP_TILING_CHECK(ptr == nullptr, - OP_LOGE(K_INNER_DEBUG, "epWorldSize is nullptr."), return GRAPH_FAILED); - bool isValidEpWorldSize = std::find(std::begin(VALID_EP_WORLD_SIZE), - std::end(VALID_EP_WORLD_SIZE), - *ptr) != std::end(VALID_EP_WORLD_SIZE); + OP_TILING_CHECK(ptr == nullptr, OP_LOGE(K_INNER_DEBUG, "epWorldSize is nullptr."), return GRAPH_FAILED); + bool isValidEpWorldSize = std::find(std::begin(VALID_EP_WORLD_SIZE), std::end(VALID_EP_WORLD_SIZE), *ptr) != + std::end(VALID_EP_WORLD_SIZE); OP_TILING_CHECK(!isValidEpWorldSize, - OP_LOGE(K_INNER_DEBUG, "epWorldSize should be one of {2, 4, 8, 16, 32, 64}, but got %ld.", - *ptr), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "epWorldSize should be one of {2, 4, 8, 16, 32, 64}, but got %ld.", *ptr), + return GRAPH_FAILED); return ge::GRAPH_SUCCESS; } static ge::graphStatus CheckMaxRecvTokenNumAttr(const int64_t *ptr) { - OP_TILING_CHECK(ptr == nullptr, - OP_LOGE(K_INNER_DEBUG, "maxRecvTokenNum is nullptr."), return GRAPH_FAILED); - OP_TILING_CHECK(*ptr < 0, - OP_LOGE(K_INNER_DEBUG, "maxRecvTokenNum is invalid, should be >= 0, but got %ld.", *ptr), - return GRAPH_FAILED); + OP_TILING_CHECK(ptr == nullptr, OP_LOGE(K_INNER_DEBUG, "maxRecvTokenNum is nullptr."), return GRAPH_FAILED); + OP_TILING_CHECK(*ptr < 0, OP_LOGE(K_INNER_DEBUG, "maxRecvTokenNum is invalid, should be >= 0, but got %ld.", *ptr), + return GRAPH_FAILED); return ge::GRAPH_SUCCESS; } // 校验 dispatch_quant_mode 和 dispatch_quant_out_dtype，两者都依赖 weight1 的 dataType -static ge::graphStatus CheckDispatchQuantAttrs(gert::TilingContext *context, - const int64_t *dispatchQuantModePtr, const int64_t *dispatchQuantOutDtypePtr) +static ge::graphStatus CheckDispatchQuantAttrs(gert::TilingContext *context, const int64_t *dispatchQuantModePtr, + const int64_t *dispatchQuantOutDtypePtr) { - OP_TILING_CHECK(dispatchQuantModePtr == nullptr, - OP_LOGE(K_INNER_DEBUG, "dispatchQuantMode is null."), return GRAPH_FAILED); - OP_TILING_CHECK(dispatchQuantOutDtypePtr == nullptr, - OP_LOGE(K_INNER_DEBUG, "dispatchQuantOutDtype is null."), return GRAPH_FAILED); + OP_TILING_CHECK(dispatchQuantModePtr == nullptr, OP_LOGE(K_INNER_DEBUG, "dispatchQuantMode is null."), + return GRAPH_FAILED); + OP_TILING_CHECK(dispatchQuantOutDtypePtr == nullptr, OP_LOGE(K_INNER_DEBUG, "dispatchQuantOutDtype is null."), + return GRAPH_FAILED); auto w1Desc = context->GetDynamicInputDesc(WEIGHT1_INDEX, 0); OP_TILING_CHECK(w1Desc == nullptr, OP_LOGE(K_INNER_DEBUG, "weight1 desc is null."), return GRAPH_FAILED); @@ -217,9 +210,10 @@ static ge::graphStatus CheckDispatchQuantAttrs(gert::TilingContext *context, expectedDispatchQuantMode = DISPATCH_QUANT_MODE_NO_QUANT; } OP_TILING_CHECK(*dispatchQuantModePtr != expectedDispatchQuantMode, - OP_LOGE(K_INNER_DEBUG, "dispatch_quant_mode is invalid, should be %ld for weight dataType %d, but got %ld.", - expectedDispatchQuantMode, static_cast(w1DataType), *dispatchQuantModePtr), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "dispatch_quant_mode is invalid, should be %ld for weight dataType %d, but got %ld.", + expectedDispatchQuantMode, static_cast(w1DataType), *dispatchQuantModePtr), + return GRAPH_FAILED); // dispatch_quant_out_dtype 校验：与 weight 数据类型对应 if (*dispatchQuantOutDtypePtr != static_cast(ge::DT_UNDEFINED)) { @@ -231,10 +225,11 @@ static ge::graphStatus CheckDispatchQuantAttrs(gert::TilingContext *context, } else if (w1DataType == ge::DT_INT4 || w1DataType == ge::DT_INT8) { expectedQuantOutType = static_cast(ge::DT_INT8); } - OP_TILING_CHECK(*dispatchQuantOutDtypePtr != expectedQuantOutType, + OP_TILING_CHECK( + *dispatchQuantOutDtypePtr != expectedQuantOutType, OP_LOGE(K_INNER_DEBUG, - "dispatch_quant_out_type is invalid, should be %ld for weight dataType %d, but got %ld.", - expectedQuantOutType, static_cast(w1DataType), *dispatchQuantOutDtypePtr), + "dispatch_quant_out_type is invalid, should be %ld for weight dataType %d, but got %ld.", + expectedQuantOutType, static_cast(w1DataType), *dispatchQuantOutDtypePtr), return GRAPH_FAILED); } return ge::GRAPH_SUCCESS; @@ -242,101 +237,99 @@ static ge::graphStatus CheckDispatchQuantAttrs(gert::TilingContext *context, static ge::graphStatus CheckCombineQuantModeAttr(const int64_t *ptr) { - OP_TILING_CHECK(ptr == nullptr, - OP_LOGE(K_INNER_DEBUG, "combineQuantMode is null."), return GRAPH_FAILED); + OP_TILING_CHECK(ptr == nullptr, OP_LOGE(K_INNER_DEBUG, "combineQuantMode is null."), return GRAPH_FAILED); OP_TILING_CHECK(*ptr != 0, - OP_LOGE(K_INNER_DEBUG, "combine_quant_mode is invalid, only support 0, but got %ld.", - *ptr), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "combine_quant_mode is invalid, only support 0, but got %ld.", *ptr), + return GRAPH_FAILED); return ge::GRAPH_SUCCESS; } // 校验 num_max_tokens_per_rank，需要从 x shape 获取 bs static ge::graphStatus CheckNumMaxTokensPerRankAttr(gert::TilingContext *context, const int64_t *ptr) { - OP_TILING_CHECK(ptr == nullptr, - OP_LOGE(K_INNER_DEBUG, "numMaxTokensPerRankPtr is null."), return GRAPH_FAILED); + OP_TILING_CHECK(ptr == nullptr, OP_LOGE(K_INNER_DEBUG, "numMaxTokensPerRankPtr is null."), return GRAPH_FAILED); const gert::StorageShape *xStorageShape = context->GetInputShape(X_INDEX); OP_TILING_CHECK(xStorageShape == nullptr, OP_LOGE(K_INNER_DEBUG, "x shape is null."), return GRAPH_FAILED); int64_t bs = xStorageShape->GetStorageShape().GetDim(0); if (*ptr > 0) { - OP_TILING_CHECK(*ptr < bs, - OP_LOGE(K_INNER_DEBUG, "num_max_tokens_per_rank is invalid, should be >= bs(%ld), but got %ld.", - bs, *ptr), return GRAPH_FAILED); + OP_TILING_CHECK( + *ptr < bs, + OP_LOGE(K_INNER_DEBUG, "num_max_tokens_per_rank is invalid, should be >= bs(%ld), but got %ld.", bs, *ptr), + return GRAPH_FAILED); OP_TILING_CHECK(*ptr < MIN_BS || *ptr > MAX_BS, - OP_LOGE(K_INNER_DEBUG, "num_max_tokens_per_rank is invalid, should be in [%ld, %ld], but got %ld.", - MIN_BS, MAX_BS, *ptr), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "num_max_tokens_per_rank is invalid, should be in [%ld, %ld], but got %ld.", MIN_BS, + MAX_BS, *ptr), + return GRAPH_FAILED); } return ge::GRAPH_SUCCESS; } static ge::graphStatus CheckActivationAttr(const char *ptr) { - OP_TILING_CHECK(ptr == nullptr, - OP_LOGE(K_INNER_DEBUG, "activation is null."), return GRAPH_FAILED); + OP_TILING_CHECK(ptr == nullptr, OP_LOGE(K_INNER_DEBUG, "activation is null."), return GRAPH_FAILED); std::string activationStr(ptr); OP_TILING_CHECK(activationStr != "swiglu", - OP_LOGE(K_INNER_DEBUG, "activation is invalid, only support 'swiglu', but got '%s'.", - ptr), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "activation is invalid, only support 'swiglu', but got '%s'.", ptr), + return GRAPH_FAILED); return ge::GRAPH_SUCCESS; } static ge::graphStatus CheckActivationClampAttr(const float *ptr) { - OP_TILING_CHECK(ptr == nullptr, - OP_LOGE(K_INNER_DEBUG, "activationClamp is null."), return GRAPH_FAILED); - OP_TILING_CHECK(*ptr < 0 || std::isnan(*ptr), - OP_LOGE(K_INNER_DEBUG, "activation_clamp is invalid, should be >= 0 and not NAN, but got %f.", - *ptr), return GRAPH_FAILED); + OP_TILING_CHECK(ptr == nullptr, OP_LOGE(K_INNER_DEBUG, "activationClamp is null."), return GRAPH_FAILED); + OP_TILING_CHECK( + *ptr < 0 || std::isnan(*ptr), + OP_LOGE(K_INNER_DEBUG, "activation_clamp is invalid, should be >= 0 and not NAN, but got %f.", *ptr), + return GRAPH_FAILED); return ge::GRAPH_SUCCESS; } static ge::graphStatus CheckActivationOutDtypeAttr(const int64_t *ptr) { - OP_TILING_CHECK(ptr == nullptr, - OP_LOGE(K_INNER_DEBUG, "activationOutDtypePtr is null."), return GRAPH_FAILED); + OP_TILING_CHECK(ptr == nullptr, OP_LOGE(K_INNER_DEBUG, "activationOutDtypePtr is null."), return GRAPH_FAILED); OP_TILING_CHECK(*ptr != static_cast(ge::DT_UNDEFINED), - OP_LOGE(K_INNER_DEBUG, "activation_out_dtype is invalid, should be default value %ld, but got %ld.", - static_cast(ge::DT_UNDEFINED), *ptr), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "activation_out_dtype is invalid, should be default value %ld, but got %ld.", + static_cast(ge::DT_UNDEFINED), *ptr), + return GRAPH_FAILED); return ge::GRAPH_SUCCESS; } static ge::graphStatus CheckTransposeWeightAttr(const bool *w1, const bool *w2) { - OP_TILING_CHECK(w1 == nullptr, - OP_LOGE(K_INNER_DEBUG, "transposeWeight1 is null."), return GRAPH_FAILED); - OP_TILING_CHECK(w2 == nullptr, - OP_LOGE(K_INNER_DEBUG, "transposeWeight2 is null."), return GRAPH_FAILED); + OP_TILING_CHECK(w1 == nullptr, OP_LOGE(K_INNER_DEBUG, "transposeWeight1 is null."), return GRAPH_FAILED); + OP_TILING_CHECK(w2 == nullptr, OP_LOGE(K_INNER_DEBUG, "transposeWeight2 is null."), return GRAPH_FAILED); OP_TILING_CHECK(*w1 != *w2, - OP_LOGE(K_INNER_DEBUG, - "transpose_weight1 and transpose_weight2 must be the same, " - "transpose_weight1 = %d, transpose_weight2 = %d.", - static_cast(*w1), static_cast(*w2)), return GRAPH_FAILED); - OP_TILING_CHECK(*w1 == true, - OP_LOGE(K_INNER_DEBUG, - "Current soc do not support transpose weight now."), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "transpose_weight1 and transpose_weight2 must be the same, " + "transpose_weight1 = %d, transpose_weight2 = %d.", + static_cast(*w1), static_cast(*w2)), + return GRAPH_FAILED); + OP_TILING_CHECK(*w1 == true, OP_LOGE(K_INNER_DEBUG, "Current soc do not support transpose weight now."), + return GRAPH_FAILED); return ge::GRAPH_SUCCESS; } // 校验 weight1_interleave，需要从 weight1 tensor 获取 N static ge::graphStatus CheckWeight1InterleaveAttr(gert::TilingContext *context, const int64_t *ptr) { - OP_TILING_CHECK(ptr == nullptr, - OP_LOGE(K_INNER_DEBUG, "weight1Interleave is null."), return GRAPH_FAILED); + OP_TILING_CHECK(ptr == nullptr, OP_LOGE(K_INNER_DEBUG, "weight1Interleave is null."), return GRAPH_FAILED); OP_TILING_CHECK(*ptr < 0, - OP_LOGE(K_INNER_DEBUG, "weight1_interleave is invalid, should be >= 0, but got %ld.", - *ptr), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "weight1_interleave is invalid, should be >= 0, but got %ld.", *ptr), + return GRAPH_FAILED); auto w1Tensor = context->GetDynamicInputTensor(WEIGHT1_INDEX, 0); OP_TILING_CHECK(w1Tensor == nullptr, OP_LOGE(K_INNER_DEBUG, "weight1 tensor is null."), return GRAPH_FAILED); uint32_t w1TensorDims = w1Tensor->GetOriginShape().GetDimNum(); uint32_t N = w1Tensor->GetStorageShape().GetDim(w1TensorDims - 1); uint32_t maxInterleave = N / 2; OP_TILING_CHECK(static_cast(*ptr) > maxInterleave, - OP_LOGE(K_INNER_DEBUG, "weight1_interleave is invalid, should be in [0, %u], but got %ld.", - maxInterleave, *ptr), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "weight1_interleave is invalid, should be in [0, %u], but got %ld.", + maxInterleave, *ptr), + return GRAPH_FAILED); return ge::GRAPH_SUCCESS; } -static ge::graphStatus MegaMoeA2A3CheckAttrAndSetTiling(gert::TilingContext *context, MegaMoeA2A3TilingData& info) +static ge::graphStatus MegaMoeA2A3CheckAttrAndSetTiling(gert::TilingContext *context, MegaMoeA2A3TilingData &info) { auto attrs = context->GetAttrs(); OP_TILING_CHECK(attrs == nullptr, OP_LOGE(K_INNER_DEBUG, "attrs is null."), return ge::GRAPH_FAILED); @@ -359,85 +352,72 @@ static ge::graphStatus MegaMoeA2A3CheckAttrAndSetTiling(gert::TilingContext *con // 1. moe_expert_num OP_TILING_CHECK(CheckMoeExpertNumAttr(moeExpertNumPtr) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckMoeExpertNumAttr failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckMoeExpertNumAttr failed."), return GRAPH_FAILED); info.moeExpertNum = static_cast(*moeExpertNumPtr); // 2. ep_world_size OP_TILING_CHECK(CheckEpWorldSizeAttr(epWorldSizePtr) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckEpWorldSizeAttr failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckEpWorldSizeAttr failed."), return GRAPH_FAILED); info.epWorldSize = static_cast(*epWorldSizePtr); // 3. ccl_buffer_size - OP_TILING_CHECK(cclBufferSizePtr == nullptr, - OP_LOGE(K_INNER_DEBUG, "ccl_buffer_size is nullptr."), return GRAPH_FAILED); + OP_TILING_CHECK(cclBufferSizePtr == nullptr, OP_LOGE(K_INNER_DEBUG, "ccl_buffer_size is nullptr."), + return GRAPH_FAILED); info.cclBufferSize = static_cast(*cclBufferSizePtr); - + // 4. max_recv_token_num OP_TILING_CHECK(CheckMaxRecvTokenNumAttr(maxRecvTokenNumPtr) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckMaxRecvTokenNumAttr failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckMaxRecvTokenNumAttr failed."), return GRAPH_FAILED); info.maxRecvTokenNum = static_cast(*maxRecvTokenNumPtr); // max_recv_token_num 为 0 时赋值为最大值； if (info.maxRecvTokenNum == 0U) { - info.maxRecvTokenNum = static_cast(info.M) * info.epWorldSize * - std::min(info.topK, info.expertPerRank); - OP_LOGD(K_INNER_DEBUG, - "maxRecvTokenNum auto-calculated to %u (bs=%u * ep=%u * min(topK=%u, expertPerRank=%u))", + info.maxRecvTokenNum = + static_cast(info.M) * info.epWorldSize * std::min(info.topK, info.expertPerRank); + OP_LOGD(K_INNER_DEBUG, "maxRecvTokenNum auto-calculated to %u (bs=%u * ep=%u * min(topK=%u, expertPerRank=%u))", info.maxRecvTokenNum, info.M, info.epWorldSize, info.topK, info.expertPerRank); } // 5-6. dispatch_quant_mode && dispatch_quant_out_dtype - OP_TILING_CHECK(CheckDispatchQuantAttrs(context, dispatchQuantModePtr, - dispatchQuantOutDtypePtr) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckDispatchQuantAttrs failed."), - return GRAPH_FAILED); + OP_TILING_CHECK(CheckDispatchQuantAttrs(context, dispatchQuantModePtr, dispatchQuantOutDtypePtr) != + ge::GRAPH_SUCCESS, + OP_LOGE(K_INNER_DEBUG, "CheckDispatchQuantAttrs failed."), return GRAPH_FAILED); info.dispatchQuantMode = static_cast(*dispatchQuantModePtr); info.dispatchQuantOutDtype = static_cast(*dispatchQuantOutDtypePtr); // 7. combine_quant_mode OP_TILING_CHECK(CheckCombineQuantModeAttr(combineQuantModePtr) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckCombineQuantModeAttr failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckCombineQuantModeAttr failed."), return GRAPH_FAILED); info.combineQuantMode = static_cast(*combineQuantModePtr); // 8. num_max_token_per_rank OP_TILING_CHECK(CheckNumMaxTokensPerRankAttr(context, numMaxTokensPerRankPtr) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckNumMaxTokensPerRankAttr failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckNumMaxTokensPerRankAttr failed."), return GRAPH_FAILED); info.numMaxTokensPerRank = static_cast(*numMaxTokensPerRankPtr); // 9. activation OP_TILING_CHECK(CheckActivationAttr(activationPtr) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckActivationAttr failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckActivationAttr failed."), return GRAPH_FAILED); // 10. activation_clamp OP_TILING_CHECK(CheckActivationClampAttr(activationClampPtr) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckActivationClampAttr failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckActivationClampAttr failed."), return GRAPH_FAILED); info.activationClamp = *activationClampPtr; // 11. activation_out_dtype OP_TILING_CHECK(CheckActivationOutDtypeAttr(activationOutDtypePtr) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckActivationOutDtypeAttr failed."), - return GRAPH_FAILED); - info.activationOutDtype = activationOutDtypePtr != nullptr ? - static_cast(*activationOutDtypePtr) : - static_cast(ge::DT_UNDEFINED); + OP_LOGE(K_INNER_DEBUG, "CheckActivationOutDtypeAttr failed."), return GRAPH_FAILED); + info.activationOutDtype = activationOutDtypePtr != nullptr ? static_cast(*activationOutDtypePtr) : + static_cast(ge::DT_UNDEFINED); // 12. transpose_weight1/2 OP_TILING_CHECK(CheckTransposeWeightAttr(transposeWeight1Ptr, transposeWeight2Ptr) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckTransposeWeightAttr failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckTransposeWeightAttr failed."), return GRAPH_FAILED); info.isTransposeW1 = *transposeWeight1Ptr ? 1U : 0U; info.isTransposeW2 = *transposeWeight2Ptr ? 1U : 0U; // 13. weight1_interleave OP_TILING_CHECK(CheckWeight1InterleaveAttr(context, weight1InterleavePtr) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckWeight1InterleaveAttr failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckWeight1InterleaveAttr failed."), return GRAPH_FAILED); info.weight1Interleave = static_cast(*weight1InterleavePtr); info.worldSize = info.epWorldSize; @@ -464,58 +444,63 @@ static ge::graphStatus MegaMoeA2A3CheckAttrAndSetTiling(gert::TilingContext *con static uint32_t GetDynamicInputTensorListLen(gert::TilingContext *context, uint32_t inputIndex) { uint32_t listLen = 0; - while (context->GetDynamicInputTensor(inputIndex, ++listLen) != nullptr) {} + while (context->GetDynamicInputTensor(inputIndex, ++listLen) != nullptr) { + } return listLen; } // 校验 context 输入 -static ge::graphStatus CheckContextInput(gert::TilingContext *context, - const gert::StorageShape *contextStorageShape) +static ge::graphStatus CheckContextInput(gert::TilingContext *context, const gert::StorageShape *contextStorageShape) { auto contextDesc = context->GetInputDesc(CONTEXT_INDEX); - OP_TILING_CHECK(contextDesc == nullptr, - OP_LOGE(K_INNER_DEBUG, "context desc is null."), return GRAPH_FAILED); + OP_TILING_CHECK(contextDesc == nullptr, OP_LOGE(K_INNER_DEBUG, "context desc is null."), return GRAPH_FAILED); OP_TILING_CHECK(contextDesc->GetDataType() != ge::DT_INT32, - OP_LOGE(K_INNER_DEBUG, "context dataType is invalid, should be INT32, but got %d.", - static_cast(contextDesc->GetDataType())), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "context dataType is invalid, should be INT32, but got %d.", + static_cast(contextDesc->GetDataType())), + return GRAPH_FAILED); OP_TILING_CHECK(static_cast(ge::GetPrimaryFormat(contextDesc->GetStorageFormat())) != ge::FORMAT_ND, - OP_LOGE(K_INNER_DEBUG, "context format is invalid, should be FORMAT_ND."), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "context format is invalid, should be FORMAT_ND."), return GRAPH_FAILED); OP_TILING_CHECK(contextStorageShape->GetStorageShape().GetDimNum() != ONE_DIM, - OP_LOGE(K_INNER_DEBUG, "contextShape dims must be 1, but current dim num is %zu.", - contextStorageShape->GetStorageShape().GetDimNum()), return ge::GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "contextShape dims must be 1, but current dim num is %zu.", + contextStorageShape->GetStorageShape().GetDimNum()), + return ge::GRAPH_FAILED); return ge::GRAPH_SUCCESS; } // 校验 x 输入，返回 bs 和 hiddenSize -static ge::graphStatus CheckXInput(gert::TilingContext *context, - const gert::StorageShape *xStorageShape, - int64_t &outBs, int64_t &outHiddenSize) +static ge::graphStatus CheckXInput(gert::TilingContext *context, const gert::StorageShape *xStorageShape, + int64_t &outBs, int64_t &outHiddenSize) { - OP_TILING_CHECK(xStorageShape->GetStorageShape().GetDimNum() != TWO_DIMS, - OP_LOGE(K_INNER_DEBUG, "x must be 2-dimension, but got %lu dim.", - xStorageShape->GetStorageShape().GetDimNum()), return GRAPH_FAILED); + OP_TILING_CHECK( + xStorageShape->GetStorageShape().GetDimNum() != TWO_DIMS, + OP_LOGE(K_INNER_DEBUG, "x must be 2-dimension, but got %lu dim.", xStorageShape->GetStorageShape().GetDimNum()), + return GRAPH_FAILED); int64_t bs = xStorageShape->GetStorageShape().GetDim(0); int64_t hiddenSize = xStorageShape->GetStorageShape().GetDim(1); - OP_TILING_CHECK(bs < MIN_BS || bs > MAX_BS, - OP_LOGE(K_INNER_DEBUG, "x's dim0(bs) is invalid, should be in [%ld, %ld], but got %ld.", - MIN_BS, MAX_BS, bs), return GRAPH_FAILED); + OP_TILING_CHECK( + bs < MIN_BS || bs > MAX_BS, + OP_LOGE(K_INNER_DEBUG, "x's dim0(bs) is invalid, should be in [%ld, %ld], but got %ld.", MIN_BS, MAX_BS, bs), + return GRAPH_FAILED); OP_TILING_CHECK(hiddenSize < MIN_HIDDEN_SIZE || hiddenSize > MAX_HIDDEN_SIZE, - OP_LOGE(K_INNER_DEBUG, "x's dim1(hidden_size) is invalid, should be in [%ld, %ld], but got %ld.", - MIN_HIDDEN_SIZE, MAX_HIDDEN_SIZE, hiddenSize), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "x's dim1(hidden_size) is invalid, should be in [%ld, %ld], but got %ld.", + MIN_HIDDEN_SIZE, MAX_HIDDEN_SIZE, hiddenSize), + return GRAPH_FAILED); OP_TILING_CHECK(hiddenSize % HIDDEN_SIZE_ALIGN != 0, - OP_LOGE(K_INNER_DEBUG, "x's dim1(hidden_size) should be %ld aligned, but got %ld.", - HIDDEN_SIZE_ALIGN, hiddenSize), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "x's dim1(hidden_size) should be %ld aligned, but got %ld.", + HIDDEN_SIZE_ALIGN, hiddenSize), + return GRAPH_FAILED); auto xDesc = context->GetInputDesc(X_INDEX); OP_TILING_CHECK(xDesc == nullptr, OP_LOGE(K_INNER_DEBUG, "x desc is null."), return GRAPH_FAILED); ge::DataType xDataType = xDesc->GetDataType(); - OP_TILING_CHECK(xDataType != ge::DT_BF16, - OP_LOGE(K_INNER_DEBUG, "x dataType is invalid, should be BF16, but got %d.", - static_cast(xDataType)), return GRAPH_FAILED); + OP_TILING_CHECK( + xDataType != ge::DT_BF16, + OP_LOGE(K_INNER_DEBUG, "x dataType is invalid, should be BF16, but got %d.", static_cast(xDataType)), + return GRAPH_FAILED); OP_TILING_CHECK(static_cast(ge::GetPrimaryFormat(xDesc->GetStorageFormat())) != ge::FORMAT_ND, - OP_LOGE(K_INNER_DEBUG, "x format is invalid, should be FORMAT_ND."), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "x format is invalid, should be FORMAT_ND."), return GRAPH_FAILED); outBs = bs; outHiddenSize = hiddenSize; @@ -523,31 +508,35 @@ static ge::graphStatus CheckXInput(gert::TilingContext *context, } // 校验 topk_ids 输入，返回 topK -static ge::graphStatus CheckTopkIdsInput(gert::TilingContext *context, - const gert::StorageShape *topkIdsStorageShape, int64_t bs, - int64_t &outTopK) +static ge::graphStatus CheckTopkIdsInput(gert::TilingContext *context, const gert::StorageShape *topkIdsStorageShape, + int64_t bs, int64_t &outTopK) { OP_TILING_CHECK(topkIdsStorageShape->GetStorageShape().GetDimNum() != TWO_DIMS, - OP_LOGE(K_INNER_DEBUG, "topk_ids must be 2-dimension, but got %lu dim.", - topkIdsStorageShape->GetStorageShape().GetDimNum()), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "topk_ids must be 2-dimension, but got %lu dim.", + topkIdsStorageShape->GetStorageShape().GetDimNum()), + return GRAPH_FAILED); int64_t topkIdsDim0 = topkIdsStorageShape->GetStorageShape().GetDim(0); int64_t topK = topkIdsStorageShape->GetStorageShape().GetDim(1); OP_TILING_CHECK(topkIdsDim0 != bs, - OP_LOGE(K_INNER_DEBUG, "topk_ids's dim0 not equal to x's dim0, topk_ids's dim0 = %ld, x's dim0 = %ld.", - topkIdsDim0, bs), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "topk_ids's dim0 not equal to x's dim0, topk_ids's dim0 = %ld, x's dim0 = %ld.", + topkIdsDim0, bs), + return GRAPH_FAILED); OP_TILING_CHECK(topK < MIN_TOPK || topK > MAX_TOPK, - OP_LOGE(K_INNER_DEBUG, "topk_ids's dim1(topK) is invalid, should be in [%ld, %ld], but got %ld.", - MIN_TOPK, MAX_TOPK, topK), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "topk_ids's dim1(topK) is invalid, should be in [%ld, %ld], but got %ld.", + MIN_TOPK, MAX_TOPK, topK), + return GRAPH_FAILED); auto topkIdsDesc = context->GetInputDesc(TOPK_IDS_INDEX); OP_TILING_CHECK(topkIdsDesc == nullptr, OP_LOGE(K_INNER_DEBUG, "topk_ids desc is null."), return GRAPH_FAILED); OP_TILING_CHECK(topkIdsDesc->GetDataType() != ge::DT_INT32, - OP_LOGE(K_INNER_DEBUG, "topk_ids dataType is invalid, should be INT32, but got %d.", - static_cast(topkIdsDesc->GetDataType())), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "topk_ids dataType is invalid, should be INT32, but got %d.", + static_cast(topkIdsDesc->GetDataType())), + return GRAPH_FAILED); OP_TILING_CHECK(static_cast(ge::GetPrimaryFormat(topkIdsDesc->GetStorageFormat())) != ge::FORMAT_ND, - OP_LOGE(K_INNER_DEBUG, "topk_ids format is invalid, should be FORMAT_ND."), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "topk_ids format is invalid, should be FORMAT_ND."), return GRAPH_FAILED); outTopK = topK; return ge::GRAPH_SUCCESS; @@ -555,73 +544,82 @@ static ge::graphStatus CheckTopkIdsInput(gert::TilingContext *context, // 校验 topk_weights 输入 static ge::graphStatus CheckTopkWeightsInput(gert::TilingContext *context, - const gert::StorageShape *topkWeightsStorageShape, int64_t bs, int64_t topK) + const gert::StorageShape *topkWeightsStorageShape, int64_t bs, + int64_t topK) { OP_TILING_CHECK(topkWeightsStorageShape->GetStorageShape().GetDimNum() != TWO_DIMS, - OP_LOGE(K_INNER_DEBUG, "topk_weights must be 2-dimension, but got %lu dim.", - topkWeightsStorageShape->GetStorageShape().GetDimNum()), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "topk_weights must be 2-dimension, but got %lu dim.", + topkWeightsStorageShape->GetStorageShape().GetDimNum()), + return GRAPH_FAILED); int64_t topkWeightsDim0 = topkWeightsStorageShape->GetStorageShape().GetDim(0); int64_t topkWeightsDim1 = topkWeightsStorageShape->GetStorageShape().GetDim(1); OP_TILING_CHECK(topkWeightsDim0 != bs, - OP_LOGE(K_INNER_DEBUG, "topk_weights's dim0 not equal to x's dim0, topk_weights's dim0 = %ld, x's dim0 = %ld.", - topkWeightsDim0, bs), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "topk_weights's dim0 not equal to x's dim0, topk_weights's dim0 = %ld, x's dim0 = %ld.", + topkWeightsDim0, bs), + return GRAPH_FAILED); OP_TILING_CHECK(topkWeightsDim1 != topK, - OP_LOGE(K_INNER_DEBUG, - "topk_weights's dim1 not equal to topk_ids's dim1, " - "topk_weights's dim1 = %ld, topk_ids's dim1 = %ld.", - topkWeightsDim1, topK), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "topk_weights's dim1 not equal to topk_ids's dim1, " + "topk_weights's dim1 = %ld, topk_ids's dim1 = %ld.", + topkWeightsDim1, topK), + return GRAPH_FAILED); auto topkWeightsDesc = context->GetInputDesc(TOPK_WEIGHTS_INDEX); - OP_TILING_CHECK(topkWeightsDesc == nullptr, - OP_LOGE(K_INNER_DEBUG, "topk_weights desc is null."), return GRAPH_FAILED); + OP_TILING_CHECK(topkWeightsDesc == nullptr, OP_LOGE(K_INNER_DEBUG, "topk_weights desc is null."), + return GRAPH_FAILED); ge::DataType topkWeightsDataType = topkWeightsDesc->GetDataType(); OP_TILING_CHECK(topkWeightsDataType != ge::DT_FLOAT && topkWeightsDataType != ge::DT_BF16, - OP_LOGE(K_INNER_DEBUG, "topk_weights dataType is invalid, should be FLOAT or BF16, but got %d.", - static_cast(topkWeightsDataType)), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "topk_weights dataType is invalid, should be FLOAT or BF16, but got %d.", + static_cast(topkWeightsDataType)), + return GRAPH_FAILED); OP_TILING_CHECK(static_cast(ge::GetPrimaryFormat(topkWeightsDesc->GetStorageFormat())) != ge::FORMAT_ND, - OP_LOGE(K_INNER_DEBUG, "topk_weights format is invalid, should be FORMAT_ND."), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "topk_weights format is invalid, should be FORMAT_ND."), + return GRAPH_FAILED); return ge::GRAPH_SUCCESS; } // 校验 x_active_mask 输入（可选），返回 isActiveMask -static ge::graphStatus CheckXActiveMaskInput(gert::TilingContext *context, int64_t bs, - bool &outIsActiveMask) +static ge::graphStatus CheckXActiveMaskInput(gert::TilingContext *context, int64_t bs, bool &outIsActiveMask) { const gert::StorageShape *xActiveMaskStorageShape = context->GetOptionalInputShape(X_ACTIVE_MASK_INDEX); bool isActiveMask = (xActiveMaskStorageShape != nullptr); if (isActiveMask) { int64_t xActiveMaskDimNums = xActiveMaskStorageShape->GetStorageShape().GetDimNum(); OP_TILING_CHECK(xActiveMaskDimNums != ONE_DIM, - OP_LOGE(K_INNER_DEBUG, - "x_active_mask must be 1-dimension, but got %ld dim. " - "2-dimension is not supported now.", - xActiveMaskDimNums), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "x_active_mask must be 1-dimension, but got %ld dim. " + "2-dimension is not supported now.", + xActiveMaskDimNums), + return GRAPH_FAILED); int64_t xActiveMaskDim0 = xActiveMaskStorageShape->GetStorageShape().GetDim(0); OP_TILING_CHECK(xActiveMaskDim0 != bs, - OP_LOGE(K_INNER_DEBUG, - "x_active_mask's dim0 not equal to x's dim0, " - "x_active_mask's dim0 = %ld, x's dim0 = %ld.", - xActiveMaskDim0, bs), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "x_active_mask's dim0 not equal to x's dim0, " + "x_active_mask's dim0 = %ld, x's dim0 = %ld.", + xActiveMaskDim0, bs), + return GRAPH_FAILED); auto xActiveMaskDesc = context->GetOptionalInputDesc(X_ACTIVE_MASK_INDEX); - OP_TILING_CHECK(xActiveMaskDesc == nullptr, - OP_LOGE(K_INNER_DEBUG, "x_active_mask desc is null."), return GRAPH_FAILED); + OP_TILING_CHECK(xActiveMaskDesc == nullptr, OP_LOGE(K_INNER_DEBUG, "x_active_mask desc is null."), + return GRAPH_FAILED); OP_TILING_CHECK(xActiveMaskDesc->GetDataType() != ge::DT_INT8, - OP_LOGE(K_INNER_DEBUG, "x_active_mask dataType is invalid, should be INT8, but got %d.", - static_cast(xActiveMaskDesc->GetDataType())), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "x_active_mask dataType is invalid, should be INT8, but got %d.", + static_cast(xActiveMaskDesc->GetDataType())), + return GRAPH_FAILED); OP_TILING_CHECK( static_cast(ge::GetPrimaryFormat(xActiveMaskDesc->GetStorageFormat())) != ge::FORMAT_ND, - OP_LOGE(K_INNER_DEBUG, "x_active_mask format is invalid, should be FORMAT_ND."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "x_active_mask format is invalid, should be FORMAT_ND."), return GRAPH_FAILED); } outIsActiveMask = isActiveMask; return ge::GRAPH_SUCCESS; } // 校验 weight1 动态输入，返回 N / expertPerRank / w1DataType / w1Format -static ge::graphStatus CheckWeight1Input(gert::TilingContext *context, int64_t hiddenSize, - uint32_t &outN, uint32_t &outExpertPerRank, ge::DataType &outW1DataType, ge::Format &outW1Format) +static ge::graphStatus CheckWeight1Input(gert::TilingContext *context, int64_t hiddenSize, uint32_t &outN, + uint32_t &outExpertPerRank, ge::DataType &outW1DataType, + ge::Format &outW1Format) { auto w1Tensor = context->GetDynamicInputTensor(WEIGHT1_INDEX, 0); OP_TILING_CHECK(w1Tensor == nullptr, OP_LOGE(K_INNER_DEBUG, "weight1 tensor is null."), return GRAPH_FAILED); @@ -630,16 +628,18 @@ static ge::graphStatus CheckWeight1Input(gert::TilingContext *context, int64_t h OP_TILING_CHECK(w1Desc == nullptr, OP_LOGE(K_INNER_DEBUG, "weight1 desc is null."), return GRAPH_FAILED); ge::DataType w1DataType = w1Desc->GetDataType(); OP_TILING_CHECK(w1DataType != ge::DT_BF16 && w1DataType != ge::DT_INT8 && w1DataType != ge::DT_INT4, - OP_LOGE(K_INNER_DEBUG, "weight1 dataType is invalid, should be BF16/INT8/INT4, but got %d.", - static_cast(w1DataType)), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "weight1 dataType is invalid, should be BF16/INT8/INT4, but got %d.", + static_cast(w1DataType)), + return GRAPH_FAILED); ge::Format w1Format = static_cast(ge::GetPrimaryFormat(w1Desc->GetStorageFormat())); - OP_TILING_CHECK( - ((w1Format != ge::FORMAT_ND && w1DataType == ge::DT_BF16) || - (w1Format != ge::FORMAT_FRACTAL_NZ && (w1DataType == ge::DT_INT8 || w1DataType == ge::DT_INT4))), - OP_LOGE(K_INNER_DEBUG, "weight1 format is invalid. " - "Expected: (FORMAT_ND with BF16) or (FORMAT_FRACTAL_NZ with INT8/INT4). " - "Actual: format=%d, dataType=%d", w1Format, w1DataType), - return GRAPH_FAILED); + OP_TILING_CHECK(((w1Format != ge::FORMAT_ND && w1DataType == ge::DT_BF16) || + (w1Format != ge::FORMAT_FRACTAL_NZ && (w1DataType == ge::DT_INT8 || w1DataType == ge::DT_INT4))), + OP_LOGE(K_INNER_DEBUG, + "weight1 format is invalid. " + "Expected: (FORMAT_ND with BF16) or (FORMAT_FRACTAL_NZ with INT8/INT4). " + "Actual: format=%d, dataType=%d", + w1Format, w1DataType), + return GRAPH_FAILED); uint32_t w1TensorDims = w1Tensor->GetOriginShape().GetDimNum(); uint32_t N = w1Tensor->GetStorageShape().GetDim(w1TensorDims - 1); @@ -647,17 +647,19 @@ static ge::graphStatus CheckWeight1Input(gert::TilingContext *context, int64_t h uint32_t expertPerRank = GetDynamicInputTensorListLen(context, WEIGHT1_INDEX); for (uint32_t i = 0; i < expertPerRank; i++) { auto wTensorI = context->GetDynamicInputTensor(WEIGHT1_INDEX, i); - OP_TILING_CHECK(wTensorI == nullptr, - OP_LOGE(K_INNER_DEBUG, "weight1[%u] tensor is null.", i), return GRAPH_FAILED); + OP_TILING_CHECK(wTensorI == nullptr, OP_LOGE(K_INNER_DEBUG, "weight1[%u] tensor is null.", i), + return GRAPH_FAILED); OP_TILING_CHECK(wTensorI->GetOriginShape().GetDimNum() != TWO_DIMS, - OP_LOGE(K_INNER_DEBUG, "weight1[%u] must be 2-dimension, but got %lu dim.", - i, wTensorI->GetOriginShape().GetDimNum()), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "weight1[%u] must be 2-dimension, but got %lu dim.", i, + wTensorI->GetOriginShape().GetDimNum()), + return GRAPH_FAILED); int64_t w1IDim0 = wTensorI->GetStorageShape().GetDim(0); OP_TILING_CHECK(w1IDim0 != hiddenSize, - OP_LOGE(K_INNER_DEBUG, - "weight1[%u]'s dim0 not equal to hidden_size, " - "weight1[%u]'s dim0 = %ld, hidden_size = %ld.", - i, i, w1IDim0, hiddenSize), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "weight1[%u]'s dim0 not equal to hidden_size, " + "weight1[%u]'s dim0 = %ld, hidden_size = %ld.", + i, i, w1IDim0, hiddenSize), + return GRAPH_FAILED); } outN = N; @@ -668,10 +670,9 @@ static ge::graphStatus CheckWeight1Input(gert::TilingContext *context, int64_t h } // 校验 weight2 动态输入，与 weight1 做交叉校验 -static ge::graphStatus CheckWeight2Input(gert::TilingContext *context, int64_t hiddenSize, - uint32_t N, uint32_t expertPerRank, - ge::DataType w1DataType, ge::Format w1Format, - ge::DataType &outW2DataType) +static ge::graphStatus CheckWeight2Input(gert::TilingContext *context, int64_t hiddenSize, uint32_t N, + uint32_t expertPerRank, ge::DataType w1DataType, ge::Format w1Format, + ge::DataType &outW2DataType) { auto w2Tensor = context->GetDynamicInputTensor(WEIGHT2_INDEX, 0); OP_TILING_CHECK(w2Tensor == nullptr, OP_LOGE(K_INNER_DEBUG, "weight2 tensor is null."), return GRAPH_FAILED); @@ -680,58 +681,70 @@ static ge::graphStatus CheckWeight2Input(gert::TilingContext *context, int64_t h OP_TILING_CHECK(w2Desc == nullptr, OP_LOGE(K_INNER_DEBUG, "weight2 desc is null."), return GRAPH_FAILED); ge::DataType w2DataType = w2Desc->GetDataType(); OP_TILING_CHECK(w2DataType != ge::DT_BF16 && w2DataType != ge::DT_INT8 && w2DataType != ge::DT_INT4, - OP_LOGE(K_INNER_DEBUG, "weight2 dataType is invalid, should be BF16/INT8/INT4, but got %d.", - static_cast(w2DataType)), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "weight2 dataType is invalid, should be BF16/INT8/INT4, but got %d.", + static_cast(w2DataType)), + return GRAPH_FAILED); // weight1 和 weight2 数据类型必须一致 OP_TILING_CHECK(w1DataType != w2DataType, - OP_LOGE(K_INNER_DEBUG, - "weight1 and weight2 must have the same dataType, " - "weight1 dataType = %d, weight2 dataType = %d.", - static_cast(w1DataType), static_cast(w2DataType)), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "weight1 and weight2 must have the same dataType, " + "weight1 dataType = %d, weight2 dataType = %d.", + static_cast(w1DataType), static_cast(w2DataType)), + return GRAPH_FAILED); ge::Format w2Format = static_cast(ge::GetPrimaryFormat(w2Desc->GetStorageFormat())); OP_TILING_CHECK(w2Format != ge::FORMAT_ND && w2Format != ge::FORMAT_FRACTAL_NZ, - OP_LOGE(K_INNER_DEBUG, "weight2 format is invalid, should be FORMAT_ND or FORMAT_FRACTAL_NZ."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "weight2 format is invalid, should be FORMAT_ND or FORMAT_FRACTAL_NZ."), + return GRAPH_FAILED); OP_TILING_CHECK(w1Format != w2Format, - OP_LOGE(K_INNER_DEBUG, - "weight1 and weight2 must have the same format, " - "weight1 format = %d, weight2 format = %d.", - static_cast(w1Format), static_cast(w2Format)), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "weight1 and weight2 must have the same format, " + "weight1 format = %d, weight2 format = %d.", + static_cast(w1Format), static_cast(w2Format)), + return GRAPH_FAILED); uint32_t w2ExpertPerRank = GetDynamicInputTensorListLen(context, WEIGHT2_INDEX); OP_TILING_CHECK(w2ExpertPerRank != expertPerRank, - OP_LOGE(K_INNER_DEBUG, - "weight2's listLen not equal to weight1's listLen, " - "weight2's listLen = %u, weight1's listLen = %u.", - w2ExpertPerRank, expertPerRank), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "weight2's listLen not equal to weight1's listLen, " + "weight2's listLen = %u, weight1's listLen = %u.", + w2ExpertPerRank, expertPerRank), + return GRAPH_FAILED); uint32_t n2 = N / 2; OP_TILING_CHECK(n2 % HIDDEN_SIZE_ALIGN != 0, - OP_LOGE(K_INNER_DEBUG, "weight2's dim0(intermediate_hidden) should be %ld aligned, but got %ld.", - HIDDEN_SIZE_ALIGN, n2), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "weight2's dim0(intermediate_hidden) should be %ld aligned, but got %ld.", + HIDDEN_SIZE_ALIGN, n2), + return GRAPH_FAILED); OP_TILING_CHECK(n2 < MIN_INTERMEDIATE_HIDDEN || n2 > MAX_INTERMEDIATE_HIDDEN, - OP_LOGE(K_INNER_DEBUG, "weight2's dim0(intermediate_hidden) is invalid, should be in [%ld, %ld], but got %ld.", - MIN_INTERMEDIATE_HIDDEN, MAX_INTERMEDIATE_HIDDEN, n2), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "weight2's dim0(intermediate_hidden) is invalid, should be in [%ld, %ld], but got %ld.", + MIN_INTERMEDIATE_HIDDEN, MAX_INTERMEDIATE_HIDDEN, n2), + return GRAPH_FAILED); for (uint32_t i = 0; i < expertPerRank; i++) { auto wTensorI = context->GetDynamicInputTensor(WEIGHT2_INDEX, i); - OP_TILING_CHECK(wTensorI == nullptr, - OP_LOGE(K_INNER_DEBUG, "weight2[%u] tensor is null.", i), return GRAPH_FAILED); + OP_TILING_CHECK(wTensorI == nullptr, OP_LOGE(K_INNER_DEBUG, "weight2[%u] tensor is null.", i), + return GRAPH_FAILED); OP_TILING_CHECK(wTensorI->GetOriginShape().GetDimNum() != TWO_DIMS, - OP_LOGE(K_INNER_DEBUG, "weight2[%u] must be 2-dimension, but got %lu dim.", - i, wTensorI->GetOriginShape().GetDimNum()), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "weight2[%u] must be 2-dimension, but got %lu dim.", i, + wTensorI->GetOriginShape().GetDimNum()), + return GRAPH_FAILED); int64_t w2IDim0 = wTensorI->GetStorageShape().GetDim(0); int64_t w2IDim1 = wTensorI->GetStorageShape().GetDim(1); OP_TILING_CHECK(w2IDim0 != n2, - OP_LOGE(K_INNER_DEBUG, "weight2[%u]'s dim0 not equal to intermediate_hidden, weight2[%u]'s dim0 = %ld, " - "intermediate_hidden = %u.", i, i, w2IDim0, n2), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "weight2[%u]'s dim0 not equal to intermediate_hidden, weight2[%u]'s dim0 = %ld, " + "intermediate_hidden = %u.", + i, i, w2IDim0, n2), + return GRAPH_FAILED); OP_TILING_CHECK(w2IDim1 != hiddenSize, - OP_LOGE(K_INNER_DEBUG, - "weight2[%u]'s dim1 not equal to hidden_size, " - "weight2[%u]'s dim1 = %ld, hidden_size = %ld.", - i, i, w2IDim1, hiddenSize), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "weight2[%u]'s dim1 not equal to hidden_size, " + "weight2[%u]'s dim1 = %ld, hidden_size = %ld.", + i, i, w2IDim1, hiddenSize), + return GRAPH_FAILED); } outW2DataType = w2DataType; @@ -739,79 +752,86 @@ static ge::graphStatus CheckWeight2Input(gert::TilingContext *context, int64_t h } // 校验 weight_scales1/weight_scales2 输入（可选），quant 权重时必选 -static ge::graphStatus CheckWeightScaleInput(gert::TilingContext *context, uint32_t inputIndex, - const char *inputName, uint32_t expertPerRank, int64_t dim1Expected, const char *dim1Name) +static ge::graphStatus CheckWeightScaleInput(gert::TilingContext *context, uint32_t inputIndex, const char *inputName, + uint32_t expertPerRank, int64_t dim1Expected, const char *dim1Name) { auto wScaleTensor = context->GetDynamicInputTensor(inputIndex, 0); OP_TILING_CHECK(wScaleTensor == nullptr, - OP_LOGE(K_INNER_DEBUG, "%s is required when weight1/weight2 is INT8 or INT4.", inputName), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "%s is required when weight1/weight2 is INT8 or INT4.", inputName), + return GRAPH_FAILED); auto wScaleDesc = context->GetDynamicInputDesc(inputIndex, 0); - OP_TILING_CHECK(wScaleDesc == nullptr, - OP_LOGE(K_INNER_DEBUG, "%s desc is null.", inputName), return GRAPH_FAILED); + OP_TILING_CHECK(wScaleDesc == nullptr, OP_LOGE(K_INNER_DEBUG, "%s desc is null.", inputName), return GRAPH_FAILED); OP_TILING_CHECK(wScaleDesc->GetDataType() != ge::DT_UINT64, - OP_LOGE(K_INNER_DEBUG, "%s dataType is invalid, should be UINT64, but got %d.", - inputName, static_cast(wScaleDesc->GetDataType())), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "%s dataType is invalid, should be UINT64, but got %d.", inputName, + static_cast(wScaleDesc->GetDataType())), + return GRAPH_FAILED); OP_TILING_CHECK(static_cast(ge::GetPrimaryFormat(wScaleDesc->GetStorageFormat())) != ge::FORMAT_ND, - OP_LOGE(K_INNER_DEBUG, "%s format is invalid, should be FORMAT_ND.", inputName), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "%s format is invalid, should be FORMAT_ND.", inputName), + return GRAPH_FAILED); uint32_t scaleListLen = GetDynamicInputTensorListLen(context, inputIndex); OP_TILING_CHECK(scaleListLen != expertPerRank, - OP_LOGE(K_INNER_DEBUG, - "%s's listLen not equal to weight1's listLen, %s's listLen = %u, weight1's listLen = %u.", - inputName, inputName, scaleListLen, expertPerRank), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "%s's listLen not equal to weight1's listLen, %s's listLen = %u, weight1's listLen = %u.", + inputName, inputName, scaleListLen, expertPerRank), + return GRAPH_FAILED); // 多个scale tensor，逐个检查 for (uint32_t i = 0; i < scaleListLen; i++) { auto sTensorI = context->GetDynamicInputTensor(inputIndex, i); - OP_TILING_CHECK(sTensorI == nullptr, - OP_LOGE(K_INNER_DEBUG, "%s[%u] tensor is null.", inputName, i), return GRAPH_FAILED); + OP_TILING_CHECK(sTensorI == nullptr, OP_LOGE(K_INNER_DEBUG, "%s[%u] tensor is null.", inputName, i), + return GRAPH_FAILED); OP_TILING_CHECK(sTensorI->GetOriginShape().GetDimNum() != ONE_DIM, - OP_LOGE(K_INNER_DEBUG, "%s[%u] must be 1-dimension, but got %lu dim.", - inputName, i, sTensorI->GetOriginShape().GetDimNum()), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "%s[%u] must be 1-dimension, but got %lu dim.", inputName, i, + sTensorI->GetOriginShape().GetDimNum()), + return GRAPH_FAILED); } return ge::GRAPH_SUCCESS; } // 校验 bias1/bias2 输入, 当且仅当 INT4 权重时存在 bias -static ge::graphStatus CheckBiasInput(gert::TilingContext *context, uint32_t inputIndex, - const char *inputName, uint32_t expertPerRank, int64_t dim0Expected, const char *dim0Name) +static ge::graphStatus CheckBiasInput(gert::TilingContext *context, uint32_t inputIndex, const char *inputName, + uint32_t expertPerRank, int64_t dim0Expected, const char *dim0Name) { const gert::StorageShape *biasStorageShape = context->GetDynamicInputShape(inputIndex, 0); OP_TILING_CHECK(biasStorageShape == nullptr, - OP_LOGE(K_INNER_DEBUG, "%s is required when weight1/weight2 is INT4.", inputName), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "%s is required when weight1/weight2 is INT4.", inputName), + return GRAPH_FAILED); auto biasTensor = context->GetDynamicInputTensor(inputIndex, 0); - OP_TILING_CHECK(biasTensor == nullptr, - OP_LOGE(K_INNER_DEBUG, "%s tensor is null.", inputName), return GRAPH_FAILED); + OP_TILING_CHECK(biasTensor == nullptr, OP_LOGE(K_INNER_DEBUG, "%s tensor is null.", inputName), + return GRAPH_FAILED); auto biasDesc = context->GetDynamicInputDesc(inputIndex, 0); - OP_TILING_CHECK(biasDesc == nullptr, - OP_LOGE(K_INNER_DEBUG, "%s desc is null.", inputName), return GRAPH_FAILED); + OP_TILING_CHECK(biasDesc == nullptr, OP_LOGE(K_INNER_DEBUG, "%s desc is null.", inputName), return GRAPH_FAILED); OP_TILING_CHECK(biasDesc->GetDataType() != ge::DT_FLOAT, - OP_LOGE(K_INNER_DEBUG, "%s dataType is invalid, should be FP32, but got %d.", - inputName, static_cast(biasDesc->GetDataType())), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "%s dataType is invalid, should be FP32, but got %d.", inputName, + static_cast(biasDesc->GetDataType())), + return GRAPH_FAILED); OP_TILING_CHECK(static_cast(ge::GetPrimaryFormat(biasDesc->GetStorageFormat())) != ge::FORMAT_ND, - OP_LOGE(K_INNER_DEBUG, "%s format is invalid, should be FORMAT_ND.", inputName), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "%s format is invalid, should be FORMAT_ND.", inputName), + return GRAPH_FAILED); uint32_t biasListLen = GetDynamicInputTensorListLen(context, inputIndex); OP_TILING_CHECK(biasListLen != expertPerRank, - OP_LOGE(K_INNER_DEBUG, - "%s's listLen not equal to weight1's listLen, %s's listLen = %u, weight1's listLen = %u.", - inputName, inputName, biasListLen, expertPerRank), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "%s's listLen not equal to weight1's listLen, %s's listLen = %u, weight1's listLen = %u.", + inputName, inputName, biasListLen, expertPerRank), + return GRAPH_FAILED); for (uint32_t i = 0; i < biasListLen; i++) { auto bTensorI = context->GetDynamicInputTensor(inputIndex, i); - OP_TILING_CHECK(bTensorI == nullptr, - OP_LOGE(K_INNER_DEBUG, "%s[%u] tensor is null.", inputName, i), return GRAPH_FAILED); + OP_TILING_CHECK(bTensorI == nullptr, OP_LOGE(K_INNER_DEBUG, "%s[%u] tensor is null.", inputName, i), + return GRAPH_FAILED); OP_TILING_CHECK(bTensorI->GetOriginShape().GetDimNum() != ONE_DIM, - OP_LOGE(K_INNER_DEBUG, "%s[%u] must be 1-dimension, but got %lu dim.", - inputName, i, bTensorI->GetOriginShape().GetDimNum()), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "%s[%u] must be 1-dimension, but got %lu dim.", inputName, i, + bTensorI->GetOriginShape().GetDimNum()), + return GRAPH_FAILED); uint32_t bDim0 = bTensorI->GetStorageShape().GetDim(0); OP_TILING_CHECK(bDim0 != static_cast(dim0Expected), - OP_LOGE(K_INNER_DEBUG, "%s[%u]'s dim0 not equal to %s, %s[%u]'s dim0 = %u, %s = %ld.", - inputName, i, dim0Name, inputName, i, bDim0, dim0Name, dim0Expected), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "%s[%u]'s dim0 not equal to %s, %s[%u]'s dim0 = %u, %s = %ld.", + inputName, i, dim0Name, inputName, i, bDim0, dim0Name, dim0Expected), + return GRAPH_FAILED); } return ge::GRAPH_SUCCESS; } @@ -823,62 +843,55 @@ static ge::graphStatus MegaMoeA2A3CheckShapeAndSetTiling(gert::TilingContext *co const gert::StorageShape *xStorageShape = context->GetInputShape(X_INDEX); const gert::StorageShape *topkIdsStorageShape = context->GetInputShape(TOPK_IDS_INDEX); const gert::StorageShape *topkWeightsStorageShape = context->GetInputShape(TOPK_WEIGHTS_INDEX); - OP_TILING_CHECK(contextStorageShape == nullptr, - OP_LOGE(K_INNER_DEBUG, "context shape is null."), return GRAPH_FAILED); - OP_TILING_CHECK(xStorageShape == nullptr, - OP_LOGE(K_INNER_DEBUG, "x shape is null."), return GRAPH_FAILED); - OP_TILING_CHECK(topkIdsStorageShape == nullptr, - OP_LOGE(K_INNER_DEBUG, "topk_ids shape is null."), return GRAPH_FAILED); - OP_TILING_CHECK(topkWeightsStorageShape == nullptr, - OP_LOGE(K_INNER_DEBUG, "topk_weights shape is null."), return GRAPH_FAILED); + OP_TILING_CHECK(contextStorageShape == nullptr, OP_LOGE(K_INNER_DEBUG, "context shape is null."), + return GRAPH_FAILED); + OP_TILING_CHECK(xStorageShape == nullptr, OP_LOGE(K_INNER_DEBUG, "x shape is null."), return GRAPH_FAILED); + OP_TILING_CHECK(topkIdsStorageShape == nullptr, OP_LOGE(K_INNER_DEBUG, "topk_ids shape is null."), + return GRAPH_FAILED); + OP_TILING_CHECK(topkWeightsStorageShape == nullptr, OP_LOGE(K_INNER_DEBUG, "topk_weights shape is null."), + return GRAPH_FAILED); // ==================== 2. context 输入校验 ==================== OP_TILING_CHECK(CheckContextInput(context, contextStorageShape) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckContextInput failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckContextInput failed."), return GRAPH_FAILED); // ==================== 3. x 输入校验 ==================== int64_t bs = 0; int64_t hiddenSize = 0; OP_TILING_CHECK(CheckXInput(context, xStorageShape, bs, hiddenSize) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckXInput failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckXInput failed."), return GRAPH_FAILED); // ==================== 4. topk_ids 输入校验 ==================== int64_t topK = 0; OP_TILING_CHECK(CheckTopkIdsInput(context, topkIdsStorageShape, bs, topK) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckTopkIdsInput failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckTopkIdsInput failed."), return GRAPH_FAILED); // ==================== 5. topk_weights 输入校验 ==================== OP_TILING_CHECK(CheckTopkWeightsInput(context, topkWeightsStorageShape, bs, topK) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckTopkWeightsInput failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckTopkWeightsInput failed."), return GRAPH_FAILED); // ==================== 6. weight1 输入校验 ==================== uint32_t N = 0; uint32_t expertPerRank = 0; ge::DataType w1DataType = ge::DT_UNDEFINED; ge::Format w1Format = ge::FORMAT_RESERVED; - OP_TILING_CHECK(CheckWeight1Input(context, hiddenSize, - N, expertPerRank, w1DataType, w1Format) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckWeight1Input failed."), - return GRAPH_FAILED); + OP_TILING_CHECK(CheckWeight1Input(context, hiddenSize, N, expertPerRank, w1DataType, w1Format) != ge::GRAPH_SUCCESS, + OP_LOGE(K_INNER_DEBUG, "CheckWeight1Input failed."), return GRAPH_FAILED); info.isQuantRouting = (w1DataType == ge::DT_FLOAT16 || w1DataType == ge::DT_BF16) ? 0U : 1U; info.isW4A8 = (w1DataType == ge::DT_INT32 || w1DataType == ge::DT_INT4) ? 1U : 0U; // expertPerRank 范围校验 OP_TILING_CHECK(expertPerRank < MIN_EXPERT_PER_RANK || expertPerRank > MAX_EXPERT_PER_RANK, - OP_LOGE(K_INNER_DEBUG, "expertPerRank is invalid, should be in [%ld, %ld], but got %u.", - MIN_EXPERT_PER_RANK, MAX_EXPERT_PER_RANK, expertPerRank), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "expertPerRank is invalid, should be in [%ld, %ld], but got %u.", + MIN_EXPERT_PER_RANK, MAX_EXPERT_PER_RANK, expertPerRank), + return GRAPH_FAILED); // ==================== 7. weight2 输入校验 ==================== ge::DataType w2DataType = ge::DT_UNDEFINED; - OP_TILING_CHECK(CheckWeight2Input(context, hiddenSize, - N, expertPerRank, w1DataType, w1Format, w2DataType) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckWeight2Input failed."), - return GRAPH_FAILED); + OP_TILING_CHECK(CheckWeight2Input(context, hiddenSize, N, expertPerRank, w1DataType, w1Format, w2DataType) != + ge::GRAPH_SUCCESS, + OP_LOGE(K_INNER_DEBUG, "CheckWeight2Input failed."), return GRAPH_FAILED); // ==================== 量化权重必须提供 scale 校验 ==================== // 当 weight1 和 weight2 的类型都为 INT8 或 INT4 时，GMM1/GMM2 需要进行反量化 @@ -889,10 +902,10 @@ static ge::graphStatus MegaMoeA2A3CheckShapeAndSetTiling(gert::TilingContext *co // 支持两种场景： // 1. 传入一个 TensorList，list 包含一个 2 维 tensor，shape 为 (num_experts_per_rank, N) // 2. 传入一个 TensorList，list 包含 num_experts_per_rank 个 tensor，每个 tensor 的 shape 为 (N,) - OP_TILING_CHECK(CheckWeightScaleInput(context, WEIGHT_SCALES1_INDEX, "weight_scales1", - expertPerRank, N, "N") != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckWeightScaleInput for weight_scales1 failed."), - return GRAPH_FAILED); + OP_TILING_CHECK(CheckWeightScaleInput(context, WEIGHT_SCALES1_INDEX, "weight_scales1", expertPerRank, N, "N") != + ge::GRAPH_SUCCESS, + OP_LOGE(K_INNER_DEBUG, "CheckWeightScaleInput for weight_scales1 failed."), + return GRAPH_FAILED); // ==================== 9. weight_scales2 输入校验（可选） ==================== // weight_scales2：可选输入，weight1/weight2 为 INT8 或 INT4 时需要，GMM2 反量化参数 @@ -900,69 +913,64 @@ static ge::graphStatus MegaMoeA2A3CheckShapeAndSetTiling(gert::TilingContext *co // 1. 传入一个 TensorList，list 包含一个 2 维 tensor，shape 为 (num_experts_per_rank, hidden_size) // 2. 传入一个 TensorList，list 包含 num_experts_per_rank 个 tensor， // 每个 tensor 的 shape 为 (hidden_size,) - OP_TILING_CHECK(CheckWeightScaleInput(context, WEIGHT_SCALES2_INDEX, "weight_scales2", - expertPerRank, hiddenSize, "hidden_size") != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckWeightScaleInput for weight_scales2 failed."), - return GRAPH_FAILED); + OP_TILING_CHECK(CheckWeightScaleInput(context, WEIGHT_SCALES2_INDEX, "weight_scales2", expertPerRank, + hiddenSize, "hidden_size") != ge::GRAPH_SUCCESS, + OP_LOGE(K_INNER_DEBUG, "CheckWeightScaleInput for weight_scales2 failed."), + return GRAPH_FAILED); } else { uint32_t weight1ListLen = GetDynamicInputTensorListLen(context, WEIGHT_SCALES1_INDEX); uint32_t weight2ListLen = GetDynamicInputTensorListLen(context, WEIGHT_SCALES2_INDEX); // 若tensor为空placeholder（shape为0，由CreateEmptyTensor创建），视为未提供 - bool hasWeight1Scale = (weight1ListLen > 0U) && - (context->GetDynamicInputDesc(WEIGHT_SCALES1_INDEX, 0) != nullptr) && + bool hasWeight1Scale = + (weight1ListLen > 0U) && (context->GetDynamicInputDesc(WEIGHT_SCALES1_INDEX, 0) != nullptr) && (context->GetDynamicInputShape(WEIGHT_SCALES1_INDEX, 0)->GetStorageShape().GetDim(0) > 0); - bool hasWeight2Scale = (weight2ListLen > 0U) && - (context->GetDynamicInputDesc(WEIGHT_SCALES2_INDEX, 0) != nullptr) && + bool hasWeight2Scale = + (weight2ListLen > 0U) && (context->GetDynamicInputDesc(WEIGHT_SCALES2_INDEX, 0) != nullptr) && (context->GetDynamicInputShape(WEIGHT_SCALES2_INDEX, 0)->GetStorageShape().GetDim(0) > 0); OP_TILING_CHECK(hasWeight1Scale || hasWeight2Scale, - OP_LOGE(K_INNER_DEBUG, "weight_scale is only supported for INT8/INT4 data type, " - "but got w1DataType=%d, got weight1ListLen=%d, weight2ListLen=%d", - w1DataType, weight1ListLen, weight2ListLen), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "weight_scale is only supported for INT8/INT4 data type, " + "but got w1DataType=%d, got weight1ListLen=%d, weight2ListLen=%d", + w1DataType, weight1ListLen, weight2ListLen), + return GRAPH_FAILED); } if (w1DataType == ge::DT_INT4) { // ==================== 10. bias1 输入校验 ==================== // 当 weight1 和 weight2 数据类型都为 INT4 时，bias1 必须存在 - OP_TILING_CHECK(CheckBiasInput(context, BIAS1_INDEX, "bias1", - expertPerRank, N, "N") != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckBiasInput for bias1 failed."), - return GRAPH_FAILED); + OP_TILING_CHECK(CheckBiasInput(context, BIAS1_INDEX, "bias1", expertPerRank, N, "N") != ge::GRAPH_SUCCESS, + OP_LOGE(K_INNER_DEBUG, "CheckBiasInput for bias1 failed."), return GRAPH_FAILED); // ==================== 11. bias2 输入校验 ==================== // 当 weight1 和 weight2 数据类型都为 INT4 时，bias2 必须存在 - OP_TILING_CHECK(CheckBiasInput(context, BIAS2_INDEX, "bias2", - expertPerRank, hiddenSize, "hiddenSize") != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckBiasInput for bias2 failed."), - return GRAPH_FAILED); + OP_TILING_CHECK(CheckBiasInput(context, BIAS2_INDEX, "bias2", expertPerRank, hiddenSize, "hiddenSize") != + ge::GRAPH_SUCCESS, + OP_LOGE(K_INNER_DEBUG, "CheckBiasInput for bias2 failed."), return GRAPH_FAILED); } else { uint32_t bias1ListLen = GetDynamicInputTensorListLen(context, BIAS1_INDEX); uint32_t bias2ListLen = GetDynamicInputTensorListLen(context, BIAS2_INDEX); // 若tensor为空placeholder（shape为0，由CreateEmptyTensor创建），视为未提供 - bool hasBias1 = (bias1ListLen > 0U) && - (context->GetDynamicInputDesc(BIAS1_INDEX, 0) != nullptr) && - (context->GetDynamicInputShape(BIAS1_INDEX, 0)->GetStorageShape().GetDim(0) > 0); - bool hasBias2 = (bias2ListLen > 0U) && - (context->GetDynamicInputDesc(BIAS2_INDEX, 0) != nullptr) && - (context->GetDynamicInputShape(BIAS2_INDEX, 0)->GetStorageShape().GetDim(0) > 0); + bool hasBias1 = (bias1ListLen > 0U) && (context->GetDynamicInputDesc(BIAS1_INDEX, 0) != nullptr) && + (context->GetDynamicInputShape(BIAS1_INDEX, 0)->GetStorageShape().GetDim(0) > 0); + bool hasBias2 = (bias2ListLen > 0U) && (context->GetDynamicInputDesc(BIAS2_INDEX, 0) != nullptr) && + (context->GetDynamicInputShape(BIAS2_INDEX, 0)->GetStorageShape().GetDim(0) > 0); OP_TILING_CHECK(hasBias1 || hasBias2, - OP_LOGE(K_INNER_DEBUG, "bias is only supported for INT4 data type, " - "but got w1DataType=%d, got bias1ListLen=%d, bias2ListLen=%d", - w1DataType, bias1ListLen, bias2ListLen), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "bias is only supported for INT4 data type, " + "but got w1DataType=%d, got bias1ListLen=%d, bias2ListLen=%d", + w1DataType, bias1ListLen, bias2ListLen), + return GRAPH_FAILED); } // ==================== 12. x_active_mask 输入校验（可选） ==================== bool isActiveMask = false; OP_TILING_CHECK(CheckXActiveMaskInput(context, bs, isActiveMask) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "CheckXActiveMaskInput failed."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "CheckXActiveMaskInput failed."), return GRAPH_FAILED); // ==================== 13. scales 输入校验（可选，预留） ==================== const gert::StorageShape *scalesStorageShape = context->GetOptionalInputShape(SCALES_INDEX); OP_TILING_CHECK(scalesStorageShape != nullptr, - OP_LOGE(K_INNER_DEBUG, "scales is not supported yet, please pass nullptr."), - return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "scales is not supported yet, please pass nullptr."), return GRAPH_FAILED); // ==================== 14. 设置 tiling data ==================== info.M = static_cast(bs); @@ -983,7 +991,7 @@ static ge::graphStatus MegaMoeA2A3CheckShapeAndSetTiling(gert::TilingContext *co return ge::GRAPH_SUCCESS; } -static ge::graphStatus MegaMoeA2A3GetPlatformInfoAndSetTiling(gert::TilingContext *context, MegaMoeA2A3TilingData& info) +static ge::graphStatus MegaMoeA2A3GetPlatformInfoAndSetTiling(gert::TilingContext *context, MegaMoeA2A3TilingData &info) { auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); uint32_t aivNum = ascendcPlatform.GetCoreNumAiv(); @@ -1004,8 +1012,7 @@ static ge::graphStatus MegaMoeA2A3CommAlg(const gert::TilingContext *context) auto attrs = context->GetAttrs(); OP_TILING_CHECK(attrs == nullptr, OP_LOGE(K_INNER_DEBUG, "attrs is null."), return ge::GRAPH_FAILED); auto commAlg = attrs->GetAttrPointer(static_cast(ATTR_COMM_ALG_INDEX)); - OP_TILING_CHECK(commAlg == nullptr, - OP_LOGE(K_INNER_DEBUG, "commAlg is null."), return ge::GRAPH_FAILED); + OP_TILING_CHECK(commAlg == nullptr, OP_LOGE(K_INNER_DEBUG, "commAlg is null."), return ge::GRAPH_FAILED); if (strlen(commAlg) > 0) { OP_LOGE(K_INNER_DEBUG, "Attr commAlg is invalid, current version only supports \"\", but got \"%s\".", commAlg); @@ -1015,14 +1022,13 @@ static ge::graphStatus MegaMoeA2A3CommAlg(const gert::TilingContext *context) return ge::GRAPH_SUCCESS; } -static ge::graphStatus MegaMoeA2A3CheckHcclBuffSize(const gert::TilingContext *context, - MegaMoeA2A3TilingData &info) +static ge::graphStatus MegaMoeA2A3CheckHcclBuffSize(const gert::TilingContext *context, MegaMoeA2A3TilingData &info) { const char *nodeName = K_INNER_DEBUG; // info.cclBufferSize是winIn和winOut空间之和（两者相等），这里只通过winIn校验 int64_t cclBufferSize = info.cclBufferSize / 2; - OP_LOGD(nodeName, "cclBufferSize = %ld Bytes (%ld MB).", - cclBufferSize, ops::CeilDiv(cclBufferSize, static_cast(MB_SIZE))); + OP_LOGD(nodeName, "cclBufferSize = %ld Bytes (%ld MB).", cclBufferSize, + ops::CeilDiv(cclBufferSize, static_cast(MB_SIZE))); int64_t h = info.K; int64_t maxRecvTokenNum = info.maxRecvTokenNum; @@ -1034,21 +1040,20 @@ static ge::graphStatus MegaMoeA2A3CheckHcclBuffSize(const gert::TilingContext *c std::string socVersion = mc2tiling::GetSocVersion(context); bool isA3 = (socVersion == "Ascend910_93"); - int64_t leastCclBufferSize = CalcLeastCclBufferSize(maxRecvTokenNum, h, - epWorldSize, expertPerRank, isQuantRouting, isW4A8, isA3, - static_cast(info.M), static_cast(info.topK)); + int64_t leastCclBufferSize = + CalcLeastCclBufferSize(maxRecvTokenNum, h, epWorldSize, expertPerRank, isQuantRouting, isW4A8, isA3, + static_cast(info.M), static_cast(info.topK)); OP_TILING_CHECK(cclBufferSize < leastCclBufferSize, - OP_LOGE(nodeName, "cclBufferSize(%ld Bytes, %ld MB) should be >= leastCclBufferSize(%ld Bytes, %ld MB). " - "maxRecvTokenNum=%ld, h=%ld, epWorldSize=%ld, expertPerRank=%ld, " - "isQuantRouting=%d, isW4A8=%d, isA3=%d. " - "Suggestion: set hccl buffsize to %ld MB.", - cclBufferSize, cclBufferSize / MB_SIZE, - leastCclBufferSize, leastCclBufferSize / MB_SIZE, - maxRecvTokenNum, h, epWorldSize, - expertPerRank, isQuantRouting, isW4A8, isA3, - (leastCclBufferSize + MB_SIZE) / MB_SIZE), - return ge::GRAPH_FAILED); + OP_LOGE(nodeName, + "cclBufferSize(%ld Bytes, %ld MB) should be >= leastCclBufferSize(%ld Bytes, %ld MB). " + "maxRecvTokenNum=%ld, h=%ld, epWorldSize=%ld, expertPerRank=%ld, " + "isQuantRouting=%d, isW4A8=%d, isA3=%d. " + "Suggestion: set hccl buffsize to %ld MB.", + cclBufferSize, cclBufferSize / MB_SIZE, leastCclBufferSize, leastCclBufferSize / MB_SIZE, + maxRecvTokenNum, h, epWorldSize, expertPerRank, isQuantRouting, isW4A8, isA3, + (leastCclBufferSize + MB_SIZE) / MB_SIZE), + return ge::GRAPH_FAILED); OP_LOGD(nodeName, "cclBufferSize is %ld, leastCclBufferSize is %ld", cclBufferSize, leastCclBufferSize); return ge::GRAPH_SUCCESS; @@ -1058,33 +1063,33 @@ static ge::graphStatus MegaMoeA2A3CheckOutputTensor(gert::TilingContext *context { // ==================== 1. y 输出校验 ==================== const gert::StorageShape *yStorageShape = context->GetOutputShape(OUTPUT_Y_INDEX); - OP_TILING_CHECK(yStorageShape == nullptr, - OP_LOGE(K_INNER_DEBUG, "y shape is null."), return GRAPH_FAILED); + OP_TILING_CHECK(yStorageShape == nullptr, OP_LOGE(K_INNER_DEBUG, "y shape is null."), return GRAPH_FAILED); // 维度数校验：必须为 2 维 - OP_TILING_CHECK(yStorageShape->GetStorageShape().GetDimNum() != TWO_DIMS, - OP_LOGE(K_INNER_DEBUG, "y must be 2-dimension, but got %lu dim.", - yStorageShape->GetStorageShape().GetDimNum()), return GRAPH_FAILED); + OP_TILING_CHECK( + yStorageShape->GetStorageShape().GetDimNum() != TWO_DIMS, + OP_LOGE(K_INNER_DEBUG, "y must be 2-dimension, but got %lu dim.", yStorageShape->GetStorageShape().GetDimNum()), + return GRAPH_FAILED); int64_t yDim0 = yStorageShape->GetStorageShape().GetDim(0); int64_t yDim1 = yStorageShape->GetStorageShape().GetDim(1); // Shape 一致性校验：y 的 shape 应该与 x 的 shape 一致 [M, K] OP_TILING_CHECK(yDim0 != static_cast(info.M), - OP_LOGE(K_INNER_DEBUG, "y's dim0 not equal to M, y's dim0 = %ld, M = %u.", - yDim0, info.M), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "y's dim0 not equal to M, y's dim0 = %ld, M = %u.", yDim0, info.M), + return GRAPH_FAILED); OP_TILING_CHECK(yDim1 != static_cast(info.K), - OP_LOGE(K_INNER_DEBUG, "y's dim1 not equal to K, y's dim1 = %ld, K = %u.", - yDim1, info.K), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "y's dim1 not equal to K, y's dim1 = %ld, K = %u.", yDim1, info.K), + return GRAPH_FAILED); // 数据类型校验 auto yDesc = context->GetOutputDesc(OUTPUT_Y_INDEX); - OP_TILING_CHECK(yDesc == nullptr, - OP_LOGE(K_INNER_DEBUG, "y desc is null."), return GRAPH_FAILED); + OP_TILING_CHECK(yDesc == nullptr, OP_LOGE(K_INNER_DEBUG, "y desc is null."), return GRAPH_FAILED); ge::DataType yDataType = yDesc->GetDataType(); OP_TILING_CHECK(yDataType != ge::DT_FLOAT16 && yDataType != ge::DT_BF16, - OP_LOGE(K_INNER_DEBUG, "y dataType is invalid, should be FLOAT16 or BF16, but got %d.", - static_cast(yDataType)), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "y dataType is invalid, should be FLOAT16 or BF16, but got %d.", + static_cast(yDataType)), + return GRAPH_FAILED); OP_LOGD(K_INNER_DEBUG, "y dim0 = %ld", yDim0); OP_LOGD(K_INNER_DEBUG, "y dim1 = %ld", yDim1); @@ -1092,30 +1097,33 @@ static ge::graphStatus MegaMoeA2A3CheckOutputTensor(gert::TilingContext *context // ==================== 2. expert_token_nums 输出校验 ==================== const gert::StorageShape *expertTokenNumsStorageShape = context->GetOutputShape(OUTPUT_EXPERT_TOKEN_NUMS_INDEX); - OP_TILING_CHECK(expertTokenNumsStorageShape == nullptr, - OP_LOGE(K_INNER_DEBUG, "expert_token_nums shape is null."), return GRAPH_FAILED); + OP_TILING_CHECK(expertTokenNumsStorageShape == nullptr, OP_LOGE(K_INNER_DEBUG, "expert_token_nums shape is null."), + return GRAPH_FAILED); // 维度数校验：必须为 1 维 OP_TILING_CHECK(expertTokenNumsStorageShape->GetStorageShape().GetDimNum() != ONE_DIM, - OP_LOGE(K_INNER_DEBUG, "expert_token_nums must be 1-dimension, but got %lu dim.", - expertTokenNumsStorageShape->GetStorageShape().GetDimNum()), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "expert_token_nums must be 1-dimension, but got %lu dim.", + expertTokenNumsStorageShape->GetStorageShape().GetDimNum()), + return GRAPH_FAILED); int64_t expertTokenNumsDim0 = expertTokenNumsStorageShape->GetStorageShape().GetDim(0); // Shape 校验：expert_token_nums 的 dim0 应该等于 expertPerRank（本卡专家数量） OP_TILING_CHECK(expertTokenNumsDim0 != static_cast(info.expertPerRank), - OP_LOGE(K_INNER_DEBUG, - "expert_token_nums's dim0 not equal to expertPerRank, " - "expert_token_nums's dim0 = %ld, expertPerRank = %u.", - expertTokenNumsDim0, info.expertPerRank), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, + "expert_token_nums's dim0 not equal to expertPerRank, " + "expert_token_nums's dim0 = %ld, expertPerRank = %u.", + expertTokenNumsDim0, info.expertPerRank), + return GRAPH_FAILED); // 数据类型校验 auto expertTokenNumsDesc = context->GetOutputDesc(OUTPUT_EXPERT_TOKEN_NUMS_INDEX); - OP_TILING_CHECK(expertTokenNumsDesc == nullptr, - OP_LOGE(K_INNER_DEBUG, "expert_token_nums desc is null."), return GRAPH_FAILED); + OP_TILING_CHECK(expertTokenNumsDesc == nullptr, OP_LOGE(K_INNER_DEBUG, "expert_token_nums desc is null."), + return GRAPH_FAILED); OP_TILING_CHECK(expertTokenNumsDesc->GetDataType() != ge::DT_INT32, - OP_LOGE(K_INNER_DEBUG, "expert_token_nums dataType is invalid, should be INT32, but got %d.", - static_cast(expertTokenNumsDesc->GetDataType())), return GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "expert_token_nums dataType is invalid, should be INT32, but got %d.", + static_cast(expertTokenNumsDesc->GetDataType())), + return GRAPH_FAILED); OP_LOGD(K_INNER_DEBUG, "expert_token_nums dim0 = %ld", expertTokenNumsDim0); OP_LOGD(K_INNER_DEBUG, "expert_token_nums dataType = %d", static_cast(expertTokenNumsDesc->GetDataType())); @@ -1131,30 +1139,24 @@ static ge::graphStatus MegaMoeA2A3TilingFuncImpl(gert::TilingContext *context) // 1. tilingData MegaMoeTilingDataQuant *tilingData = context->GetTilingData(); - OP_TILING_CHECK(tilingData == nullptr, OP_LOGE(K_INNER_DEBUG, "tilingData is nullptr."), - return ge::GRAPH_FAILED); + OP_TILING_CHECK(tilingData == nullptr, OP_LOGE(K_INNER_DEBUG, "tilingData is nullptr."), return ge::GRAPH_FAILED); OP_LOGI(K_INNER_DEBUG, "MegaMoeA2A3 get tilingData."); - MegaMoeA2A3TilingData& info = tilingData->common; + MegaMoeA2A3TilingData &info = tilingData->common; OP_LOGI(K_INNER_DEBUG, "MegaMoeA2A3 get tilingData info."); OP_TILING_CHECK(MegaMoeA2A3CommAlg(context) != ge::GRAPH_SUCCESS, - VECTOR_INNER_ERR_REPORT_TILING(K_INNER_DEBUG, "MegaMoeA2A3 CheckCommAlg Failed"), - return ge::GRAPH_FAILED); + VECTOR_INNER_ERR_REPORT_TILING(K_INNER_DEBUG, "MegaMoeA2A3 CheckCommAlg Failed"), + return ge::GRAPH_FAILED); OP_TILING_CHECK(MegaMoeA2A3CheckShapeAndSetTiling(context, info) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "MegaMoeA2A3 CheckShapeAndSetTiling Failed"), - return ge::GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "MegaMoeA2A3 CheckShapeAndSetTiling Failed"), return ge::GRAPH_FAILED); OP_TILING_CHECK(MegaMoeA2A3CheckAttrAndSetTiling(context, info) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "MegaMoeA2A3 CheckAttrAndSetTiling Failed"), - return ge::GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "MegaMoeA2A3 CheckAttrAndSetTiling Failed"), return ge::GRAPH_FAILED); OP_TILING_CHECK(MegaMoeA2A3CheckOutputTensor(context, info) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "MegaMoeA2A3 CheckOutputTensor Failed"), - return ge::GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "MegaMoeA2A3 CheckOutputTensor Failed"), return ge::GRAPH_FAILED); OP_TILING_CHECK(MegaMoeA2A3GetPlatformInfoAndSetTiling(context, info) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "MegaMoeA2A3 GetPlatformInfoAndSetTiling Failed"), - return ge::GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "MegaMoeA2A3 GetPlatformInfoAndSetTiling Failed"), return ge::GRAPH_FAILED); OP_TILING_CHECK(MegaMoeA2A3CheckHcclBuffSize(context, info) != ge::GRAPH_SUCCESS, - OP_LOGE(K_INNER_DEBUG, "MegaMoeA2A3 CheckHcclBuffSize Failed"), - return ge::GRAPH_FAILED); + OP_LOGE(K_INNER_DEBUG, "MegaMoeA2A3 CheckHcclBuffSize Failed"), return ge::GRAPH_FAILED); // 2. set blockDim uint32_t blockDim = 1U; @@ -1165,32 +1167,29 @@ static ge::graphStatus MegaMoeA2A3TilingFuncImpl(gert::TilingContext *context) context->SetBlockDim(blockDim); // 3. set tiling key - uint32_t quantMode = info.isQuantRouting != 0U - ? MEGA_MOE_QUANT_MODE_PER_TENSOR - : MEGA_MOE_QUANT_MODE_NO_QUANT; + uint32_t quantMode = info.isQuantRouting != 0U ? MEGA_MOE_QUANT_MODE_PER_TENSOR : MEGA_MOE_QUANT_MODE_NO_QUANT; uint32_t quantOutType = MEGA_MOE_QUANT_OUT_TYPE_UNDEFINED; if (info.dispatchQuantOutDtype == static_cast(ge::DT_INT8)) { quantOutType = MEGA_MOE_QUANT_OUT_TYPE_INT8; } if (quantMode == MEGA_MOE_QUANT_MODE_NO_QUANT) { - OP_TILING_CHECK(quantOutType != MEGA_MOE_QUANT_OUT_TYPE_UNDEFINED, - OP_LOGE(K_INNER_DEBUG, "dispatch_quant_out_type must be UNDEFINED in non-quant mode, got %u", - quantOutType), return ge::GRAPH_FAILED); + OP_TILING_CHECK( + quantOutType != MEGA_MOE_QUANT_OUT_TYPE_UNDEFINED, + OP_LOGE(K_INNER_DEBUG, "dispatch_quant_out_type must be UNDEFINED in non-quant mode, got %u", quantOutType), + return ge::GRAPH_FAILED); } else { - OP_TILING_CHECK(quantOutType != MEGA_MOE_QUANT_OUT_TYPE_INT8, - OP_LOGE(K_INNER_DEBUG, "dispatch_quant_out_type must be INT8 in quant mode, got %u", - quantOutType), return ge::GRAPH_FAILED); + OP_TILING_CHECK( + quantOutType != MEGA_MOE_QUANT_OUT_TYPE_INT8, + OP_LOGE(K_INNER_DEBUG, "dispatch_quant_out_type must be INT8 in quant mode, got %u", quantOutType), + return ge::GRAPH_FAILED); } // archCode 在910B时为0, 910_93时为1 std::string socVersion = mc2tiling::GetSocVersion(context); uint32_t archCode = socVersion == "Ascend910B" ? SOC_ASCEND910B : SOC_ASCEND910_93; - uint64_t tilingKey = GET_TPL_TILING_KEY( - static_cast(info.isTransposeW1), - static_cast(info.isTransposeW2), - static_cast(quantMode), - static_cast(quantOutType), - static_cast(archCode)); + uint64_t tilingKey = GET_TPL_TILING_KEY(static_cast(info.isTransposeW1), + static_cast(info.isTransposeW2), static_cast(quantMode), + static_cast(quantOutType), static_cast(archCode)); context->SetTilingKey(tilingKey); int64_t inuptXDtypeSize = sizeof(int16_t); @@ -1213,10 +1212,9 @@ static ge::graphStatus MegaMoeA2A3TilingFuncImpl(gert::TilingContext *context) quantTilingData->common = info; MoeInitRoutingQuantV2TilingBase moeInitRoutingQuantV2TilingBase; - moeInitRoutingQuantV2TilingBase.DoTiling(info.M, info.K, info.topK, - expertCapacity, expertNum, activeNum, dropPadMode, - expertTokensCountOrCumsumFlag, expertTokensBeforeCapacityFlag, - inuptXDtypeSize, quantModeRouting, scaleDim0, aivNum, ubSize); + moeInitRoutingQuantV2TilingBase.DoTiling( + info.M, info.K, info.topK, expertCapacity, expertNum, activeNum, dropPadMode, expertTokensCountOrCumsumFlag, + expertTokensBeforeCapacityFlag, inuptXDtypeSize, quantModeRouting, scaleDim0, aivNum, ubSize); initRoutingQuantTilingKey = moeInitRoutingQuantV2TilingBase.tilingKey_; initRoutingWorkspace = moeInitRoutingQuantV2TilingBase.workspaceSize_; quantTilingData->moeInitRoutingQuantV2TilingData = moeInitRoutingQuantV2TilingBase.quantTilingData; @@ -1228,35 +1226,33 @@ static ge::graphStatus MegaMoeA2A3TilingFuncImpl(gert::TilingContext *context) nonQuantTilingData->common = info; MoeInitRoutingV2TilingBase moeInitRoutingV2TilingBase; - moeInitRoutingV2TilingBase.DoTiling(info.M, info.K, info.topK, - expertCapacity, expertNum, activeNum, dropPadMode, - expertTokensCountOrCumsumFlag, expertTokensBeforeCapacityFlag, - inuptXDtypeSize, quantModeRouting, scaleDim0, aivNum, ubSize); + moeInitRoutingV2TilingBase.DoTiling(info.M, info.K, info.topK, expertCapacity, expertNum, activeNum, + dropPadMode, expertTokensCountOrCumsumFlag, expertTokensBeforeCapacityFlag, + inuptXDtypeSize, quantModeRouting, scaleDim0, aivNum, ubSize); initRoutingQuantTilingKey = moeInitRoutingV2TilingBase.tilingKey_; initRoutingWorkspace = moeInitRoutingV2TilingBase.workspaceSize_; nonQuantTilingData->moeInitRoutingV2TilingData = moeInitRoutingV2TilingBase.moeInitRoutingTilingData; nonQuantTilingData->common.initRoutingQuantTilingKey = initRoutingQuantTilingKey; } - OP_LOGD(K_INNER_DEBUG, "initRoutingQuantTilingKey=%lu (isQuantRouting=%u)", - initRoutingQuantTilingKey, info.isQuantRouting); + OP_LOGD(K_INNER_DEBUG, "initRoutingQuantTilingKey=%lu (isQuantRouting=%u)", initRoutingQuantTilingKey, + info.isQuantRouting); // 4. workspace size_t *workSpaces = context->GetWorkspaceSizes(1); - OP_TILING_CHECK(workSpaces == nullptr, OP_LOGE(K_INNER_DEBUG, "workSpaces is nullptr."), - return ge::GRAPH_FAILED); + OP_TILING_CHECK(workSpaces == nullptr, OP_LOGE(K_INNER_DEBUG, "workSpaces is nullptr."), return ge::GRAPH_FAILED); uint32_t n2 = info.K; uint32_t k2 = info.N / 2; uint64_t megeMoeWorkspace = 0; if (archCode == SOC_ASCEND910_93) { - megeMoeWorkspace = (info.M + 256 - 1) / 256 * 256 * info.topK * sizeof(int32_t) + // expandedRowIdx - info.worldSize * info.worldSize * info.expertPerRank * sizeof(int32_t) + // cumsum - info.maxRecvTokenNum * std::max(info.N, n2) * sizeof(int16_t) + // GMM1&2 Out - info.maxRecvTokenNum * std::max(info.K, k2) * sizeof(int16_t) + // GMM1&2 input - (info.worldSize * - (info.expertPerRank + 16 - 1) / 16 * 16 * sizeof(int32_t)) + // SumBeforeRank - info.worldSize * sizeof(int32_t) * 16; // sync + megeMoeWorkspace = + (info.M + 256 - 1) / 256 * 256 * info.topK * sizeof(int32_t) + // expandedRowIdx + info.worldSize * info.worldSize * info.expertPerRank * sizeof(int32_t) + // cumsum + info.maxRecvTokenNum * std::max(info.N, n2) * sizeof(int16_t) + // GMM1&2 Out + info.maxRecvTokenNum * std::max(info.K, k2) * sizeof(int16_t) + // GMM1&2 input + (info.worldSize * (info.expertPerRank + 16 - 1) / 16 * 16 * sizeof(int32_t)) + // SumBeforeRank + info.worldSize * sizeof(int32_t) * 16; // sync if (info.isQuantRouting == 1U) { megeMoeWorkspace += info.maxRecvTokenNum * sizeof(float) * 2; // perTokenScale GMM1&2 } @@ -1267,12 +1263,13 @@ static ge::graphStatus MegaMoeA2A3TilingFuncImpl(gert::TilingContext *context) } else if (archCode == SOC_ASCEND910B) { uint64_t swigluSize = info.isQuantRouting ? 1UL : 2UL; uint64_t paddedExpertNumAligned = (info.worldSize * info.expertPerRank + 1 + 128 - 1) / 128 * 128; - megeMoeWorkspace = (info.M + 256 - 1) / 256 * 256 * info.topK * sizeof(int32_t) + // expandedRowIdx - // cumsum + ptrSumBeforeRankForDispatch + ptrSumBeforeRankForCombine - paddedExpertNumAligned * (info.worldSize + 2UL) * sizeof(int32_t) + - info.maxRecvTokenNum * std::max(info.N, n2) * sizeof(int16_t) + // GMM1&2 Out - info.maxRecvTokenNum * std::max(info.K, k2) * swigluSize + // swiglu out & quantizedToken - info.worldSize * sizeof(int32_t) * 16UL; // sync + megeMoeWorkspace = (info.M + 256 - 1) / 256 * 256 * info.topK * + sizeof(int32_t) + // expandedRowIdx + // cumsum + ptrSumBeforeRankForDispatch + ptrSumBeforeRankForCombine + paddedExpertNumAligned * (info.worldSize + 2UL) * sizeof(int32_t) + + info.maxRecvTokenNum * std::max(info.N, n2) * sizeof(int16_t) + // GMM1&2 Out + info.maxRecvTokenNum * std::max(info.K, k2) * swigluSize + // swiglu out & quantizedToken + info.worldSize * sizeof(int32_t) * 16UL; // sync if (info.isQuantRouting) { megeMoeWorkspace += info.maxRecvTokenNum * sizeof(float) * 2UL; // perTokenScale GMM1&2 } @@ -1287,7 +1284,7 @@ static ge::graphStatus MegaMoeA2A3TilingFuncImpl(gert::TilingContext *context) return ge::GRAPH_SUCCESS; } -static ge::graphStatus MegaMoeA2A3TilingFunc(gert::TilingContext* context) +static ge::graphStatus MegaMoeA2A3TilingFunc(gert::TilingContext *context) { return MegaMoeA2A3TilingFuncImpl(context); } @@ -1299,9 +1296,7 @@ ge::graphStatus TilingParseForMegaMoeA2A3(gert::TilingParseContext *context) return ge::GRAPH_SUCCESS; } -IMPL_OP_OPTILING(MegaMoe) - .Tiling(MegaMoeA2A3TilingFunc) - .TilingParse(TilingParseForMegaMoeA2A3); +IMPL_OP_OPTILING(MegaMoe).Tiling(MegaMoeA2A3TilingFunc).TilingParse(TilingParseForMegaMoeA2A3); #if RUNTIME_VERSION_NUM >= EXCEPTION_DUMP_SUPPORT_VERSION && METADEF_VERSION_NUM >= EXCEPTION_DUMP_SUPPORT_VERSION inline void MegaMoeExceptionImplWrapper(aclrtExceptionInfo *args, void *userdata) @@ -1325,13 +1320,12 @@ __attribute__((constructor)) void RegisterMegaMoeExceptionFunc() if (runtimeVersionNum < EXCEPTION_DUMP_SUPPORT_VERSION || metadefVersionNum < EXCEPTION_DUMP_SUPPORT_VERSION) { OP_LOGW("MegaMoe", - "The runtime(%d) or metadata(%d) version is lower than the version(%d) supporting exception func.", - runtimeVersionNum, metadefVersionNum, EXCEPTION_DUMP_SUPPORT_VERSION); + "The runtime(%d) or metadata(%d) version is lower than the version(%d) supporting exception func.", + runtimeVersionNum, metadefVersionNum, EXCEPTION_DUMP_SUPPORT_VERSION); return; } - IMPL_OP(MegaMoe) - .ExceptionDumpParseFunc(MegaMoeExceptionImplWrapper); + IMPL_OP(MegaMoe).ExceptionDumpParseFunc(MegaMoeExceptionImplWrapper); } #endif } // namespace MegaMoeA2A3Tiling diff --git a/mc2/mega_moe/op_kernel/arch22/utils/hccl_shmem.hpp b/mc2/mega_moe/op_kernel/arch22/utils/hccl_shmem.hpp index 469c8d3..ce382ff 100644 --- a/mc2/mega_moe/op_kernel/arch22/utils/hccl_shmem.hpp +++ b/mc2/mega_moe/op_kernel/arch22/utils/hccl_shmem.hpp @@ -99,25 +99,21 @@ FORCE_INLINE_AICORE void gm_signal_wait_until_ne(__gm__ int32_t *sig_addr, int32 return; } -FORCE_INLINE_AICORE void AIVRDMAPostSend( - GM_ADDR srcDmaAddr, - GM_ADDR destDmaAddr, - uint64_t destRankId, - uint64_t messageLen, - __gm__ HcclAiRMAInfo* QpInfo, - AscendC::LocalTensor &ubLocal, - AscendC::LocalTensor &ubLocalHead) +FORCE_INLINE_AICORE void AIVRDMAPostSend(GM_ADDR srcDmaAddr, GM_ADDR destDmaAddr, uint64_t destRankId, + uint64_t messageLen, __gm__ HcclAiRMAInfo *QpInfo, + AscendC::LocalTensor &ubLocal, + AscendC::LocalTensor &ubLocalHead) { auto qpNum = QpInfo->qpNum; - auto qp_ctx_entry = (__gm__ HcclAiRMAWQ*)(QpInfo->sqPtr + - destRankId * qpNum * static_cast(QpInfo->sizeOfAiRMAWQ)); + auto qp_ctx_entry = + (__gm__ HcclAiRMAWQ *)(QpInfo->sqPtr + destRankId * qpNum * static_cast(QpInfo->sizeOfAiRMAWQ)); auto mem_info_table = QpInfo->memPtr; auto sizeof_memdetail = QpInfo->sizeOfAiRMAMem; auto sqBaseAddr = qp_ctx_entry->bufAddr; auto wqeSize = qp_ctx_entry->wqeSize; auto curHardwareHead = qp_ctx_entry->headAddr; - cacheWriteThrough((__gm__ uint8_t*)curHardwareHead, 8); - uint64_t curHead = *(__gm__ uint32_t*)(curHardwareHead); + cacheWriteThrough((__gm__ uint8_t *)curHardwareHead, 8); + uint64_t curHead = *(__gm__ uint32_t *)(curHardwareHead); auto curHardwareTailAddr = qp_ctx_entry->tailAddr; uint64_t shift = 15U; auto QP_DEPTH = qp_ctx_entry->depth; @@ -125,39 +121,42 @@ FORCE_INLINE_AICORE void AIVRDMAPostSend( AscendC::PipeBarrier(); while (true) { - cacheWriteThrough((__gm__ uint8_t*)curHardwareTailAddr, 8); - if ((curHead - *(__gm__ uint32_t*)(curHardwareTailAddr)) < QP_DEPTH - 1) { + cacheWriteThrough((__gm__ uint8_t *)curHardwareTailAddr, 8); + if ((curHead - *(__gm__ uint32_t *)(curHardwareTailAddr)) < QP_DEPTH - 1) { break; } int64_t systemCycleAfter = AscendC::GetSystemCycle(); (void)systemCycleAfter; } - __gm__ uint8_t* wqeAddr = (__gm__ uint8_t*)(sqBaseAddr + wqeSize * (curHead % QP_DEPTH)); + __gm__ uint8_t *wqeAddr = (__gm__ uint8_t *)(sqBaseAddr + wqeSize * (curHead % QP_DEPTH)); uint64_t ownBit = (curHead >> shift) & 1U; uint32_t byte_4 = 3U; byte_4 |= ((~ownBit) << 7U) & (1U << 7U); byte_4 |= 1U << 8U; - *(__gm__ uint32_t*)(wqeAddr) = byte_4; - *(__gm__ uint32_t*)(wqeAddr + 4) = messageLen; - *(__gm__ uint32_t*)(wqeAddr + 8) = 0; - *(__gm__ uint32_t*)(wqeAddr + 12) = 1U << 24U; - *(__gm__ uint32_t*)(wqeAddr + 16) = 0; - __gm__ HcclAiRMAMemInfo* memDetail = - (__gm__ HcclAiRMAMemInfo*)(mem_info_table + sizeof_memdetail * destRankId); - *(__gm__ uint32_t*)(wqeAddr + 20) = ((__gm__ MemDetails*)(memDetail->memDetailPtr + - memDetail->sizeOfMemDetails * static_cast(HcclAiRMAMemType::REMOTE_INPUT)))->key; - *(__gm__ uint64_t*)(wqeAddr + 24) = (uint64_t)(destDmaAddr); + *(__gm__ uint32_t *)(wqeAddr) = byte_4; + *(__gm__ uint32_t *)(wqeAddr + 4) = messageLen; + *(__gm__ uint32_t *)(wqeAddr + 8) = 0; + *(__gm__ uint32_t *)(wqeAddr + 12) = 1U << 24U; + *(__gm__ uint32_t *)(wqeAddr + 16) = 0; + __gm__ HcclAiRMAMemInfo *memDetail = (__gm__ HcclAiRMAMemInfo *)(mem_info_table + sizeof_memdetail * destRankId); + *(__gm__ uint32_t *)(wqeAddr + 20) = + ((__gm__ MemDetails *)(memDetail->memDetailPtr + + memDetail->sizeOfMemDetails * static_cast(HcclAiRMAMemType::REMOTE_INPUT))) + ->key; + *(__gm__ uint64_t *)(wqeAddr + 24) = (uint64_t)(destDmaAddr); // Setup SGE and write to GM - __gm__ uint8_t* sgeAddr = wqeAddr + sizeof(struct hns_roce_rc_sq_wqe); - *(__gm__ uint32_t*)(sgeAddr) = messageLen; - memDetail = (__gm__ HcclAiRMAMemInfo*)(mem_info_table + sizeof_memdetail * destRankId); - *(__gm__ uint32_t*)(sgeAddr + sizeof(uint32_t)) = ((__gm__ MemDetails*)(memDetail->memDetailPtr + - memDetail->sizeOfMemDetails * static_cast(HcclAiRMAMemType::LOCAL_OUTPUT)))->key; - *(__gm__ uint64_t*)(sgeAddr + 2 * sizeof(uint32_t)) = (uint64_t)(srcDmaAddr); + __gm__ uint8_t *sgeAddr = wqeAddr + sizeof(struct hns_roce_rc_sq_wqe); + *(__gm__ uint32_t *)(sgeAddr) = messageLen; + memDetail = (__gm__ HcclAiRMAMemInfo *)(mem_info_table + sizeof_memdetail * destRankId); + *(__gm__ uint32_t *)(sgeAddr + sizeof(uint32_t)) = + ((__gm__ MemDetails *)(memDetail->memDetailPtr + + memDetail->sizeOfMemDetails * static_cast(HcclAiRMAMemType::LOCAL_OUTPUT))) + ->key; + *(__gm__ uint64_t *)(sgeAddr + 2 * sizeof(uint32_t)) = (uint64_t)(srcDmaAddr); cacheWriteThrough(wqeAddr, sizeof(struct hns_roce_rc_sq_wqe) + sizeof(struct hns_roce_lite_wqe_data_seg)); AscendC::PipeBarrier(); @@ -169,7 +168,7 @@ FORCE_INLINE_AICORE void AIVRDMAPostSend( doorBellInfo |= (curHead % 65536UL) << 32UL; doorBellInfo |= static_cast(qp_ctx_entry->sl) << 48UL; - __gm__ uint64_t* doorBellAddr = (__gm__ uint64_t*)(qp_ctx_entry->dbAddr); + __gm__ uint64_t *doorBellAddr = (__gm__ uint64_t *)(qp_ctx_entry->dbAddr); AscendC::PipeBarrier(); ubLocal.SetValue(0, doorBellInfo); @@ -182,7 +181,7 @@ FORCE_INLINE_AICORE void AIVRDMAPostSend( ubLocalHead.SetValue(0, static_cast(curHead)); AscendC::GlobalTensor HeadGlobalTensor; - HeadGlobalTensor.SetGlobalBuffer((__gm__ uint32_t*)curHardwareHead); + HeadGlobalTensor.SetGlobalBuffer((__gm__ uint32_t *)curHardwareHead); AscendC::DataCopyExtParams copyParamsHead{1, 1 * sizeof(uint32_t), 0, 0, 0}; AscendC::PipeBarrier(); AscendC::DataCopyPad(HeadGlobalTensor, ubLocalHead, copyParamsHead); @@ -192,8 +191,7 @@ FORCE_INLINE_AICORE void AIVRDMAPostSend( template class HcclShmem { public: - std::conditional_t WinContext_{nullptr}; + std::conditional_t WinContext_{nullptr}; AscendC::LocalTensor ub; FORCE_INLINE_AICORE HcclShmem() @@ -286,7 +284,8 @@ public: } else { return (GM_ADDR)((rankId == m_rank) ? WinContext_->localWindowsIn : - ((HcclRankRelationResV2 *)(WinContext_->remoteRes[rankId].nextDevicePtr))->windowsIn) + offset; + ((HcclRankRelationResV2 *)(WinContext_->remoteRes[rankId].nextDevicePtr))->windowsIn) + + offset; } } @@ -404,12 +403,13 @@ public: // i.e. needs at least 3 * UB_ALIGN = 96 bytes of UB and must be free for the // duration of this call (matches how CrossRankSyncV2Set lays out its scratch). FORCE_INLINE_AICORE - void CrossRankSync(AscendC::LocalTensor ctrBuffer) { - uint64_t flag_offset_i64 = ((m_segmentSize - m_tailReservedSize)) / sizeof(int64_t); + void CrossRankSync(AscendC::LocalTensor ctrBuffer) + { + uint64_t flag_offset_i64 = ((m_segmentSize - m_tailReservedSize)) / sizeof(int64_t); uint64_t flag_offset_bytes = ((m_segmentSize - m_tailReservedSize)); - __gm__ int64_t* sync_counter = (__gm__ int64_t*)(*this)() + flag_offset_i64; - __gm__ int64_t* sync_base = (__gm__ int64_t*)(*this)() + flag_offset_i64 + m_rankSize * 8; - gm_dcci((__gm__ uint8_t*)sync_base); + __gm__ int64_t *sync_counter = (__gm__ int64_t *)(*this)() + flag_offset_i64; + __gm__ int64_t *sync_base = (__gm__ int64_t *)(*this)() + flag_offset_i64 + m_rankSize * 8; + gm_dcci((__gm__ uint8_t *)sync_base); int64_t count = gm_load(sync_base) + 1; int vec_id = AscendC::GetBlockIdx(); int vec_size = AscendC::GetBlockNum() * AscendC::GetTaskRation(); @@ -418,11 +418,10 @@ public: AscendC::LocalTensor rdmaUbLocalHead; AscendC::GlobalTensor gmCrossServerPayload; if constexpr (IS_A2) { - rdmaUbLocal = ctrBuffer[UB_ALIGN / sizeof(int32_t)].template ReinterpretCast(); + rdmaUbLocal = ctrBuffer[UB_ALIGN / sizeof(int32_t)].template ReinterpretCast(); rdmaUbLocalHead = ctrBuffer[2 * UB_ALIGN / sizeof(int32_t)].template ReinterpretCast(); - gmCrossServerPayload.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>( - windowsOutAddr() + flag_offset_bytes - + static_cast(m_rank) * 8 * sizeof(int64_t))); + gmCrossServerPayload.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>( + windowsOutAddr() + flag_offset_bytes + static_cast(m_rank) * 8 * sizeof(int64_t))); // Only this core's stride may include cross-server targets; skip UB→GM staging if none. bool needRdmaStage = false; int32_t const localServerId = m_rank / SERVER_RANK_SIZE_A2; @@ -450,25 +449,20 @@ public: // may already have arrived while our own cross-server WQEs are still // sitting unposted, and the producer side never makes forward progress). for (int i = vec_id; i < m_rankSize; i += vec_size) { - __gm__ int64_t* sync_remote = - (__gm__ int64_t*)((*this)(i)) + flag_offset_i64 + m_rank * 8; + __gm__ int64_t *sync_remote = (__gm__ int64_t *)((*this)(i)) + flag_offset_i64 + m_rank * 8; if constexpr (IS_A2) { int32_t localServerId = m_rank / SERVER_RANK_SIZE_A2; - int32_t dstServerId = i / SERVER_RANK_SIZE_A2; + int32_t dstServerId = i / SERVER_RANK_SIZE_A2; if (dstServerId != localServerId) { - AIVRDMAPostSend( - (GM_ADDR)gmCrossServerPayload.GetPhyAddr(), - (GM_ADDR)sync_remote, - static_cast(i), - 4 * sizeof(int64_t), - reinterpret_cast<__gm__ HcclAiRMAInfo*>(WinContext_->aiRMAInfo), - rdmaUbLocal, - rdmaUbLocalHead); + AIVRDMAPostSend((GM_ADDR)gmCrossServerPayload.GetPhyAddr(), (GM_ADDR)sync_remote, + static_cast(i), 4 * sizeof(int64_t), + reinterpret_cast<__gm__ HcclAiRMAInfo *>(WinContext_->aiRMAInfo), rdmaUbLocal, + rdmaUbLocalHead); continue; } } gm_store(sync_remote, count); - gm_dcci((__gm__ uint8_t*)sync_remote); + gm_dcci((__gm__ uint8_t *)sync_remote); } // Phase 2: wait for the same set of source ranks to land their count in @@ -482,7 +476,7 @@ public: AscendC::SyncAll(); gm_store(sync_base, count); - gm_dcci((__gm__ uint8_t*)sync_base); + gm_dcci((__gm__ uint8_t *)sync_base); } FORCE_INLINE_AICORE @@ -538,24 +532,20 @@ public: } for (uint32_t dstEpIdx = vec_id; dstEpIdx < m_rankSize; dstEpIdx += vec_size) { AscendC::GlobalTensor gmDstStates; - gmDstStates.SetGlobalBuffer((__gm__ int32_t*)((*this)(flag_offset, dstEpIdx))); + gmDstStates.SetGlobalBuffer((__gm__ int32_t *)((*this)(flag_offset, dstEpIdx))); if constexpr (IS_A2) { int32_t localServerId = m_rank / SERVER_RANK_SIZE_A2; int32_t dstServerId = dstEpIdx / SERVER_RANK_SIZE_A2; if (dstServerId != localServerId) { AscendC::GlobalTensor gmLocalOutputState; gmLocalOutputState.SetGlobalBuffer( - reinterpret_cast<__gm__ int32_t*>(windowsOutAddr() + flag_offset)); + reinterpret_cast<__gm__ int32_t *>(windowsOutAddr() + flag_offset)); AscendC::DataCopy(gmLocalOutputState, ctrBuffer, 8); AscendC::PipeBarrier(); - AIVRDMAPostSend( - (GM_ADDR)gmLocalOutputState.GetPhyAddr(), - (GM_ADDR)gmDstStates.GetPhyAddr(), - static_cast(dstEpIdx), - 8 * sizeof(int32_t), - reinterpret_cast<__gm__ HcclAiRMAInfo*>(WinContext_->aiRMAInfo), - rdmaUbLocal, - rdmaUbLocalHead); + AIVRDMAPostSend((GM_ADDR)gmLocalOutputState.GetPhyAddr(), (GM_ADDR)gmDstStates.GetPhyAddr(), + static_cast(dstEpIdx), 8 * sizeof(int32_t), + reinterpret_cast<__gm__ HcclAiRMAInfo *>(WinContext_->aiRMAInfo), rdmaUbLocal, + rdmaUbLocalHead); continue; } } @@ -606,8 +596,7 @@ public: while ((sumOfFlag < minTarget) || (sumOfFlag > maxTarget)) { AscendC::DataCopy( - statusTensor, epStatusSpaceGlobalTensor_[startRankId_ * stateOffset_ / sizeof(float)], - intriParams); + statusTensor, epStatusSpaceGlobalTensor_[startRankId_ * stateOffset_ / sizeof(float)], intriParams); AscendC::SetFlag(EVENT_ID0); AscendC::WaitFlag(EVENT_ID0); @@ -651,38 +640,28 @@ public: // 模板参数 ExceptionDumpT 需提供： // void Dump(GM_ADDR dumpAddr, size_t blockCount, size_t blockLen, size_t srcStride) template - FORCE_INLINE_AICORE - void DumpSyncRegions(ExceptionDumpT &exceptionDump) + FORCE_INLINE_AICORE void DumpSyncRegions(ExceptionDumpT &exceptionDump) { uint64_t flagOffsetBytes = m_segmentSize - m_tailReservedSize; GM_ADDR base = (*this)() + flagOffsetBytes; // Region 0: CrossRankSync sync_counter[rank]，数据类型 int64_t // 访问模式：[rank*64B, rank*64B+8B)，共 m_rankSize 个 block，stride=64B - exceptionDump.Dump(base, - static_cast(m_rankSize), - sizeof(int64_t), - 64U); + exceptionDump.Dump(base, static_cast(m_rankSize), sizeof(int64_t), 64U); // Region 1: CrossRankSync sync_base，单 int64_t // 地址 = base + m_rankSize * 64（sync_counter 之后紧接着 sync_base） - exceptionDump.Dump(base + static_cast(m_rankSize) * 64U, - sizeof(int64_t)); + exceptionDump.Dump(base + static_cast(m_rankSize) * 64U, sizeof(int64_t)); // Region 2: CrossRankSyncV2 per-rank status，每 rank 32B（V2Set 写 8 int32_t） // 访问模式：[rank*512B, rank*512B+32B)，共 m_rankSize 个 block，stride=STATE_OFFSET - exceptionDump.Dump(base, - static_cast(m_rankSize), - UB_ALIGN, - static_cast(STATE_OFFSET)); + exceptionDump.Dump(base, static_cast(m_rankSize), UB_ALIGN, static_cast(STATE_OFFSET)); // Region 3: CrossRankSyncV2 selfStatus，本 rank 各 AIV 核的 init 状态 // InitStatusTargetSum 用 GetBlockIdx() 索引，范围 [0, get_block_num()) // selfStatusTensor[coreIdx * UB_ALIGN] 中 UB_ALIGN 是 int32_t 元素索引， // 实际字节间距 = UB_ALIGN * sizeof(int32_t) = 128B - exceptionDump.Dump(base + SELF_STATE_OFFSET, - static_cast(GetBlockNum()), - sizeof(int32_t), + exceptionDump.Dump(base + SELF_STATE_OFFSET, static_cast(GetBlockNum()), sizeof(int32_t), static_cast(UB_ALIGN * sizeof(int32_t))); } @@ -691,7 +670,7 @@ private: int32_t m_rank; int32_t m_rankSize; size_t m_segmentSize; - size_t m_tailReservedSize{0}; // 尾部 CrossRankSync 保留空间 + size_t m_tailReservedSize{0}; // 尾部 CrossRankSync 保留空间 float sumTarget_{0.0}; int32_t epStateValue_; }; diff --git a/torch_extension/cann_ops_transformer/ops/csrc/mega_moe.cpp b/torch_extension/cann_ops_transformer/ops/csrc/mega_moe.cpp index 1b1b1d3..ed9790c 100644 --- a/torch_extension/cann_ops_transformer/ops/csrc/mega_moe.cpp +++ b/torch_extension/cann_ops_transformer/ops/csrc/mega_moe.cpp @@ -16,31 +16,22 @@ namespace op_api { using npu_utils = at_npu::native::NpuUtils; const int DIM_TWO = 2; -std::tuple npu_mega_moe( - const at::Tensor &context, const at::Tensor &x, const at::Tensor &topk_ids, - const at::Tensor &topk_weights, const std::vector &weight1, - const std::vector &weight2, int64_t moe_expert_num, int64_t ep_world_size, - int64_t ccl_buffer_size, - const c10::optional> &weight_scales1, - const c10::optional> &weight_scales2, - const c10::optional> &bias1, - const c10::optional> &bias2, - const c10::optional &x_active_mask, - int64_t max_recv_token_num, - int64_t dispatch_quant_mode, - int64_t combine_quant_mode, - std::string comm_alg, - int64_t num_max_tokens_per_rank, - std::string activation, - c10::optional activation_clamp, - c10::optional dispatch_quant_out_dtype, - c10::optional weight1_type, - c10::optional weight2_type) +std::tuple +npu_mega_moe(const at::Tensor &context, const at::Tensor &x, const at::Tensor &topk_ids, const at::Tensor &topk_weights, + const std::vector &weight1, const std::vector &weight2, int64_t moe_expert_num, + int64_t ep_world_size, int64_t ccl_buffer_size, + const c10::optional> &weight_scales1, + const c10::optional> &weight_scales2, + const c10::optional> &bias1, const c10::optional> &bias2, + const c10::optional &x_active_mask, int64_t max_recv_token_num, int64_t dispatch_quant_mode, + int64_t combine_quant_mode, std::string comm_alg, int64_t num_max_tokens_per_rank, std::string activation, + c10::optional activation_clamp, c10::optional dispatch_quant_out_dtype, + c10::optional weight1_type, c10::optional weight2_type) { TORCH_CHECK((ep_world_size > 0), "The ep_world_sizes should be greater than 0, current is: ", ep_world_size); TORCH_CHECK((x.dim() == DIM_TWO) && (topk_ids.dim() == DIM_TWO), "The x and topk_ids should be 2D"); TORCH_CHECK(((x.scalar_type() == at::kBFloat16) || (x.scalar_type() == at::kHalf)) && - (topk_ids.scalar_type() == at::kInt), + (topk_ids.scalar_type() == at::kInt), "dtype of x should be bfloat16, float16, dtype of topk_ids should be int."); at::TensorList weight1_ref = weight1; @@ -71,28 +62,26 @@ std::tuple npu_mega_moe( bias2_ref = at::TensorList(); } - aclDataType weight1_ref_dtype = weight1_type.has_value() ? GetAclDataType(weight1_type.value()) - : ConvertToAclDataType(weight1_ref[0].scalar_type()); + aclDataType weight1_ref_dtype = weight1_type.has_value() ? GetAclDataType(weight1_type.value()) : + ConvertToAclDataType(weight1_ref[0].scalar_type()); aclDataType weight_scales1_dtype; - if (weight1_ref_dtype == aclDataType::ACL_FLOAT8_E5M2 || - weight1_ref_dtype == aclDataType::ACL_FLOAT8_E4M3FN || + if (weight1_ref_dtype == aclDataType::ACL_FLOAT8_E5M2 || weight1_ref_dtype == aclDataType::ACL_FLOAT8_E4M3FN || weight1_ref_dtype == aclDataType::ACL_FLOAT4_E2M1) { weight_scales1_dtype = aclDataType::ACL_FLOAT8_E8M0; } else { weight_scales1_dtype = aclDataType::ACL_UINT64; } - aclDataType weight2_ref_dtype = weight2_type.has_value() ? GetAclDataType(weight2_type.value()) - : ConvertToAclDataType(weight2_ref[0].scalar_type()); + aclDataType weight2_ref_dtype = weight2_type.has_value() ? GetAclDataType(weight2_type.value()) : + ConvertToAclDataType(weight2_ref[0].scalar_type()); aclDataType weight_scales2_dtype; - if (weight2_ref_dtype == aclDataType::ACL_FLOAT8_E5M2 || - weight2_ref_dtype == aclDataType::ACL_FLOAT8_E4M3FN || + if (weight2_ref_dtype == aclDataType::ACL_FLOAT8_E5M2 || weight2_ref_dtype == aclDataType::ACL_FLOAT8_E4M3FN || weight2_ref_dtype == aclDataType::ACL_FLOAT4_E2M1) { weight_scales2_dtype = aclDataType::ACL_FLOAT8_E8M0; } else { weight_scales2_dtype = aclDataType::ACL_UINT64; } - + auto x_size = x.sizes(); auto topk_ids_size = topk_ids.sizes(); int64_t bs = x_size[0]; @@ -101,9 +90,8 @@ std::tuple npu_mega_moe( if ((dispatch_quant_out_dtype.has_value()) && (dispatch_quant_out_dtype.value() == static_cast(DType::FLOAT4_E2M1))) { - TORCH_CHECK(h % 2 == 0, - "The last dim input shape must be divisible by 2 if " - "dispatch quant output type is torch_npu.float4_e2m1"); + TORCH_CHECK(h % 2 == 0, "The last dim input shape must be divisible by 2 if " + "dispatch quant output type is torch_npu.float4_e2m1"); } int64_t local_moe_expert_num = 1; @@ -119,9 +107,9 @@ std::tuple npu_mega_moe( float activation_clamp_value = activation_clamp.value_or(std::numeric_limits::max()); - int64_t dispatch_quant_result_type = dispatch_quant_out_dtype.has_value() - ? static_cast(GetAclDataType(dispatch_quant_out_dtype.value())) - : 28; + int64_t dispatch_quant_result_type = dispatch_quant_out_dtype.has_value() ? + static_cast(GetAclDataType(dispatch_quant_out_dtype.value())) : + 28; at::Tensor y; y = at::empty({bs, h}, topk_ids.options().dtype(x.scalar_type())); @@ -133,12 +121,11 @@ std::tuple npu_mega_moe( TensorListWrapper bias1_wrapper = {bias1_ref, aclDataType::ACL_FLOAT}; TensorListWrapper bias2_wrapper = {bias2_ref, aclDataType::ACL_FLOAT}; - ACLNN_CMD(aclnnMegaMoe, context, x, topk_ids, topk_weights, weight1_wrapper, - weight2_wrapper, weight_scales1_wrapper, weight_scales2_wrapper, bias1_wrapper, bias2_wrapper, - x_active_mask, moe_expert_num, ep_world_size, ccl_buffer_size, max_recv_token_num, - dispatch_quant_mode, dispatch_quant_result_type, combine_quant_mode, - comm_alg_ptr, num_max_tokens_per_rank, activation_ptr, activation_clamp_value, - y, expert_token_nums); + ACLNN_CMD(aclnnMegaMoe, context, x, topk_ids, topk_weights, weight1_wrapper, weight2_wrapper, + weight_scales1_wrapper, weight_scales2_wrapper, bias1_wrapper, bias2_wrapper, x_active_mask, + moe_expert_num, ep_world_size, ccl_buffer_size, max_recv_token_num, dispatch_quant_mode, + dispatch_quant_result_type, combine_quant_mode, comm_alg_ptr, num_max_tokens_per_rank, activation_ptr, + activation_clamp_value, y, expert_token_nums); return std::tie(y, expert_token_nums); } @@ -158,25 +145,23 @@ int64_t CeilAlign(int64_t val, int64_t align) // A2 minimum buffer size (MB). // Matches tiling_arch22.cpp CalcLeastCclBufferSize with isA3=false. -int64_t CalcLeastCclBufferSizeA2(int64_t maxRecvTokenNum, int64_t h, - int64_t epWorldSize, bool isQuantRouting, int64_t bs, int64_t topK) +int64_t CalcLeastCclBufferSizeA2(int64_t maxRecvTokenNum, int64_t h, int64_t epWorldSize, bool isQuantRouting, + int64_t bs, int64_t topK) { // Data block 1: TokenPerExpert // EP × CeilAlign(EP × MAX_EXPERTS_PER_RANK_A2A3 + 1, 128) × 4B - int64_t offsetTokenPerExpert = epWorldSize * - CeilAlign(epWorldSize * MAX_EXPERTS_PER_RANK_A2A3 + 1, ALIGN_128) * - static_cast(sizeof(int32_t)); + int64_t offsetTokenPerExpert = epWorldSize * CeilAlign(epWorldSize * MAX_EXPERTS_PER_RANK_A2A3 + 1, ALIGN_128) * + static_cast(sizeof(int32_t)); // Data block 2: tensors // ===== winIn ===== - int64_t offsetAAfterDispatch = maxRecvTokenNum * - (isQuantRouting ? (h + ALIGN_512) : h * static_cast(sizeof(int16_t))); + int64_t offsetAAfterDispatch = + maxRecvTokenNum * (isQuantRouting ? (h + ALIGN_512) : h * static_cast(sizeof(int16_t))); int64_t offsetD = bs * topK * h * static_cast(sizeof(int16_t)); int64_t winInTensorSize = offsetAAfterDispatch + offsetD; // ===== winOut ===== - int64_t offsetA = bs * topK * - (!isQuantRouting ? h * static_cast(sizeof(int16_t)) : (h + ALIGN_512)); + int64_t offsetA = bs * topK * (!isQuantRouting ? h * static_cast(sizeof(int16_t)) : (h + ALIGN_512)); int64_t offsetC = maxRecvTokenNum * h * static_cast(sizeof(int16_t)); int64_t winOutTensorSize = offsetA + offsetC; int64_t offsetTensor = std::max(winInTensorSize, winOutTensorSize); @@ -185,27 +170,25 @@ int64_t CalcLeastCclBufferSizeA2(int64_t maxRecvTokenNum, int64_t h, } // Data block 3: sync flags - int64_t offsetFlag = epWorldSize * ALIGN_512; // CrossRankSync - offsetFlag += epWorldSize * MAX_EXPERTS_PER_RANK_A2A3 * 64LL; // DispatchFlag - offsetFlag += epWorldSize * 64LL; // AllGatherFlag + int64_t offsetFlag = epWorldSize * ALIGN_512; // CrossRankSync + offsetFlag += epWorldSize * MAX_EXPERTS_PER_RANK_A2A3 * 64LL; // DispatchFlag + offsetFlag += epWorldSize * 64LL; // AllGatherFlag return (offsetTokenPerExpert + offsetTensor + offsetFlag + RESERVED_SPACE_SIZE + MB_SIZE) / MB_SIZE; } // A3 minimum buffer size (MB). // Matches tiling_arch22.cpp CalcLeastCclBufferSize with isA3=true. -int64_t CalcLeastCclBufferSizeA3(int64_t h, - int64_t epWorldSize, bool isQuantRouting, int64_t bs, int64_t topK) +int64_t CalcLeastCclBufferSizeA3(int64_t h, int64_t epWorldSize, bool isQuantRouting, int64_t bs, int64_t topK) { // Data block 1: TokenPerExpert // EP × CeilAlign(EP × MAX_EXPERTS_PER_RANK_A2A3 + 1, 128) × 4B - int64_t offsetTokenPerExpert = epWorldSize * - CeilAlign(epWorldSize * MAX_EXPERTS_PER_RANK_A2A3 + 1, ALIGN_128) * - static_cast(sizeof(int32_t)); + int64_t offsetTokenPerExpert = epWorldSize * CeilAlign(epWorldSize * MAX_EXPERTS_PER_RANK_A2A3 + 1, ALIGN_128) * + static_cast(sizeof(int32_t)); // Data block 2: tensors (winIn only, no winOut) - int64_t offsetAAfterDispatch = bs * topK * - (isQuantRouting ? (h + ALIGN_512) : h * static_cast(sizeof(int16_t))); + int64_t offsetAAfterDispatch = + bs * topK * (isQuantRouting ? (h + ALIGN_512) : h * static_cast(sizeof(int16_t))); int64_t offsetD = bs * topK * h * static_cast(sizeof(int16_t)); int64_t offsetTensor = offsetAAfterDispatch + offsetD; if (isQuantRouting) { @@ -219,8 +202,8 @@ int64_t CalcLeastCclBufferSizeA3(int64_t h, } // A5 half-buffer minimum size (MB). Ported 1:1 from the original Python implementation. -int64_t CalcHalfBufferSizeMBA5(int64_t epWorldSize, int64_t moeExpertNum, - int64_t numMaxTokensPerRank, int64_t numTopk, int64_t hidden) +int64_t CalcHalfBufferSizeMBA5(int64_t epWorldSize, int64_t moeExpertNum, int64_t numMaxTokensPerRank, int64_t numTopk, + int64_t hidden) { int64_t expertPerRank = moeExpertNum / epWorldSize; @@ -248,29 +231,27 @@ int64_t CalcHalfBufferSizeMBA5(int64_t epWorldSize, int64_t moeExpertNum, } } // namespace -int64_t GetMegaMoeCclBufferSize(int64_t epWorldSize, int64_t moeExpertNum, - int64_t numMaxTokensPerRank, int64_t numTopk, int64_t hidden, - int64_t maxRecvTokenNum, - int64_t dispatchQuantMode, c10::optional dispatchQuantOutDtype, - int64_t combineQuantMode, std::string commAlg) +int64_t GetMegaMoeCclBufferSize(int64_t epWorldSize, int64_t moeExpertNum, int64_t numMaxTokensPerRank, int64_t numTopk, + int64_t hidden, int64_t maxRecvTokenNum, int64_t dispatchQuantMode, + c10::optional dispatchQuantOutDtype, int64_t combineQuantMode, + std::string commAlg) { const char *socName = aclrtGetSocName(); bool isA2 = (socName != nullptr && std::strstr(socName, "Ascend910B") != nullptr); bool isA3 = (socName != nullptr && std::strstr(socName, "Ascend910_93") != nullptr); if (isA2 || isA3) { - TORCH_CHECK(epWorldSize == 2 || epWorldSize == 4 || epWorldSize == 8 || - epWorldSize == 16 || epWorldSize == 32 || epWorldSize == 64, - "ep_world_size only support {2, 4, 8, 16, 32, 64} on A2/A3, but got ", epWorldSize); + TORCH_CHECK(epWorldSize == 2 || epWorldSize == 4 || epWorldSize == 8 || epWorldSize == 16 || + epWorldSize == 32 || epWorldSize == 64, + "ep_world_size only support {2, 4, 8, 16, 32, 64} on A2/A3, but got ", epWorldSize); TORCH_CHECK(hidden >= 1024 && hidden <= 8192 && hidden % 512 == 0, - "hidden only support [1024, 8192] and hidden % 512 == 0 on A2/A3, but got ", hidden); + "hidden only support [1024, 8192] and hidden % 512 == 0 on A2/A3, but got ", hidden); TORCH_CHECK(numMaxTokensPerRank >= 1 && numMaxTokensPerRank <= 4096, - "num_max_tokens_per_rank only support [1, 4096] on A2/A3, but got ", numMaxTokensPerRank); + "num_max_tokens_per_rank only support [1, 4096] on A2/A3, but got ", numMaxTokensPerRank); TORCH_CHECK(moeExpertNum >= 1 && moeExpertNum <= 2048, - "moe_expert_num only support [1, 2048] on A2/A3, but got ", moeExpertNum); - TORCH_CHECK(numTopk >= 1 && numTopk <= 16, - "num_topk only support [1, 16] on A2/A3, but got ", numTopk); + "moe_expert_num only support [1, 2048] on A2/A3, but got ", moeExpertNum); + TORCH_CHECK(numTopk >= 1 && numTopk <= 16, "num_topk only support [1, 16] on A2/A3, but got ", numTopk); TORCH_CHECK(dispatchQuantMode == 0 || dispatchQuantMode == 2 || dispatchQuantMode == 4, - "dispatch_quant_mode only support {0, 2, 4} on A2/A3, but got ", dispatchQuantMode); + "dispatch_quant_mode only support {0, 2, 4} on A2/A3, but got ", dispatchQuantMode); bool isQuantRouting = (dispatchQuantMode == 4); // max_recv_token_num 为 0 时自动计算为 bs * epWorldSize * min(topK, expertPerRank)， @@ -279,32 +260,27 @@ int64_t GetMegaMoeCclBufferSize(int64_t epWorldSize, int64_t moeExpertNum, maxRecvTokenNum = numMaxTokensPerRank * epWorldSize * std::min(numTopk, expertPerRank); } if (isA3) { - return CalcLeastCclBufferSizeA3( - hidden, epWorldSize, isQuantRouting, - numMaxTokensPerRank, numTopk); + return CalcLeastCclBufferSizeA3(hidden, epWorldSize, isQuantRouting, numMaxTokensPerRank, numTopk); } - return CalcLeastCclBufferSizeA2( - maxRecvTokenNum, hidden, epWorldSize, isQuantRouting, - numMaxTokensPerRank, numTopk); + return CalcLeastCclBufferSizeA2(maxRecvTokenNum, hidden, epWorldSize, isQuantRouting, numMaxTokensPerRank, + numTopk); } // A5 / 950 - TORCH_CHECK(epWorldSize >= 2 && epWorldSize <= 768, - "ep_world_size only support in [2, 768], but got ", epWorldSize); - TORCH_CHECK(hidden >= 1024 && hidden <= 8192, - "hidden only support in [1024, 8192], but got ", hidden); + TORCH_CHECK(epWorldSize >= 2 && epWorldSize <= 768, "ep_world_size only support in [2, 768], but got ", + epWorldSize); + TORCH_CHECK(hidden >= 1024 && hidden <= 8192, "hidden only support in [1024, 8192], but got ", hidden); TORCH_CHECK(numMaxTokensPerRank >= 1 && numMaxTokensPerRank <= 512, - "num_max_tokens_per_rank only support in [1, 512], but got ", numMaxTokensPerRank); - TORCH_CHECK(moeExpertNum >= 1 && moeExpertNum <= 1024, - "moe_expert_num only support in [1, 1024], but got ", moeExpertNum); - TORCH_CHECK(numTopk >= 1 && numTopk <= 16, - "num_topk only support in [1, 16], but got ", numTopk); - - return CalcHalfBufferSizeMBA5(epWorldSize, moeExpertNum, - numMaxTokensPerRank, numTopk, hidden); + "num_max_tokens_per_rank only support in [1, 512], but got ", numMaxTokensPerRank); + TORCH_CHECK(moeExpertNum >= 1 && moeExpertNum <= 1024, "moe_expert_num only support in [1, 1024], but got ", + moeExpertNum); + TORCH_CHECK(numTopk >= 1 && numTopk <= 16, "num_topk only support in [1, 16], but got ", numTopk); + + return CalcHalfBufferSizeMBA5(epWorldSize, moeExpertNum, numMaxTokensPerRank, numTopk, hidden); } -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ m.def("npu_mega_moe", &npu_mega_moe, "npu_mega_moe"); m.def("get_mega_moe_ccl_buffer_size", &GetMegaMoeCclBufferSize, "get_mega_moe_ccl_buffer_size"); }