MemGPT: Towards LLMs as Operating Systems

论文:2个版本论文不太一样

  1. v1 https://arxiv.org/pdf/2310.08560v1
  2. v2 https://yiyibooks.cn/arxiv/2310.08560v2/index.html

代码:https://github.com/cpacker/MemGPT

核心亮点:

  1. 在有限的上下文窗口(8k)里,支持无线上下文的“错觉”,属于长文的一种解决方案
  2. 对智能体友好,能够提供超长记忆力,以及个性化的智能体体验

1. 产品场景

微软的类似产品能力:https://microsoft.github.io/autogen/docs/ecosystem/mem0

千帆 app-builder 类似的产品能力:

2. 论文背景

由于 Transformer 架构的自注意力机制,直接扩展 Transformer 的上下文长度会导致计算时间和内存成本成倍增加,这使得新的长上下文架构的设计成为紧迫的研究挑战(戴等人,2019; Kitaev 等人, 2020; Beltagy 等人, 2020)。虽然开发更长的模型是一个活跃的研究领域(Dong等人,2023),即使我们能够克服上下文缩放的计算挑战,最近的研究表明长上下文模型很难利用额外的上下文实际上(Liu等人,2023a)。因此,考虑到训练最先进的大语言模型所需的大量资源以及上下文扩展的收益递减,迫切需要替代技术来支持长上下文。
在本文中,我们研究如何在继续使用固定上下文模型的同时提供无限上下文的错觉。我们的方法借鉴了虚拟内存分页的理念,该理念通过在主内存和磁盘之间进行数据分页,使应用程序能够处理远远超出可用内存的数据集。
在MemGPT中,我们将上下文窗口视为受限内存资源,并为大语言模型设计了类似于传统操作系统中使用的内存层的内存层次结构(Patterson等人,1988)。传统操作系统中的应用程序与虚拟内存交互,通过将溢出数据分页到磁盘,并在应用程序访问时(通过页面错误)将数据检索回内存,提供了一种内存资源多于物理内存(即主内存)实际可用资源的假象。为了提供更长上下文长度(类似于虚拟内存)的类似错觉,我们允许大语言模型通过 “大语言模型操作系统”(我们称之为 MemGPT)来管理放置在自己上下文中的内容(类似于物理内存)。MemGPT 使大语言模型能够检索上下文中丢失的相关历史数据,并且还可以将不太相关的数据从上下文中驱逐到外部存储系统中。图3说明了MemGPT的组件。


Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving

一、摘要

MoonCake设计了KV Cache分离式架构,该架构下prefill和decode分离为2个集群。同时,存算分离KV Cache利用了GPU集群的CPU、MEM、SSD资源来构建。
1、MoonCake关键核心:以KV Cache为中心的调度,该调度核心价值:最大化吞吐、同时不打破TTFT和TBT。
2、MoonCake应对超高负载:设计了提前预估的拒绝策略
3、MoonCake效果:模拟场景中,MoonCake相比基线提升了5倍多的吞吐(满足延迟前提下)
4、MoonCake实际应用:MoonCake使得Kimi处理的请求量提升了75%

二、结论

为了高效LLM推理,尤其是在长上下文和超载场景,我们设计了MoonCake。文章中讨论了该架构的必要性、挑战、针对最大化吞吐和保障延迟间做的设计折中。

三、介绍

1、研发分离式KV Cache背景和动机
  • prefill和decode具有不同计算特性,prefill偏计算,decode偏显存和带宽,拆分有利于提升资源利用率和吞吐
  • 最大化吞吐的核心办法:prefill阶段最大程度reuse KV Cache(本地Cache+远程Cache)、推理阶段在一个batch中最大化token数量(提升MFU,模型浮点计算)
  • 核心问题:最大化吞吐,通过远程Cache复用会增大TTFT,增大batch的size过大会导致TBT变大
  • 关键解法:设计一个分布式KV Cache,并以他为中心设计核心调度和优化
2、关键流程
  • 关键核心流程设计:对于一个请求,全局调度器Conductor根据最大化Cache和工作负载相结合原则,挑选prefill实例及decode实例
    • 1)传输尽可能多的可复用KV Cache到被选择的prefill实例
    • 2)按层分块并行完成prefill stage,并且prefill stage过程中按层分块流式传输KV Cache到对应decode实例
    • 3)decode实例load 被写入本地的KV Cache,然后继续流程完成推理
3、核心挑战
  • prefill阶段(最大化复用KV Cache)
    • 内存/SSD KV Cache load时延保证:分布式KV Cache核心存储在内存和SSD,如何控制加载时间,不影响TTFT
    • 短时间高并发大量KV Cache传输时延保证:存在KV Cache需要远程传输到本地,可能会出现网络拥塞,进一步加剧延迟
    • 调度优化和保障:Conductor需要预测未来KV Cache的使用情况,并提前执行swap或多副本策略
        • 最热的一些block Cache,需要复制到多个节点,避免网络拥塞
    • 内存限制:单机内存大部分被KV Cache Pool预留,留给prefill调度的内存有限,影响调度效果?
  • decode阶段(最大化增加batch size,增加token数量)
    • batch过大,TBT SLO可能被打破
    • prefill实例节点显存有限,影响size数量
  • 负载调度
    • 请求提前拒绝:prefill后,需要预测是否有decode槽位(空闲decode实例),如果没有则需要尽早拒绝,节省资源给其他请求
    • 粗暴拒绝会导致负载波动剧烈:需要预估生成的长度,然后做短期内的全局的负载预估,支撑更好的拒绝策略
    • 优先差异调度:区分不同请求优先级,实现更好的优先级调度
4、关键策略设计
  • CPP:Chunk Pipline Parallelism,将一个请求跨多节点并行处理,并且优化了网络占用,以应对长文本分布的动态变化,这样可以有效降低TTFT
  • 中心式KV Cache请求调度算法:基于启发式的自动化热点迁移方案,自动复制热的KV Cache块,无需依赖 精确预估未来KV Cache使用。能显著降低TTFT
  • 请求拒绝:基于实例系统负载,决定是接受请求还是拒绝请求

四、相关工作

五、核心设计

5.1 整体架构