我们很高兴宣布在 TorchAO 中添加了低比特权重（1-8 比特）的嵌入算子和具有 8 比特动态量化和低比特权重（1-8 比特）的线性算子，适用于 Arm CPU。这些算子可以在所有 PyTorch 表面上无缝工作，包括 eager、torch.compile、AOTI 和 ExecuTorch，并在 torchchat 中可用。
在开发这些线性算子时，我们的重点是 PyTorch 和 ExecuTorch 之间的代码共享，以及高级算子和低级内核之间的明确边界。这种设计允许第三方供应商轻松地更换他们自己的内核。我们还致力于创建一个实验新 CPU 量化和测试这些想法在 PyTorch 生态系统中的基础设施。
通用低比特内核
没有硬件支持低比特算术。在我们所说的通用内核中，我们明确地将解包低比特值到 int8 值的逻辑和 int8 GEMV 内核逻辑以模块化的方式分开。我们从一个 8 比特内核开始，例如这个 1x8 8 比特 GEMV 内核，它使用了 Arm neondot 指令。在 8 比特内核中，我们调用内联的解包例程将低比特值转换为 int8 值。这个解包例程是强制内联的，并且针对某些低比特值进行了模板化。我们的实验表明，使用单独的强制内联解包例程和直接内联解包代码在性能上没有差异。
这种模块化设计的优势在于提高了开发速度和代码可维护性。在编写了 8 位内核之后，我们通过编写简单的位打包例程迅速实现了全低位覆盖。实际上，从事位打包例程的开发者不需要成为 GEMV/GEMM 内核编写的专家。我们还重用了嵌入内核中的线性内核的相同位打包例程。在未来，我们可以重用相同的位打包例程用于通用 GEMM 内核或基于 fma 或 i8mm 指令的内核。
PyTorch 和 ExecuTorch 之间的共享代码
为了实现 PyTorch 和 ExecuTorch 之间的共享代码，我们使用原始指针而不是 PyTorch 张量编写了内核。此外，我们在一个头文件中实现了线性算子，该头文件包含在独立的 PyTorch 和 ExecuTorch 算子注册代码中。通过仅使用 ATen 和 ExecuTorch 张量共有的功能，我们确保了两个框架之间的兼容性。对于多线程计算，我们引入了 torchao::parallel_1d，它根据编译时标志编译为 at::parallel_for 或 ExecuTorch 的 threadpool。
可交换的内核
我们为高级多线程线性算子设计的方案对底层单线程内核是透明的，允许第三方供应商替换他们自己的实现。算子与内核之间的接口由 ukernel 配置定义，该配置指定了用于准备激活数据、准备权重数据和运行内核的内核函数指针。负责分块和调度的算子仅通过此配置与内核交互。
性能
下表展示了在配备 32GB RAM 的 M1 Macbook Pro 上使用 6 个 CPU 线程时，Llama3.1 8B token 生成性能。
| 位宽 x | torch.compile（解码令牌/秒） | ExecuTorch (解码令牌/秒) | ExecuTorch PTE 大小（GiB） |
| 1 | 24.18 | 17.86 | 1.46 |
| 2 | 27.02 | 19.65 | 2.46 |
| 3 | 21.01 | 22.25 | 3.46 |
| 4 | 19.51 | 19.47 | 4.47 |
| 5 | 14.78 | 16.34 | 5.47 |
| 6 | 12.80 | 13.61 | 6.47 |
| 7 | 8.16 | 11.73 | 7.48 |
结果在配备 8 个性能核心和 2 个效率核心的 M1 Macbook Pro（32GB RAM，6 线程）上使用 torchchat 运行。每个测试中，生成了最大序列长度为 128 个令牌的结果。对于每个位宽 x，嵌入层被分组量化为 x 位，分组大小为 32。在线性层中，激活项按令牌动态量化为 8 位，权重被分组量化为 x 位，分组大小为 256。我们在这里关注性能，不报告准确率或困惑度数值。根据模型的不同，较低的位宽可能需要量化感知训练、混合位宽的模型量化或调整分组大小以获得可接受的准确率。
尝试它们并做出贡献！
如果你想看到新的低比特内核的实际效果，请尝试设置 torchchat，并使用内核在本地运行LLM进行量化。
如果你想贡献力量，可以考虑添加以下领域的支持：
- 为 Arm CPU 添加通用的低比特 GEMM 内核，重用通用的 GEMV 内核的相同位打包例程。
- 改进基于 ISA、打包格式和激活形状的 ukernel 配置的运行时选择。
- 为其他 CPU 指令集架构如 x86 添加低比特内核。
- 将第三方库如 KleidiAI 与操作框架集成。
