Jenga: Effective Memory Management for Serving LLM with Heterogeneity

大规模语言模型（LLMs）被广泛使用，但运行成本高昂，尤其是在推理工作负载不断增加的情况下。为了降低成本，通过高效管理GPU内存以最大化请求批处理大小变得至关重要。尽管最近提出的分页注意力机制（PagedAttention）旨在提高内存管理的效率，但我们发现现代LLM架构中嵌入维度、注意力机制和访问模式的日益异构化为内存分配带来了新的挑战。 在本文中，我们提出了Jenga，一种专为LLM中的异构嵌入设计的新型内存分配框架。Jenga解决了两个关键问题：(1) 在管理不同尺寸的嵌入时最小化内存碎片化，以及 (2) 提供灵活的缓存和驱逐策略，以适应各层特定的标记依赖模式。Jenga采用两级内存分配器，利用嵌入尺寸的最小公倍数（LCM）优化内存使用，并提供API来表达特定层的缓存逻辑，从而增强内存复用。 我们在最先进的LLM推理引擎vLLM上实现了Jenga，并使用多种LLM、数据集和GPU配置对其进行了评估。评估结果显示，Jenga可将GPU内存利用率提升高达79.6%，并将服务吞吐量提升至最高4.92倍（平均1.80倍）。

该论文试图解决大型语言模型（LLMs）在推理过程中因嵌入维度、注意力机制和访问模式的异构性导致的GPU内存管理效率低下问题。这是一个重要但尚未完全解决的问题，尤其是在多模型、多任务部署场景中。

论文提出了一种名为Jenga的新型内存分配框架，通过两层内存分配器优化异构嵌入的管理。其核心创新在于：1）利用嵌入大小的最小公倍数（LCM）减少内存碎片化；2）提供灵活的缓存和驱逐策略以适应不同层的标记依赖模式。相比现有方法，Jenga更适配现代LLM架构的复杂性和多样性。

实验结果表明，Jenga能够显著提升GPU内存利用率（最高79.6%），并大幅增加服务吞吐量（最高4.92倍，平均1.80倍）。论文在vLLM推理引擎上实现了Jenga，并测试了多种LLM模型、数据集和GPU配置，验证了其泛化能力。此外，研究还强调了API设计的重要性，允许开发者根据具体需求定制缓存逻辑。目前暂未提及代码是否开源，但这一框架为未来异构内存管理的研究提供了方向。

近期相关研究包括PagedAttention（用于高效管理长上下文序列）、FlexGen（针对稀疏激活的内存优化）以及VPGAE（虚拟页全局地址映射）。其他值得注意的工作有《Memory-Efficient Inference for Large Language Models》和《Optimizing GPU Memory Management for Transformers》。这些研究主要集中在降低内存占用或加速推理，而Jenga则进一步探索了异构嵌入管理的具体挑战。