CANN AMCT hifloat8_cast 算子解析FP16/BF16 与 HiFloat8 双向转换的 LUT 查表实现与工程实践【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct本文围绕 CANN AMCT 模型压缩工具仓中的amct_ops/hifloat8_cast算子展开系统讲解 HiFloat8 数据类型在昇腾 NPU 上实现 FP16/BF16 ↔ HiFloat8 双向转换的功能设计、Python/C 双层接口、LUT 查表内核原理、跨平台A2/A3/A5编译方法以及完整的性能与精度验证方案。读完本文你将掌握该算子的调用方式、底层查表与半空间 LUT 优化的实现细节并能够在 Atlas A2/A3 系列产品上独立完成算子编译、部署与测试。算子功能与四种转换模式hifloat8_cast是 AMCT 提供的 HiFloat8 数据类型转换算子支持 FP16/BF16 ↔ HiFloat8 双向转换。它面向大模型压缩与推理场景将 FP16/BF16 权重或激活编码为 1 字节的 HiFloat8 格式以降低内存占用和带宽开销需要时再解码回原精度。算子底层统一由设备端内核类KernelHiFloat8CastLut实现见 hifloat8_cast_kernel.cpp通过castMode区分 4 种转换模式castMode转换底层 LUT0FP16 → HiFloat832768 项半空间 UB-LUT32 KB1BF16 → HiFloat832768 项半空间 UB-LUT32 KB2HiFloat8 → FP16256 项 UB-LUT512 B3HiFloat8 → BF16256 项 UB-LUT512 B其中 castMode 的枚举定义可在 hifloat8_cast_tiling.h 中找到FP16_TO_HIF8 0、BF16_TO_HIF8 1、HIF8_TO_FP16 2、HIF8_TO_BF16 3。Python 接口会根据输入 dtype 自动选择 castMode无需手动指定dtype 为 float16 时走模式 0bfloat16 时走模式 1解码时依据目标 dtype 自动选择模式 2 或 3。设计上编码方向使用 32768 项大 LUT 覆盖 FP16/BF16 的全部幅值位型解码方向仅需 256 项 LUT 覆盖全部 HiFloat8 编码8 bit 全空间。编码 LUT 之所以采用半空间32768 项而非 65536 项利用的是 FP16/BF16 编码关于符号位的对称性后文原理章节会展开说明。接口说明算子对外暴露两个 Python 函数实现见 ops.py同时也可通过torch.ops.amct直接调用底层算子。encode_to_hifloat8(x: Tensor) - Tensor将 FP16/BF16 张量编码为 HiFloat8 格式。参数x输入张量NPU 设备dtype 为torch.float16或torch.bfloat16任意非空 shape返回torch.uint8张量shape 与输入相同异常RuntimeErrordtype 不是 float16 / bfloat16或输入不在 NPU 上示例y encode_to_hifloat8(x_fp16) y encode_to_hifloat8(x_bf16)decode_from_hifloat8(x: Tensor, dtype: dtype torch.bfloat16) - Tensor将 HiFloat8 格式解码为 FP16/BF16 张量。参数xHiFloat8 编码张量torch.uint8任意非空 shapedtype输出类型默认torch.bfloat16支持float16/bfloat16返回解码后的张量shape 与输入相同异常RuntimeError输入 dtype 不是 uint8或输出 dtype 不支持示例z decode_from_hifloat8(y) # 默认输出 BF16 z decode_from_hifloat8(y, torch.float16) # 指定输出 FP16从源码看这两个 Python 函数只是对torch.ops.amct.encode_to_hifloat8/torch.ops.amct.decode_from_hifloat8的薄封装dtype 校验、NPU 设备校验、castMode 选择以及真正的内核调度都发生在 C 侧。注册逻辑位于 register.cpp通过TORCH_LIBRARY_FRAGMENT(amct, m)定义 schema并通过TORCH_LIBRARY_IMPL(amct, PrivateUse1, m)注册 NPUPrivateUse1 后端实现、TORCH_LIBRARY_IMPL(amct, Meta, m)注册仅做 shape 推断的 meta 实现供 torch 的 fake tensor / shape 计算机制使用如EncodeMeta直接返回与输入同 shape 的 uint8 空张量。这也解释了为什么 Python 侧不做显式类型检查也能给出准确的异常提示。目录结构hifloat8_cast/ ├── op_kernel/ │ ├── hifloat8_cast_kernel.cpp # device 端 kernelLUT encode decode半空间 LUT 优化 │ └── hifloat8_cast_tiling.h # TilingData 结构体定义tileLength 由 host 运行时填入 ├── op_extension/ │ ├── hifloat8_cast_torch.cpp # PyTorch host 实现CPU 端 LUT 预计算与缓存调用 ASC 生成的 host stub │ ├── ops.h # PyTorch C 扩展函数声明 │ └── register.cpp # TORCH_LIBRARY 注册 ├── python/ │ └── hifloat8_cast/ │ ├── __init__.py # Python 包入口加载 .so、注册算子 │ └── ops.py # Python 接口实现 ├── CMakeLists.txt # CMake 构建配置 └── README.md # 算子说明文档各部分的职责边界很清晰op_kernelAscendC 设备侧实现包括内核入口hifloat8_cast_kernel_luthifloat8_cast_kernel.cpp与 TilingData 结构定义hifloat8_cast_tiling.h。TilingData 中的blockNumblock 数、totalLength总元素数、numPerCore每核元素数、tailNumLastCore末核尾部元素数、castMode、tileLength每 tile 元素数运行时根据平台 UB 大小确定由 host 侧填充。op_extensionPyTorch host 侧实现负责输入校验、LUT 的 CPU 端预计算与设备端缓存、tiling 参数推导并调用 ASC 编译器生成的内核启动 stub。pythonPython 包入口init.py 在导入时通过torch.ops.load_library加载libhifloat8_cast_ops.so并重新导出encode_to_hifloat8、decode_from_hifloat8两个函数。环境依赖SOC平台标称 UB / coreCANN 实测可用 UBascend910bAtlas A2 系列产品256 KB~192 KBascend910_93Atlas A3 系列产品512 KB~384 KBascend950Ascend950PR/Ascend950DT512 KB—CANN 运行时会保留部分 UBtileLength在每次调用时由GetCoreMemSize()查询实际可用量并计算无需手动指定。其余依赖CANN 9.0.0Python ≥ 3.9PyTorch torch_npu适配对应 CANN 版本从 host 实现可以印证上述机制GetUbSizeBytes()通过platform_ascendc::PlatformAscendCManager的GetCoreMemSize(CoreMemType::UB, ub)查询实际可用 UB见 hifloat8_cast_torch.cpp查询失败时才回退默认值 256 KBGetAivCoreNum()则查询可用的 AIV 核数默认 32 核见同文件 L261-L267。编译方式一统一打包推荐cd amct_ops bash ops_build.sh [--soc soc] [hifloat8_cast]--soc平台--npu-arch默认ascend910bA2910B1/B2/B3UB 256 KBdav-2201✓ascend910_93A3910_93UB 512 KBdav-2201ascend950A5UB 512 KBdav-3510A2/A3 共用同一 ISAdav-2201编译产物相同UB 大小差异由运行时平台 API 区分自动选择最优tileLength。 A5 构建要求当前 CANN 编译器支持dav-3510。如果bisheng报Unsupported NPU architecture or soc需要更换支持 A5 目标的 CANN 编译环境。bash ops_build.sh # 全部算子默认平台 bash ops_build.sh --soc ascend910_93 # 全部算子指定平台 bash ops_build.sh hifloat8_cast # 指定算子默认平台 bash ops_build.sh --soc ascend950 hifloat8_cast # 指定算子指定平台 pip install dist/amct_ops-*.whl方式二单算子独立编译开发调试针对单个算子进行独立编译适用于开发调试场景。通过-DNPU_ARCH指定目标平台默认 A2平台-DNPU_ARCHAtlas A2 系列产品dav-2201默认Atlas A3 系列产品dav-2201同 A2Ascend950PR/Ascend950DTdav-3510cd amct_ops/hifloat8_cast source $ASCEND_HOME_PATH/set_env.sh mkdir -p build cd build cmake .. -DNPU_ARCHdav-2201 # A2 / A3默认 # cmake .. -DNPU_ARCHdav-3510 # A5 make -j8编译产物build/libhifloat8_cast_ops.so从 CMakeLists.txt 可以看到构建的关键约束内核文件op_kernel/hifloat8_cast_kernel.cpp以 ASC 语言编译set_source_files_properties(... PROPERTIES LANGUAGE ASC)并通过--npu-arch${NPU_ARCH}指定目标指令集同时要求环境变量ASCEND_HOME_PATH已设置source set_env.sh即为此目的并自动探测 host 架构目录x86_64-linux或aarch64-linux。链接依赖包含torch_npu、tiling_api、ascendcl、ascendc_runtime等 CANN 与 torch_npu 库这正是-DNPU_ARCH之外还必须正确配置 CANN 环境的原因。已知 CMake 告警使用 pip 安装的 PyTorch 构建时find_package(Torch)可能输出如下告警static library kineto_LIBRARY-NOTFOUND not found.该告警来自 PyTorch 自带的TorchConfig.cmake表示未找到 Kineto profiler 的静态库。hifloat8_cast不依赖 PyTorch profiler / Kineto 能力只要 CMake configure、编译和链接成功该告警可以忽略。使用示例import torch import torch_npu from amct_ops.hifloat8_cast import encode_to_hifloat8, decode_from_hifloat8 # 编码FP16/BF16 → HiFloat8 x torch.randn(1024, 256, dtypetorch.bfloat16, devicenpu) y encode_to_hifloat8(x) # shape [1024, 256], dtype uint8 # 解码HiFloat8 → FP16/BF16 z decode_from_hifloat8(y, torch.bfloat16) # shape [1024, 256], dtype bfloat16也可直接通过torch.ops.amct调用底层算子y torch.ops.amct.encode_to_hifloat8(x) z torch.ops.amct.decode_from_hifloat8(y, torch.bfloat16)注意使用前必须先import torch_npuPython 包init.py 内部已引入torch_npu以注册 PrivateUse1 后端。调用链整体为Python 薄封装 →torch.ops.amct.*经TORCH_LIBRARY分发到 PrivateUse1 实现→AscendKernel::Hifloat8CastTorch校验、LUT 获取、tiling 构建、启动内核最终落到设备端KernelHiFloat8CastLut。源码级实现原理LUT 查表策略以内存换指令编码方向FP16/BF16 → HiFloat8本质上是位型到位型的非线性映射涉及指数段切换、舍入与饱和直接标量计算指令开销大。算子选择在 host 侧CPU一次性构建 32768 项编码 LUT 并缓存在 device 侧设备端每条数据只需几次查表与位运算即可完成转换encodecastMode 0/1每元素约 5 条指令AND 取幅值 LUT 查表 SHR 取符号 compare OR 叠加符号位对应ComputeEncode16hifloat8_cast_kernel.cppdecodecastMode 2/3每元素仅 1 条 LUT 查表指令256 项 LUT 直接消除所有分支对应ComputeDecode同文件 L184-L194。半空间 LUT 优化FP16/BF16 编码关于符号位对称encode(-x) encode(x) | 0x80当encode(x) ≠ 0时因此 LUT 只存储正半空间32768 项幅值编码键为去掉符号位的 15 bit 幅值位型设备端 compute 时先剥离符号、查表、再叠加符号位。这一设计带来两个收益见内核文件头部注释 L40-L44节省 32 KB UB32 KB 的 LUT 数据可通过一次DataCopyPad完成搬入无需分两次。注意两个边界细节内核注释与实现均有体现下溢/零时mag 0不叠加符号位保持0x00encode(-0) 0x00。此外A2 平台编译器不支持Castuint32_t, uint16_t及uint16_t/uint32_t的移位指令无法用 Gather 向量化因此当前 BF16/FP16 encode 为标量循环——但每步只需约 5 条 UB 查表位运算指令开销可控。Host 侧 LUT 预计算与缓存32768 项编码 LUT 与 256 项解码 LUT 都在 CPU 端构建逻辑位于 hifloat8_cast_torch.cppBuildLut16CpuL221-L257把 FP16/BF16 的 15 bit 幅值位型先规整到 FP32再调用HostFp32MagnitudeToHif8得到 HiFloat8 幅值编码。FP16 的 subnormal 通过__builtin_clz归一化处理exp 全 1Inf/NaN显式映射到 FP32 exp255避免丢特殊值语义超出 HiFloat8 表示范围时饱和到最大正有限值0x6E过小值归零denormal 边界值保留最小非零编码0x01。HostFp32MagnitudeToHif8L164-L216实现 HiFloat8 的动态指数段dot value映射HostGetHif8BitsNumL139-L162根据指数幅值选择 1/2/3/4 位指数 3/3/2/1 位尾数的组合舍入时多取 1 bit尾数溢出进位到下一个指数段。BuildLut8CpuL123-L135对 256 个 HiFloat8 编码逐个用HostHif8ToFpBits计算 FP16/BF16 位型其中0x80映射为 NaN、0x6F/0xEF映射为 ±Inf、0x00映射为 0。GetOrBuildLutOnDeviceL353-L376以(deviceIndex, castMode)为键做进程级缓存首次使用时构建并搬到 device此后直接复用避免重复构建开销。Tiling 与核数自适应tileLength由 host 在每次调用时根据平台实际 UB 大小动态计算ComputeMaxTileLengthL278-L295每个元素在 UB 中占 3 字节encode 输入 2B 输出 1Bdecode 输入 1B 输出 2B单缓冲TQue depth1、InitBuffer num1——因为 compute 为标量循环与 MTE 流水重叠收益为 0释放队列另一半 UB 给 tilemaxTile (ubBytes - lutBytes) / 3上限 65536≥ 32768 时对齐到 32768 的整数倍使CopyIn大块搬运走{n, 32768}分支避免DataCopyParams.blockLenuint16_t最大 65535溢出否则对齐到 32向量指令粒度核数分配遵循核数优先策略ComputeNumBlocksL297-L303按每核最少处理量encode 4096 元素、decode 2048 元素计算 block 数小数据时不会让所有核都去加载 LUT避免过大的 LUT 加载开销。设备端CopyIn依据上述对齐特性实现了三种DataCopyParams拆分策略hifloat8_cast_kernel.cppbyteCount ≤ 65535用{1, byteCount}可整除 32768 用{byteCount/32768, 32768}其余偶数用{2, byteCount/2}保证blockLen始终不越界。性能验证测试平台Ascend 910B3A2ascend910bCANN 9.0.0吞吐量定义(输入字节 输出字节) / 耗时100 次迭代均值10 次预热NPU synchronize 计时推荐场景数据量 ≥ 256K 时 NPU 优势显著BF16 ↔ HiFloat8数据大小Encode (ms)吞吐 (MB/s)Decode (ms)吞吐 (MB/s)1K0.20714.90.22813.54K0.21557.20.21956.016K0.215228.40.255192.964K0.215914.00.290677.6256K0.2892723.90.2932679.81M0.8673627.20.8633645.74M3.1334016.73.1114044.316M12.2364113.412.1524141.7FP16 ↔ HiFloat8数据大小Encode (ms)吞吐 (MB/s)Decode (ms)吞吐 (MB/s)1K0.20814.80.22413.74K0.21656.90.22355.116K0.215228.60.257191.664K0.217904.10.289680.9256K0.2882734.80.2932687.81M0.8673628.10.8633645.24M3.1304019.63.1214032.216M12.2364113.412.1484143.1RoundtripEncode Decode数据大小FP16 (ms)吞吐 (MB/s)BF16 (ms)吞吐 (MB/s)1K0.38910.50.39210.54K0.39241.80.39341.716K0.394166.20.390168.164K0.388674.80.389673.7256K0.5481911.90.5541891.61M1.7142446.71.7162444.94M6.2462686.26.2492684.816M24.3932751.124.4012750.3小数据 256K吞吐较低但已通过核数优先策略优化避免过多核加载 LUT 开销。 大数据≥ 4Mencode/decode 吞吐均超过 4 GB/s接近 HBM 带宽上限。 tileLength 由运行时根据平台实际 UB 大小和核数动态计算。精度验证精度验证通过以下方式进行NPU API 综合测试类型推断、边界值、特殊值、非规格化数、全 256 个 HiFloat8 decode 完备性见 test_hifloat8_cast.py边界/特殊值±0、±Inf、NaN、最大值、最小非规格化数等边界情况。测试场景结果随机 FP16/BF16 roundtrip✓ 相对误差在 HiFloat8 精度范围内边界值 roundtrip✓ 覆盖 0、正负数、小数、较大值特殊值±0 / ±Inf / NaN✓ 编码字节和解码语义符合规范全 256 个 HiF8 值解码FP16 BF16✓ 仅 0x80 为 NaN 编码注FP16 subnormal 编解码含 HiFloat8 非规格化值映射到 FP16 非规格化值已全部验证正确。测试用例与实现细节可以互相印证例如test_special_values断言0/-0均编码为0x00、Inf为0x6F、-Inf为0xEF、NaN为0x80与内核的零不下溢叠加符号、HostHif8ToFpBits的特殊值映射一致test_fp16_subnormal_rounding_regression验证 FP16 非规格化值2.7180e-05编码为0x7E即HIF8_POSITIVE_MIN_NORMAL定义的 HiFloat8 最小正规格化数test_full_256_decode验证全 256 个编码中仅0x80解码为 NaNtest_type_inference/test_decode_output_type则覆盖了输入 dtype 校验与输出 dtype 推断。测试方法# 在仓库根目录执行。先构建 amct_ops再通过 staging 运行测试。 bash amct_ops/ops_build.sh hifloat8_cast # Python API 综合测试类型推断、边界值、特殊值、非规格化数 PYTHONPATHamct_ops/staging python3 -m unittest tests.amct_ops.test_hifloat8_cast也可以先安装 wheel 后再执行测试pip install amct_ops/dist/amct_ops-*.whl python3 -m unittest tests.amct_ops.test_hifloat8_cast更多测试执行说明见 tests/amct_ops/README.md。小结hifloat8_cast以host 预计算 LUT device 查表的思路将 FP16/BF16 与 HiFloat8 的双向转换做到每元素仅 15 条指令并通过半空间 LUT、单缓冲 tile 化、运行时 UB/核数自适应 tiling、进程级 LUT 缓存等手段在 A2/A3/A5 多平台下取得接近 HBM 带宽上限≥ 4M 数据时超 4 GB/s的吞吐表现。其接口设计兼顾易用性dtype 自动选择 castMode与灵活性torch.ops.amct底层直调并配套了覆盖类型推断、边界值、特殊值、非规格化数与全 256 编码解码完备性的 NPU 实测用例。对于在昇腾 NPU 上做 HiFloat8 低比特压缩的开发者该算子既是一个可直接复用的转换工具也是一份理解 LUT 查表算子在 AscendC 上落地优化的完整参考实现。【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
