通用人工智能时代，大模型推理面临存储瓶颈，上海交大研究团队提出MILLION框架，基于乘积量化实现键值缓存压缩和推理加速，有望突破长上下文应用限制。
在通用人工智能（Artificial General Intelligence, AGI）的曙光初现之际，大型语言模型（Large Language Models, LLMs）正以前所未有的速度渗透到各个领域，驱动着复杂任务的智能化进程。它们在自然语言处理、机器翻译、文本生成、代码编写等方面的卓越表现，充分展现了其巨大的潜力和商业价值。然而，LLMs的广泛应用也带来了严峻的挑战，尤其是在计算和存储资源方面。这些模型动辄拥有数十亿甚至数千亿的参数，对硬件设施提出了极高的要求。
其中，基于Transformer架构的大模型，在长上下文场景下，键值缓存（Key-Value Cache）成为了制约其发展的关键瓶颈。键值缓存作为一种“以存代算”的优化策略，旨在通过存储中间计算结果来加速推理过程。然而，随着上下文长度的增加，键值缓存所占用的存储空间也呈指数级增长，使得有限的硬件资源捉襟见肘。
为了突破这一瓶颈，上海交通大学计算机学院蒋力教授与刘方鑫助理教授带领的IMPACT课题组，联合华为2012实验室和上海期智研究院，共同研发了一种名为MILLION的创新框架。该框架基于乘积量化（Product Quantization, PQ）技术，实现了键值缓存的有效压缩和推理加速，为大模型在长上下文场景下的应用提供了新的解决方案。这项研究成果已被电子设计自动化领域的顶级会议DAC 2025接收，标志着中国科研团队在该领域取得了重要突破。
大模型推理的存储困境：键值缓存的挑战
Transformer模型是当前大语言模型的主流架构。在推理过程中，Transformer模型需要重复计算注意力机制，这部分计算量巨大。为了避免重复计算，键值缓存被引入，用于存储先前计算得到的键（Key）和值（Value）向量。在后续的推理步骤中，模型可以直接从缓存中读取这些向量，而无需重新计算，从而显著加速推理速度。
然而，随着上下文长度的增加，键值缓存的存储需求也急剧增长。以半精度的LLaMA-2-7B模型为例，其模型权重约为14GB。当上下文长度达到128K时，键值缓存将占用高达64GB的存储空间，总和接近高端显卡NVIDIA A100的80GB显存容量上限。这意味着，即使使用最先进的硬件设备，也难以支持LLMs在长上下文场景下的高效推理。
此外，键值缓存的存储需求还会受到模型大小、量化精度等因素的影响。更大的模型和更高的量化精度通常意味着更大的键值缓存。因此，如何有效地压缩键值缓存，降低其存储需求，成为了大模型推理面临的关键挑战。
MILLION框架：乘积量化赋能键值压缩
为了解决上述问题，上海交大研究团队提出了MILLION框架，该框架的核心思想是利用乘积量化技术对键值缓存进行压缩。乘积量化是一种矢量量化方法，它将高维向量空间分解成多个低维子空间，然后在每个子空间中进行量化。通过这种方式，可以将原始向量压缩成一组码本索引，从而显著降低存储空间的需求。
MILLION框架的具体实现步骤如下：
- 键值缓存划分： 首先，将键值缓存划分为多个子集，每个子集包含一定数量的键和值向量。
- 子空间分解： 对于每个子集，将其对应的向量空间分解成多个低维子空间。
- 码本训练： 在每个子空间中，使用聚类算法（如K-means）训练一个码本，码本包含一组代表性的向量。
- 向量量化： 对于每个向量，找到其在每个子空间中最接近的码本向量，并用码本索引代替原始向量。
- 索引存储： 将所有码本索引存储起来，作为压缩后的键值缓存。
在推理过程中，当需要访问键值缓存时，首先根据码本索引从码本中检索出对应的向量，然后将这些向量组合起来，得到原始向量的近似值。虽然这种方法会引入一定的量化误差，但可以通过调整码本大小和子空间数量来控制误差的大小。
MILLION框架的优势在于：
- 高压缩率： 乘积量化能够有效地压缩键值缓存，显著降低存储需求。
- 低计算复杂度： 码本检索和向量组合的计算复杂度较低，不会显著增加推理时间。
- 可控的量化误差： 可以通过调整码本大小和子空间数量来控制量化误差的大小，从而保证模型性能。
实验结果：性能与效率的双重提升
为了验证MILLION框架的有效性，研究团队在多个基准数据集上进行了实验。实验结果表明，MILLION框架能够在显著降低键值缓存存储需求的同时，保持较高的模型性能。
具体而言，与传统的全精度存储相比，MILLION框架可以将键值缓存的存储需求降低4倍甚至更多，而模型性能的下降幅度仅在可接受范围内。此外，MILLION框架还能够加速推理速度，这是因为压缩后的键值缓存可以更快地从内存中读取，从而减少了数据传输的延迟。
这些实验结果充分证明了MILLION框架在键值缓存压缩和推理加速方面的优势。该框架为大模型在长上下文场景下的应用提供了新的可能性，有望突破存储瓶颈，推动通用人工智能的发展。
开源框架：定义模型量化推理新范式
为了促进MILLION框架的广泛应用，研究团队将其开源，并发布了相应的代码和文档。开源链接为：https://github.com/ZongwuWang/MILLION。
MILLION框架的开源，不仅为研究人员提供了一个便捷的工具，可以用于研究和开发新的键值缓存压缩算法，也为工程师提供了一个高效的解决方案，可以用于优化大模型的推理性能。
通过开源，MILLION框架有望成为模型量化推理的新范式，推动整个领域的发展。研究团队也希望能够与更多的研究者和开发者合作，共同完善和改进MILLION框架，使其能够更好地服务于通用人工智能的发展。
未来展望：探索键值压缩的更多可能性
尽管MILLION框架已经取得了显著的成果，但研究团队认为，键值压缩领域仍有许多值得探索的方向。
- 自适应量化： 目前，MILLION框架使用固定的量化参数。未来可以研究自适应量化方法，根据键值缓存的特性动态调整量化参数，从而进一步提高压缩率和模型性能。
- 混合量化： 可以将不同的量化方法结合起来，例如将乘积量化与标量量化相结合，以获得更好的压缩效果。
- 硬件加速： 可以设计专门的硬件加速器，用于加速乘积量化的计算过程，从而进一步提高推理速度。
此外，还可以将MILLION框架应用于其他类型的模型和任务，例如图像识别、语音识别等。通过不断探索和创新，键值压缩技术有望在通用人工智能领域发挥更大的作用。
结语
上海交大研究团队提出的MILLION框架，为大模型推理的存储瓶颈问题提供了一种有效的解决方案。该框架基于乘积量化技术，实现了键值缓存的有效压缩和推理加速，有望突破长上下文应用限制，推动通用人工智能的发展。
MILLION框架的开源，不仅为研究人员提供了一个便捷的工具，也为工程师提供了一个高效的解决方案。我们期待MILLION框架能够在未来得到广泛应用，并推动模型量化推理领域的发展。
在通用人工智能的浪潮下，我们相信，通过不断的技术创新和开放合作，我们一定能够克服各种挑战，实现人工智能的真正潜力。
参考文献：
- 汪宗武, 许鹏, 等. MILLION: A Product Quantization Framework for Key-Value Cache Compression and Inference Acceleration. arXiv preprint arXiv:2504.03661 (2025). https://arxiv.org/abs/2504.03661
致谢：
感谢上海交通大学计算机学院蒋力教授与刘方鑫助理教授的指导，以及华为2012实验室和上海期智研究院的支持。感谢博士生汪宗武与硕士生许鹏的辛勤工作。
Views: 1