TorchAO 为 Arm CPU 带来了高性能低位线性算子和嵌入算子。在本次更新中，我们很高兴分享三项重大改进：动态内核选择、与 Arm 的 KleidiAI 库集成，以及对量化绑定嵌入（Quantized Tied Embeddings）的支持——所有这些改进旨在提升 PyTorch 中低位推理的性能并扩大其覆盖范围，包括ExecuTorch（PyTorch 用于高效端侧执行的解决方案）。
事实上，借助 KleidiAI 内核，我们在 M1 Mac 上运行 4 位量化 Llama1B 模型时，预填充（prefill）性能提升了超过 2 倍（达到 373 tokens/秒）！
动态内核选择
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 集成到了我们的动态内核选择系统中。在支持的情况下（例如，具有 4 位权重的 8 位动态激活），KleidiAI 内核会自动注册并使用。当出现覆盖盲区时（例如非 4 位权重或绑定嵌入层），我们会回退到我们自研的 GEMV neondot 内核——无需任何配置。这一切都是使用前面讨论过的打包权重格式完成的。
这种混合方法让我们兼得两者的优点：来自 KleidiAI 的顶级性能，以及 torchao 内核带来的广泛算子支持。
通过 KleidiAI，我们观察到解码性能与现有的 torchao 内核相当。然而，由于我们没有自研的 GEMM 内核，使用 KleidiAI 显著提升了预填充性能——在 M1 Mac 上使用 ExecuTorch 时实现了超过 2 倍的加速，速度超过 373 tokens/秒！
量化绑定嵌入与 lm_head 内核
绑定嵌入（Tied Embeddings）广泛应用于 LLaMA 3.2 和 Phi-4 Mini 等紧凑型大模型中。在绑定嵌入中，嵌入层的权重与计算 Logits 的最终线性层（lm_head）的权重是共享的：
现代模型通常使用超过 100,000 个标记的词汇表大小，而在较小的模型中，嵌入层和 lm_head 层可能占据模型总大小的很大一部分。
然而，在移动设备上，这些权重有时不得不被重复存储。这是因为高效的线性内核和嵌入内核要求权重以不同的格式进行打包，导致在不牺牲性能的情况下进行权重共享变得不切实际。
为了解决这个问题，我们为嵌入操作和线性操作开发了高效的量化内核，这些内核使用统一的权重格式。我们通过新的 SharedEmbeddingQuantizer 来实现这一点，它允许相同的量化权重同时用于输入嵌入和输出 lm_head。这可以在不影响性能的情况下减少模型大小。该量化器支持多种配置，包括：
- 8 位动态激活
- x 位权重（x 的范围从 1 到 8）
快来尝试并参与贡献吧！
所有这些增强功能现已通过 torchao 的量化 API 提供，并且已经集成到 ExecuTorch 中，用于将量化模型高效部署到移动设备和边缘设备。
我们非常期待您的尝试、反馈和贡献！