CANN 鸿蒙端侧 GatherDequantInt8 自定义算子：基于 Ascend C 的 INT8 Embedding 查表与反量化融合实现

CANN 鸿蒙端侧 GatherDequantInt8 自定义算子基于 Ascend C 的 INT8 Embedding 查表与反量化融合实现【免费下载链接】cann-recipes-harmony-infer本项目为鸿蒙开发者提供基于CANN平台的业务实践案例方便开发者参考实现端云能力迁移及端侧推理部署。项目地址: https://gitcode.com/cann/cann-recipes-harmony-infer导读GatherDequantInt8是 CANN 开源仓库 cann-recipes-harmony-infer 中面向鸿蒙端侧Kirin 9020 处理器提供的 Ascend C 自定义算子样例它把「INT8 embedding 查表」与「逐行per-token非对称反量化」融合为单个算子直接在 AI Core 上以 uint8 读表、Cast 到 fp16 并逐行反量化解决端侧模型 embedding 无法以 INT8 内置进图的痛点。读完本文你将掌握该算子的数学语义、算子规格、Host/Device 两侧源码实现原理、ONNX 前端适配方式以及从环境准备、编译安装到单算子测试的完整实战流程。应用背景为什么需要自定义 Gather 算子在端侧部署带 embedding 的模型如标点恢复、ASR 等时为压缩模型物理尺寸通常希望把 embedding 权重以 INT8 形式内置进计算图。然而 Kirin 9020 工具链的框架 Gather 算子GatherV2D的data输入仅支持 fp32/fp16/int32不接受 uint8/int8导致图内 INT8 embedding 查表无法用标准算子组合表达。GatherDequantInt8正是针对该场景设计的融合算子在 AI Core 上直接以 uint8 读表、Cast 到 half、逐行反量化使图内 embedding 可压缩到 INT8约为 fp32 的 1/4与外置反量化 bin 的体积对齐数值与 fp16 反量化路径一致在标点模型端到端部署中已验证与 fp32 embedding 的标点预测 argmax 100% 一致见 应用场景说明。该样例还提供 ONNX 框架适配插件可通过 ATC/OMG 将含该自定义节点的 ONNX 模型转换为端侧离线模型从而在开发者自研 INT8 量化 embedding 的部署场景下实现图内 INT8 embedding 查表 反量化。数学语义与算子规格数学表达式算子对应的数学表达式为y[i, :] (half(table[indices[i], :]) - zero_point[indices[i]]) * scale[indices[i]]table逐行非对称量化的 uint8 embedding 表indices待查表的 token idscale/zero_point逐行per-token的缩放因子与零点y反量化后的 fp16 embedding。算子规格表名称角色数据类型形状说明table输入uint8[V, E]逐行非对称量化的 embedding 表indices输入int32[...]token id按元素查表支持任意 shapescale输入fp16[V]逐行 scalezero_point输入fp16[V]逐行 zero_pointy输出fp16[..., E]反量化后的 embedding其中 V 为词表大小E 为 embedding 维度输出 shape indices.shape [E]dtype 固定为 fp16。逐行量化公式与测试数据生成一致本样例test/gen_data.py采用的 per-row 非对称 uint8 量化关系为scale_v (max(W_v) - min(W_v)) / 255 zp_v -min(W_v) / scale_v q_v round(W_v / scale_v zp_v)即在生成测试数据时先对原始 fp32 embedding 逐行求 min/max 得到 scale 与 zero_point再量化到 uint8golden 输出则严格按 fp16 计算路径(half(q) - zp) * scale生成用于与 device kernel 对齐验证见 gen_data.py。支持的产品型号Kirin 9020 处理器系列产品如需适配 Kirin X90 / 9030需同步修改CMakePresets.json的ASCEND_COMPUTE_UNIT与 op_host/gather_dequant_int8.cpp 中AICore().AddConfig(...)。算子工程目录结构gather_dequant_int8_custom ├── build_and_install.sh # 编译安装脚本 ├── CMakeLists.txt ├── CMakePresets.json # 编译配置ASCEND_CANN_PACKAGE_PATH / 算力型号 ├── framework │ ├── CMakeLists.txt │ └── onnx_plugin │ ├── CMakeLists.txt │ └── gather_dequant_int8_plugin.cpp # ONNX 前端适配插件 ├── op_host │ ├── CMakeLists.txt │ ├── gather_dequant_int8.cpp # 原型注册 / InferShape / InferDataType / Tiling │ └── gather_dequant_int8_tiling.h # TilingData 定义 ├── op_kernel │ ├── CMakeLists.txt │ └── gather_dequant_int8.cpp # 核函数实现 └── test ├── create_onnx.py # 生成单算子 ONNX 测试模型 └── gen_data.py # 生成输入与 golden 数据源码级实现原理该算子的实现遵循 Ascend C 自定义算子工程的标准三段式结构Host 侧原型注册、Shape/DType 推导、Tiling、Device 侧核函数、框架适配ONNX 插件。Host 侧原型注册与 Tilingop_host/gather_dequant_int8.cpp 承担三部分职责算子原型注册GatherDequantInt8类中声明 4 个 REQUIRED 输入tableuint8 ND、indicesint32 ND、scalefp16 ND、zero_pointfp16 ND与 1 个输出yfp16 ND并通过AICore().SetTiling(...)绑定 Tiling 函数、AICore().AddConfig(kirin9020)声明算力配置。InferShape输出y的 shape 由indices的 rank 与table最后一维拼接得到即y.shape indices.shape [E]其中E table.shape[-1]。InferDataType输出固定为 fp16与 scale / zero_point 一致。Tiling 逻辑中从输入 shape 解析出三个关键参数并写入 TilingDatavocabVtable 第 0 维embDimEtable 最后一维numIndicesindices 全元素乘积即总查表次数。同时设置SetBlockDim(1)即 Kirin 9020 AI Core 单核执行且 workspace 大小为 0。TilingData 结构定义于 op_host/gather_dequant_int8_tiling.h通过BEGIN_TILING_DATA_DEF声明numIndices、embDim、vocab三个 uint32 字段并以REGISTER_TILING_DATA_CLASS完成注册供核函数侧GET_TILING_DATA读取。Device 侧核函数实现op_kernel/gather_dequant_int8.cpp 实现核函数KernelGatherDequantInt8其设计要点查表数据驻留 UBindices/scale/zero_point一次性载入统一缓冲区UB供按 idx 随机访问通过DataCopyPad处理非 32B 对齐场景并用PipeBarrierPIPE_ALL()保证 MTE2 搬运完成后标量单元GetValue才能读取。逐 token 流水对每个 index依次执行CopyInuint8 行→Castuint8 转 half→Adds减 zero_point→Muls乘 scale→CopyOuthalf 行写回其中Adds/Muls对应数学表达式中的(q - zp) * scale。Double buffer 并行通过BUFFER_NUM 2的队列配置让搬运与计算流水并行。架构细节由于 dav_l310kirin9020不允许标量 half 算术代码将 zero_point 取负运算放到 float 上进行再转回 half同时idxBuf/scaleBuf/zpBuf/calcBuf使用VECCALC位置输入/输出行队列使用VECIN/VECOUT。核函数入口gather_dequant_int8通过GET_TILING_DATA读取 Host 侧下发的 tiling 数据后初始化并执行Process()。ONNX 前端适配插件framework/onnx_plugin/gather_dequant_int8_plugin.cpp 将 ONNX 图中op_typeGatherDequantInt8的自定义节点映射到 GE 自定义算子使omg --framework5能解析并入图通过REGISTER_CUSTOM_OP(GatherDequantInt8)与FrameworkType(ONNX)声明框架类型OriginOpType同时覆盖裸算子名与多种 domain 前缀如custom::GatherDequantInt8、ai.onnx::1::GatherDequantInt8直至ai.onnx::16::GatherDequantInt8该节点无属性输入/输出顺序与 GE 原型一致故直接使用AutoMappingByOpFn自动映射。操作步骤1. 环境准备参考 环境准备 完成环境搭建核心前提如下Python 3.7.0、gcc 7.3.0、cmake 3.16.0建议使用 Ubuntu 22.04 以上环境依赖 glibc 2.34安装鸿蒙社区版 CANN 开发套件包Ascend-cann-toolkit_${cann_version}_linux-${arch}-mobile-station.run安装命令形如chmod x Ascend-cann-toolkit_${cann_version}_linux-${arch}-mobile-station.run ./Ascend-cann-toolkit_${cann_version}_linux-${arch}-mobile-station.run --install --force --install-path${install_path}配置环境变量source /usr/local/Ascend/cann-${cann_version}/set_env.sh编译前确认 CMakePresets.json 中ASCEND_CANN_PACKAGE_PATH指向正确的 toolkit 安装路径一般为${install_path}/cann。该文件还集中定义了ASCEND_COMPUTE_UNIT当前为kirin9020、vendor_namecustomize、ENABLE_TEST、ENABLE_CROSS_COMPILE等编译选项适配新算力型号时主要修改ASCEND_COMPUTE_UNIT一项。2. 编译安装在算子工程目录下执行chmod x build_and_install.sh ./build_and_install.shbuild_and_install.sh 内部会依次完成设置ASCEND_HOME_PATH与setenv.bash环境、将CMakePresets.json中的默认 CANN 路径替换为实际安装路径、以defaultpreset 配置并编译binary与package目标最后执行生成的custom_opp_${OS_ID}_${arch}.run --quiet完成安装。编译产物为自定义算子 run 包并自动安装到packages/vendors/customize/下。详细流程还可参考 算子工程编译安装指南。3. 单算子测试cd test python3 create_onnx.py # 生成 GatherDequantInt8.onnx python3 gen_data.py # 生成 table/indices/scale/zero_point.bin 与 golden output.bincreate_onnx.py 使用默认规格 V1024、E256、N30SEQ_LEN构造含单自定义节点GatherDequantInt8的 ONNX 模型输入/输出类型与顺序和 GE 原型保持一致table uint8、indices int32、scale/zero_point fp16、y fp16opset 版本设为 11。gen_data.py 以固定随机种子生成量化后的 table、随机 indices、逐行 scale/zero_point并按 fp16 计算路径生成 goldenoutput.bin供精度比对使用。随后可通过 ATC 工具转换测试模型参考 ATC 工具使用指南调用鸿蒙维测接口完成单算子的精度与性能验证。总结GatherDequantInt8以一个自定义算子替代标准算子无法表达的 INT8 embedding 查表 反量化为切入点完整展示了鸿蒙端侧 Ascend C 算子开发的通用范式Host 侧原型注册与 Tiling 下发、Device 侧 double-buffer 流水核函数、ONNX 前端插件适配以及配套的单算子 ONNX 测试与 golden 数据生成链路。对需要在 Kirin 9020 上以 INT8 压缩 embedding 并保持 fp16 精度路径一致的开发者而言本样例既是一份可直接复用的算子实现也是学习端侧自定义算子全流程的最佳参考。【免费下载链接】cann-recipes-harmony-infer本项目为鸿蒙开发者提供基于CANN平台的业务实践案例方便开发者参考实现端云能力迁移及端侧推理部署。项目地址: https://gitcode.com/cann/cann-recipes-harmony-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考