vllm 学习笔记

Appearance

将不活跃的 KV Cache 从 GPU 显存卸载到 CPU 内存或磁盘，在有限显存下支持更长上下文或更大 batch。

为什么需要 KV Cache Offloading ​

KV Cache 随序列长度和 batch size 线性增长，很快耗尽 GPU 显存。在长上下文场景（128K+ tokens）中，KV Cache 可能占用数十 GB 显存，远超模型权重本身。KV Cache Offloading 将暂时不活跃的 KV block 搬移到 CPU 内存（或磁盘），在需要时再搬回 GPU，以时间换空间。

核心原理 ​

  • Swap 机制：当 GPU 显存不足时，调度器选择部分 sequence 的 KV Cache swap 到 CPU pinned memory。
  • Recompute 机制：也可选择丢弃 KV Cache 并在需要时从原始 token 重新计算，适用于重计算成本较低的场景。
  • 预取（Prefetch）：在 sequence 即将被调度前异步将其 KV Cache 从 CPU 预取到 GPU，掩盖传输延迟。
  • 层级管理：Hot KV 在 GPU、Warm KV 在 CPU、Cold KV 可选择丢弃，形成多级缓存层次。

在源码中的实现 ​

  • vllm/v1/kv_offload/cpu/ — CPU 层 swap 实现（swap_blocks_triton.py、gpu_worker.py），执行 GPU-CPU 数据搬运。
  • vllm/v1/kv_offload/base.py — OffloadingWorker 抽象（原 OffloadingHandler，运行于 worker 进程）与 LookupResult 枚举（HIT / MISS / HIT_PENDING / RETRY）；CPU pinned memory 注册（cudaHostRegister）在 worker 初始化时同步完成（mmap 区域经 pin_mmap_region()、其余 CPU 张量经 PyTorch pin_memory=True）。注：PR #45850 曾把该注册挪到后台 CPUTensorPinThread 异步执行以避免启动阻塞，但因引入竞态（传输在 pin 完成前就开始）仅 5 天后即被 #46958 回退，当前为同步实现。
  • vllm/v1/core/sched/scheduler.py — Scheduler 在显存不足时触发 preempt 决策，并协调 connector 的异步 load/store。
  • vllm/v1/core/block_pool.py — 物理块池，追踪哪些 block 在 GPU、哪些被卸载。
  • vllm/v1/kv_offload/tiering/ — 多级卸载编排器：TieringOffloadingManager 以 CPU 为主层级（gateway，唯一可直接访问 GPU），其下挂磁盘（fs/）、对象存储（obj/）等 secondary tier，secondary 层须经 CPU 中转。AsyncLookupManager 用后台线程批量做存在性检查。
  • vllm/v1/kv_offload/tiering/fs/ — 文件系统分层：FileSystemTierManager + DualQueueThreadPool 实现磁盘级 KV 缓存。
  • vllm/v1/kv_offload/tiering/p2p/ — P2P/PD 分离 secondary tier（v0.26 新增，PR #42285 + #48021）：P2PSecondaryTierManager 让 decoder 经 RDMA 从 prefiller 的 CPU 缓存拉块。控制面用 ZMQ（control/ZmqTransport，默认 VLLM_P2P_SIDE_CHANNEL_HOST:PORT=localhost:5710，端口按 DP rank 偏移），数据面用 NIXL RDMA（data/NixlTransport，backends 默认 ["UCX"]，可选 MOONCAKE/GDS_MT/LIBFABRIC，选择逻辑与主 NixlConnector 一致）。两种角色：remote_prefiller（经典 PD，立即 HIT）、remote_kv_source（对称 P2P，先 RETRY 再批量探测）。tiering/base.py 新增 ParentManager ABC 供 secondary tier 反查主层级。握手还通告 get_none_hash_seed()（NONE_HASH 种子跨实例对齐，#51875）。
  • vllm/v1/kv_cache_interface.py — AttentionSpec.page_size_bytes 拆为 unpadded_page_size_bytes（含 per-token-head 量化 scale）+ padding（PR #48411），修正了 per-token-head 量化下卸载页传输尺寸的低估。
  • vllm/v1/simple_kv_offload/ — Simple offload 的 manager/worker；命中计算尊重 skip_reading_prefix_cache（#54998），DMA region 注册支持同 group 混合 page size（#54756，按 block_stride 注册、区分 layer/block-outermost）。
  • vllm/distributed/kv_transfer/kv_connector/v1/offloading/ — OffloadingConnectorWorker（封装 OffloadingWorker）与 OffloadingConnector，通过 spec_name 选择 CPUOffloadingSpec（单 CPU 层）或 TieringOffloadingSpec（多级）；调度器接入 prefix_cache_retention_interval（#51886）并在 SWA 命中查找时放大首串窗口防欠覆盖（#55712）。
  • vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/ — MooncakeStoreConnector 分布式 KV 缓存连接器（支持多组 KV cache、流水线并行 PD 分离、混合 DCP 前缀命中 #53324、Mamba align 边界状态 #51358/#53663）；KV load 由 VLLM_MOONCAKE_LOAD_RECV_THREADS 控制多线程并行接收。
  • vllm/distributed/kv_transfer/kv_connector/v1/nixl/ — NIXL 连接器，拆分为 NixlPullConnector（D 端 READ 拉取）与 NixlPushConnector（P 端 WRITE 直接写入 D 端显存）。
  • vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py — Mamba1/Mamba2/GDN 的卷积状态按 TP rank 拆解，使状态化 SSM 模型也能做 NIXL P/D 分离。
  • vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py — Simple CPU Offload 后端（支持 DSV4 混合注意力、PCP/DCP、reset_cache()）。
  • vllm/distributed/ec_transfer/ec_connector/ — Encoder cache 传输（EPD）：mooncake/ 的 ECMooncakeConnector（#41567，约 3200 行）与 CPU P2P 版（ZMQ 控制面 + NIXL 数据面，#47941），把多模态 encoder 输出转移到代理侧。

相关概念 ​

  • kv-cache — 被卸载的核心数据结构
  • paged-attention — Block 粒度的管理使 swap 高效（以 block 为单位搬运）
  • continuous-batching — 调度器的 swap 决策与 continuous batching 深度耦合
  • prefix-caching — 常用前缀的 KV Cache 可常驻 GPU，非活跃部分卸载到 CPU

反向链接