vllm分析(六)——KV cache offload

KV cache offload到外部存储（比如本机cpu内存），依然使用connector接口。不管是pd分离场景的KV cache 传输]还是KV cache offload， KV cache需要有两个过程：store(存储到外部) 和 load（从外部加载）。为了处理load和store，vllm的scheduler和worker之间传递信息。pd分离场景的kv cache的处理，参考上一篇博客。

OffloadingConnector 整体架构

text 复制代码

┌─────────────────────────────────────────────────────────────────┐
│                     vLLM Scheduler Process                       │
│  ┌───────────────────────────────────────────────────────────┐  │
│  │              OffloadingConnector (SCHEDULER)               │  │
│  │  ┌─────────────────────────────────────────────────────┐  │  │
│  │  │           OffloadingConnectorScheduler               │  │  │
│  │  │  - get_num_new_matched_tokens() → 可加载 token 数    │  │  │
│  │  │  - update_state_after_alloc() → 创建 load 作业       │  │  │
│  │  │  - build_connector_meta() → 构造 store 作业          │  │  │
│  │  │  - update_connector_output() → 处理 worker 完成信息  │  │  │
│  │  │  - _req_status, _jobs, _block_id_to_pending_jobs     │  │  │
│  │  └─────────────────────────────────────────────────────┘  │  │
│  └───────────────────────────┬───────────────────────────────┘  │
└──────────────────────────────┼──────────────────────────────────┘
                               │ OffloadingConnectorMetadata
                               │ (load_jobs, store_jobs, jobs_to_flush)
                               ▼
┌─────────────────────────────────────────────────────────────────┐
│                      vLLM Worker Process                         │
│  ┌───────────────────────────────────────────────────────────┐  │
│  │              OffloadingConnector (WORKER)                  │  │
│  │  ┌─────────────────────────────────────────────────────┐  │  │
│  │  │           OffloadingConnectorWorker                  │  │  │
│  │  │  - register_kv_caches() → 注册 GPU KV cache          │  │  │
│  │  │  - start_kv_transfers() → 提交 load 作业             │  │  │
│  │  │  - prepare_store_kv() → 延迟 store 作业              │  │  │
│  │  │  - get_finished() → 收集完成作业 → 返回元数据        │  │  │
│  │  │  - _load_jobs, _unsubmitted_store_jobs               │  │  │
│  │  └───────────────────────────┬─────────────────────────┘  │  │
│  │                              │                              │  │
│  │  ┌───────────────────────────▼─────────────────────────┐  │  │
│  │  │              OffloadingWorker                        │  │  │
│  │  │  - 路由 (src_medium, dst_medium) → Handler          │  │  │
│  │  │  - transfer_async() / get_finished() / wait()       │  │  │
│  │  └───────────────────────────┬─────────────────────────┘  │  │
│  │                              │                              │  │
│  │  ┌───────────────────────────▼─────────────────────────┐  │  │
│  │  │  SingleDirectionOffloadingHandler (CPU offload)      │  │  │
│  │  │  - GPU→CPU 或 CPU→GPU                                │  │  │
│  │  │  - CUDA stream + events 控制异步拷贝                 │  │  │
│  │  │  - swap_blocks_batch 批量传输                        │  │  │
│  │  └─────────────────────────────────────────────────────┘  │  │
│  └───────────────────────────────────────────────────────────┘  │
└─────────────────────────────────────────────────────────────────┘
                               │ KVConnectorOutput
                               │ (completed jobs)
                               ▼
                         (回到调度器)

OffloadingConnector

python 复制代码

class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
    def __init__(
        self,
        vllm_config: VllmConfig,
        role: KVConnectorRole,
        kv_cache_config: KVCacheConfig,
    ):
        super().__init__(vllm_config, role, kv_cache_config)

        spec = OffloadingSpecFactory.create_spec(vllm_config, kv_cache_config)

        self.connector_scheduler: OffloadingConnectorScheduler | None = None
        self.connector_worker: OffloadingConnectorWorker | None = None
        if role == KVConnectorRole.SCHEDULER:
            self.connector_scheduler = OffloadingConnectorScheduler(spec)
        elif role == KVConnectorRole.WORKER:
            self.connector_worker = OffloadingConnectorWorker(spec)

涉及的类：
OffloadingConnectorScheduler
OffloadingConnectorWorker

load 作业的处理流程

Scheduler 向 Worker 下发 load作业

text 复制代码

