在设计 KVCache 的时候,需要考虑我们能接触到的硬件的性能,在作出设计,包含显卡、网卡、HBM、DRAM 和 SSD。
基础知识
大模型生成回答时,会不断用到前面的上下文。在常规 Transformer 中,attention 会为历史 token 保存 Key 和 Value 两类向量,也就是 KVCache。保留这些向量后,生成下一个 token 时就不用把历史 token 的 K/V 重新计算一遍。代价是上下文越长、同时处理的请求越多,KVCache 占用的内存也越大。
为了理解后面的取舍,可以把一次推理分成两个阶段。Prefill 是处理输入提示词、建立其 KVCache 的阶段。Decode 是利用已有 KVCache,逐步生成后续 token 的阶段。对于读取全部历史上下文的注意力计算,也就是本文所说的 dense attention,历史 KV 会在后续 decode 中反复被访问。
既然要反复读,最直接的办法就是把活跃请求的 KV 放在 GPU 的 HBM 中。HBM 是 High Bandwidth Memory,也就是高带宽内存,通过封装内的宽接口连接 GPU。它很快,但容量有限,而且模型权重、计算临时数据也要占用它。因此,HBM 用完之后,系统通常会考虑把暂时不用的 KV 搬到 CPU 内存,或者保存到其他机器上,以便之后再取回来。
本文把 CPU 一侧的普通系统内存简称为 DRAM,也会称它为 Host 内存。严格说,HBM 本身也是一种 DRAM 技术,这里的区分是为了说明数据放在 GPU 旁边,还是放在 CPU 一侧。NIC 就是网卡,CX-7 和 CX-8 是 NVIDIA ConnectX 系列网卡的两个代际。跨机器共享 KV 时,网卡负责把这些数据运过去。
由此会出现两个不同的问题。第一个是“有多少空间可以保存 KV”,这是容量问题。第二个是“GPU 要用的时候,能不能及时把 KV 搬回来”,这是带宽和延迟问题。增加 DRAM 可以解决一部分容量压力,但不会自动缩短搬运时间。
后文提到的 KV 恢复,泛指把一个推理请求要用的 KV 从本地或远端 DRAM 加载到 GPU。它可能是新请求复用了缓存前缀,也可能是暂停的请求重新开始执行。一个推理请求可以分布在多张 GPU 上,一次恢复也可以拆成很多次传输。
一张大表看架构
| 资源 | 容量 | 带宽数值与口径 | 同时双向 / 混合读写如何计算 |
|---|---|---|---|
| CX-7,400 Gb/s 配置 | — | 单向 50 GB/s,网卡线速 | 全双工:理论上可同时发送 50、接收 50 GB/s;双向合计 100 GB/s |
| CX-8,800 Gb/s 总速率配置 | — | 单向合计 100 GB/s,整卡网络标称总速率 | 全双工:理论上可同时发送合计 100、接收合计 100 GB/s;双向合计 200 GB/s;Ethernet 的 2×400GbE 配置需同时利用两个端口 |
| PCIe 4.0 ×16 | — | 单向约 31.5 GB/s,编码后理论值 | 全双工:两个方向理论上各约 31.5 GB/s;双向合计约 63 GB/s |
| PCIe 5.0 ×16 | — | 单向约 63 GB/s,编码后理论值 | 全双工:两个方向理论上各约 63 GB/s;双向合计约 126 GB/s |
| PCIe 6.0 ×16 | — | 单向标称 128 GB/s,原始带宽,未扣除协议开销 | 全双工:两个方向标称各 128 GB/s;双向合计 256 GB/s;有效载荷吞吐低于此值 |
| H200 SXM HBM | 141 GB/卡 | 最高 4.8 TB/s,单张 GPU 所有 HBM 通道合计峰值 | 纯读理论上限为 4.8 TB/s,无需除以 2;混合读写共享总带宽,不能同时读 4.8、写 4.8 TB/s |
| B200 SXM HBM | 180 GB/卡 | 最高 8 TB/s,单张 GPU 所有 HBM 通道合计峰值 | 纯读理论上限为 8 TB/s,无需除以 2;混合读写共享总带宽,不能同时读 8、写 8 TB/s |
速度差距
当 KV 已经在 GPU 的 HBM 中时,GPU 通过本地内存接口读取它。当 KV 在本机 CPU 内存中时,需要通过 PCIe 搬到 GPU。PCIe 是连接 CPU、GPU、网卡等设备的高速互联。如果 KV 在另一台机器,还需要先跨越网络,再经过目标侧的设备连接到达 GPU。
这几种路径的物理条件不同。HBM 紧邻 GPU,可以使用大量并行的短距离连接。PCIe 要连接板卡和主机,网络还要连接不同机器,接口规模、传输距离和可靠性要求都不同。因此,GPU 读取本地 HBM 的速度远高于从外部获取数据,是很自然的结果。
详细说明
用这张表估算 KV 搬运时,最重要的是选对路径和方向。比如,从 CPU 内存往 GPU 写入一份 KV,只能使用 PCIe 的一个方向,不能拿双向合计的 126 GB/s 来算。数据已经在 B200 的 HBM 中之后,GPU 才能通过本地内存接口反复读取它。此时的标称上限是 8 TB/s = 8,000 GB/s,与从外部搬进来的速度属于两条不同路径。
表里的“峰值”也不等于应用实际能拿到的速度。接口上传输的内容除了 KV 本身,还有报头、控制信息和用于保证可靠性的冗余数据。设备排队、访问模式以及软件提交是否及时,也会影响最后测到的吞吐。因此,做容量规划时可以参考规格,估算请求延迟时则应尽量使用实测的有效带宽,也就是每秒真正搬完了多少 KV 字节。
如果想进一步理解为什么 PCIe 的数字不是整齐的 32、64、128,可以从 PCIe 5.0 ×16 看起。×16 表示这条连接有 16 条并行 lane,也就是传输通道。Gen5 是 PCIe 5.0 的简称,每条通道的原始单向速率是 32 Gb/s,但它采用 128b/130b 编码,即每 128 bit 内容需要占用 130 bit 传输空间。因此,扣除这部分开销后得到:
1 | 32 Gb/s × 16 条通道 ÷ 8 × 128/130 ≈ 63.0 GB/s |
PCIe 4.0 每条通道的速率是 Gen5 的一半,所以同样计算得到约 31.5 GB/s。这些数值还没有扣除 TLP(Transaction Layer Packet,事务层数据包) 的开销。TLP 是 PCIe 用来表达内存读写等操作的数据包,其中除了数据,还要携带地址、操作类型等信息。也就是说,表里的 63 GB/s 依然不是应用有效吞吐。
PCIe 6.0 在更高速度下还要兼顾传输可靠性,因此引入了固定大小的数据打包和纠错机制。FLIT(Flow Control Unit,流控单元) 是链路使用的固定大小传输单元,PCIe 6.0 将其定为 256 byte。FEC(Forward Error Correction,前向纠错) 则通过附加冗余信息,让接收方能直接纠正一部分传输错误,减少重传。一个 FLIT 中也要容纳控制和校验信息,所以 256 byte 不会全部用于承载应用数据。这就是表中 128 GB/s 被标为“原始带宽”的原因。理解 KV 搬运时,不必记住它们的每个字段,只需要知道可靠传输会占用一部分链路能力。PCI-SIG 对 FLIT 和 FEC 的解释
关于 GDR
跨机器读取 KV 时,如果每一批数据都需要 CPU 执行网络收发、拷贝和转交,CPU 开销和额外内存访问就可能拖慢过程。RDMA(Remote Direct Memory Access,远程直接内存访问) 允许网卡直接访问事先授权和注册的内存,从而减少 CPU 对数据搬运的参与。这里的“注册”是把指定内存区域及其访问权限交给网卡,并不是把数据复制一份。
GPUDirect RDMA 把这条直接访问路径延伸到 GPU 显存。在硬件和驱动支持的情况下,远端数据可以由目标网卡直接写入 GPU,不必先落到目标机器的 CPU 内存,再进行一次 CPU 内存到 GPU 的复制。下面将它与另外两条路径对比,其中本地复制可以通过 CUDA,也就是 NVIDIA 的 GPU 编程与运行时平台来发起。
1 | 本机内存中已有 KV,使用本地 CUDA 拷贝: |
GPUDirect RDMA 省掉的是目标侧的内存中转,并没有省掉网卡到 GPU 的 PCIe 连接。远端取数、网络传输和目标 GPU 接收数据仍然要花时间,所以整体速度取决于这条路径中最慢的一段。NVIDIA GPUDirect RDMA 文档
实际服务器还需要考虑设备接在哪里。多路 CPU 机器的内存通常分属不同 CPU,本地访问与跨 CPU 访问的代价不同,这就是 NUMA(非统一内存访问)。同样,多个 GPU 和网卡也可能共用 PCIe 交换芯片及其上行连接。即使每个设备单独测试都很快,并发时仍可能争用同一条通道。后面计算多 GPU 的总带宽时,必须说明是否存在这种共享瓶颈。
上图第一条是可选的本地拷贝方案,并不等于 KVCache 当前一定这么实现。
关于网卡
升级网卡的直接收益是提高网络这段的传输能力。CX-7 的 400G 配置单向为 50 GB/s,CX-8 的 800G 总速率配置则为 100 GB/s。只有原来的瓶颈确实在网络上,而且源端取数和目标端写入都能跟上时,KV 搬运才可能获得接近两倍的提升。
先看端口。InfiniBand 和 Ethernet 是两种不同的网络体系,RDMA 可以在它们上面实现,但可用的端口速率不一定相同。当前官方手册列出的 CX-8 C8180 支持单端口 800G XDR InfiniBand,XDR 是这里的 InfiniBand 速率代际名称。它的 Ethernet 配置包括两个 400GbE 端口,C8240 的默认配置也是两个 400GbE 端口。因此,在这类 Ethernet 配置下,单向 100 GB/s 是两个端口加起来的能力,每个端口仍是单向 50 GB/s。搬运软件必须把工作分配到两个端口,才能利用整卡总带宽。CX-8 官方端口配置
再看网卡接入主机的通道。CX-8 支持 PCIe 6.0 ×16。如果服务器只有 PCIe 5.0,适用的 Socket Direct 型号可以通过辅助卡和线缆,额外接入一条 PCIe ×16 连接,让网卡使用两条 Gen5 路径的总 I/O 能力。它解决的是单条 Gen5 ×16 带宽不足的问题,需要对应的网卡、插槽和配件。只插一条 Gen5 ×16 时,经过该连接的数据仍受约 63 GB/s 的单向理论上限约束。CX-8 接口与 Socket Direct 说明
最后还要看 GPU 怎样接收数据。即使网卡已经接在 Gen6 上,如果某个请求的 KV 全部要写入同一张、仅通过 Gen5 ×16 接收数据的 GPU,这条 GPU 连接仍会限制速度。多个 GPU 分担传输时,可能利用上更高的网卡总带宽,但前提是这些 GPU 的连接没有再汇聚到一条更窄的共享通道上。因此,换成 CX-8 后不能直接把所有恢复耗时除以 2。
KVCache 的架构比较
比较实现方案时,先问“这份 KV 接下来会怎样被使用”通常比先选传输库更有用。如果它马上要被连续读取几十次,放在 HBM 的价值很大。如果它只是可能被未来请求复用,放在更便宜的 DRAM 中通常更合适。
在这个前提下,各种方案可以理解为对“容量、第一次等待时间、后续重复访问成本”的不同取舍。
| 方案 | 为什么会选择它 | 为此要付出什么 |
|---|---|---|
| 活跃 KV 留在 HBM | 后续每个生成步骤都能直接访问本地内存,适合对响应速度要求高的请求。 | 需要为当前请求保留显存,能够同时容纳的上下文和请求数有限。 |
| 本机 DRAM 保存 KV,用时再搬入 HBM | 利用本机较大的内存扩充缓存,命中后只需经过本地设备连接。 | 加载速度受 PCIe 和内存位置影响,也不能天然与其他机器共享。 |
| 远端 DRAM 保存 KV,直接写入目标 HBM | 多个推理实例可以复用同一批历史 KV,并省掉目标侧内存中转。 | 请求开始计算前要等待网络搬运,还可能与其他 GPU 通信争用网卡。 |
| 远端 KV 先放本机 DRAM,再搬入 HBM | 可以提前从远端取回数据,并在本机留下后续可复用的副本。 | 第一次使用需要两段搬运,本机内存也要额外读写一遍。 |
| Prefill GPU 直接把 KV 交给 Decode GPU | 把处理提示词和生成回答交给不同 GPU,以便分别调度这两类工作。 | 源 GPU 必须等传输完成才能释放对应数据,跨节点交接仍受网络限制。 |
| 每步计算前从 DRAM 读取所需 KV | 即使 GPU 装不下全部历史 KV,也能继续运行较长上下文。 | 同一份历史 KV 可能被反复搬运,每个 token 的生成时间容易受外部带宽限制。 |
| 更冷的 KV 放到 NVMe SSD | 对很久以后才可能复用的数据,用 SSD 换取更大的低成本容量。 | SSD 还引入存储 I/O 等待,需要更早预取并单独评估恢复时间。 |
其中,把 prefill 和 decode 分给不同 GPU 的方式通常叫 P/D 分离。把 GPU 内存中的数据转移到 CPU 内存或存储设备,通常叫 offload,也就是卸载。表中的 NVMe 是访问固态硬盘的协议,NVMe SSD 在这里承担较慢但容量更大的存储层。这些设计可以组合使用,例如活跃 KV 留在 HBM,历史前缀保存在共享 DRAM,更冷的数据再下沉到 SSD。
这里还要区分两种收益。前缀缓存是在模型、位置等计算条件兼容时,复用已有提示词前缀的 KV,从而省掉这一段 prefill。它并不会自动减少后续 dense decode 对历史 KV 的读取,所以前缀命中率提高,通常首先影响的是请求开始生成之前的等待时间。vLLM 前缀缓存说明
已有系统也遵循这种思路。SGLang HiCache 把 GPU、Host 内存和共享存储组织成多层缓存,并支持预先取回后面要用的数据。Mooncake 则围绕共享 KV 组织存储和推理调度。它们提供了可参考的实现方式,但具体收益仍要放回自己的请求复用模式和硬件路径中判断。HiCache 设计、Mooncake 论文