远端 KV Cache 传输
大模型推理正在逐渐从单机内的 Prefill-Decode 连续执行,扩展到 Prefill-Decode 分离、分层 KV Cache 和跨节点 Prefix Cache 复用。随着计算阶段和缓存空间被拆分到不同设备,KV Cache 也开始频繁跨越 GPU、CPU 和服务器边界。
在 Prefill-Decode 分离架构中,这一问题尤其直接。Prefill 节点计算 Prompt,并在本地 GPU 上生成 KV Cache;Decode 节点负责后续 Token 生成。Decode 开始执行之前,需要先获得 Prefill 阶段已经生成的 KV Cache。
因此,一次推理请求的数据路径从:
Prefill → Decode
进一步展开为:
Prefill 计算 → KV Cache 生成 → 远端传输 → Decode 加载 → Token 生成。
KV Cache 的远端传输由此进入推理关键路径。
对于长上下文模型,这部分数据量并不小。随着上下文长度和模型规模增长,单个请求产生的 KV Cache 可以达到数百 MB 甚至数 GB。传输延迟过高时,会显著影响推理的性能。
如何在不同 GPU 之间快速移动 KV Cache,因此成为分离式推理系统中的重要问题。
1. 远端 KV Cache 的基本路径
当前推理系统中的远端 KV Cache 大致存在两类数据路径。
第一类出现在 Prefill-Decode 分离中。KV Cache 刚刚由 Prefill GPU 产生,随后立即交给 Decode GPU 使用,数据具有明确的源端和目标端。
第二类出现在分层缓存和共享 Prefix Cache 中。GPU 上产生的 KV Cache 可以进入 CPU、SSD 或远端缓存池,后续请求根据 Prefix 命中情况重新加载。
两类路径可以统一表示为:
对于分布式推理而言,底层最终都要解决相似的数据搬运问题:
- KV Cache 位于哪块内存;
- 目标节点应当接收到哪些 KV Cache;
- 数据应该写入目标 GPU 的哪个位置;
- 如何建立跨节点内存访问;
- 如何异步执行传输并判断传输完成。
vLLM 和 LMCache 分别从不同层次提供了相应机制。
2. vLLM 与 LMCache 的传输方案
vLLM 使用 KVConnector 抽象外部 KV Cache 数据交换。不同 Connector 可以连接不同的远端缓存或数据传输系统。
其中,与高性能节点间 KV Cache 传输关系最直接的方案包括:
| 方案 | 主要用途 | 典型数据路径 |
|---|---|---|
| NixlConnector | PD 分离 KV Cache 传输 | GPU → GPU |
| MooncakeConnector | PD 分离 KV Cache 传输 | GPU → GPU |
| LMCacheConnector | 分层/远端 KV Cache | GPU ↔ CPU/远端存储 |
| Mooncake Store 等 | 共享 KV Cache 池 | GPU ↔ KV Store |
LMCache 提供了更完整的 KV Cache 存储层,可以将 KV Cache 扩展到 CPU DRAM、本地存储和远端缓存。其节点间 P2P Cache Sharing 同样需要高性能的数据搬运能力,当前 P2P transfer engine 采用 NIXL。
从整个软件栈来看,NIXL 因此出现在两个典型位置:
一方面,vLLM 可以直接使用 NixlConnector 在 Prefill 与 Decode 之间传输 KV Cache;
另一方面,LMCache 可以在其缓存管理机制下使用 NIXL 完成节点间的数据搬运。
3. 为什么需要 NIXL
从功能上看,将 KV Cache 从一台服务器传到另一台服务器并不困难。
最直接的方法可以经过 CPU 和传统网络协议栈:
GPU → CPU → Network → CPU → GPU
这种路径包含额外的 GPU/CPU 数据复制,同时占用 PCIe、CPU DRAM 和 CPU Core。对于位于推理关键路径上的大规模 KV Cache,这些额外开销会迅速积累。
高性能分离式推理更希望建立 GPU 到 GPU 的直接数据通路:
GPUDirect RDMA 等技术可以减少 CPU 中转,使网卡直接访问 GPU Memory。
但推理框架若直接操作 RDMA,还需要管理大量底层机制,例如:
- GPU Memory Registration;
- Remote Address;
- Memory Key;
- Endpoint;
- 网络 backend;
- RDMA READ/WRITE;
- 异步 Completion;
- 多 NIC 与不同 Fabric。
这些机制与 vLLM 的 Request、Block Table、Prefix Cache 等推理语义处于不同层次。
NIXL 在二者之间提供了一层统一的数据搬运接口。
推理系统向 NIXL 描述需要访问的内存区域和传输方向,NIXL 再通过 UCX、Libfabric 等 backend 执行具体传输。
于是软件栈形成了清晰的分层关系:
4. NIXL 在推理框架中的定位
NIXL,全称 NVIDIA Inference Xfer Library,其设计目标就是为分布式推理系统提供统一的数据移动层。
从推理框架向下看,NIXL 屏蔽了不同通信 backend 和不同内存介质之间的差异。
从通信系统向上看,NIXL 提供了适合推理框架使用的 Memory Descriptor、Remote Metadata 和异步 Transfer 接口。
可以把一次 KV Cache 数据移动抽象成三个步骤:
- 描述源端和目标端内存;
- 提交 READ 或 WRITE;
- 异步等待传输完成。
NIXL 官方架构中的几个核心对象也围绕这一过程组织:
-
Memory Section 用于描述可访问的内存区域;
-
Metadata 用于描述远端节点及其已注册内存;
-
Transfer Backend 负责完成具体的数据移动;
-
Transfer Handle 用于跟踪异步操作状态。
一个 NIXL descriptor 更接近:
地址 + 长度 + 内存类型 + Device
上层推理系统负责在此前完成 KV Cache 语义到内存地址的转换。
这一点决定了 NIXL 在整个远端 KV Cache 系统中的位置:
KV Cache 管理位于 NIXL 之上,NIXL 将已经确定的数据移动需求转换成具体的跨节点内存传输。
5. vLLM 中的 NixlConnector
vLLM 并不直接让 Scheduler 调用 NIXL API。两者之间存在一层 NixlConnector。 NixlConnector 继承自 vLLM 的 KV Connector 架构,负责将 vLLM 的 Request 和 KV Block 信息转换成 NIXL 可以执行的数据传输操作。
当前实现可以从逻辑上划分为两个部分:
- Scheduler 侧 Connector;
- Worker 侧 Connector。
Scheduler 仍然工作在 vLLM 的 KV Cache 语义层。
它知道当前请求需要哪些 KV Cache、哪些 block 已经存在、哪些 block 需要从远端获得,以及什么时候可以释放 block。
Worker 更接近实际数据。
它负责:
- 注册 GPU KV Cache;
- 维护 NIXL Agent;
- 获取远端 Metadata;
- 将 block ID 映射为实际地址;
- 创建 NIXL descriptor;
- 提交 READ 或 WRITE;
- 检查 transfer completion。
因此,一次数据请求会逐层完成下面的转换:
这里存在一个很重要的分界点。
Scheduler 处理的主要对象是 Block ID。
NIXL 最终处理的主要对象是 Memory Descriptor。
6. KV Block 如何变成 Descriptor
假设 Prefill 节点生成了一个请求对应的 KV Cache。
其中一部分逻辑 block 为:
| Prefill Block | Decode Block |
|---|---|
| 17 | 3 |
| 23 | 8 |
| 41 | 9 |
| 55 | 14 |
这些编号分别属于 Prefill 和 Decode 自己的 KV Cache 地址空间。
因此,Block 17 → Block 3 仍然只是 vLLM 层面的描述。
Worker 需要继续计算两者对应的实际 GPU Memory。
对于 packed KV Cache,vLLM 通常将多个 KV block 放置在一块连续 GPU allocation 中。已知 allocation 的起始地址和每个 block 的 stride 后,就可以计算任意 block 的位置。
从概念上看:
Block Address = KV Cache Base Address + Block ID × Block Stride
随后,Worker 将这些地址和长度组织成 NIXL descriptor。
于是原来的:
Prefill Block 17 → Decode Block 3
在 NIXL 层会进一步变成:
Remote Memory Descriptor → Local Memory Descriptor
NIXL 到这里已经不再看到 Block 17 或 Block 3 的推理语义。
它只需要按照 descriptor 描述执行数据移动。
这也解释了 vLLM 与 NIXL 所使用“传输粒度”的区别。
vLLM 在调度时按照 KV block 选择数据。
NIXL 在执行时按照对应的 memory descriptor 搬运数据。
7. Memory Registration
完成地址映射之后,还有一个 RDMA 数据通路必须解决的问题:Memory Registration。
RDMA 通信需要提前注册可访问内存。对于 GPU Memory,同样需要建立相应的 registration 信息,使通信 backend 能够执行直接 DMA。
因此,vLLM 初始化 KV Cache 后,NixlConnectorWorker 会将相应 GPU Memory 注册到 NIXL。
这一过程包含:
KV Cache allocation → 获取 GPU 地址和大小 → 构造 registration descriptor → NIXL register_memory
对于 packed KV Cache,一块连续 allocation 可以作为一个完整 memory region 注册。
后续每个 KV block 只是这块 region 中不同的地址范围,因此不需要在每个请求到来时重新注册 GPU Memory。
这种设计非常适合推理场景。
KV Cache Memory Pool 通常在 Engine 初始化时就已经创建,其生命周期远长于单个 Request。提前注册整块 KV Cache,使后续请求能够直接构造 descriptor 并提交传输。
这样,一次请求真正进入远端 KV Cache 数据路径时,主要剩下:
Block Mapping → Descriptor → Transfer
昂贵的 Memory Registration 已经从关键路径中移出。
8. Metadata 与远端内存访问
本地 GPU Memory 注册完成后,另一端还需要知道如何访问这些内存。
NIXL 使用 Metadata 描述远端 Agent、Memory Region 和 backend 相关信息。
因此,在真正执行 KV Cache READ/WRITE 之前,Prefill 和 Decode Worker 会首先完成控制信息交换。
从系统角度看,可以将整个通信过程拆成两条路径:
Metadata 路径的数据量很小,主要负责建立远端访问关系。
KV Cache Data Path 承担真正的大规模数据传输。
这种分离使运行时不需要在每次 KV Cache 搬运之前重新建立完整连接。远端节点和 Memory Region 的信息可以缓存并持续复用。
到这里,一次高性能远端传输所需要的准备工作已经完成:
- GPU KV Cache 已分配;
- GPU Memory 已注册;
- 双方已经交换 Metadata;
- Worker 可以根据 KV block 构造 descriptor。
接下来只需要选择实际的数据传输方式。
9. Pull:Decode 发起 READ
vLLM 的 NIXL 数据路径支持 Pull 和 Push 两种模式。
Pull 使用 NIXL READ。
典型过程从 Prefill 开始。
Prefill Scheduler 为请求分配 KV block,Worker 执行模型计算,并将 Prompt 对应的 K/V 写入本地 GPU KV Cache。
Prefill 完成后,相关 block 会继续保留一段时间,等待 Decode 读取。
Decode 收到请求后,在自己的 KV Cache 中分配目标 block。随后,Decode Worker 获得:
- Prefill 的远端 block;
- Decode 的本地 block;
- Prefill 的 NIXL Metadata。
Worker 分别计算本地和远端 descriptor,然后发起 NIXL READ。
完整过程如下:
在这一模式中,KV Cache 的数据方向始终为:
Prefill GPU → Decode GPU
Decode 是传输操作的 initiator。
从 Decode 的视角看,它执行的是:
READ remote Prefill KV into local Decode KV。
vLLM 当前 Pull Worker 中的 _read_blocks() 最终会构造 NIXL READ transfer,并保存异步 handle,后续通过 completion 状态判断对应 KV Cache 是否已经准备完成。
10. Push:Prefill 发起 WRITE
Push 模式将传输控制权放在 Prefill 一侧。
Decode 首先为即将接收的 KV Cache 分配目标 block,并将目标信息发送给 Prefill。
此时 Prefill 在完成 Prompt 计算后,已经知道:
- 自己的源 KV block;
- Decode 的目标 KV block;
- Decode 的远端 NIXL Metadata。
因此 Prefill 可以直接发起 WRITE,将 KV Cache 写入 Decode GPU。
Push 与 Pull 的 KV Cache 数据方向保持一致:
Prefill → Decode
控制方向有所不同。
| 模式 | 发起方 | NIXL Operation |
|---|---|---|
| Pull | Decode | READ |
| Push | Prefill | WRITE |
Push 可以让 Decode 提前准备目标空间,同时让 Prefill 在计算完成后立即启动数据传输。
vLLM 当前已经分别实现了 NixlPullConnector 和 NixlPushConnector,使这两种执行模式可以在同一套 NIXL 数据面上工作。
11. 为什么同时需要 READ 和 WRITE
READ 与 WRITE 都可以完成 Prefill 到 Decode 的 KV Cache 搬运。
差异来自 initiator 的视角。
在 Pull 模式中,Decode 持有自己的 Local Buffer,并从 Prefill 的 Remote Buffer 中读取:
Remote Prefill → Local Decode
对应 NIXL READ。
在 Push 模式中,Prefill 持有 Local Buffer,并向 Decode 的 Remote Buffer 写入:
Local Prefill → Remote Decode
对应 NIXL WRITE。
这种设计与 RDMA one-sided communication 十分契合。
一旦双方 Memory Region 已经注册,并完成远端 Metadata 交换,initiator 就可以直接操作远端内存,不需要为每一批 KV Cache 建立传统的 send/recv 配对。
这使推理框架可以根据自己的调度方式决定由哪一侧驱动传输。
Decode 调度逻辑更适合主动拉取时,可以使用 READ。
Prefill 能够提前获得目标地址并希望完成后立即发送时,可以使用 WRITE。
NIXL 的数据面保持一致,上层只需要改变 transfer operation 和 descriptor 的 local/remote 关系。
12. NIXL 落到底层 RDMA
NIXL 本身仍然位于 RDMA 之上。一次典型的 vLLM KV Cache 传输继续向下展开后,可以表示为:
各层承担不同职责。
vLLM Scheduler 维护 Request、Prefix Cache 和 KV Block。
NixlConnector 将 KV block 转换为对应的 Memory Descriptor。
NIXL 管理 Memory Registration、Remote Metadata、READ/WRITE 请求和异步状态。
UCX 或 Libfabric 实现具体通信 backend。
RDMA 和 NIC 最终完成跨节点 DMA。
13. NIXL 带来的抽象边界
当系统进入 NixlConnector 之前,数据仍然具有丰富的推理语义。
例如:
- Request ID;
- Token;
- Prefix;
- KV Block;
- Cache Hit;
- Block Table。
进入 NIXL 后,数据逐渐变成通信系统能够处理的形式:
- Local Address;
- Remote Address;
- Length;
- Memory Type;
- Device;
- READ/WRITE。
可以将这条语义变化总结如下:
| 层次 | 核心对象 |
|---|---|
| vLLM Scheduler | Request / Token |
| KV Cache Manager | Logical KV Block |
| NixlConnector | Block Mapping |
| NIXL | Memory Descriptor |
| UCX / Libfabric | Transport Endpoint |
| NIC | DMA / Network Packet |
这个边界对于后续 KV Cache 优化尤其重要。
Prefix-aware 传输、按需加载以及 KV Cache 命中判断,需要访问 Request 和 Block 层的信息,因此更接近 vLLM 或 LMCache。
数据压缩、descriptor 聚合、传输调度以及 SmartNIC/DPU 数据处理,则更加接近 NixlConnector、NIXL 和底层 Transport。
随着数据沿软件栈向下移动,KV Cache 的模型语义逐渐消失,最终成为一组连续或离散的内存区域。
14. 总结
远端 KV Cache 已经成为分布式大模型推理中的基础数据通路。vLLM 通过 KVConnector 管理不同形式的远端 KV Cache 交换;LMCache 进一步提供多级缓存和跨节点 KV Cache 管理。在需要低延迟、高带宽 GPU 数据搬运的路径上,NIXL 提供了面向推理系统的数据传输抽象。
在 vLLM 中,一次 NIXL KV Cache 传输可以归纳为:
-
Scheduler 确定需要传输的 KV Block;
-
NixlConnector 将 Block 映射为 GPU Memory Descriptor;
-
NIXL 根据远端 Metadata 构造 READ 或 WRITE;
-
UCX、Libfabric 等 backend 将请求映射到底层 RDMA;
-
NIC 最终完成 GPU 之间的数据搬运。
因此,NIXL 位于两个具有明显不同语义的系统之间:
-
上层是以 Request、Prefix 和 KV Block 为核心的推理框架;
-
下层是以 Memory Region、DMA 和网络传输为核心的通信系统。
NixlConnector 和 NIXL 共同完成了这两层之间的转换。