CANN ascend-transformer-boost MultiLatentAttentionOperation C Demo 实战指南【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost导读本文以 CANN ascend-transformer-boost 开源仓库中的 example/op_demo/multi_latent_attention/README_en.md 为骨架系统讲解如何基于加速库ATBAscend Transformer Boost在华为昇腾 Atlas A2/A3 系列产品上编译、运行并理解 MultiLatentAttentionOperation 的 C 调用示例。文章不仅完整覆盖 README 中的环境配置、构建命令与参数规格还深入算子参数定义、参数校验逻辑与 runner 实现源码帮助读者掌握 MLA 融合算子在 INT8 量化 NZ Cache 与 KROPE 分离 Cache 两种典型场景下的数据组织方式与调用流程。背景为什么需要 MultiLatentAttentionOperationMLAMulti-head Latent Attention多潜变量注意力是 DeepSeek 等大模型采用的高效注意力机制其核心思想是将 KV 压缩到低秩潜变量空间显著降低推理时的 KV Cache 显存占用。ascend-transformer-boost 通过MultiLatentAttentionOperation将 MLA 的完整计算链融合为一个高性能算子上层业务只需传入 query、KV Cache、blockTables 等张量即可完成注意力计算省去逐算子拼接的低效路径。从源码结构看该算子位于 src/ops/ops_infer/multi_latent_attention/包含算子主体multi_latent_attention_operation.cpp/.h、执行 runnermulti_latent_attention_ops_runner.cpp、multi_latent_attention_ops_runner_prefill.cpp、参数解析param.cpp/.h以及针对不同 Cache 模式的实现atb_acl_mla.cpp。环境准备Source CANN 与 NNAL 安装路径README 明确要求运行 demo 前先加载 CANN 与 NNAL加速库两个软件包的环境变量这是 demo 能够链接到libatb.so与libascendcl.so的前提。加载 CANN 工具包环境source [CANN安装路径]/set_env.sh # 默认路径示例 source /usr/local/Ascend/ascend-toolkit/set_env.sh加载 NNAL加速库环境source [NNAL安装路径]/set_env.sh # 默认路径示例 source /usr/local/Ascend/nnal/atb/set_env.sh如果直接使用仓库源码编译产物则改为 source 编译输出目录下的环境脚本source [加速库源码路径]/output/atb/set_env.sh # 例如 source ./ascend-transformer-boost/output/atb/set_env.sh从仓库中的 build.sh 可以看出编译时实际依赖ATB_HOME_PATH与ASCEND_HOME_PATH两个环境变量分别定位加速库与 CANN 的头文件和库目录上述set_env.sh正是负责注入这些变量。编译与运行 demo一键构建在 multi_latent_attention 示例目录 下执行bash build.shCXX ABI 注意事项CANN 与加速库在编译期使用了特定的 C ABI 规范demo 编译时必须保持一致否则会出现链接错误。README 给出了两种配置使用cxx_abi0默认时设置D_GLIBCXX_USE_CXX11_ABI为0g -D_GLIBCXX_USE_CXX11_ABI0 -I ...使用cxx_abi1时更改D_GLIBCXX_USE_CXX11_ABI为1g -D_GLIBCXX_USE_CXX11_ABI1 -I ...实际 build.sh 会通过 Python 检测 PyTorch 的 ABI 设置自动决定cxx_abi取值并拼装完整编译命令g -D_GLIBCXX_USE_CXX11_ABI$cxx_abi -I ${ATB_HOME_PATH}/include -I ${ASCEND_HOME_PATH}/include \ -L ${ATB_HOME_PATH}/lib -L ${ASCEND_HOME_PATH}/lib64 \ mlapa_demo.cpp ../demo_util.h -l atb -l ascendcl -o mlapa_demo ./mlapa_demo替换编译目标README 特别说明提供的build.sh仅用于编译并运行mlapa_demo.cpp。如需编译其他 demo如mlapa_ds_demo.cpp需要将脚本中的mlapa_demo替换为对应的.cpp文件名。Demo 场景与数据规格详解README 明确该算子提供的 demo 仅支持在 Atlas A2/A3 系列产品上运行并给出两个典型场景。场景一mlapa_demo.cppINT8 量化 NZ Cache参数设置名称取值headNum128qkScale1/sqrt(576)kvHeadNum1maskTypeUNDEFINEDcalcTypeCALC_TYPE_UNDEFINEDcacheModeINT8_NZCACHE注意qkScale的取值基于 MLA 做 RoPE 投影前的 headSize即512原始 64投影 576。这一细节在 mlapa_demo.cpp 中有直接对应实现param.qkScale 1 / sqrt(512 64)。数据规格Tensor数据类型数据格式维度cpu/npuqNopeint8nd[4, 128, 512]npuqRopefloat16nd[4, 128, 64]npuctKVint8nz[48, 16, 128, 32]npukRopefloat16nz[48, 4, 128, 16]npublockTablesint32nd[4, 12]npucontextLensint32nd[4]cpuqkDescalefloatnd[128]npupvDescalefloatnd[128]npuattenOutfloat16nd[4, 128, 512]npu该场景对应INT8_NZCACHE高性能分离 Cache模式query 的 Nope 部分与 KV CachectKV以 int8 量化存储并额外提供qkDescale、pvDescale反量化尺度张量。ctKV采用 FRACTAL_NZ 格式int8 的 NZ 分形块基元为 32因此维度 [48, 16, 128, 32] 中 51216×32对应每个 block 的 latent 维度kRope同理644×16为 RoPE 投影后的头维度。场景二mlapa_ds_demo.cppKROPE_CTKV 分离 Cache参数设置名称取值headNum128qkScale0.1352667747812271kvHeadNum1maskTypeUNDEFINEDcalcTypeCALC_TYPE_UNDEFINEDcacheModeKROPE_CTKV数据规格Tensor数据类型数据格式维度cpu/npuqNopefloat16nd[32, 128, 512]npuqRopefloat16nd[7168, 128, 64]npuctKVfloat16nd[160, 128, 1, 512]npukRopefloat16nd[160, 128, 1, 64]npublockTablesint32nd[32, 5]npucontextLensint32nd[32]cpuattenOutfloat16nd[32, 128, 512]npu该场景对应KROPE_CTKV分离 Cache也是参数默认值模式全部张量为 float16 ND 格式无需反量化尺度张量。与场景一相比qRope 的 token 维度7168 32×224大于 qNope 的 token 数32可以推断该模式支持对 RoPE 部分单独组织更长/更灵活的 token 序列且ctKV/kRope维度中显式保留了kvHeadNum1这一维[160, 128, 1, 512]。此处qkScale采用直接的十进制数值0.1352667747812271与场景一的1/sqrt(576)含义相同但表示方式不同。源码级解读demo 的调用主流程两个 demo 遵循相同的 ATB 算子调用范式下面以 mlapa_demo.cpp 为主线拆解。1. 初始化环境与创建 Contextmain函数首先初始化 ACL 运行时、设置 Device、创建 ATB Context 并绑定执行流CHECK_STATUS(aclInit(nullptr)); CHECK_STATUS(aclrtSetDevice(DEVICE_ID)); // DEVICE_ID 1 CHECK_STATUS(atb::CreateContext(context)); CHECK_STATUS(aclrtCreateStream(stream)); CHECK_STATUS(context-SetExecuteStream(stream));demo 同时支持命令行传参控制数据类型与规模argv[1]为dtypeStrbf16时切换到ACL_BF16argv[2]、argv[3]、argv[4]分别对应tokenNum、headNum、kSeqLen。默认值为tokenNum4、headNum128、kSeqLen1500。2. 创建算子通过参数结构体实例化atb::infer::MultiLatentAttentionParam param; param.headNum headNum; param.qkScale 1 / sqrt(512 64); param.kvHeadNum 1; param.cacheMode atb::infer::MultiLatentAttentionParam::CacheMode::INT8_NZCACHE; return atb::CreateOperation(param, mlaOp);参数结构体定义位于 include/atb/infer_op_params.h核心字段包括字段类型默认值说明headNumint32_t0query 头大小qkScalefloat1.0算子 torch 值在 Q×K^T 后乘kvHeadNumint32_t0kv 头数量maskTypeMaskTypeUNDEFINEDmask 类型calcTypeCalcTypeCALC_TYPE_UNDEFINED计算类型cacheModeCacheModeKVCACHEcache 类型windowSizeuint32_t0滑动窗口大小maskUseStatusTypeMaskUseStatusTypeMASK_USE_STATUS_TYPE_UNDEFINED是否按 batch 控制 mask 使用rsvuint8_t[32]{0}预留参数各枚举取值均来自同一头文件MaskTypeUNDEFINED默认全 0 mask、MASK_TYPE_SPECqSeqLen1 时的 mask、MASK_TYPE_MASK_FREEmask free、MASK_TYPE_CAUSAL_MASK内部生成 mask、MASK_TYPE_SWA_NORM。CalcTypeCALC_TYPE_UNDEFINED默认、CALC_TYPE_SPEC支持传入大于 1 的 qSeqLen、CALC_TYPE_RINGringAttention、CALC_TYPE_SPEC_AND_RING大 qSeqLen ringAttention、CALC_TYPE_PREFILL全量场景。CacheModeKVCACHE拼接 cache、KROPE_CTKV分离 cache默认值、INT8_NZCACHE高性能分离 cache、NZCACHE非量化 NZ cache。3. 准备输入张量VariantPackPrepareInTensor通过CreateTensorFromVector封装在 example/op_demo/demo_util.h 中依次创建输入张量并放入variantPack.inTensors。注意以下几个关键点blockNum由tokenNum与kSeqLen按块大小 128 推导maxBlockNumPerSeq (kSeqLen blockSize - 1) / blockSizeblockNum tokenNum * maxBlockNumPerSeq。默认参数下kSeqLen1500得到每序列 12 块、总计 48 块与数据规格中blockTables的 [4, 12] 与ctKV的 48 一致。contextLens是唯一位于 host 侧cpu的输入通过contextLens.hostData contextLensHost.data()直接挂接 host 内存不需要设备侧分配。INT8_NZCACHE模式会额外附加qkDescale、pvDescale两个 shape 为 [headNum] 的 float 张量。输入张量的组织顺序可以从 runner 实现 multi_latent_attention_ops_runner.cpp 得到印证基础 6 个输入依次为 queryqNope、queryRopeqRope、kvCachectKV、kvCacheRopekRope、blockTables、contextLens随后按配置可选追加 mask、qLens、qkDescale/pvDescale仅 INT8_NZ、maskUseStatus。基础输出为attenOutring 模式还会追加第二个输出 LSE。4. Setup 计算 workspace 并循环 Executeuint64_t workspaceSize 0; CHECK_STATUS(mlaOp-Setup(variantPack, workspaceSize, context)); uint8_t *workspacePtr nullptr; if (workspaceSize 0) { CHECK_STATUS(aclrtMalloc(workspacePtr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST)); } for (size_t i 0; i 2; i) { mlaOp-Execute(variantPack, workspacePtr, workspaceSize, context); CHECK_STATUS(aclrtSynchronizeStream(stream)); // 流同步等待 device 侧任务计算完成 }Setup阶段完成 shape 推导与 workspace 大小估算Execute真正下发任务demo 循环执行 2 次以验证算子可重复调用。执行完毕后依次释放输入/输出张量与 workspace最后atb::DestroyOperation(mlaOp)销毁算子对象。5. 输出与资源释放输出attenOutshape 为 [tokenNum, headNum, 512]数据类型与qNope一致mlapa_demo 中为 float16。释放顺序上demo 遵循先释放 deviceData、再销毁 Operation、最后销毁 Context 与流的原则避免资源生命周期交叉。参数校验规则哪些组合合法从 multi_latent_attention_operation.cpp 的参数校验逻辑可以提炼出使用该算子必须遵守的约束这些约束对读者自行修改 demo 参数至关重要headNum取值范围为 [1, 128]。qkScale必须满足0 qkScale 1。kvHeadNum必须为 1仅支持 MQAMulti-Query Attention当前实现不支持KVCACHE拼接模式CacheMode::KVCACHE会直接报错。maskType、calcType、cacheMode均不能超出各自枚举范围。仅CALC_TYPE_SPEC/CALC_TYPE_SPEC_AND_RINGMTP 场景支持非UNDEFINED的 maskType且 SPEC 系列不支持MASK_TYPE_SWA_NORM如需 SWA_NORM 请改用CALC_TYPE_PREFILL或MASK_FREE。INT8_NZCACHE ringCALC_TYPE_RING/CALC_TYPE_SPEC_AND_RING组合下不支持 mask。maskUseStatusType仅在CALC_TYPE_SPEC_AND_RING下支持MASK_USE_STATUS_TYPE_BATCH_MASK。MASK_TYPE_SWA_NORM模式下windowSize必须大于 0。若calcType CALC_TYPE_PREFILL全量场景cacheMode必须为KROPE_CTKV且对 maskType 有更严格的组合要求相关校验见ParamPrefillCheck与 atb_acl_mla.cpp 中 Prefill 场景的 mask 处理逻辑。测试用例与数据生成参考README 特别声明demo 中生成的数据不代表实际场景输出。需要真实场景的数据组织与 golden 计算参考时应查看仓库根目录下的 Python 用例目录tests/apitest/opstest/python/operations/multi_latent_attention/test_multi_latent_attention.py基础 MLA 用例内含完整的分组矩阵乘、softmax、反量化de-scale的 PyTorch golden 实现同目录下的test_multi_latent_attention_int8nz.py、test_multi_latent_attention_lse.py、test_multi_latent_attention_mtp.py、test_multi_latent_attention_prefill.py、test_multi_latent_attention_ring.py则分别覆盖 INT8 NZ 量化、LSE 输出、MTPSPEC、全量 Prefill 与 ringAttention 等扩展场景可作为理解各cacheMode/calcType组合下张量语义的权威参考。常见问题与排查建议链接报错未定义的atb::符号多为 CXX ABI 不匹配请核对D_GLIBCXX_USE_CXX11_ABI与所用包一致或环境未正确 sourceATB_HOME_PATH/ASCEND_HOME_PATH为空。运行报kvHeadNum should be 1该算子当前仅支持 MQA请勿将kvHeadNum配为大于 1 的值。cacheModeKVCACHE报不支持当前实现未支持拼接 Cache 模式请改用KROPE_CTKV、INT8_NZCACHE或NZCACHE。INT8_NZCACHE模式下输入张量缺失需要同时提供qkDescale与pvDescale两个 [headNum] 的 float 张量且ctKV/kRope需按 FRACTAL_NZ 格式组织。精度不符预期demo 数据为占位数据请以 Python 用例目录 的数据生成与精度对比流程为准。硬件限制该算子 demo 仅支持 Atlas A2/A3 系列产品其他硬件上无法运行。小结通过本文可以完整掌握 ascend-transformer-boost 中 MultiLatentAttentionOperation C demo 的环境配置、编译运行方法与两种典型 Cache 模式的参数/张量规格并借助算子参数定义infer_op_params.h、参数校验multi_latent_attention_operation.cpp与 runner 实现multi_latent_attention_ops_runner.cpp理解其底层约束。在此基础上读者可以仿照 mlapa_demo.cpp 的调用范式将 MLA 算子接入自己的推理服务或结合 Python 测试用例 扩展出适配自身模型配置的调用代码。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考