┌─────────────────────────────────────────────────────────────────────────────┐
│                          Scheduler 侧 (调度器)                               │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
        Step 1: 调度决策 (schedule)
            ├─ get_num_new_matched_tokens → 确定需要加载的 token 数
            ├─ update_state_after_alloc → 创建 load_jobs
            └─ build_connector_meta → 封装到 SchedulerOutput.kv_connector_metadata
                                    │
        Step 2: 发送 SchedulerOutput 给 Worker (通过执行器)
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                          Worker 侧 (模型执行前)                               │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
        Step 3: GPUModelRunner.execute_model 调用 pre_forward
                                    │
            ┌─────────────────────────────────────────────────────────────┐
            │  ActiveKVConnector.pre_forward(scheduler_output)           │
            │                                                             │
            │  1. 从 scheduler_output 取出 kv_connector_metadata          │
            │  2. 调用 self.kv_connector.handle_preemptions(metadata)    │
            │  3. 调用 self.kv_connector.bind_connector_metadata(metadata)│
            │  4. 调用 self.kv_connector.start_load_kv(forward_context)  │
            └─────────────────────────────────────────────────────────────┘
                                    │
            ┌─────────────────────────────────────────────────────────────┐
            │  OffloadingConnector.start_load_kv (继承自 KVConnectorBase) │
            │    调用 self.connector_worker.start_kv_transfers(metadata)  │
            │    → 提交所有 load_jobs 到 OffloadingWorker                 │
            └─────────────────────────────────────────────────────────────┘
                                    │
        Step 4: OffloadingWorker 路由到 Handler，执行异步传输
            （cpu_to_gpu_handler.transfer_async 发起异步拷贝）
                                    │
        Step 5: 模型执行 (model forward)
            （load 作业在后台传输，不阻塞计算）

start_kv_transfers 中触发 store 和 load 操作，数据传输是异步的。 _unsubmitted_store_jobs存储的是上一轮提交的store任务。

python 复制代码

# https://github.com/vllm-project/vllm/blob/v0.21.0/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py#L89
class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
        assert self.connector_worker is not None
        assert isinstance(self._connector_metadata, OffloadingConnectorMetadata)
        self.connector_worker.start_kv_transfers(self._connector_metadata)
# https://github.com/vllm-project/vllm/blob/v0.21.0/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py#L295
class OffloadingConnectorWorker:
    def start_kv_transfers(self, metadata: OffloadingConnectorMetadata):
        for job_id, transfer_spec in self._unsubmitted_store_jobs:
            success = self.worker.transfer_async(job_id, transfer_spec)
            assert success
        self._unsubmitted_store_jobs.clear()

        for job_id, entry in metadata.load_jobs.items():
            self._load_jobs[job_id] = entry.req_id
            success = self.worker.transfer_async(job_id, entry.transfer_spec)
            assert success

Worker向Scheduler返回load 任务完成状态

text 复制代码

┌─────────────────────────────────────────────────────────────────────────────┐
│                    Worker 侧 (模型执行后)                                     │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
        Step 6: GPUModelRunner.execute_model 调用 post_forward
                                    │
            ┌─────────────────────────────────────────────────────────────┐
            │  ActiveKVConnector.post_forward(scheduler_output,           │
            │                            wait_for_save=True)              │
            │                                                             │
            │  1. 如果 wait_for_save: self.kv_connector.wait_for_save()   │
            │  2. 调用 self.kv_connector.get_finished(finished_req_ids)   │
            │  3. 获取 get_block_ids_with_load_errors() 等                │
            │  4. 返回 KVConnectorOutput                                  │
            └─────────────────────────────────────────────────────────────┘
                                    │
            ┌─────────────────────────────────────────────────────────────┐
            │  OffloadingConnector.get_finished (worker 侧)               │
            │                                                             │
            │  1. 调用 self.connector_worker.prepare_store_kv(metadata)   │
            │     → 将 metadata.store_jobs 添加到 _unsubmitted_store_jobs │
            │     （延迟提交 store 作业）                                  │
            │  2. 调用 self.connector_worker.get_finished(finished_req_ids)│
            │     → 返回 finished_sending, finished_recving               │
            │     （finished_recving 包含已完成 load 的请求 ID）           │
            └─────────────────────────────────────────────────────────────┘
                                    │
        Step 7: 返回 KVConnectorOutput 给 Scheduler
            （包含 finished_recving、completed_jobs 等）
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                          Scheduler 侧 (调度器)                               │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
        Step 8: Scheduler.update_from_output 处理 KVConnectorOutput
                                    │
            ┌─────────────────────────────────────────────────────────────┐
            │  _update_from_kv_xfer_finished(kv_connector_output)         │
            │     - 调用 connector.update_connector_output()              │
            │     - 遍历 finished_recving:                                │
            │        将 req_id 加入 finished_recving_kv_req_ids           │
            │     - 遍历 finished_sending: 释放块                         │
            └─────────────────────────────────────────────────────────────┘
                                    │
        Step 9: 下一轮 schedule 中提升请求
            _try_promote_blocked_waiting_request 检查 finished_recving_kv_req_ids
            → 调用 _update_waiting_for_remote_kv 缓存块，状态改为 WAITING

