Star1281
1394
代码
代码
Issues94
Pull Requests192
Actions
项目讨论
Wiki
分析
Star1281
1394
cann-robotcann-robotdocs目录基本概念md文件名改为英文
51815064创建于 7月22日历史提交

aclnnWeightQuantPreprocess

📄 查看源码

产品支持情况

  • Ascend 950PR/Ascend 950DT：支持
  • Atlas A3 训练系列产品/Atlas A3 推理系列产品：不支持
  • Atlas A2 训练系列产品/Atlas A2 推理系列产品：不支持
  • Atlas 200I/500 A2 推理产品：不支持
  • Atlas 推理系列产品：不支持
  • Atlas 训练系列产品：不支持

功能说明

完成伪量化Matmul（包括QuantBatchMatmulV5、GroupedMatmul-伪量化）的参数预处理：主要将weight从ND格式转换为FRACTAL_NZ格式，并在需要时对weightScale、weightOffsetOptional、biasOptional进行同步处理。

函数原型

每个算子分为两段式接口：

  1. 调用aclnnWeightQuantPreprocessGetWorkspaceSize获取workspace大小及执行器；
  2. 调用aclnnWeightQuantPreprocess执行计算。

注意：用户需自行构造输出张量，参考约束说明中的shape计算公式。

aclnnStatus aclnnWeightQuantPreprocessGetWorkspaceSize(
    const aclTensor *weight,
    const aclTensor *weightScale,
    const aclTensor *weightOffsetOptional,
    const aclTensor *biasOptional,
    aclDataType      xDtype,
    aclDataType      xScaleDtype,
    int64_t          kGroupSize,
    aclTensor       *outWeight,
    aclTensor       *outWeightScale,
    aclTensor       *outWeightOffsetOptional,
    aclTensor       *outBiasOptional,
    uint64_t        *workspaceSize,
    aclOpExecutor  **executor)

aclnnStatus aclnnWeightQuantPreprocess(
    void          *workspace,
    uint64_t       workspaceSize,
    aclOpExecutor *executor,
    aclrtStream    stream)

aclnnWeightQuantPreprocessGetWorkspaceSize

  • 参数说明

    参数名 输入/输出 描述 使用说明 数据类型 数据格式 维度（shape） 非连续Tensor
    weight(aclTensor *) 输入 Matmul的权重矩阵 不支持空tensor float4_e2m1 ND 2-3 仅转置场景支持
    weightScale(aclTensor *) 输入 权重的反量化scale参数 不支持空tensor float8_e8m0 ND/NCL/NCHW 3-4 仅转置场景支持
    weightOffsetOptional(aclTensor *) 可选输入 权重的反量化offset参数 当前MM_MX_A8W4/GMM_MX_A8W4数据流不支持，必须为nullptr - ND 1-2 仅转置场景支持
    biasOptional(aclTensor *) 可选输入 Matmul的偏置矩阵 不支持空tensor，必须contiguous float16/bfloat16 ND 1-2 不支持
    xDtype(aclDataType) 输入 Matmul的激活矩阵的数据类型 - aclDataType - - -
    xScaleDtype(aclDataType) 输入 激活的量化scale参数的数据类型 - aclDataType - - -
    kGroupSize(int64_t) 输入 权重在per-group量化时K维度的group的大小 - int64 - - -
    outWeight(aclTensor *) 输出 预处理后的weight - int8/int4/fp8_e4m3/hif8/fp4_e2m1 NZ 2-5 仅转置场景支持
    outWeightScale(aclTensor *) 输出 预处理后的weightScale - float16/bfloat16/fp8_e8m0 ND/NCL/NCHW 3-4 仅转置场景支持
    outWeightOffsetOptional(aclTensor *) 输出 预处理后的weightOffset 当前MM_MX_A8W4/GMM_MX_A8W4数据流不支持，必须为nullptr float16/bfloat16 ND 1-2 仅转置场景支持
    outBiasOptional(aclTensor *) 输出 预处理后的bias 必须contiguous float16/bfloat16 ND 1-2 不支持
    workspaceSize(uint64_t *) 输出 计算所需的workspace大小 - uint64* - - -
    executor(aclOpExecutor **) 输出 包含算子计算流程的执行器 - aclOpExecutor** - - -

  • 返回值

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

    第一段接口完成入参校验，出现以下场景时报错：

    返回值 错误码 描述
    ACLNN_ERR_PARAM_NULLPTR 161001 weight、weightScale、outWeight或outWeightScale是空指针；或biasOptional非空但outBiasOptional是空指针。
    ACLNN_ERR_PARAM_INVALID 161002 输入的数据类型组合不支持，无法匹配当前支持的MM_MX_A8W4/GMM_MX_A8W4数据流。
    weight、weightScale、outWeight或outWeightScale是空tensor；或biasOptional/outBiasOptional在提供时为空tensor。
    weight、weightScale、biasOptional、outWeight、outWeightScale或outBiasOptional的数据类型和数据格式不在支持的范围之内。
    weight、weightScale、biasOptional、outWeight、outWeightScale或outBiasOptional的shape或storage shape不满足校验条件。
    weight或weightScale的stride不满足转置要求，或biasOptional/outBiasOptional在提供时不连续。
    weightOffsetOptional或outWeightOffsetOptional非空，或kGroupSize不等于32。
    ACLNN_ERR_RUNTIME_ERROR 361001 产品型号不支持。
    ACLNN_ERR_INNER_CREATE_EXECUTOR 561101 内部错误，执行器创建失败。
    ACLNN_ERR_INNER_NULLPTR 561103 workspaceSize或executor是空指针，或API内部构图接口返回空指针。

