昇腾社区首页
中文
注册
开发者
下载
支持

aclnnGroupedMatmulFinalizeRouting

支持的产品型号

  • Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件
  • Atlas A3 训练系列产品/Atlas A3 推理系列产品

功能说明

算子功能：GroupedMatmul和MoeFinalizeRouting的融合算子，GroupedMatmul计算后的输出按照索引做combine动作

函数原型

每个算子分为两段式接口，必须先调用“aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnGroupedMatmulFinalizeRouting”接口执行计算。

  • aclnnStatus aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize(const aclTensor *x, aclTensor *w, const aclTensor *scaleOptional, const aclTensor* biasOptional, const aclTensor *pertokenScaleOptional, const aclTensor *groupListOptional, const aclTensor *sharedInputOptional, const aclTensor* logitOptional, const aclTensor *rowIndexOptional, int64_t dtype, float sharedInputWeight, int64_t sharedInputOffset, bool transposeX, bool transposeW, int64_t groupListType, aclTensor *y, uint64_t *workspaceSize, aclOpExecutor **executor)
  • aclnnStatus aclnnGroupedMatmulFinalizeRouting(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)

aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize

  • 参数说明：

    • x（aclTensor*，计算输入）：Device侧的aclTensor，输入x(左矩阵)，数据格式支持ND，数据类型支持INT8，shape支持2维，维度为(m, k)，维度m的取值范围为[1,16*1024*8]，k支持2048。
    • w（aclTensor*，计算输入）：Device侧的aclTensor，输入weight(右矩阵)，数据格式支持ND，数据类型支持INT32（传入数据为INT32，实际按照8个INT4来处理），shape支持三维，当输入为INT32时维度为(e, k, n / 8)，输入转为INT4时维度为(e, k, n)，e取值范围[1,256]，k支持2048，n支持7168。
    • scaleOptional（aclTensor*，可选计算输入）可选参数，Device侧的aclTensor，量化参数中的缩放因子，数据格式支持ND，数据类型支持INT64，shape支持三维，维度为(e, 1, n)，e、n和w的e、n一致。
    • biasOptional（aclTensor*，可选计算输入）可选参数，Device侧的aclTensor，矩阵的偏移，数据格式支持ND，数据类型支持FLOAT32，shape支持二维，维度为(e, n)，e、n和w的e、n一致。
    • pertokenScaleOptional（aclTensor*，可选计算输入）可选参数，Device侧的aclTensor，矩阵计算的反量化参数，数据格式支持ND，数据类型支持FLOAT32，shape支持一维，维度为(m)，m和x的m一致。
    • groupListOptional（aclTensor*，可选计算输入）：可选参数，Device侧的aclTensor类型，代表输入和输出分组轴方向的matmul大小分布，数据格式支持ND，数据类型支持INT64，shape支持一维，维度为(e)，e和w的e一致。
    • sharedInputOptional（aclTensor*，可选计算输入）：可选参数，Device侧的aclTensor类型，moe计算中共享专家的输出，需要与moe专家的输出进行combine操作数据格式支持ND，数据类型支持BFLOAT16，shape支持二维，维度(batch/dp,n)，batch/dp取值范围[1,2*1024]，batch取值范围[1,16*1024]。
    • logitOptional（aclTensor*，可选计算输入）：可选参数，Device侧的aclTensor类型，moe专家对各个token的logit大小，矩阵乘的计算输出与该logit做乘法，然后索引进行combine，数据格式支持ND，数据类型支持FLOAT32，shape支持一维，维度为(m)，m和x的m一致。
    • rowIndexOptional（aclTensor*，可选计算输入）：可选参数，Device侧的aclTensor类型，moe专家输出按照该rowIndex进行combine，其中的值即为combine做scatter add的索引，数据格式支持ND，数据类型支持FLOAT32，shape支持一维，维度为(m)，m和x的m一致。
    • dtype（int64_t，计算输入）：GroupedMatmul计算的输出类型，0：FLOAT32；1：FLOAT16；2：BFLOAT16。仅支持取0。
    • sharedInputWeight（float，计算输入）：共享专家与moe专家进行combine的系数，sharedInput先与该参数乘，然后在和moe专家结果累加。
    • sharedInputOffset（int64_t，计算输入）：共享专家输出的在总输出中的偏移。
    • transposeX（bool，计算输入）：左矩阵是否转置。
    • transposeW（bool，计算输入）：右矩阵是否转置。
    • groupListType(int64_t，计算输入)：GroupedMatmul的分组模式：配置为0：cumsum模式，即为前缀和；配置为1：count模式。
    • y（aclTensor*，计算输出）：数据格式支持ND。
    • workspaceSize（uint64_t*，出参）：返回需要在Device侧申请的workspace大小。
    • executor（aclOpExecutor**，出参）：返回op执行器，包含了算子计算流程。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

    第一段接口完成入参校验，若出现以下错误码，则对应原因为：