要点：

post_forward 中的 get_finished 会触发 prepare_store_kv，从而将 store 作业存入 _unsubmitted_store_jobs 队列，但此时不启动传输。
store 作业的实际传输将在下一次 start_kv_transfers（即下一个 step 的 pre_forward）开始时提交，实现与模型计算的重叠。

get_num_new_matched_tokens

text 复制代码

get_num_new_matched_tokens(request, num_computed_tokens)
    │
    ├─ 1. 获取或创建 RequestOffloadState
    │      - 若新请求，创建并保存到 _req_status
    │      - 否则复用现有状态，并清空之前的 block_ids
    │
    ├─ 2. 更新 offload keys（调用 req_status.update_offload_keys()）
    │      - 基于 request.block_hashes 生成该请求的所有 offload keys
    │      - 每个 offload key 标识一个外部存储块
    │
    ├─ 3. 记录本地已计算 token 数
    │      req_status.num_locally_computed_tokens = num_computed_tokens
    │
    ├─ 4. 调用 _lookup(req_status) 获取命中 token 数
    │      
    │
    ├─ 5. 若为新请求，更新命中的 offload 块数
    │      req_status.update_num_hit_blocks(num_computed_tokens + (hit_tokens or 0))
    │
    ├─ 6. 调用 _touch(req_status) 更新缓存访问热度
    │      - 全注意力组：touch 所有 offload keys
    │      - 滑动窗口组：只 touch 最近窗口内的 keys
    │
    └─ 7. 返回 (hit_tokens, bool(hit_tokens))

_lookup负责实际查询外部缓存中存在且可读的连续块，并返回可加载的 token 总数。
_touch更新缓存管理器的访问热度（用于 LRU/ARC 淘汰策略。

store 作业的处理流程

text 复制代码

┌─────────────────────────────────────────────────────────────────────────────┐
│                          Scheduler 侧 (调度器)                               │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
        Step 1: 构建 Store 作业 (_build_store_jobs)
                                    │
            ┌─────────────────────────────────────────────────────────────┐
            │  遍历本轮调度的所有请求:                                      │
            │    - 计算 num_offloadable_tokens (本轮后 token 总数)         │
            │    - 找出需要存储的新 offload 块 (从 next_stored_block_idx)  │
            │    - 过滤跳过块 (block_id == 0)                              │
            │    - 调用 manager.prepare_store() 处理缓存驱逐和分配         │
            │    - 生成 job_id，记录作业状态 (pending_count = world_size)  │
            │    - 区分滑动窗口块和非滑动窗口块，加入监视集合               │
            │    - 返回 store_jobs 字典 {job_id: TransferJob}              │
            └─────────────────────────────────────────────────────────────┘
                                    │
        Step 2: 封装到 ConnectorMetadata (build_connector_meta)
            metadata.store_jobs = store_jobs
                                    │
        Step 3: 发送给 Worker (通过 SchedulerOutput)
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                          Worker 侧 (工作器)                                  │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
        Step 4: 预处理阶段 (pre_forward)
            handle_preemptions(metadata)
                └─ 若有抢占，立即提交 _unsubmitted_store_jobs 中所有作业
                   (确保数据不丢失)
                                    │
        Step 5: 模型前向计算期间
            （GPU 生成新 token，store 作业暂未启动）
                                    │
        Step 6: 后处理阶段 (post_forward → get_finished)
                                    │
            ┌─────────────────────────────────────────────────────────────┐
            │  OffloadingConnector.get_finished()                         │
            │                                                             │
            │  1. 调用 self.connector_worker.prepare_store_kv(metadata)   │
            │     → 遍历 metadata.store_jobs                              │
            │     → 将 (job_id, transfer_spec) 追加到                     │
            │        _unsubmitted_store_jobs 列表                         │
            │     （仅暂存，不启动传输）                                   │
            │                                                             │
            │  2. 调用 self.connector_worker.get_finished()               │
            │     → 返回已完成传输的作业信息（load 和已完成 store）        │
            └─────────────────────────────────────────────────────────────┘
                                    │
        Step 7: 返回 KVConnectorOutput 给 Scheduler
            （包含 completed_jobs、finished_sending 等）
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                   下一个 Step 的 Worker 侧 (延迟提交)                        │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
        Step 8: 下一轮 pre_forward → start_kv_transfers
                                    │
            ┌─────────────────────────────────────────────────────────────┐
            │  OffloadingConnectorWorker.start_kv_transfers(metadata)     │
            │                                                             │
            │  1. 首先处理 _unsubmitted_store_jobs 中的延迟作业:           │
            │     for (job_id, spec) in _unsubmitted_store_jobs:          │
            │         worker.transfer_async(job_id, spec)                 │
            │     _unsubmitted_store_jobs.clear()                         │
            │                                                             │
            │  2. 再处理本轮新的 load_jobs (如果有)                        │
            └─────────────────────────────────────────────────────────────┘
                                    │
        Step 9: OffloadingWorker 路由到 Handler (GPU→CPU)
                                    │
            ┌─────────────────────────────────────────────────────────────┐
            │  gpu_to_cpu_handler.transfer_async(job_id, spec)            │
            │     - 解析 TransferSpec (src: GPU 块, dst: CPU 块)          │
            │     - 计算指针，在独立 CUDA stream 上异步拷贝               │
            │     - 记录 Transfer 到队列，等待完成                        │
            └─────────────────────────────────────────────────────────────┘
                                    │
        Step 10: 后台异步传输 (与下一轮模型计算重叠)
                                    │
                                    ▼
        Step 11: 完成查询 (后续的 get_finished 轮询)
            handler.get_finished() 检测到 end_event 完成
            → TransferResult 上报给 OffloadingWorker
            → OffloadingConnectorWorker 记录到 _connector_worker_meta
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                          Scheduler 侧 (调度器)                               │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
        Step 12: 调度器接收完成信息 (update_connector_output)
            - 遍历 completed_jobs，更新 pending_count
            - 当 pending_count == 0:
                * manager.complete_store() (标记块为 ready)
                * 清理 _block_id_to_pending_jobs 中的监视
                * 删除作业记录
            - 对于 finished_sending 中的 req_id，释放 GPU 块

