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 的 Query（Q）张量是前向过程中的瞬态数据，常规 KV Cache 层不会保存它。LMCache 的Q Ring Buffer通过在 vLLM 工作进程内部按“分页 KV Cache 的布局”将每一层的 Q 暂存进 GPU 环形缓冲区，再复用既有的 STORE 通路把整块数据异步落盘到 LMCache MP 服务器，从而让下游 SDK（qcache）能够按 token id 检索、编辑并回存 Query 张量（典型场景如 token-dropping 等中间张量分析）。读完本文，你将掌握 Q Ring Buffer 的组件划分、行到 token 的归属算法、五阶段生命周期、缓存寻址规则与配置开关，以及当前实现的功能边界。

一、设计目标：让瞬态的 Q 张量可被持久化复用

Attention 的 Query 张量天然是瞬态的：它只在一次前向中出现，计算完 K、V 点积后即被丢弃。LMCache 的目标是让这类张量也能像 KV Cache 一样被保存、检索与复用（检索侧见 LMCache SDK 文档 中的qcache）。为此，设计采用了以下关键决策：

• 在生产者侧捕获：Q 的捕获发生在 vLLM（生产者）内部，而非消费端 SDK，SDK 用户不需要也不应该直接调用这里的任何接口。
• 分页化暂存：每一层的 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] = { f"lmcache_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(f"cannot 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 reserved

2.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）被调用，完成以下工作：

1. 挑选注意力层：通过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 或没有注意力层时打警告并跳过。
2. 计算几何信息：从vllm_config.model_config获取num_q_heads（考虑parallel_config）与head_size，从cache_config.block_size获取块大小，dtype 从模型配置推导。
3. 计算 ring 大小，优先级如下：
   • 显式指定：lmcache.q.ring_blocks（extra config），取max(1, int(explicit))；
   • 否则按深度估算：depth = max(1, lmcache.q.ring_depth)（默认 2），max_batched = scheduler_config.max_num_batched_tokens（缺省 8192），则num_ring_blocks = max(1, ceil(max_batched / block_size) * depth)。
4. 调用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 存储计划”（_QStepState：ring_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_salt=q_store.cache_salt)

即每个请求提交一个STORE_Q；本步没有捕获计划或事件缺失时静默跳过，且状态被重置，不会泄漏到下一步。

四、行到 token 的归属（Row-to-Token Attribution）：不做位置假设的匹配算法

这是整个设计中最微妙的部分。在 continuous batching 下，一步的 Q 张量把多个请求的行拼接在一起，且存在三重“不对齐”：

1. 行数与存储 token 数不等：请求在某一步的调度 token 数（行数）未必等于其 store 操作的按 chunk 对齐 token 数——prompt 尾部越过最后一个 chunk 边界的部分仍然产生行；
2. 顺序不一致：批内行的排列顺序未必与连接器元数据中的请求顺序一致；
3. 部分存在：某请求的 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_size（op 的块列表已预先切片到其[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；
• 分配QRingBuffer（hidden_dim = num_q_heads * head_size），构造EngineGroupInfo(engine_group_id=0, layer_indices=tuple(range(num_layers)), tokens_per_block=block_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_name=self.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阶段被调用：轮询每个 future，store 完成后立即归还 ring 块；服务器不健康时批量归还并清空。这保证了环形缓冲区容量能在异步落盘完成后及时回收复用。

4.reregister_q_ring()与shutdown_q_ring()

• reregister_q_ring：服务器恢复后重发REGISTER_Q_CACHE（ring 从未构建时为空操作），由 Dispatcher 的 reregister 在心跳恢复流程中调用；
• shutdown_q_ring：teardown 时发送UNREGISTER_Q_CACHE（req_client.unregister_q_cache(instance_id)），超时仅打警告、继续关闭流程。

六、五阶段生命周期

1. 注册（Register）：setup_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，不重建上下文。
2. 捕获（Capture）：每个注意力层的save_q_layer把该层 Q 按本步计划 scatter 进为 store 请求预留的 ring 块。
3. 存储（Store）：batched_submit_qstore_requests按请求逐个发送STORE_Q；服务器端QStoreModule.store_q像处理 KV store 一样把 ring 块从 GPU 拷贝到 CPU（TransferDirection.D2H），走reserve_write(..., "new")实现 chunk 级去重，并以MP_STORE_SUBMITTED / MP_STORE_START / MP_STORE_END事件发布可观测性指标。
4. 回收（Reclaim）：reclaim_finished_q_stores在get_finished中随 store 完成归还块。
5. 关闭（Shutdown）：shutdown_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 不足以覆盖所有 chunk（len(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.py）；
• Q ring 以 worker 自身的instance_id（与其 KV cache 相同）注册，靠模型名区分两套对象；
• 键的其余字段（token chunk 哈希、kv_rank、cache_salt）与 KV 完全一致，因此 Q 的检索键可直接由 SDK 侧按相同规则构造。

消费端（LMCache SDK 文档）以kind=LMCacheSDKCacheKind.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_QUERY而服务器未声明支持，init_dispatcher会抛ValueError；
• Q ring 必须在 KV 传输上下文建立之后注册（register_q_ring对未建立上下文直接抛RuntimeError）；
• 捕获依赖 vLLM 前向传入intermediate_tensors（含q/query）与attn_metadata.slot_mapping，两者缺失时相应层/步会安全跳过。

九、当前限制

1. 仅支持 CUDA / lmcache-driven 传输：register_q/submit_q_store只在LMCacheDrivenTransferContext中实现（见 worker_transfer.py）；engine-driven（CPU）传输路径未实现，因此 CPU-only 场景无法使用 Q 捕获。
2. 仅捕获 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