aclnnWeightQuantPreprocess

  • 参数说明

    参数名 输入/输出 描述
    workspace 输入 在Device侧申请的workspace内存地址。
    workspaceSize 输入 在Device侧申请的workspace大小，由第一段接口aclnnWeightQuantPreprocessGetWorkspaceSize获取。
    executor 输入 op执行器，包含了算子计算流程。
    stream 输入 指定执行任务的Stream。

  • 返回值

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

约束说明

  • 确定性说明：aclnnWeightQuantPreprocess默认确定性实现。
MM_MX_A8W4数据流（MM表示Matmul；MX_A8W4表示x的数据类型为FLOAT8_E4M3FN，weight的数据类型为FLOAT4_E2M1，Mx量化模式）

  • weight

    • 数据类型：FLOAT4_E2M1
    • 格式：ND
    • K % kGroupSize == 0
    • view shape：2-D {K, N}
    • storage shape：{N, K}（transposed）
    • stride：[1, K]（最后两维transposed）
    • 不支持空tensor

  • weightScale

    • 数据类型：FLOAT8_E8M0
    • 格式：ND/NCL
    • view shape：3-D {ceildiv(K, 64), N, 2}
    • storage shape：{N, ceildiv(K, 64), 2}（transposed）
    • stride：[2, 2*ceildiv(K,64), 1]（维度0和1交换）
    • 不支持空tensor

  • weightOffsetOptional

    • 当前不支持，必须为nullptr
    • outWeightOffsetOptional也必须为nullptr

  • biasOptional

    • 数据类型：float16/bfloat16
    • 格式：ND
    • 必须为contiguous
    • 不支持空tensor（若提供）

  • kGroupSize

    • 必须等于 32

  • xDtype

    • FLOAT8_E4M3FN

  • xScaleDtype

    • FLOAT8_E8M0

  • outWeight

    • 数据类型：与weight相同
    • 格式：FRACTAL_NZ_C0_32
    • view shape：与weight view shape相同 {K, N}
    • storage shape：4-D {ceildiv(K, 32), ceildiv(N, 16), 16, 32}

  • outWeightScale

    • 数据类型：与weightScale相同
    • 格式：ND
    • view shape：与weightScale view shape相同
    • storage shape：与weightScale storage shape相同

  • outBiasOptional

    • 数据类型：与biasOptional相同
    • 格式：ND
    • 必须为contiguous
    • view shape：与biasOptional相同
    • storage shape：与biasOptional相同
