十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

LMCache Q Ring Buffer 深度解析:基于分页 KV 机制捕获与持久化 Attention Query 张量

LMCache Q Ring Buffer 深度解析:基于分页 KV 机制捕获与持久化 Attention Query 张量 LMCache Q Ring Buffer 深度解析基于分页 KV 机制捕获与持久化 Attention Query 张量【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache导读Attention 的 QueryQ张量是前向过程中的瞬态数据常规 KV Cache 层不会保存它。LMCache 的Q Ring Buffer通过在 vLLM 工作进程内部按“分页 KV Cache 的布局”将每一层的 Q 暂存进 GPU 环形缓冲区再复用既有的 STORE 通路把整块数据异步落盘到 LMCache MP 服务器从而让下游 SDKqcache能够按 token id 检索、编辑并回存 Query 张量典型场景如 token-dropping 等中间张量分析。读完本文你将掌握 Q Ring Buffer 的组件划分、行到 token 的归属算法、五阶段生命周期、缓存寻址规则与配置开关以及当前实现的功能边界。一、设计目标让瞬态的 Q 张量可被持久化复用Attention 的 Query 张量天然是瞬态的它只在一次前向中出现计算完 K、V 点积后即被丢弃。LMCache 的目标是让这类张量也能像 KV Cache 一样被保存、检索与复用检索侧见 LMCache SDK 文档 中的qcache。为此设计采用了以下关键决策在生产者侧捕获Q 的捕获发生在 vLLM生产者内部而非消费端 SDKSDK 用户不需要也不应该直接调用这里的任何接口。分页化暂存每一层的 Q 先被写入一个 GPU环形缓冲区ring buffer其逻辑布局刻意设计成与分页 KV Cache 一致[num_layers, num_blocks, block_size, hidden_dim]从而直接复用 LMCache 现有的分页 KV 传输内核无需为 Q 单独开发新的传输路径。复用 STORE 通路环形缓冲区中的整块数据通过既有的 STORE 消息路径提交给 LMCache MP 服务器仅在键名上使用 query 专属模型名model##query与 KV 对象做隔离。功能开关门控整条捕获链路由 vLLM MP 连接器的transfer_intermediate_tensors标志驱动默认关闭。从源码结构看该功能被组织为“实验性中间张量”特性连接器与 worker 适配器只感知一个统一的 Dispatcher由FEATURE_REGISTRY {TRANSFER_QUERY: QTensorFeature}把QTensorFeature注册进去将register / save_kv_layer / wait_for_save / reclaim / reregister / shutdown等生命周期钩子扇出fan-out给具体特性实现。这样在连接器层增加新的中间张量类型时不需要改动连接器与适配器本身。二、核心组件一QRingBuffer—— 分页化的 GPU 环形缓冲区QRingBuffer实现在 lmcache/sdk/qringbuffer.py是一个 GPU 张量逻辑形状为[num_layers, num_blocks, block_size, hidden_dim]其中num_layers参与捕获的注意力层数num_blocks环形缓冲区中的块block总数block_size每块包含的 token 数与 vLLM 的cache_config.block_size一致hidden_dim等于num_q_heads * head_size即单层 Q 展平后的宽度。构造时对四个维度都做了正数校验num_layers 0 or num_blocks 0 ...会抛出ValueError。每层维护独立的torch.empty((num_blocks, block_size, hidden_dim))张量并通过字典暴露为self.tensors: dict[str, torch.Tensor] { flmcache_q_layer_{i}: self._layer_tensors[i] for i in range(num_layers) }空闲块用self._free_blocks: list[int]维护初始为range(num_blocks)。核心方法如下1.allocate(n)—— 申请空闲块从空闲块列表尾部弹出n个块 ID 返回若空闲块不足则返回None调用方据此决定跳过本步捕获n为负数时抛出ValueError。def allocate(self, n: int) - list[int] | None: if n 0: raise ValueError(fcannot allocate a negative block count: {n}) if n len(self._free_blocks): return None reserved self._free_blocks[-n:] if n else [] del self._free_blocks[len(self._free_blocks) - n:] return reserved2.free(block_ids)—— 归还块对越界block_id 0 or block_id num_blocks或重复释放的块 ID 记录警告日志后跳过合法块追加回空闲列表。该接口的容错设计保证了在服务器健康检查失败、store 失败等异常路径下也能安全回收块。3.num_free_blocks()—— 查询空闲块数直接返回len(self._free_blocks)供规划阶段判断本步是否有足够容量。4.scatter(layer_index, query, ring_slots)—— 写入一层 Q将某一层展平后的 query 行按给定槽位写入环形缓冲区flat_q query.reshape(query.shape[0], -1) # [num_tokens, num_q_heads * head_size] ring[slots // self.block_size, slots % self.block_size] flat_q[valid].to(ring.dtype)关键语义ring_slots是int64张量形状为[num_tokens]每个 token 映射到一个 ring 槽位 ID槽位为-1的行被丢弃不写入这正好对应“本行没有对应 KV 写操作”或“CUDA-graph padding 行”的情况。若flat_q的宽度与hidden_dim不一致会抛出ValueError提示检查配置的 query head 数 / head size。写入前会做 dtype 转换to(ring.dtype)保证环形缓冲区张量与模型 dtype 一致。三、核心组件二QRingBufferCapture—— 挂钩连接器前向生命周期QRingBufferCapture负责把 Q 捕获接入连接器的 forward 生命周期包含三个方法1.setup_q_ring(kv_caches, kv_cache_config, vllm_config)—— 初始化与注册在 KV 注册阶段即连接器的register_kv_caches被调用完成以下工作挑选注意力层通过attention_layer_names_from_vllm(kv_cache_config, kv_caches)获得注意力层名列表。该函数读取 vLLM 的kv_cache_groups只保留kv_cache_spec为AttentionSpec的层若 vLLM 未提供 groups 则退化为list(kv_caches.keys())。没有任何 KV cache 或没有注意力层时打警告并跳过。计算几何信息从vllm_config.model_config获取num_q_heads考虑parallel_config与head_size从cache_config.block_size获取块大小dtype 从模型配置推导。计算 ring 大小优先级如下显式指定lmcache.q.ring_blocksextra config取max(1, int(explicit))否则按深度估算depth max(1, lmcache.q.ring_depth)默认 2max_batched scheduler_config.max_num_batched_tokens缺省 8192则num_ring_blocks max(1, ceil(max_batched / block_size) * depth)。调用self.q_ring_adapter.register_q_ring(...)完成 GPU 张量分配与服务器端注册。2.save_q_layer(layer_name, metadata, **kwargs)—— 每层 scatter每个注意力层的 KV 保存钩子都会触发本方法非 KV 写者is_kv_writer为假或本步已禁用时直接返回从kwargs[intermediate_tensors]中按[q, query]依次查找 Q 张量取不到则返回在本步的第一个注意力层上调用_build_q_step_state(...)构建“本步 Q 存储计划”_QStepStatering_slots行槽映射 每请求一个_QLayerStore后续层直接复用该计划保证所有层把同一批 token 写入相同的 ring 块计划构建失败例如无 STORE 请求、缺slot_mapping时置q_step_disabled True本步剩余层全部跳过。3.batched_submit_qstore_requests(event)—— forward 出口批量提交在连接器的wait_for_save阶段模型推理步完成、event已记录后被调用state self.q_step_state self.q_step_state None # 消费即重置 self.q_step_disabled False if state is None or event is None: return for q_store in state.stores: self.q_ring_adapter.submit_q_store_request( q_store.request_id, q_store.op, q_store.ring_block_ids, event, cache_saltq_store.cache_salt)即每个请求提交一个STORE_Q本步没有捕获计划或事件缺失时静默跳过且状态被重置不会泄漏到下一步。四、行到 token 的归属Row-to-Token Attribution不做位置假设的匹配算法这是整个设计中最微妙的部分。在 continuous batching 下一步的 Q 张量把多个请求的行拼接在一起且存在三重“不对齐”行数与存储 token 数不等请求在某一步的调度 token 数行数未必等于其 store 操作的按 chunk 对齐 token 数——prompt 尾部越过最后一个 chunk 边界的部分仍然产生行顺序不一致批内行的排列顺序未必与连接器元数据中的请求顺序一致部分存在某请求的 token 可能只有部分出现在本步例如在更早的 chunked-prefill 迭代中已计算的部分没有 Q 行。因此计划绝不按位置给行分配 token而是通过attn_metadata.slot_mapping匹配行r会把它的 KV 写到 GPU 槽位slot_mapping[r]store 操作的 tokeni位于 GPU 槽位block_ids[i // block_size] * block_size i % block_sizeop 的块列表已预先切片到其[start, end)范围由于每个 GPU 槽位每一步最多被写一次因此两个集合的“完全交集”就是 op 的 token 与行之间的一一对应bijection。算法实现_build_q_step_state大致为row_slots slot_mapping[:n_rows] # 行 → GPU KV 槽位 op_slots block_tensor[:, None] * block_size offsets[None, :] # op token → GPU 槽位 sorted_slots, token_order torch.sort(op_slots) pos torch.searchsorted(sorted_slots, row_slots) hit (row_slots 0) (sorted_slots[pos_clamped] row_slots)然后对每个命中的行回填 ring 槽位ring_slots[:n_rows][hit] ring_slot_by_token[matched_token_idx]由此得到的关键性质逐请求独立跳过若某 op 的 token 只有部分在本步例如更早 chunked-prefill 迭代已算过该请求被单独跳过打警告不影响本步其他请求的捕获非 STORE 请求直接忽略metadata.requests中direction ! STORE的请求不参与计划不会像旧实现那样让整个 step 的捕获失效严格校验op 的 token 数必须是block_size的整数倍且 op 块布局必须完整覆盖其 token 范围len(gpu_blocks) * block_size num_tokens否则跳过该请求——这是防止传输内核越界读 GPU 内存的安全闸门CUDA-graph padding 行行数超过slot_mapping长度的部分与槽位为-1的行直接丢弃ring_slots初始全为-1。五、核心组件三QRingBufferAdapter—— 与 LMCache 服务器的交互边界QRingBufferAdapter拥有 ring 与 LMCache 之间的全部交互包括1.register_q_ring(...)要求传输上下文已建立否则抛RuntimeError提示先调用register_kv_caches()计算块大小block_size lmcache_tokens_per_chunk // blocks_in_chunk分配QRingBufferhidden_dim num_q_heads * head_size构造EngineGroupInfo(engine_group_id0, layer_indicestuple(range(num_layers)), tokens_per_blockblock_size, sw_size_tokens-1)调用transfer_ctx.register_q(...)内部走req_client.register_q_cache对应协议REGISTER_Q_CACHE把q_ring.tensors与q_model_name一并注册超时则抛ConnectionError。2.submit_q_store_request(request_id, op, ring_block_ids, event, cache_salt)先_ensure_heartbeat_started()并检查健康状态不健康时直接q_ring.free(ring_block_ids)归还块op.token_ids is None时同样跳过并归还块用_create_key(op.token_ids, op.start, op.end, request_id..., cache_salt...)构建键再replace(key, model_nameself.q_model_name)把模型名替换为 query 专属名调用transfer_ctx.submit_q_store(...)对应协议STORE_Q并把(future, ring_block_ids)与 event 按自增序号记入q_store_futures/q_store_events。3.reclaim_finished_q_stores()在连接器/适配器的get_finished阶段被调用轮询每个 futurestore 完成后立即归还 ring 块服务器不健康时批量归还并清空。这保证了环形缓冲区容量能在异步落盘完成后及时回收复用。4.reregister_q_ring()与shutdown_q_ring()reregister_q_ring服务器恢复后重发REGISTER_Q_CACHEring 从未构建时为空操作由 Dispatcher 的 reregister 在心跳恢复流程中调用shutdown_q_ringteardown 时发送UNREGISTER_Q_CACHEreq_client.unregister_q_cache(instance_id)超时仅打警告、继续关闭流程。六、五阶段生命周期注册Registersetup_q_ring→register_q_ring→REGISTER_Q_CACHE服务器端的QStoreModule.register_q_cache为该实例构建 Q 缓存上下文GPUCacheContext、布局描述符并登记layout_desc_registry。注意REGISTER_Q_CACHE在 MQ 主循环上被 SYNC 串行化是_q_contexts的唯一插入点重复注册如恢复期 worker 首次 PING 时的再注册只刷新last_seen不重建上下文。捕获Capture每个注意力层的save_q_layer把该层 Q 按本步计划 scatter 进为 store 请求预留的 ring 块。存储Storebatched_submit_qstore_requests按请求逐个发送STORE_Q服务器端QStoreModule.store_q像处理 KV store 一样把 ring 块从 GPU 拷贝到 CPUTransferDirection.D2H走reserve_write(..., new)实现 chunk 级去重并以MP_STORE_SUBMITTED / MP_STORE_START / MP_STORE_END事件发布可观测性指标。回收Reclaimreclaim_finished_q_stores在get_finished中随 store 完成归还块。关闭Shutdownshutdown_q_ring发送UNREGISTER_Q_CACHE服务器释放该实例的上下文与设备内存含torch_dev.empty_cache()与ipc_collect()。服务器端模块 qstore.py 还实现了实例活性跟踪reap_stale_instances区分已 PING 证明的实例与从未 PING 的实例分别用reap_timeout_s与更宽松的registration_grace_s判定过期、report_status状态上报暴露registered_q_ids与每实例的q_ring_layout以及store_q的 fail-closed 语义只要任一 LMCache group 的块 ID 不足以覆盖所有 chunklen(group_block_ids) num_chunks * bpc整次 store 被跳过且不提交任何内容后续 retrieve 只会 miss 并触发重算绝不会缓存部分或越界数据。七、缓存寻址model##query与 KV 永不冲突Q 与 KV 共享同一条 store 路径但通过 query 专属模型名隔离worker 端构造q_model_name LMCacheSDKCacheKind.QUERY.server_model_name(model_name)即model##query详见 cache_kind.py 与 dispatcher.pyQ ring 以 worker 自身的instance_id与其 KV cache 相同注册靠模型名区分两套对象键的其余字段token chunk 哈希、kv_rank、cache_salt与 KV 完全一致因此 Q 的检索键可直接由 SDK 侧按相同规则构造。消费端LMCache SDK 文档以kindLMCacheSDKCacheKind.QUERY模块 lmcache/sdk/qcache.py连接时同样使用model##query后缀的模型名进行握手与检索其 Q 布局正是由 vLLM worker 的 Q ring 通过REGISTER_Q_CACHE注册的而非 KV 的REGISTER_KV_CACHE。SDK 侧按“world_size 1、单一非 hybrid kernel group”约束从/status读取布局完成[2, num_layers, hit_tokens, hidden_dim]连续张量的 retrieve / store / close可参考端到端示例 e2e_kv_edit.ipynb。八、配置开关与依赖前提配置项位置默认值作用transfer_intermediate_tensorslmcache.mp.transfer_intermediate_tensorsFalse总开关开启后连接器才把TRANSFER_QUERY特性加入 dispatcher见 lmcache_mp_connector.pylmcache.q.ring_blocksvLLMkv_transfer_configextra config无按深度估算显式指定 ring 块总数max(1, int(...))lmcache.q.ring_depth同上2未显式指定块数时用ceil(max_batched_tokens / block_size) * depth估算 ring 容量依赖前提该特性要求连接器与服务器就实验特性达成一致若连接器请求了TRANSFER_QUERY而服务器未声明支持init_dispatcher会抛ValueErrorQ ring 必须在 KV 传输上下文建立之后注册register_q_ring对未建立上下文直接抛RuntimeError捕获依赖 vLLM 前向传入intermediate_tensors含q/query与attn_metadata.slot_mapping两者缺失时相应层/步会安全跳过。九、当前限制仅支持 CUDA / lmcache-driven 传输register_q/submit_q_store只在LMCacheDrivenTransferContext中实现见 worker_transfer.pyengine-drivenCPU传输路径未实现因此 CPU-only 场景无法使用 Q 捕获。仅捕获 prefill 步的 Q当某一步allocate无法预留足够的块时典型发生在 decode 阶段该步的 Q 捕获被跳过——块被归还而非排队等待。当前这一行为可接受因为 SDK 以离线方式使用。十、小结Q Ring Buffer 是 LMCache 在“KV 之外的中间张量持久化”方向上的一个实现样本通过把 Query 张量包装成与分页 KV 同构的 GPU 环形缓冲区它几乎零成本地复用了既有的分页传输内核、STORE 协议与 chunk 级去重存储通过slot_mapping的槽位交集匹配它把“行”与“token”的对应从脆弱的顺序假设中解放出来天然适配 continuous batching、chunked prefill 与混合 STORE/RETRIEVE 步通过model##query的模型名隔离它与 KV 对象在同一个存储路径下共存而不冲突。对于需要在离线场景下分析、编辑或复用模型中间 Query 张量的开发者理解本机制是与lmcache.sdk.qcache配合使用、并进一步扩展其他中间张量类型的基础。【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表