ARTICLE DETAIL

资讯详情

深耕商务建站与企业官网运营的一线实战洞察。

昇腾 NPU 上的 HiFloat4 4-bit 块缩放伪量化算子:AMCT `amct_ops/hifloat4_cast` 完全解析

昇腾 NPU 上的 HiFloat4 4-bit 块缩放伪量化算子:AMCT `amct_ops/hifloat4_cast` 完全解析 昇腾 NPU 上的 HiFloat4 4-bit 块缩放伪量化算子AMCTamct_ops/hifloat4_cast完全解析【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct本技术指南以 CANN AMCT 仓库中amct_ops/hifloat4_cast算子包为核心系统讲解 HiFloat4HiF44-bit 块缩放浮点格式的原理、FP→HiF4→FP 伪量化仿真的接口用法、AscendC 设备端内核的逐级量化流水线以及面向 A2/A3/A5 平台的编译与测试方法。读完本文你将掌握hifloat4_fake_quant的完整调用方式与内部实现机制能够在自己的量化部署流程中直接复用该算子并理解其与 CPU 参考实现、测试用例之间的对应关系。HiFloat4 格式概述S1P2 元素 三级块缩放HiFloat4HiF4是一种 4-bit 块缩放浮点格式用于在保持可用精度的前提下把 FP16/BF16 权重或激活压缩到平均每元素 4 bit 的量级。其编码分为两层元素级 S1P2 表示每个元素使用 1 位符号sign 2 位尾数mantissa组内可表达的数值幅度共 8 种S1PNg其中Ng2即组内编码占 3 bit。由 hifp_impl.py 中的常量可知E6_OFFSET 48、E6M2_MBITS 2、NG 2与 README 描述的 S1P2 语义完全一致。三级缩放block scaling沿量化维每 64 个元素共享一组缩放因子L1scale_factor覆盖整个 64 元素块以 E6M2 格式存储由块内最大值推导得出L2每 8 个元素一个微指数micro exponent取值 ∈ {1, 2}L3每 4 个元素一个微指数取值 ∈ {1, 2}。需要特别强调的是该算子只实现伪量化仿真fake-quant即完成一次 FP → HiF4 → FP 的量化再反量化往返输出与输入同 shape、同 dtype 的浮点张量并不产生真实的打包 4-bit 数据。这种仿真方式适合在训练或 PTQ 标定阶段以可导/无损替换的方式评估量化误差。算子能力按输入 dtype 选择底层内核hifloat4_cast根据输入张量的 dtype 分发到不同的设备端内核两条路径在内部均以 fp32 计算差异只在转换时机与 bf16 舍入位置输入 dtype底层内核内部计算方式BF16hifx_kernel_bf16fp32加载后转 fp32仅在指定计算点scale factor、倒数、组内值做 bf16 舍入输出时转回 bf16FP16hifx_kernelfp32先在主机侧转 fp32计算完成后输出转回 fp16该分派逻辑可在 hifloat4_cast_torch.cpp 中确认cvt_fp32对 FP16 输入在主机侧完成 fp32 提升随后按in_dtype at::kBFloat16选择run_hifx_kernel_bf16否则走run_hifx_kernel。BF16 路径则由内核内的vconv_bf162f32在加载后完成转换见 hifloat4_cast_kernel.cpp。接口说明hifloat4_fake_quant(x, qdim-1)函数签名与语义hifloat4_fake_quant(x: torch.Tensor, qdim: int -1) - torch.Tensor沿qdim维将 FP16/BF16 张量按每 64 个元素一个共享 scale 分块执行一次 FP → HiF4 → FP 伪量化仿真。参数xNPU 设备上的输入张量dtype 为torch.float16或torch.bfloat16支持任意非空 shapeqdim分块维度默认-1最后一维。返回值与输入同 shape、同 dtype 的张量包含 HiFloat4 量化误差。异常RuntimeErrordtype 不是 float16 / bfloat16或输入张量不在 NPU 设备上RuntimeErrorqdim维长度不是 64 的倍数。示例from amct_ops.hifloat4_cast import hifloat4_fake_quant y hifloat4_fake_quant(x) # 沿最后一维分块 y hifloat4_fake_quant(w, qdim1) # Linear 权重 [out, in] 沿 in 维分块 y torch.ops.amct.hifloat4_fake_quant(x, -1) # 等价的 torch.ops 调用主机侧处理流程从 hifloat4_cast_torch.cpp 的实现可以看出主机侧完成以下五步维度校验qdim支持负数索引qdim 0 ? qdim : qdim nd越界会抛RuntimeErrorPermuteToLast把qdim维移动到最后一轴input.permute(perm).contiguous()使内核始终沿最后一维归约若qdim本来就是末维则仅做contiguous()64 对齐校验 512 对齐补零先校验末维长度是 64 的倍数否则抛错再通过PadLastDim以 0 补齐到 512 的倍数——由于 0 不影响块最大值补零是安全的内核调用以HIF4_BLOCK_DIM 40的 blockDim、HIF4_MANT_BIT 3调用对应内核还原slice(-1, 0, n)裁掉补零必要时把结果Unpermute回原始维度顺序。源码级原理剖析Python 包装层加载 .so 并注册算子ops.py 中hifloat4_fake_quant是一个极薄的包装直接转发给torch.ops.amct.hifloat4_fake_quant(x, qdim)init.py 在包导入时通过torch.ops.load_library加载同目录下的libhifloat4_cast_ops.so这是torch.ops.amct命名空间能够解析的前提导入前先import torch_npu以注册 PrivateUse1 后端。算子注册schema、NPU 实现与 Meta 实现register.cpp 展示了完整的注册骨架TORCH_LIBRARY_FRAGMENT(amct, m) { m.def(hifloat4_fake_quant(Tensor input, int qdim-1) - Tensor); } TORCH_LIBRARY_IMPL(amct, PrivateUse1, m) { ... } // NPU 实现dtype/device 校验 TORCH_LIBRARY_IMPL(amct, Meta, m) { ... } // 形状推导输出与输入同 shape/dtype其中PrivateUse1 实现负责设备与 dtype 校验允许 float16 / bfloat16 / float32Meta 实现声明输出形状与输入完全一致因此 shape 推断是恒等映射。设备端内核12 级计算流水线内核主体是一个手工编写的 AscendC 模板类Hifv14KernelT, E6MX其 Compute() 将一次 512 元素批BATCH 512向量指令按 64 元素段SEG_ELEMS 64操作的量化拆分为 12 个小阶段LoadInput → ZeroTail → PrepareSignAbs → HandleNonFinite → ReduceBlockMax → PoisonBlock → RoundScaleE6M2 → ComputeExp1 → Exp2PartMax → Exp2Select → ComputeMantissa → Reconstruct → StoreOutput各阶段核心语义与 hifloat4_cast_kernel.cpp 头部注释及实现一致ReduceBlockMaxL1 scalescale_factor round_E6M2(bf16(max/7))并钳制到[2^-48, 49152]vcmax归约块最大值、vbrcb广播随后通过vconv_f322bf16r/vconv_bf162f32完成 bf16 舍入ComputeExp1L2 微指数exp1 (max8 * rec 4) ? 2 : 1其中rec 1/scale经过 bf16 舍入Exp2PartMax / Exp2SelectL3 微指数先分别求 4 元素组的奇偶半组最大值vcgmax再判定exp2 (max4 / exp1 * rec 2) ? 2 : 1ComputeMantissaS1P2 组内编码mant abs / exp2 / exp1 * rec先做 bf16 舍入再通过*4 0.5后截断实现 round-half-up最终钳制到±1.75即±(2 - 2^-2)Reconstructout sign * mant * exp2 * exp1 * scale四步连乘还原。非有限值处理块级 NaN 投毒Poison BlockREADME 未展开但源码注释与测试明确强调的细节是含 NaN/±Inf 的块整体反量化为 NaN。其处理策略值得关注在HandleNonFinite中非有限元素对块最大值贡献被掩蔽为 0vsel选择 0避免污染 scale块内任意元素非有限则该块的干净标志位被置脏PoisonBlock但 NaN 替换被延迟到 Reconstruct 输出阶段执行而不是提前注入 scale——因为 NaN 一旦流入vdiv/vconv会触发 AI Core 异常测试注释中提到异常码 507057导致后续所有 launch 失败。这一点在 test_hifloat4_cast.py 的test_kernel_nan_block_poison_keeps_device_healthy回归测试中有完整覆盖含 NaN/Inf 的块输出全 NaN干净块与 CPU 参考逐位一致且后续正常输入仍可正常启动内核。与 CPU 参考实现的关系算子正确性以纯 PyTorch 参考实现为基准参考实现在 hifp_impl.py 中hif4_pack/hif4_unpackL260-L334实现真实的 2-D HiF4 打包/解包输出 scale value 两类 uint8 字节其校验规则n_cols % 64 0、仅支持 2-D与 NPU 算子保持一致hifloat4_fake_quantL347-L357是一个分发函数当输入在 NPU 且为 FP16/BF16 时优先调用amct_ops内核否则回退到纯 torch 参考实现。这意味着即使没有安装amct_ops同一接口仍可在 CPU 上完成等效仿真。目录结构amct_ops/hifloat4_cast/ ├── op_kernel/ │ ├── hifloat4_cast_kernel.cpp # 设备端内核FP→HiF4→FP 仿真往返 │ └── tensorutils.h # Ascend C 工具函数内核侧张量操作 ├── op_extension/ │ ├── hifloat4_cast_torch.cpp # PyTorch 主机实现qdim 预处理与内核调用 │ ├── ops.h # PyTorch C 扩展函数声明 │ └── register.cpp # TORCH_LIBRARY 注册 ├── python/ │ └── hifloat4_cast/ │ ├── __init__.py # Python 包入口加载 .so 并注册算子 │ └── ops.py # Python 接口实现 ├── CMakeLists.txt # CMake 构建配置 ├── README.md # 算子文档中文 └── README_en.md # 算子文档英文环境依赖CANN 9.0.0Python ≥ 3.9PyTorch torch_npu与对应 CANN 版本匹配构建时还需保证ASCEND_HOME_PATH已通过source $ASCEND_HOME_PATH/set_env.sh设置——CMakeLists 中会对其显式检查并报错见 CMakeLists.txt。编译方法方法一统一打包构建推荐在amct_ops目录下执行cd amct_ops bash ops_build.sh [--soc soc] [hifloat4_cast]平台与 NPU 架构对照--soc平台--npu-arch默认ascend910bA2910B1/B2/B3UB 256 KBdav-2201✓ascend910_93A3910_93UB 512 KBdav-2201ascend950A5UB 512 KBdav-3510A2/A3 共用同一 ISAdav-2201构建产物相同UB 差异由运行时GetCoreMemSize()解决见 CMakeLists.txt 注释。 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 hifloat4_cast # 指定算子默认平台 bash ops_build.sh --soc ascend950 hifloat4_cast # 指定算子指定平台 pip install dist/amct_ops-*.whl方法二单算子独立编译开发调试通过-DNPU_ARCH指定目标平台默认 A2平台-DNPU_ARCHAtlas A2 系列产品dav-2201默认Atlas A3 系列产品dav-2201与 A2 相同Ascend950PR / Ascend950DTdav-3510cd amct_ops/hifloat4_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/libhifloat4_cast_ops.so已知 CMake 警告使用 pip 安装的 PyTorch 构建时find_package(Torch)可能输出如下警告static library kineto_LIBRARY-NOTFOUND not found.该警告来自 PyTorch 自带的TorchConfig.cmake表示未找到 Kineto profiler 的静态库。hifloat4_cast不依赖 PyTorch profiler / Kineto 能力只要 CMake configure、编译与链接成功此警告可安全忽略。使用示例import torch import torch_npu from amct_ops.hifloat4_cast import hifloat4_fake_quant # 伪量化仿真FP16/BF16 → HiF4 → FP16/BF16 x torch.randn(1024, 256, dtypetorch.bfloat16, devicenpu) y hifloat4_fake_quant(x) # shape [1024, 256], dtype bfloat16 # 指定分块维度 w torch.randn(4096, 1024, dtypetorch.float16, devicenpu) y hifloat4_fake_quant(w, qdim1) # 沿 in_features 分块也可通过torch.ops.amct直接调用底层算子y torch.ops.amct.hifloat4_fake_quant(x, -1)精度验证精度验证方法详见 tests/amct_ops/test_hifloat4_cast.py测试场景结果随机 BF16/FP16 伪量化往返仿真✓ 与 CPU 参考实现逐元素一致max abs diff 1e-6非 64 对齐维度✓ 抛出 RuntimeErrorfake-quant 与 pack 行为一致零输入✓ 无 NaN输出全零NaN/Inf 块投毒回归✓ 非有限块整体为 NaN干净块与参考逐位一致设备健康其中逐元素一致性测试_assert_kernel_matches_reference将 NPU 内核输出与hifloat4_fake_quant纯 torch 参考在 1e-6 容差内比较CPU-only 的 shape 回归测试还覆盖了hif4_pack的 scale/value 打包形状如 2x128 → scale (2,2,4) value (2,64)并守护了 squeeze 索引错位的历史 bug。测试方法在仓库根目录执行# 先构建 amct_ops再通过 staging 运行测试 bash amct_ops/ops_build.sh hifloat4_cast # NPU 内核与 CPU 参考实现逐元素一致性验证 shape 回归测试 PYTHONPATHamct_ops/staging python3 -m unittest tests.amct_ops.test_hifloat4_cast也可以先安装 wheel 再运行测试pip install amct_ops/dist/amct_ops-*.whl python3 -m unittest tests.amct_ops.test_hifloat4_cast更多测试执行说明参见 tests/amct_ops/README.md。小结amct_ops/hifloat4_cast以“单次 FP→HiF4→FP 伪量化仿真”为核心将 S1P2 元素编码与 L1/L2/L3 三级 64 元素块缩放完整落地为 AscendC 内核并提供了从 Python 接口、torch.ops 注册、主机侧预处理到设备端 12 级量化流水线的全链路实现。无论是评估 HiFloat4 对特定层权重的量化误差、在 PTQ 标定中做误差仿真还是将其作为打包算子的精度基准hifloat4_fake_quant都是一份开箱即用、可独立编译与验证的参考实现其与 hifp_impl.py 中的 CPU 参考保持同一套语义为跨端一致性与后续训练/推理集成提供了可靠基础。【免费下载链接】amctAMCT是CANN提供的昇腾AI处理器亲和的模型压缩工具仓。项目地址: https://gitcode.com/cann/amct创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表
PREV
查看更多资讯
NEXT
返回资讯列表