- 返回161001（ACLNN_ERR_PARAM_NULLPTR）：
1. 传入的x、w、y、biasOptional是空指针。
- 返回161002（ACLNN_ERR_PARAM_INVALID）：
1.x、w、scaleOptional、biasOptional、pertokenScaleOptional、groupListOptional、shareInputOptional、logitOptional、rowIndexOptional、shareInputWeight、shareInputOffest、transposeX、transposeW、或y的数据类型和数据格式不在支持的范围内。
2.x、w、scaleOptional、biasOptional、pertokenScaleOptional、groupListOptional、shareInputOptional、logitOptional、rowIndexOptional或y的shape不满足校验条件。
3.x、w、scaleOptional、biasOptional、pertokenScaleOptional、groupListOptional、shareInputOptional、logitOptional、rowIndexOptional或y的shape是空tensor。
4.dtype、sharedInputOffset、transposeX、transposeW、groupListType的取值范围不满足条件

aclnnGroupedMatmulFinalizeRouting

  • 参数说明：

    • workspace（void*，入参）：在Device侧申请的workspace内存地址。
    • workspaceSize（uint64_t*，入参）：在Device侧申请的workspace大小，由第一段接口aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize获取。
    • executor（aclOpExecutor*，入参）：op执行器，包含了算子计算流程。
    • stream（aclrtStream，入参）：指定执行任务的AscendCL stream流。

  • 返回值：

    返回aclnnStatus状态码，具体参见aclnn返回码。

约束说明

伪量化场景支持类型 输入和输出支持以下数据类型组合：

x w scale bias pertokenScale groupList sharedInput logit rowIndex out
INT8 INT4 INT64 FLOAT32 FLOAT32 INT64 BFLOAT16 FLOAT32 INT64 FLOAT32
  • 在该场景中，x的数据格式支持ND，数据类型支持INT8，shape支持2维，维度为(m, k)，维度m的取值范围为[1,16*1024*8]，k支持2048。
  • 在该场景中，w的数据格式支持ND，数据类型支持INT32（传入数据为INT32，实际按照8个INT4来处理），shape支持三维，当输入为INT32时维度为(e, k, n / 8)，输入转为INT4后维度为(e, k, n)，e取值范围[1,256]，k支持2048，n支持7168。
  • 在该场景中，scaleOptional代表per-channel和per-group离线融合的结果，数据格式支持ND，数据类型支持INT64，shape支持三维，维度为(e, 1, n)，e、n和w的e、n一致，不支持非连续的Tensor。
  • 在该场景中，biasOptional代表离线计算的辅助结果，值要求为8*w*scale，并在第一维累加，数据格式支持ND，数据类型支持FLOAT32，shape支持二维，维度为(e, n)，e、n和w的e、n一致，不支持非连续的Tensor。
  • 在该场景中，pertokenScaleOptional数据格式支持ND，数据类型支持FLOAT32，shape支持一维，维度为(m)，m和x的m一致，不支持非连续的Tensor。
  • 在该场景中，groupListOptional数据格式支持ND，数据类型支持INT64，shape支持一维，维度为(e)，grouplist的值总和小于等于m，不支持非连续的Tensor。
  • 在该场景中，sharedInputOptional数据格式支持ND，数据类型支持BFLOAT16，shape支持二维，维度(batch/dp,n)，batch/dp取值范围[1,2*1024]，batch取值范围[1,16*1024]，不支持非连续的Tensor。
  • 在该场景中，logitOptional数据格式支持ND，数据类型支持FLOAT32，shape支持一维，维度为(m)，m和x的m一致，不支持非连续的Tensor。
  • 在该场景中，rowIndexOptional数据格式支持ND，数据类型支持INT64，shape支持一维，维度为(m)，m和x的m一致，不支持非连续的Tensor。
  • 在该场景中，dtype仅支持取0。
  • 在该场景中，transposeX仅支持false。
  • 在该场景中，transposeW仅支持false。
  • 在该场景中，y数据格式支持ND，输出的数据类型固定为FLOAT32，shape支持二维，维度为(batch, n)，不支持非连续的Tensor。

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考编译与运行样例。

#include <iostream>
#include <memory>
#include <vector>

#include "acl/acl.h"
#include "aclnnop/aclnn_permute.h"
#include "aclnnop/aclnn_grouped_matmul_finalize_routing.h"
#include "aclnnop/aclnn_trans_matmul_weight.h"

