【vllm】vLLM v1 Simple KV Offload — 系统级架构深度分析（八）

vLLM v1 Simple KV Offload — 系统级架构深度分析代码目录:vllm/vllm/v1/simple_kv_offload代码规模: 6 文件 / ~1,354 行分析时间: 2026-04-19一、整体架构概览1.1 设计思路v1/simple_kv_offload是 vLLM v1 架构中KV Cache CPU 卸载的简化实现核心目标是在 GPU 显存不足时将 KV Cache 块异步搬运到 CPU 内存并在请求重入时从 CPU 加载回来从而扩大有效 KV Cache 容量。与v1/kv_offload17 文件 / ~1.9K 行含共享内存、引用计数、LRU/ARC 驱逐策略等相比simple_kv_offload采用了极简设计维度kv_offloadsimple_kv_offload文件数176代码行数~1,900~1,354驱逐策略LRU ARCCPU BlockPool 前缀缓存共享内存mmap/dev/shm无每个 Worker 独立 CPU 缓存传输引擎swap_blocks_batch (PyTorch)cuMemcpyBatchAsync (CUDA Driver API)Worker间同步SharedOffloadRegion 引用计数WorkerMetadata.aggregate() 计数复用门控FilterReusedOffloadingManager无1.2 架构模式采用双面对称架构Scheduler-Side / Worker-Side严格分离决策与执行Scheduler-SideManager决定卸载哪些块、加载哪些块管理 CPU 侧 BlockPool 和前缀缓存匹配Worker-SideHandler决定如何搬运管理 DMA 传输、CUDA 流、事件轮询两者通过SimpleCPUOffloadMetadata/SimpleCPUOffloadWorkerMetadata进行跨侧通信完全解耦。1.3 整体运行流程┌─ 初始化 ─────────────────────────────────────────────┐ │ Scheduler: 创建 CPU KVCacheCoordinator BlockPool │ │ Worker: 注册 GPU KV caches 分配 pinned CPU 内存 │ │ Worker: 启动 DmaCopyBackend 后台线程 │ └──────────────────────────────────────────────────────┘ │ ▼ 每步循环 ┌─ Scheduler 侧 ──────────────────────────────────────┐ │ 1. update_state_after_alloc(): CPU 命中查找 → 构建加载传输对 │ │ 2. prepare_store_specs(): eager/lazy 模式 → 构建存储传输对 │ │ 3. build_connector_meta(): 生成 SimpleCPUOffloadMetadata │ └──────────────────────────────────────────────────────┘ │ metadata ▼ ┌─ Worker 侧 ─────────────────────────────────────────┐ │ 4. get_finished(): 启动 DMA 拷贝loadstore │ │ 5. 轮询完成事件 → 报告 finished_recving store_events │ └──────────────────────────────────────────────────────┘ │ worker_meta ▼ ┌─ Scheduler 侧 ──────────────────────────────────────┐ │ 6. update_connector_output(): 处理传输完成 │ │ - Load 完成: 清理请求释放 touch refs │ │ - Store 完成: 注册 CPU 前缀缓存释放 GPU refs │ └──────────────────────────────────────────────────────┘二、子模块深度分析2.1 SimpleCPUOffloadScheduler — 调度侧管理器文件:manager.py(739 行)核心作用: Scheduler 侧的 CPU 卸载决策中心负责 CPU 前缀缓存匹配、传输规划、完成处理关键类与数据结构类/数据结构作用SimpleCPUOffloadScheduler调度侧管理器主类TransferMeta传输对gpu_block_ids↔cpu_block_idsLoadRequestState加载请求状态request transfer_meta load_event finishedStoreRequestState存储请求状态eager 模式block_ids num_stored_blocks store_events关键方法方法行数作用__init__()~50初始化 CPU coordinator/block_pool, 模式选择, 事件计数器_derive_cpu_config()~20从 GPU config 推导 CPU KVCacheConfig, 按容量比例缩放 num_blocksget_num_new_matched_tokens()~15CPU 前缀缓存命中查找返回 (hit_length, is_asyncTrue)update_state_after_alloc()~60构建加载传输对查找 CPU 缓存块 → touch 防驱逐 → 存入 _reqs_to_loadbuild_connector_meta()~30每步生成 SimpleCPUOffloadMetadataprepare_store_specs()5分发到 eager/lazy store 准备_prepare_lazy_store_specs()~35游标扫描 GPU free_queue卸载接近淘汰的缓存块_prepare_eager_store_specs()~80扫描 scheduler_output卸载已确认计算的 KV 块update_connector_output()~25处理 Worker 报告的传输完成事件_process_store_completion()~12注册 CPU 前缀缓存条目释放 GPUCPU refsrequest_finished()~20请求完成处理如有传输在飞则延迟清理_cleanup_load_request()~15清理加载请求释放 touch refs移除事件映射_cleanup_store_request()~10清理存储请求元数据两种存储模式Eager Mode急切模式请求的 KV 块一旦被确认计算完成立即安排卸载通过StoreRequestState跟踪每个请求的存储进度num_stored_blocks数组作为每组的游标避免重复卸载从scheduler_output获取块 ID 信息通过yield_req_dataLazy Mode惰性模式仅在 GPU 块接近被淘汰时才卸载游标扫描gpu_block_pool.free_block_queue目标是保持_target_free个空闲块无每请求跟踪更简洁的路径2.2 Metadata — 跨侧通信协议文件:metadata.py(60 行)核心作用: 定义 Scheduler → Worker 和 Worker → Scheduler 的数据传输协议SimpleCPUOffloadMetadataScheduler → Worker字段类型作用load_eventint加载事件 ID单调递增-1无操作load_gpu_blockslist[int]加载目标 GPU 块 ID 列表load_cpu_blockslist[int]加载源 CPU 块 ID 列表load_event_to_reqsdict[int, list[str]]事件 → 请求 ID 反向映射store_eventint存储事件 IDstore_gpu_blockslist[int]存储源 GPU 块 ID 列表store_cpu_blockslist[int]存储目标 CPU 块 ID 列表need_flushbool是否有请求被抢占需同步所有传输SimpleCPUOffloadWorkerMetadataWorker → Scheduler字段类型作用completed_store_eventsdict[int, int]{event_idx: 1}每个完成的存储事件关键方法aggregate(other): 合并多个 Worker 的完成计数当count world_size时 Scheduler 才处理完成事件设计特点事件 ID 驱动Worker 通过单调递增的event_idx标识传输不暴露请求 ID请求映射在 Scheduler 侧load_event_to_reqs反向映射仅在 Scheduler 维护Worker 无需知道请求身份TP/PP 安全存储完成需所有 Worker 汇报后才处理_store_event_pending_counts2.3 SimpleCPUOffloadWorker — Worker 侧处理器文件:worker.py(305 行)核心作用: Worker 侧的传输执行器管理 GPU/CPU 缓存注册、DMA 传输启动、事件轮询关键方法方法行数作用register_kv_caches()~80注册 GPU KV 缓存 分配 pinned CPU 张量 初始化 DmaCopyBackendbind_connector_metadata()~8绑定当前步的 metadata记录 pending 事件get_finished()~40启动 DMA 拷贝 轮询完成事件 报告完成build_connector_worker_meta()~8构建 Worker → Scheduler 的存储完成报告handle_preemptions()~5抢占处理同步所有在飞传输_poll_stream_events()~12非阻塞轮询 CUDA 事件更新高水位标记_flush_and_sync_all()~8同步所有在飞事件抢占时调用register_kv_caches() — GPU/CPU 缓存注册这是 Worker 侧最复杂的方法处理多种 KV 缓存布局去重多个层可能共享同一底层存储按data_ptr去重布局分类FlashAttn/ROCm:(2, num_blocks, ...)→ K/V 外层段拆为 2 个段FlashInfer/MLA:(num_blocks, ...)→ 块在最外层1 个段int8 视图统一转换为(num_blocks, page_size_bytes)视图CPU 分配相同形状 cudaHostRegister钉住内存低优先级 CUDA 流load_streamstore_stream让 KV I/O 让位于计算延迟启动策略get_finished()中load 和 store 都在模型执行后才启动start_load_kv()和wait_for_save()均为空操作将 CPU 侧拷贝开销~5ms隐藏在 GPU 计算背后。2.4 DmaCopyBackend — 后台线程 DMA 拷贝文件:copy_backend.py(97 行)核心作用: 后台线程执行 cuMemcpyBatchAsync 批量 DMA 传输避免 GIL 和 CUDA 上下文切换类结构DmaCopyBackend ├── _store_params: BatchMemcpyParams (GPU→CPU 预建参数) ├── _load_params: BatchMemcpyParams (CPU→GPU 预建参数) ├── _load_stream / _store_stream: CUDA 流引用 ├── _queue: SimpleQueue (生产者-消费者队列) ├── _thread: Thread (后台拷贝线程) └── _shutdown: bool关键方法方法作用init()构建 store/load BatchMemcpyParams启动后台线程launch_copy()主线程入队传输请求_copy_loop()后台线程循环取任务 → copy_blocks() → 记录事件shutdown()入队 None 信标 join 线程设计决策预建参数build_params()在 init 时一次性计算src_bases/dst_bases/bpb每步只传块 ID 列表事件列表共享events_list在主线程和后台线程间共享依赖 CPythonlist.append()原子性daemon 线程后台线程标记为 daemon进程退出时自动终止2.5 CudaMemOps — CUDA Driver API 底层操作文件:cuda_mem_ops.py(153 行)核心作用: 封装cuMemcpyBatchAsyncCUDA Driver API 和cudaHostRegister内存钉住关键函数/类函数/类作用pin_tensor(tensor)通过cudaHostRegister钉住 CPU 张量绕过 PyTorch 2^n 内存对齐_resolve_batch_memcpy()通过cuGetProcAddress懒加载cuMemcpyBatchAsync函数指针BatchMemcpyParams传输参数命名元组src_bases, dst_bases, bpb, num_layers, attrs, stream_handlebuild_params()从 src/dst 缓存 CUDA 流构建 BatchMemcpyParamscopy_blocks()执行批量 DMA 传输展开块 ID → 调用 cuMemcpyBatchAsynccuMemcpyBatchAsync 优势相比 v1/kv_offload 使用的swap_blocks_batchPyTorch 层面逐块拷贝维度swap_blocks_batchcuMemcpyBatchAsync调用层PyTorch PythonCUDA Driver API ©批量性逐块循环单次调用 N×M 次传输内核启动N 次1 次驱动内调度GPU 兼容性任意需 Driver API 支持内存钉住优化PyTorch 的pin_memoryTrue通过CUDACachingHostAllocator分配会将请求大小向上取整到 2 的幂如 100GB → 128GB。cudaHostRegister直接注册已分配的内存避免此浪费。2.6 Load Flow — 加载流程核心流程命中查找get_num_new_matched_tokens()→cpu_coordinator.find_longest_cache_hit()传输构建update_state_after_alloc()→ 构建TransferMeta(gpu_ids, cpu_ids)→ touch 防驱逐元数据生成build_connector_meta()→SimpleCPUOffloadMetadata.load_*DMA 执行get_finished()→DmaCopyBackend.launch_copy(cpu→gpu)完成轮询_poll_stream_events()→finished_recving请求 ID 集合清理update_connector_output()→_cleanup_load_request()→ 释放 touch refs2.7 Store Flow — 存储流程核心流程传输规划prepare_store_specs()→ eager/lazy 模式选择块分配CPUblock_pool.get_new_blocks() 印入 block_hash → touch GPU 防释放元数据生成build_connector_meta()→SimpleCPUOffloadMetadata.store_*DMA 执行get_finished()→DmaCopyBackend.launch_copy(gpu→cpu)Worker 汇报completed_store_events[event_idx]1→aggregate()跨 Worker完成处理_process_store_completion()→ 注册 CPU 前缀缓存 释放 GPUCPU refsTP/PP 安全存储完成需所有 Workerworld_size汇报后才处理避免部分完成导致数据不一致。2.8 Data Types State Tracking核心数据结构关系TransferMeta ────── LoadRequestState ────── _reqs_to_load[req_id] gpu_ids, cpu_ids request, load_event │ transfer_meta │ _load_event_to_reqs[event] finished │ ▼ StoreRequestState ── _reqs_to_store[req_id] SimpleCPUOffloadMetadata request (eager mode only) load_event / store_event block_ids (per grp) load_gpu/cpu_blocks num_stored_blocks store_gpu/cpu_blocks store_events load_event_to_reqs finished need_flush2.9 Module Call Relations — 模块调用关系调用链Scheduler └── SimpleCPUOffloadScheduler (manager.py) ├── KVCacheCoordinator.find_longest_cache_hit() ├── BlockPool.get_new_blocks() / free_blocks() / touch() ├── build_connector_meta() → SimpleCPUOffloadMetadata └── update_connector_output() └── _process_store_completion() ├── cpu_block_pool.cached_block_hash_to_block.insert() ├── cpu_block_pool.free_blocks() └── gpu_block_pool.free_blocks() Worker └── SimpleCPUOffloadWorker (worker.py) ├── register_kv_caches() │ ├── pin_tensor() [cuda_mem_ops.py] │ └── DmaCopyBackend.init() [copy_backend.py] ├── get_finished() │ ├── DmaCopyBackend.launch_copy() │ │ └── _copy_loop() → copy_blocks() [cuda_mem_ops.py] │ └── _poll_stream_events() └── build_connector_worker_meta() → SimpleCPUOffloadWorkerMetadata三、模块间调用关系与数据流3.1 跨侧数据流Scheduler Side Worker Side ┌──────────────┐ ┌──────────────┐ │ Manager │──metadata───▶ │ Worker │ │ (739 lines) │ │ (305 lines) │ │ │◀─worker_meta─│ │ └──────┬───────┘ └──────┬───────┘ │ │ CPU BlockPool DmaCopyBackend Coordinator │ │ │ ┌─────┘ └──────┐ │ ▼ ▼ │ CudaMemOps CUDA Streams │ (copy_blocks) (load/store) │ GPU BlockPool (reference only)3.2 主流程时序Step N: Scheduler Worker ───────── ────── update_state_after_alloc() │ find CPU cache hits │ build LoadRequestState │ prepare_store_specs() │ eager: scan scheduler_output │ lazy: cursor-walk free_queue │ build_connector_meta() │ ─── SimpleCPUOffloadMetadata ──▶ │ bind_connector_metadata() │ │ [Model Execution on GPU] │ │ get_finished() │ launch DMA copies │ poll completed events │ ◀── finished_recving ──────│ │ ◀── completed_store_events ─│ │ update_connector_output() │ process load completions │ process store completions │ (register prefix cache) ▼ Step N1四、架构设计亮点与总结4.1 设计亮点cuMemcpyBatchAsync 替代 swap_blocks_batch单次 Driver API 调用完成 N×M 次传输减少内核启动开销驱动层调度更高效后台线程 DMADmaCopyBackend的生产者-消费者模式将 DMA 启动从主线程卸载避免 GIL 和 CUDA 上下文切换cudaHostRegister 绕过 PyTorch 内存对齐直接注册已分配内存避免 2^n 向上取整导致的内存浪费延迟启动策略load 和 store 在模型执行后才启动将 CPU 侧拷贝开销隐藏在 GPU 计算背后Eager/Lazy 双模式Eager 模式尽早卸载保证 CPU 缓存覆盖率Lazy 模式按需卸载减少不必要的传输事件 ID 解耦Worker 通过单调递增的 event_idx 标识传输不暴露请求 IDScheduler 侧维护反向映射TP/PP 安全完成语义存储完成需所有 Worker 汇报后才处理通过aggregate()跨 Worker 计数4.2 与 kv_offload 的定位差异维度simple_kv_offloadkv_offload定位轻量级单节点卸载多 Worker 共享内存卸载复杂度低6 文件 / 1.4K 行高17 文件 / 1.9K 行传输引擎cuMemcpyBatchAsyncswap_blocks_batch共享内存无mmap /dev/shm cudaHostRegister驱逐策略BlockPool 前缀缓存LRU ARC 双策略复用门控无FilterReusedOffloadingManager适用场景单节点 / TP 8多节点 / 大规模 TP/PP4.3 代码统计文件行数职责manager.py739Scheduler 侧管理器决策中心worker.py305Worker 侧处理器执行中心cuda_mem_ops.py153CUDA Driver API 底层操作copy_backend.py97后台线程 DMA 拷贝metadata.py60跨侧通信协议定义__init__.py0包初始化空合计1,354