设计与实现解析)
LMCache PD 异步预留式准入控制Reservation-Based Admission Control设计与实现解析【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache导读本文深入解析 LMCache 中 pd_async_reservation_design.md 所描述的核心设计——基于预留Reservation的准入控制机制。该机制专门用于解决 Prefill/Decode 分离PD场景下chunked-prefill 多请求并发时接收端物理缓冲区的死锁问题并通过发送端物理暂存缓冲区流控、Fail-Fast 协议违规检测与全量回滚构建了一套完整的并发安全与错误处理体系。读完本文你将理解 LMCache 的 PD 异步后端PDBackendAsync如何在 sender/receiver 两侧分别做流量控制、四种 ZMQ 消息如何协作、以及如何通过源码与测试验证这套机制的可靠性。问题背景chunked-prefill 下的缓冲区死锁为什么会产生死锁在 chunked-prefill 模式下一个大型 prompt 会被切分为多个顺序批次batch。以 2P1D2 个 Prefill 节点 1 个 Decode 节点部署为例Prefill 节点sender把每个 chunk 的 KV 缓存通过 RDMA 写入 Decode 节点receiver的物理缓冲区。每次写入前receiver 都要先为这批 chunk 分配物理槽位。当多个请求并发执行时这种逐批交错分配会产生经典死锁Buffer 10 chunks Req A needs 8 chunks, Req B needs 8 chunks A allocates 5 → B allocates 5 → buffer full A needs 3 more → blocked, B needs 3 more → blocked → DEADLOCK根因分析每个请求都是分批占用缓冲区的而不是一次性占满。N 个请求各自部分填满缓冲区后任何一个请求都无法凑齐自己的完整 chunk 集合缓冲区永远不会被排空因为没有请求能完成 RDMA 并释放槽位系统整体陷入僵局。这与操作系统中经典的资源部分分配导致死锁问题同构。死锁的代价死锁不仅是性能问题更是正确性问题receiver 侧的 decode 引擎需要某个请求的全部chunk 才能开始消费部分分配会让请求永远停留在半完成状态。因此在设计上必须保证要么整请求准入要么完全不占用的原子性。解决方案预留式准入控制Reservation-Based Admission核心思想在任何物理分配发生之前先为每个请求预留其total_chunks完整 chunk 数。如果缓冲区无法容纳完整预留请求就等待一旦被准入admitted该请求后续所有批次分配必然成功从而从根上消除死锁Buffer 10 chunks A requests admission (8 chunks) → reserved8, available2 B requests admission (8 chunks) → 8 2 → WAIT A completes all 8 chunks → RDMA done → release reservation → available10 B admitted → reserved8, proceeds不变式Invariant系统的关键不变式是仅 receiver 侧成立total_reserved total_chunks恒成立。即所有已准入请求的预留 chunk 总数永远不会超过缓冲区总容量。这是死锁免疫的数学保证——任何一个被准入的请求其完整 chunk 集都在预留范围内物理分配只是兑现预留。架构总览三线程模型与消息流转组件拓扑设计文档给出了一套清晰的三层架构vLLM Worker 线程、Sender 事件循环、Receiver 事件循环通过 asyncio 桥接与 ZMQ 通道协作┌─────────────────────────────────────────────┐ │ vLLM Worker Thread │ │ wait_for_save() │ │ ├─ store() × N batches ←─── allocate from_gpu per batch │ │ ├─ allocate() ←─── blocks on _staging_condition if buffer full │ │ └─ batched_submit_put_task() ←─── submits to sender loop, returns immediately └─────────────────────────────────────────────┘ │ asyncio.run_coroutine_threadsafe ▼ ┌─────────────────────────────────────────────┐ │ Sender Event Loop │ │ _async_transfer_task() (concurrent) │ │ ├─ _async_remote_allocate() ←────────── ZMQ REQ/REP to receiver │ ├─ async_batched_write() ←────────── RDMA write │ ├─ ref_count_down() ←────────── free sender staging buffer │ ├─ _notify_staging_freed() ←────────── wake allocate() waiters │ └─ _check_and_send_proxy_notif() │ │ └─ ProxyNotif via ZMQ PUSH │ └─────────────────────────────────────────────┘ │ ZMQ DEALER/ROUTER ▼ ┌─────────────────────────────────────────────┐ │ Receiver Event Loop │ │ _handle_alloc_request() │ │ ├─ CancelNotif → release keys reservation │ └─ AllocRequest → _async_allocate_and_put() │ ├─ async_try_admit() (first batch) │ │ ├─ allocate() per chunk │ │ └─ put() → register KV object │ └─────────────────────────────────────────────┘代码落点PDBackendAsync上述架构在 PDBackendAsync 中完整实现。该后端是 PD 场景默认实现pd_backend_mode默认async见 config.py。初始化时按角色分叉pd_backend_async.py L391-L434sender 角色创建_staging_condition物理暂存区等待条件、_async_alloc_locks按 receiver 串行化 ZMQ 请求、以及_completed_chunks/_req_has_last/_req_total_chunks/_sent_keys等 per-request 跟踪状态receiver 角色以_aligned_buffer_size // _chunk_size_bytes计算出缓冲区总 chunk 数total_chunks实例化ReservationManager并通过asyncio.run_coroutine_threadsafe在 receiver 事件循环内创建 asyncio 原语_create_recv_primitives。ReservationManagerreceiver 侧的唯一事实来源类设计ReservationManagerpd_backend_async.py L98-L210只在 receiver 侧使用注释与文档双重确认了这一点。其构造参数包括参数含义total_chunks缓冲区总 chunk 容量allocation_timeout等待准入的最长秒数condition_poll_interval条件等待的轮询间隔内部维护两个核心状态_reservations: dict[str, int]req_id → 预留 chunk 数与_total_reserved: int累计预留数。由于 asyncio 原语必须在事件循环内创建_async_admit_condition采用惰性初始化通过init_async_admit_condition()在 receiver 事件循环内创建L133-L138。async_try_admit预留准入准入逻辑L140-L182是一个带超时的自旋等待循环计算available total_chunks - total_reserved若available total_chunks请求所需将 req_id 写入_reservations并累加_total_reserved返回True否则在_async_admit_condition上等待每次最长等待min(remaining, condition_poll_interval)超过allocation_timeout仍未准入则返回False。async_release_reservation释放预留释放逻辑L184-L202从_reservations弹出该请求的预留数并扣减_total_reserved随后notify_all()唤醒所有等待准入的协程——这正是被阻塞的 B 请求得以继续的关键步骤。消息协议四种 ZMQ 消息PD 后端基于msgspec定义消息pd_backend_async.py L45-L95消息方向用途关键字段AllocRequestsender → receiver分配远程缓冲槽位首个批次携带total_chunks做预留keys、req_id、is_last_batch、total_chunksAllocResponsereceiver → sender返回远程缓冲区地址失败为 -1remote_indexesProxyNotifsender → proxy所有 RDMA 完成decoder 可开始消费req_idCancelNotifsender → receiver请求中止释放已分配 keysreq_id、keys值得注意的是AllocRequest.req_id允许为空字符串以兼容旧 sender此时 receiver 跳过 per-request 计数与 Fail-Fast 检测L59-L63。Receiver 侧准入控制从首批预留到整请求原子性首批准入流程在_async_allocate_and_put()pd_backend_async.py L1284-L1473中is_first_batch通过req_id not in self._req_allocated_keys判定。首批到达时依次执行Legacy sender 拒绝total_chunks 0时立即抛RuntimeErrorLegacy senders are no longer supported见 L1302-L1307容量前置校验若请求所需 chunk 数超过整个缓冲区容量直接抛错避免请求在准入等待中耗满pd_allocation_timeout_sec后以误导性的 over-subscribed 超时失败L1308-L1323预留准入调用async_try_admit(req_id, total_chunks)失败则抛RuntimeErrorL1324-L1332。后续批次兑现预留同一请求的后续批次不再重复预留直接从已有预留中兑现物理分配self.allocate()逐 chunk 分配失败时在_alloc_freed_condition上等待L1383-L1401成功则self.put(key, mem_obj)注册 KV 对象并记录到_req_allocated_keys。预留的释放时机预留的释放有三条路径设计文档 L133 明确正常完成is_last_batch True的批次分配成功后清理_req_allocated_keys并调用async_release_reservationL1467-L1469批次失败触发全请求回滚时同步释放见下文错误处理请求中止经CancelNotif释放见下文 Abort Flow。Sender 侧流控物理暂存缓冲区的背压机制与 receiver 的逻辑预留不同sender不使用 ReservationManager而是依赖物理暂存缓冲区staging buffer流控。设计文档强调这是两侧分工的本质区别L137。allocate() 阻塞与唤醒sender 的allocate()pd_backend_async.py L559-L586采用快速路径 慢速路径快速路径直接尝试分配成功立即返回慢速路径暂存区物理已满时在_staging_condition上循环等待直到_notify_staging_freed()唤醒或有剩余槽位超时_allocation_timeout则返回None。唤醒链在_async_transfer_task()中完成每批 RDMA 写完成后ref_count_down()释放暂存槽位随后_notify_staging_freed()L987-L998对_staging_condition执行notify_all()阻塞中的allocate()重试并成功。并发模型发送端允许多个请求并发分配与传输仅受物理暂存区容量限制。接收端的预留机制保证每个准入请求能完成完整 chunk 集二者共同作用既消除了死锁又最大化并发度。设计文档中的 2P1D 对比清晰展示了收益Buffer 20 chunks, P1 req A (10), P2 req B (10) Old (serial admission): P1: [ transfer A ] P2: [ wait ][ transfer B ] Total: A B New (reservation): P1: [ transfer A ] P2: [ transfer B ] ← concurrent RDMA from different peers Total: max(A, B)ProxyNotif 顺序保证由于同一请求的多个批次并发执行is_last_prefill批次可能比更早的批次更早完成 RDMA。因此ProxyNotif的发送必须满足双重条件设计文档 L85-L88completed_chunks total_chunks该请求所有 RDMA 均完成req_has_last Trueis_last_prefill批次已完成。代码中通过_req_total_chunks/_req_has_last/_completed_chunks三张 per-request 跟踪表L404-L406维护该条件并经由_proxy_send_lockasyncio.Lock串行化 ProxyNotif 发送。测试 test_sender_chunk_ordering 专门验证了乱序完成场景下 ProxyNotif 只在最后一块到达后才触发。错误处理Fail-Fast 检测与全量回滚协议违规检测Fail-Fast Overflow如果某请求的累计 chunk 数超过其声明的total_chunks说明 sender 违反协议L1334-L1356。此时 receiver 依次执行回滚之前所有批次已分配的 chunks逐个self.remove()清理请求跟踪状态_req_allocated_keys释放预留async_release_reservation(req_id)抛出带违规详情的RuntimeError。对应测试为 test_receiver_fail_fast_overflow。批次失败的全请求回滚任一批次分配失败如超时、内存耗尽时由于 decoder 需要全部 chunk 才能开始任何部分状态都不允许残留receiver 执行全量回滚L1421-L1460当前批次回滚删除失败批次已分配的 chunks先前批次回滚删除同请求此前成功批次的全部 chunks状态清理移除_req_allocated_keys跟踪并释放预留错误传播向上抛出异常通知 sender。注意异常捕获使用BaseException确保KeyboardInterrupt/SystemExit场景下也能完成清理。Abort 流程请求被中止时的调用链设计文档 L93-L100为request_finished(ABORTED) → cancel_request(req_id) # any thread → wake _staging_condition # unblocks allocate() if waiting → schedule _abort_request() # on sender loop → CancelNotif to receiver # release remote keys reservation → clear per-request state # clean up sender tracking代码中cancel_request唤醒暂存区等待者L1106-L1108_abort_request在 sender 事件循环上发送CancelNotif(req_id, keyssent_keys)L1117-L1132receiver 侧_handle_alloc_request收到CancelNotif后释放 keys 与预留。锁清单与并发模型设计文档给出的锁清单L75-L81在代码中逐一可考路径锁类型用途代码位置Sender worker 线程_staging_conditionthreading.Condition等待暂存槽位L401Sender loop_async_alloc_locks[receiver_id]asyncio.Lock串行化发往同一 receiver 的 ZMQL750-L752Sender loop_proxy_send_lockasyncio.Lock串行化 ProxyNotif 发送L685Receiver loop_async_admit_conditionasyncio.Condition异步准入等待L138Receiver loop_alloc_freed_conditionasyncio.Conditionchunk 释放时唤醒分配重试L1393配置参数与实战建议本机制相关的 PD 配置在 lmcache/v1/config.py 中统一定义均可通过环境变量注入配置项默认值说明pd_buffer_sizeNone必须设置PD 缓冲区字节数自动向下对齐到 chunk 整数倍小于单个 chunk 时报错L460-L469pd_allocation_timeout_secfloat(inf)内存分配重试的最大秒数L262-L267pd_condition_poll_interval_sec0.005Condition 等待的轮询间隔越小越灵敏、越大越省 CPUL276-L285pd_max_prefill_len0若 0初始化时校验缓冲区 token 容量必须 ≥ 该值L286-L296校验逻辑见 L484-L500pd_backend_modeasyncasync本文所述 asyncio 实现或sync旧线程实现L297-L306pd_skip_proxy_notificationFalse是否跳过 ProxyNotif 发送L307-L311实战要点pd_buffer_size必须能容纳最长 prefill 的完整 chunk 集。代码在首批准入时前置校验请求所需 chunk 数 ≤ 缓冲区总容量否则直接报错并给出调整建议增大pd_buffer_size或缩短 prefill。建议配合pd_max_prefill_len在初始化阶段就强制校验尽早暴露容量不匹配不要依赖无限超时pd_allocation_timeout_sec默认inf生产环境建议设为有限值配合 Fail-Fast 检测避免请求无限期挂在准入或分配等待上所有 sender 必须升级total_chunks 0的 legacy sender 会在首批即被 receiver 拒绝并抛RuntimeError升级前需确认所有 Prefill 节点代码版本一致。测试验证与可靠性保障测试文件 tests/v1/storage_backend/test_pd_backend_async.py 覆盖了本设计的全部关键路径test_sender_flow_control_backpressure验证 sender 暂存区满时的阻塞与唤醒test_sender_chunk_ordering验证乱序完成下 ProxyNotif 仅在最后一块到达后触发test_receiver_admission_control直接验证预留式准入——req-A 首批预留 3 chunks 后req-B 与 req-A 的后续批次可并发完成全部成功test_receiver_fail_fast_overflow验证累计 chunk 数超过声明值的协议违规检测与回滚test_receiver_is_last_batch_cleanup验证最后批次完成后清理与预留释放test_receiver_reject_legacy_sender_zero_total_chunks验证total_chunks0的 legacy sender 被拒绝。这些测试从并发正确性准入与并发推进、协议合规性违规检测与资源生命周期释放与回滚三个维度为设计文档中的每一项保证提供了可执行验证。总结LMCache 的 PD 异步预留式准入控制是一套两侧分工的并发控制方案receiver 通过ReservationManager以total_reserved total_chunks不变式保证任何被准入请求的完整 chunk 集可兑现从根源杜绝缓冲区死锁sender 则依靠物理暂存区的_staging_condition背压流控维持并发传输两者叠加配合 ProxyNotif 双重条件排序、Fail-Fast 协议违规检测、批次失败全量回滚与 CancelNotif 中止路径构成了一个具备强正确性保证、可测试、可观测的完整体系。对部署 PD 分离推理的团队而言理解该机制是正确配置pd_buffer_size、pd_allocation_timeout_sec等参数、以及在多 Prefill 节点高并发下保障系统吞吐的前提。【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考