#define CHECK_RET(cond, return_expr) \
    do {                             \
        if (!(cond)) {               \
            return_expr;             \
        }                            \
    } while (0)

#define CHECK_FREE_RET(cond, return_expr) \
    do {                                  \
        if (!(cond)) {                    \
            Finalize(deviceId, stream);   \
            return_expr;                  \
        }                                 \
    } while (0)

#define LOG_PRINT(message, ...)         \
    do {                                \
        printf(message, ##__VA_ARGS__); \
    } while (0)

int64_t GetShapeSize(const std::vector<int64_t> &shape)
{
    int64_t shapeSize = 1;
    for (auto i : shape) {
        shapeSize *= i;
    }
    return shapeSize;
}

int Init(int32_t deviceId, aclrtStream *stream)
{
    // 固定写法，AscendCL初始化
    auto ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
    ret = aclrtSetDevice(deviceId);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
    ret = aclrtCreateStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
    return 0;
}

template <typename T>
int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
                    aclDataType dataType, aclTensor **tensor)
{
    auto size = GetShapeSize(shape) * sizeof(T);
    // 调用aclrtMalloc申请device侧内存
    auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
    ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);

    // 计算连续tensor的strides
    std::vector<int64_t> strides(shape.size(), 1);
    for (int64_t i = shape.size() - 2; i >= 0; i--) {
        strides[i] = shape[i + 1] * strides[i + 1];
    }

    // 调用aclCreateTensor接口创建aclTensor
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
    return 0;
}