GMM_MX_A8W4数据流（GMM表示GroupedMatmul；MX_A8W4表示x的数据类型为FLOAT8_E4M3FN，weight的数据类型为FLOAT4_E2M1，Mx量化模式）

  • weight

    • 数据类型：FLOAT4_E2M1
    • 格式：ND
    • K % kGroupSize == 0
    • view shape：3-D {G, K, N}
    • storage shape：{G, N, K}（transposed，最后两维交换）
    • stride：[K*N, 1, K]（维度1和2 transposed）
    • 不支持空tensor

  • weightScale

    • 数据类型：FLOAT8_E8M0
    • 格式：ND/NCL/NCHW
    • view shape：4-D {G, ceildiv(K, 64), N, 2}
    • storage shape：{G, N, ceildiv(K, 64), 2}（transposed，维度2和3交换）
    • stride：[2*ceildiv(K,64)*N, 2, 2*ceildiv(K,64), 1]（维度2和3交换）
    • 不支持空tensor

  • weightOffsetOptional

    • 当前不支持，必须为nullptr
    • outWeightOffsetOptional也必须为nullptr

  • biasOptional

    • 数据类型：float16/bfloat16
    • 格式：ND
    • 必须为contiguous
    • 不支持空tensor（若提供）

  • kGroupSize

    • 必须等于 32

  • xDtype

    • FLOAT8_E4M3FN

  • xScaleDtype

    • FLOAT8_E8M0

  • outWeight

    • 数据类型：与weight相同
    • 格式：FRACTAL_NZ_C0_32
    • view shape：与weight view shape相同 {G, K, N}
    • storage shape：5-D {G, ceildiv(K, 32), ceildiv(N, 16), 16, 32}

  • outWeightScale

    • 数据类型：与weightScale相同
    • 格式：ND
    • view shape：与weightScale view shape相同
    • storage shape：与weightScale storage shape相同

  • outBiasOptional

    • 数据类型：与biasOptional相同
    • 格式：ND
    • 必须为contiguous
    • view shape：与biasOptional相同
    • storage shape：与biasOptional相同
  • 其余数据类型与shape组合为预留接口，当前调用将返回ACLNN_ERR_PARAM_INVALID

调用示例

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

注意：用户需自行计算并构造输出张量shape，参考约束说明中的公式：

  • outWeight viewShape：与weight viewShape相同
  • outWeight storageShape：{CeilDiv(K, 32), CeilDiv(N, 16), 16, 32}
  • outWeight format：ACL_FORMAT_FRACTAL_NZ_C0_32
#include <iostream>
#include <memory>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_weight_quant_preprocess.h"

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

#define CEIL_DIV(x, y) (((x) + (y) - 1) / (y))

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

class AclRuntimeGuard {
public:
    explicit AclRuntimeGuard(int32_t deviceId) : deviceId_(deviceId) {}

    ~AclRuntimeGuard()
    {
        if (stream_ != nullptr) {
            aclrtDestroyStream(stream_);
            stream_ = nullptr;
        }
        if (deviceSet_) {
            aclrtResetDevice(deviceId_);
            deviceSet_ = false;
        }
        if (aclInited_) {
            aclFinalize();
            aclInited_ = false;
        }
    }

    int Init(aclrtStream* stream)
    {
        auto ret = aclInit(nullptr);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        aclInited_ = true;

        ret = aclrtSetDevice(deviceId_);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        deviceSet_ = true;

        ret = aclrtCreateStream(stream);
        CHECK_RET(ret == ACL_SUCCESS, return ret);
        stream_ = *stream;
        return ACL_SUCCESS;
    }

private:
    int32_t deviceId_;
    aclrtStream stream_ = nullptr;
    bool aclInited_ = false;
    bool deviceSet_ = false;
};

