CANN/ascend-transformer-boost多潜在注意力算子演示

📅 2026/7/13 8:10:56 👁️ 阅读次数 📝 编程学习
CANN/ascend-transformer-boost多潜在注意力算子演示

加速库MultiLatentAttentionOperation C++ Demo

【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库,基于华为Ascend AI处理器,提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost

介绍

该目录下为加速库MultiLatentAttentionOperation C++调用示例。

使用说明

  • 首先source 对应的CANN和nnal包的安装路径

    1. source [cann安装路径]/set_env.sh 默认:source /usr/local/Ascend/ascend-toolkit/set_env.sh
    2. source [nnal安装路径]/set_env.sh
      默认:source /usr/local/Ascend/nnal/atb/set_env.sh
      ①. 如果使用加速库源码编译,source [加速库源码路径]/output/atb/set_env.sh
      例如: source ./ascend-transformer-boost/output/atb/set_env.sh
  • 运行demo

    bash build.sh

    注意

    • 使用cxx_abi=0(默认)时,设置D_GLIBCXX_USE_CXX11_ABI为0,即:

      g++ -D_GLIBCXX_USE_CXX11_ABI=0 -I ...
    • 使用cxx_abi=1时,更改D_GLIBCXX_USE_CXX11_ABI为1,即:

      g++ -D_GLIBCXX_USE_CXX11_ABI=1 -I ...
    • 提供的build脚本仅用于编译和运行mlapa_demo.cpp,如需编译其他demo,需要替换“mlapa_demo”为对应的cpp文件名

额外说明

示例中生成的数据不代表实际场景,如需数据生成参考请查看根目录下的python用例目录: tests/apitest/opstest/python/operations/multi_latent_attention/

场景说明

该算子所给demo仅支持在Atlas A2/A3系列产品上运行,demo的场景说明如下:

  • mlapa_demo.cpp

    参数设置

    成员名称取值
    headNum128
    qkScale1/sqrt(576)
    kvHeadNum1
    maskTypeUNDEFINED
    calcTypeCALC_TYPE_UNDEFINED
    cacheModeINT8_NZCACHE

    注意:qkScale设置值为MLA做rope投影前的headSize,即512(原始) + 64(投影) = 576

    数据规格

    tensor名字数据类型数据格式维度信息cpu/npu
    qNopeint8nd[4, 128, 512]npu
    qRopefloat16nd[4, 128, 64]npu
    ctKVint8nz[48, 16, 128, 32]npu
    kRopefloat16nz[48, 4, 128, 16]npu
    blockTablesint32nd[4, 12]npu
    contextLensint32nd[4]cpu
    qkDescalefloatnd[128]npu
    pvDescalefloatnd[128]npu
    attenOutfloat16nd[4, 128, 512]npu
  • mlapa_ds_demo.cpp

    参数设置

    成员名称取值
    headNum128
    qkScale0.1352667747812271
    kvHeadNum1
    maskTypeUNDEFINED
    calcTypeCALC_TYPE_UNDEFINED
    cacheModeKROPE_CTKV

    数据规格

    tensor名字数据类型数据格式维度信息cpu/npu
    qNopefloat16nd[32, 128, 512]npu
    qRopefloat16nd[7168, 128, 64]npu
    ctKVfloat16nd[160, 128, 1, 512]npu
    kRopefloat16nd[160, 128, 1, 64]npu
    blockTablesint32nd[32, 5]npu
    contextLensint32nd[32]cpu
    attenOutfloat16nd[32, 128, 512]npu

【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库,基于华为Ascend AI处理器,提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考