_build_store_jobs 过滤出需要存储的 offload keys

python 复制代码

new_offload_keys: list[OffloadKey] = []
for group_config, group_state in zip(
    self.config.kv_group_configs, req_status.group_states
):
    num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
    start_block_idx = group_state.next_stored_block_idx
    if num_blocks <= start_block_idx:
        continue
    offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
    offload_block_ids = group_state.block_ids[
        start_block_idx * block_size_factor
        + block_size_factor - 1 : num_blocks * block_size_factor : block_size_factor
    ]
    assert len(offload_keys) == len(offload_block_ids)

    for offload_key, block_id in zip(offload_keys, offload_block_ids):
        if block_id != 0:
            new_offload_keys.append(offload_key)

OffloadingWorker 与 Handler 交互流程

text 复制代码

OffloadingConnectorWorker
         │
         │ register_handler(src_cls, dst_cls, handler)
         ▼
┌────────────────────────────────────────────────────────────┐
│                    OffloadingWorker                         │
│  transfer_type_to_handler = {                               │
│      ("GPU", "CPU"): gpu_to_cpu_handler,                    │
│      ("CPU", "GPU"): cpu_to_gpu_handler,                    │
│  }                                                           │
│  handlers = {handler1, handler2}                            │
└────────────────────────────────────────────────────────────┘
         │
         │ transfer_async(job_id, (src_spec, dst_spec))
         ▼
   确定 transfer_type = (src_spec.medium(), dst_spec.medium())
         │
         ▼
   查找 handler = transfer_type_to_handler[transfer_type]
         │
         ▼
   handler.transfer_async(job_id, spec)
         │
         ▼
┌────────────────────────────────────────────────────────────┐
│          SingleDirectionOffloadingHandler                   │
│  - 解析 TransferSpec → 计算源/目标指针列表                  │
│  - 分配/复用 CUDA stream 和 events                          │
│  - ops.swap_blocks_batch() 异步拷贝                         │
│  - 将 Transfer 加入队列，记录 end_event                      │
└────────────────────────────────────────────────────────────┘
         │
         │ (后续轮询) get_finished()
         ▼
   检查队列头部 end_event.query()
         ├─ 未完成 → 返回空列表
         └─ 已完成 → 弹出 Transfer, 构造 TransferResult
                     (job_id, success, bytes, time, type)
         │
         ▼
   返回给 OffloadingWorker → 聚合 → OffloadingConnectorWorker

Reference

1\] [提速30%：vLLM推理的Swap特性实践](https://zhuanlan.zhihu.com/p/1999536171961828862) \[2\] [vllm分析(五)------pd分离kv cache的处理过程](https://blog.csdn.net/u010643777/article/details/161043777) \[3\] [Multi-tier KV offloading via the vLLM offloading connector](https://github.com/vllm-project/vllm/issues/38260)