TorchAO 将高性能低位线性算子和嵌入算子带到 Arm CPU。在本次更新中，我们很高兴分享三项重大改进：动态内核选择、与 Arm 的 KleidiAI 库集成，以及对量化绑定嵌入的支持——所有这些都旨在提升性能并扩展 PyTorch 中低位推理的覆盖范围，包括ExecuTorch，PyTorch 用于高效设备端执行的解决方案。
事实上，借助 KleidiAI 内核，我们在 M1 Mac 上对 4 位量化的 Llama1B 进行预填充时，性能提升了两倍以上（373 token/秒）！
动态内核选择
TorchAO 低位算子现在根据以下因素自动选择最佳可用内核：
- 打包权重的格式，
- CPU 特性，例如
has_arm_neon_dot和has_arm_i8mm，以及 - 激活张量的形状。
这种动态调度允许我们根据硬件和工作负载特性定制执行。
它是如何工作的？
量化模型权重可以打包成针对特定线性内核优化的格式。这些打包权重包含描述其格式的头部。当使用打包权重调用线性算子时，我们首先检查权重格式和当前的 CPU 功能。根据这些信息，我们确定一组可以在权重上操作的兼容内核，并将这些内核的函数指针缓存到注册表中。
例如，以 format1 打包的权重可能支持 GEMV 和 GEMM 内核（例如，gemv_kernel1 和gemm_kernel1），而以 format2 打包的权重可能只支持 GEMV 内核（例如，gemv_kernel2）。在这种情况下，内核注册表可能如下所示：
下次我们遇到带有format1的打包权重时，我们可以从表中快速检索兼容的内核，并根据激活的形状调度到适当的内核。如果激活形成一个向量，我们将使用gemv_kernel1；如果它们形成一个矩阵，我们将使用gemm_kernel1。
想查看哪些内核处于活动状态？设置 TORCH_CPP_LOG_LEVEL=INFO 以获取详细视图。
KleidiAI 集成
KleidiAI 是 Arm 提供的一个开源库，为 Arm CPU 提供高度优化的微内核。我们现在将 KleidiAI 集成到我们的动态内核选择系统中。在支持的情况下（例如，8 位动态激活与 4 位权重），KleidiAI 内核会被注册并自动使用。当出现覆盖空白时（例如，非 4 位权重或绑定的嵌入层），我们回退到我们内部的 GEMV neondot 内核——无需配置。所有这些都使用前面讨论的打包权重格式完成。
这种混合方法让我们两全其美：KleidiAI 提供顶级性能，而 torchao 的内核提供广泛的算子支持。
通过 KleidiAI，我们观察到解码性能与现有的 torchao 内核相当。然而，由于我们没有内部的 GEMM 内核，使用 KleidiAI 可以显著提升预填充性能——实现超过 2 倍的加速，在 M1 Mac 上使用 ExecuTorch 时的速度超过 373 token/秒！
量化绑定嵌入和 lm_head 内核
绑定嵌入广泛用于紧凑型 LLM，例如 LLaMA 3.2 和 Phi-4 Mini。在绑定嵌入中，嵌入层的权重与计算 logits 的最终线性层（lm_head）的权重共享：
最近的模型通常使用超过 100,000 个词元的词汇量，在较小的模型中，嵌入层和 lm_head 层可能占模型总大小的很大一部分。
然而，在移动设备上，这些权重有时会无论如何都被复制。这是因为高效的线性内核和嵌入内核需要以不同的格式打包权重，这使得在不牺牲性能的情况下共享权重不切实际。
为了解决这个问题，我们为嵌入和线性操作开发了高效的量化内核，这些内核使用统一的权重格式。我们通过新的SharedEmbeddingQuantizer公开了这一点，它允许相同的量化权重用于输入嵌入和输出lm_head。这在不影响性能的情况下减小了模型大小。量化器支持广泛的配置，包括：
- 8 位动态激活
- x 位权重，其中 x 范围从 1 到 8
立即尝试并贡献！
所有这些增强功能现在都可以通过 torchao 的量化 API 获得——并且它们已经集成到ExecuTorch中，用于将量化模型高效部署到移动和边缘设备。
我们希望您能尝试一下，分享反馈，并做出贡献！