int main()
{
    int32_t deviceId = 0;
    aclrtStream stream = nullptr;
    AclRuntimeGuard aclGuard(deviceId);
    auto ret = aclGuard.Init(&stream);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Init failed" << std::endl; return ret);

    // weight: FLOAT4_E2M1, transposed (MM_MX_A8W4)
    int64_t k = 64;
    int64_t n = 128;
    int64_t C0 = 32; // FLOAT4_E2M1 对应 C0=32

    std::vector<int64_t> weightViewShape = {k, n};
    std::vector<int64_t> weightStorageShape = {n, k};
    std::vector<int64_t> weightStrides = {1, k};
    int64_t weightStorageSize = GetShapeSize(weightStorageShape);
    int64_t weightBytes = weightStorageSize / 2; // FP4: 4 bits = 0.5 bytes per element

    std::vector<int8_t> weightHostData(weightBytes, 0);
    void* weightDeviceAddr = nullptr;
    ret = aclrtMalloc(&weightDeviceAddr, weightBytes, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Malloc weight failed" << std::endl; return ret);
    std::unique_ptr<void, aclError (*)(void*)> weightDeviceAddrPtr(weightDeviceAddr, aclrtFree);
    ret = aclrtMemcpy(weightDeviceAddr, weightBytes, weightHostData.data(), weightBytes, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Memcpy weight failed" << std::endl; return ret);
    aclTensor* weight = aclCreateTensor(weightViewShape.data(), weightViewShape.size(), ACL_FLOAT4_E2M1,
                                        weightStrides.data(), 0, ACL_FORMAT_ND, weightStorageShape.data(),
                                        weightStorageShape.size(), weightDeviceAddr);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> weightPtr(weight, aclDestroyTensor);
    CHECK_RET(weight != nullptr, std::cout << "Create weight tensor failed" << std::endl; return ACL_ERROR_FAILURE);

    // weightScale: FLOAT8_E8M0, 3-D transposed (MM_MX_A8W4)
    // viewShape: {ceildiv(K,64), N, 2} = {1, 128, 2}
    // storageShape: {N, ceildiv(K,64), 2} = {128, 1, 2}
    // transposed stride: {2, 2, 1} (dim0 <-> dim1)
    std::vector<int64_t> scaleViewShape = {k / 64, n, 2};
    std::vector<int64_t> scaleStorageShape = {n, k / 64, 2};
    std::vector<int64_t> scaleStrides = {2, 2, 1};
    int64_t scaleStorageSize = GetShapeSize(scaleStorageShape);
    int64_t scaleBytes = scaleStorageSize; // FP8: 1 byte per element

    std::vector<int8_t> scaleHostData(scaleBytes, 0);
    void* scaleDeviceAddr = nullptr;
    ret = aclrtMalloc(&scaleDeviceAddr, scaleBytes, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Malloc weightScale failed" << std::endl; return ret);
    std::unique_ptr<void, aclError (*)(void*)> scaleDeviceAddrPtr(scaleDeviceAddr, aclrtFree);
    ret = aclrtMemcpy(scaleDeviceAddr, scaleBytes, scaleHostData.data(), scaleBytes, ACL_MEMCPY_HOST_TO_DEVICE);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Memcpy weightScale failed" << std::endl; return ret);
    aclTensor* weightScale = aclCreateTensor(scaleViewShape.data(), scaleViewShape.size(), ACL_FLOAT8_E8M0,
                                             scaleStrides.data(), 0, ACL_FORMAT_ND, scaleStorageShape.data(),
                                             scaleStorageShape.size(), scaleDeviceAddr);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> weightScalePtr(weightScale, aclDestroyTensor);
    CHECK_RET(weightScale != nullptr, std::cout << "Create weightScale tensor failed" << std::endl;
              return ACL_ERROR_FAILURE);

    // 用户自行构造 outWeight (FRACTAL_NZ_C0_32)
    // viewShape 与 weight viewShape 相同，storageShape 按公式计算
    std::vector<int64_t> outWeightViewShape = {k, n};
    std::vector<int64_t> outWeightStorageShape = {CEIL_DIV(k, C0), CEIL_DIV(n, 16), 16, C0};
    int64_t outWeightStorageSize = GetShapeSize(outWeightStorageShape);
    int64_t outWeightBytes = outWeightStorageSize / 2; // FP4

    void* outWeightDeviceAddr = nullptr;
    ret = aclrtMalloc(&outWeightDeviceAddr, outWeightBytes, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Malloc outWeight failed" << std::endl; return ret);
    std::unique_ptr<void, aclError (*)(void*)> outWeightDeviceAddrPtr(outWeightDeviceAddr, aclrtFree);
    aclTensor* outWeight = aclCreateTensor(outWeightViewShape.data(), outWeightViewShape.size(), ACL_FLOAT4_E2M1,
                                           nullptr, 0, ACL_FORMAT_FRACTAL_NZ_C0_32, outWeightStorageShape.data(),
                                           outWeightStorageShape.size(), outWeightDeviceAddr);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> outWeightPtr(outWeight, aclDestroyTensor);
    CHECK_RET(outWeight != nullptr, std::cout << "Create outWeight tensor failed" << std::endl;
              return ACL_ERROR_FAILURE);

    // 构造 outWeightScale (viewShape 和 storageShape 都与 weightScale 相同)
    // 根据实现要求：outWeightScale 的 viewShape 和 storageShape 必须都与 weightScale 相同
    void* outScaleDeviceAddr = nullptr;
    ret = aclrtMalloc(&outScaleDeviceAddr, scaleBytes, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Malloc outWeightScale failed" << std::endl; return ret);
    std::unique_ptr<void, aclError (*)(void*)> outScaleDeviceAddrPtr(outScaleDeviceAddr, aclrtFree);
    aclTensor* outWeightScale = aclCreateTensor(scaleViewShape.data(), scaleViewShape.size(), ACL_FLOAT8_E8M0,
                                                scaleStrides.data(), 0, ACL_FORMAT_ND, scaleStorageShape.data(),
                                                scaleStorageShape.size(), outScaleDeviceAddr);
    std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> outWeightScalePtr(outWeightScale, aclDestroyTensor);
    CHECK_RET(outWeightScale != nullptr, std::cout << "Create outWeightScale tensor failed" << std::endl;
              return ACL_ERROR_FAILURE);

    aclDataType xDtype = ACL_FLOAT8_E4M3FN;
    aclDataType xScaleDtype = ACL_FLOAT8_E8M0;
    int64_t kGroupSize = 32;

    // 1. 获取 workspace 与执行器
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor = nullptr;
    ret = aclnnWeightQuantPreprocessGetWorkspaceSize(
        weight, weightScale, nullptr, nullptr, // weightOffsetOptional, biasOptional
        xDtype, xScaleDtype, kGroupSize, outWeight, outWeightScale, nullptr, nullptr, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "GetWorkspaceSize failed" << std::endl; return ret);

    void* workspaceAddr = nullptr;
    std::unique_ptr<void, aclError (*)(void*)> workspaceAddrPtr(nullptr, aclrtFree);
    if (workspaceSize > 0) {
        ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, std::cout << "Malloc workspace failed" << std::endl; return ret);
        workspaceAddrPtr.reset(workspaceAddr);
    }

    // 2. 执行计算
    ret = aclnnWeightQuantPreprocess(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Preprocess failed" << std::endl; return ret);

    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, std::cout << "Synchronize failed" << std::endl; return ret);

    // 3. 释放资源
    workspaceAddrPtr.reset();
    outWeightScalePtr.reset();
    outWeightPtr.reset();
    weightScalePtr.reset();
    weightPtr.reset();
    outScaleDeviceAddrPtr.reset();
    outWeightDeviceAddrPtr.reset();
    scaleDeviceAddrPtr.reset();
    weightDeviceAddrPtr.reset();
    return 0;
}