template <typename T>
int CreateAclTensorWeight(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
                      aclDataType dataType, aclTensor **tensor)
{
    auto size = static_cast<uint64_t>(GetShapeSize(shape));

    const aclIntArray *mat2Size = aclCreateIntArray(shape.data(), shape.size());
    auto ret = aclnnCalculateMatmulWeightSizeV2(mat2Size, dataType, &size);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnCalculateMatmulWeightSizeV2 failed. ERROR: %d\n", ret);
              return ret);
    size *= sizeof(T);

    // 调用aclrtMalloc申请device侧内存
    ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
    // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
    ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);

    // 计算连续tensor的strides
    std::vector<int64_t> strides(shape.size(), 1);
    for (int64_t i = shape.size() - 2; i >= 0; i--) {
        strides[i] = shape[i + 1] * strides[i + 1];
    }

    std::vector<int64_t> storageShape;
    storageShape.push_back(GetShapeSize(shape));

    // 调用aclCreateTensor接口创建aclTensor
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                              storageShape.data(), storageShape.size(), *deviceAddr);
    return 0;
}

  int main() {
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

    // 2. 构造输入与输出，需要根据API的接口自定义构造
    int64_t m = 8;
    int64_t k = 2048;
    int64_t n = 7168;
    int64_t e = 1;
    int64_t batch = 8;
    int64_t bsdp = 1;
    int64_t dtype = 0;
    float shareInputWeight = 1.0;
    int64_t shareInputOffest = 0;
    bool transposeX = false;
    bool transposeW = false;
    int64_t groupListType = 1;
    
    std::vector<int64_t> xShape = {m, k};
    std::vector<int64_t> wShape = {e, k, n / 8};
    std::vector<int64_t> scaleShape = {e, 1, n};
    std::vector<int64_t> biasShape = {e, n};
    std::vector<int64_t> pertokenScaleShape = {m};
    std::vector<int64_t> groupListShape = {e};
    std::vector<int64_t> sharedInputShape = {bsdp, n};
    std::vector<int64_t> logitShape = {m};
    std::vector<int64_t> rowIndexShape = {m};
    std::vector<int64_t> outShape = {batch, n};

    void *xDeviceAddr = nullptr;
    void *wDeviceAddr = nullptr;
    void *biasDeviceAddr = nullptr;
    void *scaleDeviceAddr = nullptr;
    void *pertokenScaleDeviceAddr = nullptr;
    void *groupListDeviceAddr = nullptr;
    void *sharedInputDeviceAddr = nullptr;
    void *logitDeviceAddr = nullptr;
    void *rowIndexDeviceAddr = nullptr;
    void *outDeviceAddr = nullptr;

    aclTensor* x = nullptr;
    aclTensor* w = nullptr;
    aclTensor* bias = nullptr;
    aclTensor* groupList = nullptr;
    aclTensor* scale = nullptr;
    aclTensor* pertokenScale = nullptr;
    aclTensor* sharedInput = nullptr;
    aclTensor* logit = nullptr;
    aclTensor* rowIndex = nullptr;
    aclTensor* out = nullptr;

    std::vector<int8_t> xHostData(GetShapeSize(xShape));
    std::vector<int32_t> wHostData(GetShapeSize(wShape));
    std::vector<int64_t> scaleHostData(GetShapeSize(scaleShape));
    std::vector<float> biasHostData(GetShapeSize(biasShape));
    std::vector<float> pertokenScaleHostData(GetShapeSize(pertokenScaleShape));
    std::vector<int64_t> groupListHostData(GetShapeSize(groupListShape));
    std::vector<uint16_t> sharedInputHostData(GetShapeSize(sharedInputShape));
    std::vector<int64_t> logitHostData(GetShapeSize(logitShape));
    std::vector<float> rowIndexHostData(GetShapeSize(rowIndexShape));
    std::vector<float> outHostData(GetShapeSize(outShape));  // 实际上是float16半精度方式
    // 对groupList赋值
    groupListHostData[0] = 8;
    // 创建x aclTensor
    ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_INT8, &x);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // 创建int32_t 的w aclTensor，后续转为int_4
    ret = CreateAclTensorWeight(wHostData, wShape, &wDeviceAddr, aclDataType::ACL_INT32, &w);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // 创建scale aclTensor
    ret = CreateAclTensor(scaleHostData, scaleShape, &scaleDeviceAddr, aclDataType::ACL_INT64, &scale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // 创建bias aclTensor
    ret = CreateAclTensor(biasHostData, biasShape, &biasDeviceAddr, aclDataType::ACL_FLOAT, &bias);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // 创建pertokenScale aclTensor
    ret = CreateAclTensor(pertokenScaleHostData, pertokenScaleShape, &pertokenScaleDeviceAddr, aclDataType::ACL_FLOAT, &pertokenScale);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // 创建groupList aclTensor
    ret = CreateAclTensor(groupListHostData, groupListShape, &groupListDeviceAddr, aclDataType::ACL_INT64, &groupList);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // 创建sharedInput aclTensor
    ret = CreateAclTensor(sharedInputHostData, sharedInputShape, &sharedInputDeviceAddr, aclDataType::ACL_BF16, &sharedInput);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // 创建logit aclTensor
    ret = CreateAclTensor(logitHostData, logitShape, &logitDeviceAddr, aclDataType::ACL_FLOAT, &logit);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // 创建rowIndex aclTensor
    ret = CreateAclTensor(rowIndexHostData, rowIndexShape, &rowIndexDeviceAddr, aclDataType::ACL_INT64, &rowIndex);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // 创建out aclTensor
    ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // 3. 调用CANN算子库API，需要修改为具体的Api名称
    uint64_t workspaceSize = 0;
    aclOpExecutor *executor;
    void *workspaceAddr = nullptr;

    // 调用aclnnGroupedMatmulFinalizeRouting第一段接口
    workspaceSize = 0;
    ret = aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize(x, w, scale, bias, pertokenScale, groupList, sharedInput, logit, rowIndex, dtype, shareInputWeight, shareInputOffest, transposeX, transposeW, groupListType, out, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnGroupedMatmulFinalizeRoutingGetWorkspaceSize failed. ERROR: %d\n", ret);
              return ret);
    // 根据第一段接口计算出的workspaceSize申请device内存
    if (workspaceSize > 0) {
        ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
    }
    // 调用aclnnGroupedMatmulFinalizeRouting第二段接口
    ret = aclnnGroupedMatmulFinalizeRouting(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnGroupedMatmulFinalizeRouting failed. ERROR: %d\n", ret); return ret);

    // 4. （固定写法）同步等待任务执行结束
    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

    // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
    auto size = GetShapeSize(outShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr,
                      size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret);
              return ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("result[%ld] is: %u\n", i, resultData[i]);
    }

    // 6. 释放aclTensor和aclTensor，需要根据具体API的接口定义修改
    aclDestroyTensor(x);
    aclDestroyTensor(w);
    aclDestroyTensor(scale);
    aclDestroyTensor(bias);
    aclDestroyTensor(pertokenScale);
    aclDestroyTensor(groupList);
    aclDestroyTensor(sharedInput);
    aclDestroyTensor(logit);
    aclDestroyTensor(rowIndex);
    aclDestroyTensor(out);

    // 7.释放device资源，需要根据具体API的接口定义修改
    aclrtFree(xDeviceAddr);
    aclrtFree(wDeviceAddr);
    aclrtFree(scaleDeviceAddr);
    aclrtFree(biasDeviceAddr);
    aclrtFree(pertokenScaleDeviceAddr);
    aclrtFree(groupListDeviceAddr);
    aclrtFree(sharedInputDeviceAddr);
    aclrtFree(logitDeviceAddr);
    aclrtFree(rowIndexDeviceAddr);
    aclrtFree(outDeviceAddr);

    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}