这是2025年的第32篇文章
( 本文阅读时间:15分钟 )
RTP-LLM是阿里巴巴智能引擎团队自研的大模型推理加速引擎,作为一个高性能的大模型推理解决方案,已在阿里内部众多LLM场景中得到实际应用与检验。本文探讨与分析了大模型推理引擎中P-D分离技术的意义与优势,并提出了项目自己的方案。
01
背景
在大模型的推理过程中,通常可以将任务分为两个阶段:Prefill 阶段处理所有输入的 Token,生成第一个输出 Token,并生成 KVCache。Decode 利用 KVCache 进行多轮迭代,每轮生成一个 Token。Pefill 阶段通常是计算密集型的,Decode 阶段通常是显存带宽瓶颈。
业界常见的调度器(Continuous Batching)会在每一轮调度中,剔除已经完成的请求,并且将能满足显存需求的 Prefill 请求和 Decode 请求凑批执行。Prefill 阶段运行时间较长,此时 Decode 阶段的时延受到较大影响。最终体现为只要请求出现了 P-D 请求凑批执行,那么请求的时延和 P99 时延会出现巨大波动,这个问题在线上场景时刻存在。
当然也有其他调度策略:
Prefill 优先策略,Prefill-Decode 请求不允许凑批,那么对 Decode 请求的影响更大。
Decode 优先策略,Prefill-Decode 请求不允许凑批,会使得 GPU 利用效率降低。
Chunked Prefill 技术将 Prefill 的请求拆成多个部分多轮执行,在每轮中和 Decode 请求凑批执行,可以提高 Decode 请求的交互性能,但是它的总时延还是会受到 Prefill 请求的影响。并且因为 Prefill 请求仍然长时间占用显存,导致 Decode 请求的并发受到限制。
好像并没有完美的解决方案,既然不能同时满足两个阶段的需求,干脆直接将它们拆分。
我们预计 P-D 分离可以带来一些好处:它们可以选择不同的机型,Prefill 采用高算力的 GPU,Decode 采用大显存的 GPU。它们可以选择不同的量化方式和凑批大小。这样就可以同时优化首字时间(TTFT)和Time Between Tokens(TBT)。最重要的是不同请求的 Prefill 阶段的执行不再会影响 Decode 阶段,使得请求总时延和 P99 时延稳定。
在学术界和工业界都有 P-D 分离技术的使用,总结来看,都取得了不错的效果:
Mooncake 是业界大规模使用 P-D 分离技术的例子。Mooncake 构建了以 KVCache 为中心的 P-D 分离调度集群,形成了 Prefill/Decode Pool 以及分布式异构介质 KVCache Pool,全局负载均衡技术预测了 Prefill 和 Decode 的负载,在满足 SLA 的约束下最大化吞吐,在高负载下预测请求的执行时间,提供早停能力。
Splitwise 开发了一个 GPU 集群,旨在实现三个主要目标:提升吞吐量、减少成本和降低能耗。该集群的设计将 LLM 推理过程分为两个独立的机器池,以优化不同阶段的处理效率。同时,他们还增设了第三个机器池,专门用于处理 Prefill 和 Decode 阶段的混合批处理,并且能够根据实时计算需求灵活调整其规模。
DistServe 通过将预填充和解码计算分配至不同的 GPU,从而消除了二者之间的干扰。针对应用程序的总时延和每个 Token 运行时间的需求,DistServe 为每个阶段量身定制了资源分配和并行策略的优化方案。此外,DistServe 还考虑到服务集群的带宽,将这两个阶段放置在一起,以减少由任务分解所带来的通信开销。
基于这些思考,以及其他方案的剖析,RTP-LLM 团队也提出了自己的方案,本文将介绍具体的实践过程以及踩坑经验。
02
上线效果
首先来看下成果,我们在 RTP-LLM 上实现了 P-D 分离,并且在阿里内部的 Qwen1.5-72B MaaS 场景以及机器人客服场景 Qwen2 14B 模型全量上线,都取得了不俗的效果。
QWen1.5-72B 效果如下:实例个数下降 24%,平均时延下降 48%,P99 时延下降 78%。
Qwen2-14B 上线效果:平均时延降低 32%。
2.1 整体流程
P-D分离的总体流程如上图所示:
请求到到达 Prefill 节点,检查是否可以开启 P-D 分离。
Prefill 根据自己的负载均衡策略,选择 Decode 节点。
Prefill 命令 Decode 申请显存资源。
Prefill 将请求加入本地的任务队列,等待本地调度。
Prefill 本地开始调度请求。
Prefill 命令 Decode 拉取 KVCache。
Prefill 生成 First Token,并且将 First Token 发送给 Decode。
Decode 等待 KVCache 全部传输完毕,和 First Token 一起组织成任务,进行本地调度。
Decode 流式产出 Token,返回给 Prefill,Prefill 流式返回给用户。
Prefill 和 Decode 结束请求,清理资源。
Prefill 的计算流和 RDMA 的发送之间通过 Cuda Event 来同步,从而实现了计算和传输的 Overlap。
03
Cache Store
为了支持 KVCache 传输,我们开发了 CacheStore 组件,支持 Store/Load/Delete 原语,对上层业务层屏蔽了传输和通信的细节。
3.1 KVCache传输
P-D 分离首先遇到的问题是通信传输量以及带来的 Overhead 代价。
例如 Qwen1.5-72B,Layer Number 80,Head Number 32,假设单个 Block(大模型引擎的主要显存占用为模型权重,运行时显存和 KVCache,其中 KVCache 占据很大组成部分。为了支持 Page Attention:KVCache 按照 Block 进行组织)容纳 8 个 Token,那么每层的 Block 大小为130 KB。
如果请求的 Prompt 长度为 4K,那么 Prefill 阶段产出的 KVCache 总大小约为 5G。我们需要尽可能快的将 Prefill 产生的 KVCache 传输给 Decode,降低 Overhead 代价。Mooncake [1] 采用了 Prefill 分层传输 KVCache 技术:在 Prefill 执行下一层计算的同时,将上一层的 KVCache 流式传输到 Decode。我们也使用了这种方案。
我们在 Cache Store 中实现了两种传输方式:
在 TCP 实现中,Prefill 将显存中的 KVCache 拷贝一份到内存中,通过 TCP 发送到对端。
在 RDMA 实现中,依赖通信库 ACCL-Barex,使用 GDR(GPUDirect RDMA [2])在显存之间发送数据。
在 P-D 分离下我们实现了全生命周期管理,在 TCP 复制 / RDMA 传输之时,不会将 KVCache 分配给其他请求。
3.2 TCP传输
在没有 RDMA 网络的环境中,CacheStore 支持只通过 TCP 传输,总体流程基于 ARPC 实现(ARPC 是基于 ANet 实现的 RPC 框架,提供了面向 RDMA、TCP 等不同协议的 RPC 通信能力)。TCP 下的 KVCache 传输流程如下,序号对应了上图过程:
Prefill 端按层计算产生 KVCache。
由显存拷贝到本地内存,放到 Cache Store中存储。
Decode 端通过 ARPC 向 Prefill 发起 Load Cache 请求。
Prefill 接收到请求后会在本地 CacheStore 中寻找请求对应的 Block,然后将找到的 Block 放入 ARPC的 Response 中,发回到 Decode。
Decode 接收到 Response 后从中取出 Block,并拷贝回显存待使用。
在初期原型阶段,我们使用了 TCP 方式传输 KVCache 来快速验证 P-D 分离流程。协议过程有多次拷贝,分配内存和拷贝时延高且不稳定,并且还会影响计算流,多线程拷贝还带来了竞争问题。TCP 传输的问题较多,影响较大。我们认为 TCP 传输使用在无 RDMA 卡的环境,以及简单测试环境,在线上生产环境我们使用 RDMA 传输。
3.3 RDMA传输
在 TCP 环境下,KVCache 从显存拷贝到内存需要较多的时间,而 RDMA 可以使用 GDR,直接在显存之间发送数据。GPU Direct RDMA 是 NVIDIA 推出的一项技术,允许数据在 GPU 显存之间直接传输,绕过了 CPU 和系统内存,从而显著提高了数据传输速度并降低了系统开销。RDMA 下的 KVCache 传输流程如下,序号对应了上图过程:
Prefill 端按层计算产生 KVCache。
由 Decode 向 Prefill 发起 ARPC 请求,告知对端自己的显存地址。
Prefill 接收到请求后调用 RDMA 的 Write 方法,将请求的 Block 从 Prefill 显存写入到 Decode 显存。
Prefill 在写入完成后返回 ARPC Response。
由于 RDMA 通信原语使用上的复杂性,我们基于 ACCL 通信库实现 RDMA 通信(ACCL 通信库是高性能网络团队研发的高性能通信库,支持基于服务器间的 RoCE、IB 和 TCP 网络实现高带宽和低延迟通信)。
我们在 RDMA 传输实现中,发现事情也没有想象的这么简单。
网卡模式
首先我们准备在线上机器直接使用 RDMA 卡,而线上容器只支持 VF 使用 RDMA,而 VF 的挂载使得我们启动独立容器比较复杂。我们在非 VF 模式,拉起简单容器才可以直接使用主机网络来测试 RDMA 功能。
RDMA 通信模式
RDMA 的通信方式包括 Send、Recv、Read 和 Write 四种,其中 Send、Recv 是一对类似 TCP 中的 read 和 write 的双边操作,Read 和 Write 是 RDMA 特有的单边操作(单侧网卡发起,对端用户层不感知)。在基于 RDMA 的 KVCache 传输中,数据传输可以由 Decode 发起,通过 Read 从 Prefill 传播到 Decode,也可以由 Prefill 发起,通过 Write 从 Prefill 传播到 Decode。
在我们的场景下,KVCache 的生命周期由 Prefill / Decode 各自独立管理。出于性能和安全性的考虑,我们使用 RDMA Write 模式来实现数据传输。整体流程是由 Decode 的 KVCache 分配器来分配显存,并告知 Prefill。然后 Prefill 通过 RDMA 将自身的 KVCache Write 到对端。RDMA 传输所需的 remote_addr 和 rkey 等字段,我们通过 TCP 传输,RDMA 能力集中用于 KVCache 的发送。
Register Memory Region
在发送和接收数据之前,RDMA 需要将源地址和目标地址 Register MR,这是为了让不具备 IOMMU 并且无法读取主机 MMU(内存管理单元)的网卡保存虚拟地址和物理地址的映射。因为 Register MR 比较耗时,所以我们在请求来临之前提前 Register MR。
ACCL 的显存管理有两种模式:一是ACCL 自身的显存池,管理可定制 Allocator 分配出来的显存;另一个是用户提供显存,注册到 ACCL 里。在我们的使用场景中,由 Prefill 和 Decode 自身管理显存空间,因此我们采用提前 Register User MR 的做法,在程序启动时为 Prefill 和 Decode 申请所需显存并注册到 ACCL 中。
但是我们发现,Prefill 和 Decode 通过 cudaMemcpyAsync 申请出来的显存,无论如何都会 Register MR失败,Errno 为 EFAULT(Bad Address),令人百思不得其解。最后查清楚是 Nvidia 的 GDR 还不支持cudaMallocAsync,必须使用 cudaMalloc [2],修改为 cudaMalloc 之后注册 MR 成功。
但与此同时我们发现,在实际使用中,MR 的注册总耗时极长。经过排查,确认原因是我们按照 Block 级别注册 MR,而按层分开的 Block 总数为几十万,导致注册时间达到了半个小时级别。我们优化了注册管理机制,采用注册大块显存的方式,将一整块显存,直接注册一次,注册时间降低到秒级。
ACCL 队列
为了减少块内碎片,在线上服务中,我们将 Block Size 从 512 改为 64,然而测试中出现了发送队列溢出问题。经排查是一些长请求对应的 Block 数量较多,在 ACCL 发送时,高频调用发送接口将 ACCL 的发送软队列直接打满,因此出现了提示拒绝发包的情况。我们增加了 ACCL 发送端的发送软队列和 RDMA 硬件发送队列,以及接收端的 RDMA 硬件接收队列。
但是修改后的方案,在创建 RDMA 连接时出现了大量的超时。原因是创建 RDMA QP 时,也会同时创建用于 SendRecv 的接收 Buffer,而增加接收队列长度使得接收 Buffer 的数量增加,引起建链耗时增长。在 RDMA Write 模式下,不需要基于接收 Buffer 来进行 SendRecv 通信,因此可以通过去除接收 Buffer 来优化建链耗时。
RDMA 链接管理
RDMA 的建连一般需要多次带外通信(TCP)交换 QP 信息并进行 ModifyQP 等操作,因此建连延迟较高,而高压力下建连延迟会进一步升高。因此在实际应用中,我们实现了 RDMA 连接池方便连接复用。出于充分利用网卡上联带宽和防止交换机 Hash 不均导致拥塞的考虑,每个请求会基于多条 RDMA 链接并行传输。
为了防止可能的 DirtyWrite(因非预期的 KVCache 生命周期管理,出现显存错写),目前的实现中并发的请求之间不进行连接复用,这也导致了高并发下初始化建连压力增大, 后续考虑通过 MemoryWindow、优化 KVCache 生命周期管理等机制来解决。
Prefill 和 Decode 之间的 RDMA 传输使用多链接发送,并行化数据传输过程,减少单个连接的拥塞,从而提高整体传输速度。RDMA 支持多个发送和接收队列(Queue Pair,QP),可以为每个连接配置多个发送队列,这样可以在同一连接内部进一步实现并行发送,减少队列之间的等待时间,提升并发处理能力。在链接断开和 Decode 重启后能够自动重建 RDMA 连接,使得服务不会出现中断。
04
流程优化
经过不少时间的磨砺,我们的 P-D 分离流程终于串通了,但是真正的问题才刚刚开始。
GPU降频
流程串通之后,初步测试发现 P-D 分离导致 Prefill 的 TTFT 有12%的增加,这完全不符合期望,因为我们认为 Prefill 的首字时延应该基本保持不变。一开始怀疑是 P-D 交互流程影响了时延,在复用 GRPC 链接和调优 GRPC 参数之后,此问题仍然存在。我们继续怀疑是 KVCache 的 RDMA 传输影响到了 Prefill 的计算。
经过 Nsight 抓取 Timeline,发现 Qwen-7B Prefill 的32层计算,前面 17 层计算都变慢了,但是后面 15 层,计算并没有变慢,这就排除了 RDMA 传输带来的影响。
同时我们发现只有在请求比较稀疏的情况下才会出现这个问题,最后定位是 GPU 主动降频导致的。在非 P-D 分离下,一张卡负责 Prefill/Decode 的全部计算。在 P-D 分离稀疏请求场景下,Prefill 只负责 First Token 的生成。在其他 Token 生成期间,Prefill 处于空闲状态,Decode 处于忙碌状态,那么 Prefill 比之前非 P-D 分离场景下更空闲。那么 Prefill 会更容易主动降频,而从降频状态到正常工作频率之间需要花费 170 毫秒。
我们固定 GPU SM/Graph 工作频率为最大工作频率,Prefill 不再出现这个问题。我们在和调度团队一起看,如何能将 GPU 尽快从降频状态转换到正常工作频率。并且,在线上情况,我们减少 P-D 实例的个数,使得流量尽可能的密集,节省资源。
Block管理
继续测试,我们发现在长 Prompt 场景下,P-D 分离流程增加的额外时延比较高。经过分析是因为 Decode 从 Prefill 拉取 KVCache 的时候,是利用 Token ID 进行查找,效率比较低。我们对 Token ID 按照Block 级别求得 Hash 值,这样我们就可以按照 Block Hash 快速定位到 KVCache Block。
为了复用 System Prompt 和跨请求的KVCache:对于 System Prompt,我们将这块 KVCache 固定在 Block 内;对于跨请求的 KVCache,我们在请求执行完毕之时,将请求持有的 KVCache 资源释放给 LRU Cache,它将 KVCache 按照 Block Hash 标记;两种复用方式底层都使用 Trie 树组织。
初期我们的 KVCache Match 采用的是按照 Token ID 来寻找最长前缀匹配,这样在长序列以及 KVCache Item 比较多的时候,匹配效率较低。利用上述 Block Hash,我们可以按照 Block 级别快速寻找最长前缀匹配,大大降低了 BlockCache 的 Match 时延。
但是即使采用了上述的 Block Hash 方案,但仍然存在一些代价:KVCache Block 是最小显存分配单元,Block 大小是定长的,在进程启动的时候确定,一般配置为包含 8 到 4096 个 Token。我们初期使用的 Block 大小为 8 个 Token,带来了很多问题。长 Prompt 需要的 Block 比较多,在 RDMA 传输的时候,分块计算 Block Hash 和发送带来了较大的 Overhead。
但是如果使用超大 Block,例如包含 4096 个 Token,会造成较多的内部碎片和外部碎片。好像陷入了两难的境地。最后我们选择块大小为 64,并且在 RDMA 发送的时候,将 Block 合并成大块再发送。
生命周期
在初期测试中,我们发现如果两个相同请求同时发来,那么其中一个请求的时延会出现巨大的抖动。经过排查 Decode 从 Prefill 获取 KVCache 的时候,按照 Block Hash 来查找,一旦传输完毕就可以释放。这就造成了相同的两个请求,后来的请求先获取了 KVCache 并且释放了它,而先来的请求后获取到 KVCache。乱序执行问题使得部分请求的时延大大增加。
为了解决这个问题,我们在 CacheStore 存储 KVCache 的时候,加上了 Request ID 标识。请求的结束的时候,按照 Request ID 释放 KVCache。当然相同请求的 KVCache 只占用一份,CacheStore 存储相同的 Block ID。
订阅机制
在上线接流中,我们突然发现线上出现了大面积 ARPC 队列满的情况,令人非常紧张。 经过排查是因为在 Decode 从 Prefill 拉取 KVCache 之时,此时 Prefill 可能还没有产出 KVCache。
我们采用的是 Polling 机制,不断的检查 KVCache 是否产出,而 Polling 会占住一个 RPC 线程并最终等到超时。在 RPC 线程数固定的情况下,一旦有部分请求的 KVCache 没有产出,那么 RPC 线程超时,会使得后续请求的传输出现延后,从而带来雪崩现象。
最后我们使用订阅通知机制,Decode 的 Load 请求订阅 Prefill 的 KVCache 的产出。 Prefill 在 KVCache 产出之后,通知 Decode 的 Load 过程,完成后续的传输。当然如果 Prefill 因为某些错误,永远没有产出 KVCache,那么需要把 Decode 的请求在超时的时候结束掉。
会话
在线上服务的长期运行中,我们发现出现了缓慢的内存泄漏和显存泄漏。原来是因为在一次查询期间,Prefill 和 Decode 之间存在多次命令交互。我们在初期,使用了独立的多次的 GRPC 命令来控制交互,维持请求的状态特别困难。在异常的情况下,如果没有如期发出所有命令,那么请求的状态没有进行清理,造成了内存和显存泄漏。
最后,我们使用 GRPC Stream 来维护一次请求对应的 P-D 之间的所有交互,在任何一端出现已知 / 未知问题的时候,会将两端的的会话资源都进行析构。在任何重试可成功的过程中,我们都允许进行有限制的重试。在任何大耗时的操作中,我们都加入了超时检查和取消检查,一旦发现了请求超时和取消,我们立刻结束会话。
Serverless
我们在一个 Prefill 集群下面挂载了多组异构的 Decode 集群,Prefill 会根据负载,以及请求的长短,自动选择合适的 Decode 集群来服务。Prefill-Decode 组成了多 Role 的 Serverless 集群,在 Decode 挂掉的情况下,Prefill 可以自动链接上其他 Decode 进行请求重试。Prefill-Decode 可以自由的伸缩,在机器迁移/版本升级的时候,快速 Failover,自动保障服务可用度。
健康检查
Prefill-Decode 分离的情况下,Prefill 在尚未发现 Decode 节点的情况下,Prefill 节点标记为不健康,不承接流量。Prefill 和 Decode 在停止的时候,会等待正在运行的所有查询都停止。Prefill 还实现了可选的Fallback 逻辑,在 Prefill 访问 Decode 多次重试仍然错误的情况下,Prefill 可以选择本地执行请求,提供容灾能力。在升级的时候,会保证 Prefill-Decode 同版本同比例升级,保证访问的兼容性。
服务发现
我们在P-D 分离场景中支持四种服务发现:
Local:是为了方便测试。
CM2:最初作为搜索引擎内部的名字服务,支持了节点级别的服务发现。现在还支持在大数据场景里基于节点中数据的版本、分片等,提供面向数据的服务发现能力。
Nacos:阿里巴巴开源的一款分布式服务治理和配置中心的框架,主要应用于微服务架构中的服务发现、配置管理和服务管理。
VipServer:阿里巴巴内部服务发现。
在非 Local 场景下,Prefill 定时获取 Decode 的节点状态信息并进行检测,保证健康节点的服务,提供底线容灾能力,并基于此实现了 Prefill-Decode 之间的负载均衡。目前 Prefill 选取 Decode 的负载均衡支持两种:
RR 策略,Round Robin 的方式轮询 Decode 节点。
WRR 策略,加权 RR 策略,根据 Decode 剩余可用显存分配权重,在动态变化的负载环境下维持系统的高稳定性和响应速度。
负载均衡
在逐步增大流量规模的时候,我们发现有时候 Decode 会突然出现爆显存的问题。这是因为在一个短暂的时间内,可能存在一个长请求将 Decode 显存占满的情况。Prefill 节点访问哪台 Decode 节点:我们初期采用的是 RR(Round Robin)的策略,这个策略面对突发情况无能为力。
基于这个观察,我们开发了 WRR(加权 Round Robin),Prefill 实时获取 Decode 的剩余显存,动态决定负载均衡。这个方法进一步提高了系统的鲁棒性。
KVCache 回退
引擎在执行多个请求期间,可能总显存不足,此时我们会暂停某些请求的执行,释放部分 KVCache 资源,等待下次调度的时候,重新执行。这个策略在 Decode 引起了较大的问题,因为重新执行的请求会进入 Prefill 阶段,从而违背了 P-D 分离的初衷。所以在 P-D 分离的场景下,Prefill-Decode 禁止 KVCache 资源回退。
量化
到此服务已经基本稳定,但是我们想要的更多。在线上部署的时候我们选择的是 16 个Prefill 实例对应 2 个 Decode 实例,Decode 机器承担了 8 倍的请求。我们观察到 Decode 容易出现爆显存,主要是因为并发执行的请求比较多。我们将 KVCache 量化减少 KVCache 占用,并且将 Decode 的模型权重从 FP16 替换成 INT4,使得模型权重显存占用降低到 1/4。此举大大提高了 Decode 的并行处理能力。
05
多卡
P-D 分离准备在 Qwen1.5-72B 场景上线,单卡 GPU 捉襟见肘,无法满足算力和显存的需求。所以我们开发了多卡的的 P-D 分离,而开发多卡也遇到了很多问题。
多卡控制
首先 Prefill Master 是直接通知所有 Decode Worker 去开始会话,还是只通知 Decode 的 Master。Decode 的 Master 使用 NCCL 还是 GRPC 去通知 Decode Worker 去加载 KVCache。我们判断 Decode Worker 的一切行动需要和 Decode Master 保持一致,而且 NCCL 使用起来要求较高,例如 Buffer 需要是 GPU Buffer。
所以我们最终采用了由 Decode Master 使用 GRPC 通知其他 Decode Worker Load KVCache。RTP-LLM 的多卡通信基于 NCCL 实现,Rank0 通过 NCCL Broadcast 将模型执行必要的信息同步给其他 Rank,因此只有 Rank0 会收到 GRPC 请求。但是在 P-D 分离的框架设计中,Prefill 机器的资源释放,Decode 机器的 KVCache 加载和以及单机资源申请/释放都和 GRPC 请求生命周期绑定,无法通过 NCCL 实现。
因此在多卡 P-D 分离逻辑里我们需要 Rank0 通过 GRPC 将这些信息和其他 Rank 同步,以做到状态和资源生命周期同步。
具体流程如下:
用户请求发送到 Prefill Rank0。
Prefill Rank0 开始生成 KVCache,同时通知 Decode Rank0 加载 KVCache。
Decode Rank0 并发通知其他 Decode 机器准备从 Prefill 对应 Rank 加载 KVCache。
在请求结束后,Prefill Rank0 需要并行通知其他 Prefill Rank 清理 CacheStore 资源。
异步Load
多卡继续测试,在高并发下,我们发现 Decode Load KVCache 出现了双倍甚至更高的延时。经过排查原来是同步 GRPC 造成的问题:在多卡环境下,Decode Master 命令 Decode Worker 去加载 KVCache,这是一个长耗时过程,时延在几十毫秒到几秒不等,取决于 Prefill 的 TTFT。我们初期采用的是同步的 GRPC Call,在高并发请求下,Decode Master 需要大量的线程执行同步调用,线程数不足就出现了 Load Cache 阻塞。
很自然的,我们会想到使用异步来进行改造。GRPC 提供 Async RPC Sevice,但是改造成本很大,而我们只需要 Load Cache过程是异步的,并不需要 GRPC Stream 也是异步的。最后我们使用 GRPC 的 Completetion Queue 来发起异步请求。但是测试中发现,Completetion Queue 只能承担2个异步请求,如果发起大于2个异步请求,那么会带来巨大的时延抖动,目前我们暂时采用多 Completetion Queue 来解决这个问题。
流量规划
在线上的机器中,会将一张物理网卡虚拟化成多张虚拟网卡。在线上的容器中,会将多张虚拟网卡,甚至全部的虚拟网卡,都挂载进入容器。那么此时 Prefill-Decode 实例应该使用哪些 GPU 卡,哪些虚拟网卡,哪些物理网卡。能否降低 GPU 和网卡之间传输时延,能否充分利用多物理网卡多队列的能力,这都是需要我们解决的问题。
在 Prefill-Decode 之间,RDMA 链接选择哪条最优路径,在路径出现问题的时候,如何回避到其他路径。随着 P-D 实例增多,全局通信流量增多,是否会出现局部的满载及单点故障等问题,这也是我们在大规模铺开 P-D 分离实例之前需要充分研究的问题。
这些担心并不是杞人忧天,在线上的实际部署中,我们突然发现有台机器服务异常,RDMA 流量暴涨,不仅仅影响了自己,甚至影响到了广告的服务,情况十分紧张。在故障恢复之后,我们深入排查发现当时多个 Decode 实例分配到了同一个物理机器的同一张网卡上,日常流量较高已接近打满交换机的 Buffer。
出问题的时候流量变化超过了交换机 Buffer 限制触发交换机丢包,丢包后网卡重传流量暴涨且难以自动恢复。这种情况下不仅自身服务的 RT 无法保证,甚至影响了此网卡上的其他服务。目前我们通过对优化对网卡资源的申请防止出现流量过高来保障自身与其他服务的服务能力,网络侧上也尝试通过调整配置来优化类似场景下的处置。
06
ACCL-Barex
6.1 ACCL 介绍
ACCL[3] (Alibaba Collective Communication Library) 是高性能网络团队研发的高性能通信库,是一个专门为多 GPU 之间提供集合通讯的通讯库,它具有拓扑感知的能力,提供了包括 AllReduce、Broadcast、Reduce、AllGather、ReduceScatter 等集合通信 API,也支持 send、recv 等 API 来实现各种复杂的点对点通信,可以通过服务器内的 PCIe、NVLink、NVSwitch 和服务器间的 RoCE、IB、TCP 网络实现高带宽和低延迟通信。
ACCL 提供通用的分布式多机多卡集合通信能力,并且深度结合大规模 AI 集群的互联架构和多网卡特性,实现创新的无拥塞算法保证带宽,基于 Barex 模块提供点对点通信能力,基于 C4D 提供故障诊断功能,基于 C4P 提供跨任务路径规划和 Rank 重排功能,在性能和功能上接近甚至超越业界领先的其它通信库。同时考虑到普适性,当前的 ACCL 设计全面兼容 Nvidia 的 NCCL 通信库。
在 P-D 分离场景中,使用 ACCL 通信库的点对点通信模块 ACCL-Barex,可以很好地满足 KV Cache 在 P-D 之间高性能、低时延的传输需求。基于 KV Cache 使用 WRITE 的传输模式,提供多种场景的接口支持:WriteSingle:单块数据 WRITE 到对端一块地址;WriteBatch:多块数据 WRITE 到对端多块地址;WriteBySglist:多块数据 WRITE 到对端一块地址等等。同时,ACCL-Barex 也支持 READ 的调用模式,以及 Send-Recv 等更多 RDMA 通信场景。
6.2 ACCL-RDMA 和 TCP 性能比较
从下图的对比我们可以看到,随着包的大小从 256K 涨到 32M,RDMA / TCP 传输加速比从 1.6 倍涨到了 19 倍。对于大包,RDMA 的加速优势更为明显。
6.3 网卡亲和
对于多网卡环境,往往每个 GPU 设备有自己的亲和网卡,基于亲和网卡进行 KVCache 数据的 GDR 发送,相比非亲和网卡性能更高。我们在这方面增加了网卡亲和匹配的能力,调度侧分配容器时给容器内的多个 GPU 选中并挂载亲和网卡。程序启动时,我们给每个 GPU 进程分配对应的亲和网卡,在亲和网卡缺失时,也会分配其他可用网卡,保证所有网卡带宽都被有效利用。
07
现在还存在的问题
虽然我们的部分线上业务已全量上线 P-D 分离,且效果良好稳定。但是仍然存在一些问题,需要我们进一步解决。
7.1 长序列
虽然在平均场景下,我们的 Qwen1.5 72B P-D 分离集群已经变现良好。但是在长文本场景下问题严重:首先长文本在 Prefill 执行时间比较长,在 Decode 爆显存。我们有一些想法来解决这个问题。在 Prefill 使用 CPP(Chunked Prefill Pipeline)/ PP (Pipeline Parallelism) / Ring Attention,将长文本的计算划分到多个 Prefill 节点执行。Decode 使用更大显存容量的卡,这样可以满足下长文本的 KVCache 占用。
解决 Decode 显存容量还有一个方法,那就是 Decode 选择更大的 TP Size,那么就和 Prefill 的 TP Size 不对等了。此时我们必须将 Prefill 的一个节点的 KVCache 进行划分,组织成 Decode 的 KVCache Layout。这必然带来一些代价,而且更大的 TP 意味着在计算过程出现更多的通信代价。但是更大的 TP 使得 Decode 的显存容量更大,可以并行执行更大的 Batch。那么哪种部署才是最优部署,这是一个值得探索的问题。
7.2 ReuseCache
在多轮会话的场景中,我们跨请求复用 KVCache,良好的 Cache 命中率能大大降低 Prefill 阶段的时延。目前我们在 Prefill/Decode 各自使用了 Reuse Cache 逻辑,减少 KVCache 的计算和传输。进一步的,我们将 Decode 产出的 KVCache 也回传给 Prefill,有机会提高 Prefill 在下一轮会话的 KVCache 命中率。在 Prefill 自身压力比较大时,可以将请求转发给其他 Prefill 来降低压力,当然也可以从其他 Prefill 拉取 KVCache 来进行复用。
7.3 分布式 CacheStore
我们正在实现 Prefill-Decode、Prefill-Prefill、Decode-Decode 之间拉取 KVCache,使得 CacheStore 真正成为分布式 CacheStore。
7.4 支持多设备
支持多种计算设备,例如 GPU,AMD,ARM以及它们之间的互联互通。支持多种不同 RDMA 卡型和协议。在大规模铺开 RDMA 传输的情况下,对节点/集群/选路/路径容灾进行进一步的探索的和优化。
7.5 负载均衡
在当前线上场景,服务的接入层根据 Prefill 的负载来选择 Prefill 节点发送请求,Prefill 再根据自己获取到所有 Decode 节点的剩余显存信息选择 Decode 节点来服务请求。由于接入层是一个分布式集群,因此在同一时间来的请求可能会被分配到同一台机器上,造成负载不一致。
对于 Prefill 来说,并发请求会直接增加排队时间,影响 First Token Time。同时接入层的负载均衡只考虑 Prefill 的计算速度,没有全局考虑 Prefill-Decode 对的情况,也没有考虑 KVCache 最大化复用率,更没有考虑 KVCache 传输和本地计算之间的性能对比等。集群下一些长请求在执行过程出现了超时,我们需要尽可能的避免这种计算的浪费。
08
以 KVCache 为中心的集群调度
我们可以看到 KVCache 已经成为 P-D 分离或者非 P-D 分离情况下的核心部分。它的分配效率,复用效率,传输效率,集群命中率,异构介质 Offload/Load 等,对于大模型推理引擎的性能至关重要。而这对调度提出了非常高的要求。
在集群模式下,节点之间的负载不一致,产出 Token 的能力不一致,那么集群模式下的负载均衡成为了一个新的问题。不同节点拥有的 KVCache 情况很不同,甚至固定的 KVCache 例如 SystemPrompt 或者热门的 KVCache 需要分发到整个集群。集群下的 KVCache 全局命中率也使得中央调度节点势在必行。
拥有中央调度节点,我们可以实时获取所有节点状态,并且做出当前最优决策。
当前计算选择哪些 P-D 实例来服务才能在满足 SLA 的情况下最大化吞吐?
计算节点是本地进行计算还是从其他节点拉取 KVCache?或者一边计算一部分 KVCache,一边从远端拉取?
甚至于预测节点的负载,预测请求的输出长度。如果预测出请求不能满足 SLA 约束(FFTT 约束和 TBT 约束),那么将请求早早的拒绝,节省算力。
在计算节点在产出 Token 过程中,预测出节点显存不足,一边继续产出 Token,一边迁移 KVCache 和会话状态到其他节点继续服务。
在离线高吞吐长超时的请求负载下,将冷 KVCache Offload 到主机内存/Optane/SSD,进一步提高服务的并行处理能力。
在请求尚未调度的时候,将请求对应的KVCache提前分层调入显存,并且分层执行,介质传输和模型计算之间进行 Overlap。
在离线负载下,并发请求量更多,此时如何高效调度所有异构集群,达到最大吞吐,是一个可想象的空间。
在 MaaS 服务下,请求长短不一,优先级不一样,SLA 不一致,如何在保障 SLA 的情况下优先服务重要请求。
总体上来看:中央调度节点,任务很重,想象空间很大,在线场景下可以更好的满足时间约束的请求 SLA,离线场景下可以进一步提高吞吐。
09
总结
P-D 分离成功地解决了 Prefill 和 Decode 相互影响的问题,使我们能够选择最适合的设备、量化类型和调度策略,从而提高系统的整体性能和灵活性。它的全量上线带来的效果振奋人心,验证了这个方向的正确性。目前我们也在不断地进行优化,并计划将这些改进扩展到新的场景中。
P-D 分离带来了一些转变:以 KVCache 为中心的集群调度势在必行,这将显著提升集群的服务水平(SLA)和吞吐。它确实给我们带来了相当多的惊喜和想象空间,我们会坚持不懈的继续在这条道路上努力狂奔。
最后,RTP-LLM 项目已经开源,欢迎大家与我们交流共建!
参考文献
[02] GPUDirect RDMA
https://developer.nvidia.com/blog/using-cuda-stream-ordered-memory-allocator-part-2/
[03] 高性能集合通信库(ACCL)
[04] RTP-LLM Github链接