摘要

本文探讨了在CPU上优化大型语言模型（LLM）推理性能的问题，特别是在GPU资源受限的环境中。文章提出了一种易于部署的解决方案，通过减少KV缓存大小并确保精度，以及实施分布式推理优化，来加速LLM在CPU上的推理过程。该解决方案支持多种常用LLM模型，并通过实验验证了其在CPU上的良好推理性能可扩展性。

原理

本文提出的优化方法主要包括三个方面：LLM优化、有效的KV缓存优化和分布式推理优化。在LLM优化中，文章提出了一种名为SlimAttention的新方法，通过一维分解查询和键之间的得分，减少内存使用并提高计算效率。在KV缓存优化中，文章实施了一种INT8 KV缓存方法，通过为每个令牌和每个头维护一个独特的比例，有效地减少了KV缓存的大小，同时保持了模型的输出质量。分布式推理优化则通过使用oneAPI Collective Communications Library（oneCCL）来实现，通过广播令牌ID而不是嵌入部分的值，以及在每个工作者计算top-k后进行减少，而不是直接减少所有令牌的对数，来提高分布式推理的性能。

流程

文章详细描述了SlimAttention的工作流程，首先计算水平得分，然后对得分应用softmax，并将结果与相应的值相乘，生成部分输出。这一过程在下一个块中重复使用相同的得分缓冲区。此外，文章还展示了INT8 KV缓存的工作流程，通过自定义内核支持混合数据类型的MatMul操作，动态地将INT8数据转换为FP32，以利用AVX512指令集的FMA指令。分布式推理的工作流程则包括广播令牌ID和在每个工作者计算top-k后进行减少的步骤。

应用

本文提出的优化解决方案不仅适用于当前的LLM模型，还计划扩展到更多类型的CPU，特别是那些资源受限的CPU。此外，文章还计划探索更大批量大小下的性能提升，并适应最新的模型趋势，如混合专家（MoE）模型。这些努力旨在为现有GPU解决方案提供一个实用的替代方案。