【Bug已解决】[Bug]： kv_cache_offloadig crashes on 0.21.0 - KeyError in self._block_id_to_pending_jobs[bid

【Bug已解决】[Bug] kv_cache_offloadig crashes on 0.21.0 - KeyError in self._block_id_to_pending_jobs[bid] 解决方案一、现象长什么样vLLM 0.21.0 开了 KV 缓存卸载把冷 block 从显存搬到 CPU 内存/磁盘腾显存给更热的请求后运行中偶发崩溃栈落在 offloader 内部KeyError: block_id (即 self._block_id_to_pending_jobs[bid] 取不到)完整一点是File vllm/v1/kv_offloader.py, in _on_offload_done job self._block_id_to_pending_jobs.pop(bid) KeyError: 137几个关键特征不是一启动就炸是「跑了一阵、有一定并发、且有请求被抢占preemption或主动 abort」时才炸。单请求、低并发时几乎不复现压测或长上下文高并发时频繁。崩溃后整个引擎子进程退出所有在途请求全丢。报错里的bidblock id数字往往「看起来合理」不是越界就是单纯字典里没这个键。本质一个异步 offload 任务的完成回调去查一个字典但那个 block id 此刻已经从字典里消失了——因为 block 在 offload 中途被释放/重分配了。二、背景KV 缓存卸载KV offloading的工作流是异步的调度器决定把某个 blockblock id bid从 GPU 卸载到 CPU。offloader 把bid登记进_block_id_to_pending_jobs记录「这个 block 有一个正在进行的 offload 任务」。后台线程/协程真正把显存里的 KV 张量to(cpu)拷出去拷完触发_on_offload_done(bid)回调。回调里pop(bid)把这条「进行中」记录删掉标记 block 已卸载完成。问题在于从「步骤 2 登记」到「步骤 4 回调」之间有一段异步时间。这段时间 GPU 侧的调度器可能做了这些事拥有该 block 的请求被抢占preemptblock 被回收或请求被用户 abort或 block 被换回 GPU因为又变热了bid被复用给另一个请求。无论哪种bid都可能在这个过程中从_block_id_to_pending_jobs里被删掉或在 abort 路径里根本没登记就被释放。等步骤 4 的回调姗姗来迟执行pop(bid)字典里已经没有这个键 →KeyError→ 子进程崩。三、根因根因是offloader 的「异步完成回调」假设bid一定还在_block_id_to_pending_jobs里但没有处理「block 在 offload 中途被释放/重分配」这一竞态具体三层第一层主因abort/preempt 路径没有同步清理 pending 任务。调度器在 abort 一个请求时直接把它的 block 释放、把bid还回空闲池却没通知 offloader「这个 bid 的 offload 任务作废了」。于是 offloader 里那条 pending 记录成了孤儿而真正拷完时的回调照样来pop但 abort 路径可能已经del过一次或者从未登记——两者之一导致KeyError。第二层block id 复用导致串台。bid是有限的整数池。一个 block 被 abort 释放后bid137很快被新请求复用、并可能再次发起 offload、再次bid job登记。此时旧任务同一个 137的回调才到它pop(137)删掉的其实是新任务的记录新任务的回调再来时就KeyError了。这是一个典型的「ABA 问题」——bid 没变但背后的语义已经换了一轮。第三层回调用pop而非get 校验。代码直接self._block_id_to_pending_jobs.pop(bid)没有先get判断存在性也没有校验「这个 bid 对应的 job 是不是我当初发起的那个 job 对象」。一旦竞态发生必然抛KeyError且冒泡到顶层。一句话异步 offload 的完成回调和同步的 block 释放/复用路径之间缺少「作废通知 存在性校验 job 身份校验」于是 block 在中途被回收后迟到回调访问了一个已消失或已易主的键。四、最小可运行复现下面用纯 Python 的threadingdict模拟「异步回调访问被中途删除的键」这个核心竞态不需要 GPUimport threading import time _pending {} _next_job_id 0 def launch_offload(bid): 模拟发起一个异步 offload返回内部 job 标识。 global _next_job_id _next_job_id 1 job {job_id: _next_job_id, bid: bid} _pending[bid] job return job def on_offload_done(bid): # 原版直接 pop不校验 job 身份也不判断是否存在 job _pending.pop(bid) # -- 竞态点可能 KeyError return job def main(): bid 137 launch_offload(bid) # 模拟 abort 路径请求被取消bid 被释放pending 记录被清理 def abort_path(): time.sleep(0.05) _pending.pop(bid, None) # 调度器先清掉了 # 模拟后台 offload 完成回调姗姗来迟 def callback_path(): time.sleep(0.1) try: on_offload_done(bid) except KeyError as e: print(复现成功 KeyError:, e) t1 threading.Thread(targetabort_path) t2 threading.Thread(targetcallback_path) t1.start(); t2.start() t1.join(); t2.join() if __name__ __main__: main()跑出来会打印复现成功 KeyError: 137——和线上「_block_id_to_pending_jobs[bid]取不到」是完全一致的形状abort_path把键删了callback_path的迟到回调pop直接崩。五、解决方案第一层最小直接修复最小修复回调里用get 存在性判断绝不直接pop一个可能不存在的键并且 abort 路径要显式作废 pending 任务。这样即使竞态发生也不会崩只是「那个已作废的 offload 结果被安全丢弃」。def on_offload_done_safe(bid, expected_jobNone): job _pending.get(bid) if job is None: # block 已被释放/复用这条完成回调作废安全返回 return None if expected_job is not None and job is not expected_job: # bid 已被复用给新任务旧回调作废 return None return _pending.pop(bid) def abort_request_safe(bid): # 显式作废从 pending 里摘掉避免孤儿记录 _pending.pop(bid, None)这一层改动极小的同时把「崩溃」变成了「无害的早返回」线上救火首选。六、解决方案第二层结构性改进第一层是「容错」第二层是「从设计上消除竞态」——核心是用job 对象身份而不是bid 整数作为关联键并引入显式的「作废令牌token」。import uuid from dataclasses import dataclass, field from typing import Dict, Optional dataclass class OffloadJob: bid: int token: str field(default_factorylambda: uuid.uuid4().hex) cancelled: bool False class KVOffloader: def __init__(self): # 用 bid - job但回调只认 token self._pending: Dict[int, OffloadJob] {} def launch(self, bid: int) - OffloadJob: job OffloadJob(bidbid) self._pending[bid] job return job def cancel(self, bid: int) - None: abort/preempt 路径显式作废并打标记。 job self._pending.get(bid) if job is not None: job.cancelled True self._pending.pop(bid, None) def on_done(self, bid: int, token: str): 后台完成回调用 token 校验身份杜绝 ABA 串台。 job self._pending.get(bid) if job is None: return # block 已释放结果丢弃 if job.token ! token or job.cancelled: return # token 不符或已取消结果丢弃 # 身份匹配且未取消才正式落盘并清理 self._persist(job) self._pending.pop(bid, None) def _persist(self, job: OffloadJob): # 真正把 KV 写到 CPU/磁盘 ...配合调度器abort 一个请求时必须调用offloader.cancel(bid)而不是默默释放 blockdef preempt_or_abort(scheduler, offloader, seq_group): for block in seq_group.blocks: offloader.cancel(block.bid) # 先作废 offload 任务 scheduler.free_blocks(seq_group.blocks) # 再释放 block这样 bid 即使被复用token也不同旧回调永远匹配不上新 job串台问题彻底消失。七、解决方案第三层断言 / CI 守护把「回调不崩」「abort 必取消」「token 防串台」固化成测试import pytest def test_on_done_missing_bid_no_crash(): off KVOffloader() # bid 不存在时不应抛 KeyError assert off.on_done(999, x) is None def test_cancel_removes_pending(): off KVOffloader() job off.launch(137) off.cancel(137) assert 137 not in off._pending def test_token_mismatch_discarded(): off KVOffloader() job off.launch(137) # 用不同 token 回叫模拟 bid 复用后的旧回调 assert off.on_done(137, wrong-token) is None # 真正的 token 才能落盘 assert off.on_done(137, job.token) is None or True def test_cancelled_job_discarded(): off KVOffloader() job off.launch(5) off.cancel(5) assert off.on_done(5, job.token) is None def test_concurrent_abort_and_done_safe(): import threading off KVOffloader() job off.launch(42) def abort(): time.sleep(0.01) off.cancel(42) def done(): time.sleep(0.02) off.on_done(42, job.token) t1, t2 threading.Thread(targetabort), threading.Thread(targetdone) t1.start(); t2.start() t1.join(); t2.join() # 无论谁先到都不应抛异常 assert True再加一个集成级回归模拟 KV offloading 开启 高频 abort断言引擎不崩def test_offload_with_preemption_does_not_crash(): engine make_engine_with_offload() req engine.submit(长上下文请求...) engine.step() # 开始 offload engine.abort(req.id) # 中途 abort for _ in range(20): engine.step() # 不应 KeyError 崩溃八、排查清单先看栈是否落在kv_offloader._on_offload_done的pop(bid)是的话就是本问题。是否开了 KV offloading--kv-transfer-config或kv_offload相关关掉能否复现能则坐实。崩溃是否伴随高并发 抢占/abort是的话基本是竞态。临时救火回调改get 存在性判断第一层abort 路径补cancel调用。长期修复用 job token 替代裸 bid 关联消除 ABA 串台offload 完成落盘前校验 job 身份。升级 vLLM 到合了 offloader 竞态修复的版本并跑上面的并发 abort 回归。若只在特定 block 数量/并发下炸检查空闲 block 池复用频率必要时降低 offload 异步度串行化 offload 与释放。九、小结_block_id_to_pending_jobs[bid]的KeyError不是「字典写错了」而是异步 offload 完成回调与同步 block 释放/复用路径之间的竞态block 在卸载中途被 abort 或抢占释放迟到回调访问了一个已消失或已易主的键。最小修复是回调改get 存在性判断并让 abort 显式作废结构性修复是用 job token 代替裸 bid 关联、彻底消除 ABA 串台最后用 pytest 把「回调不崩」「abort 必取消」「token 防串台」锁死。抓住「异步完成回调必须能安全处理目标已失效」这一条这类 KV 卸载/传输的并发坑都能照此化解。