/ AI Infra

[Infra-9] FlexKV：把 KV Cache 变成可分层、可搬运、可共享的资源

从 prefix 匹配、RadixTree、异步 transfer graph 到 Redis + Mooncake 的分布式复用，理解 FlexKV 如何扩展 KV Cache 的生命周期，以及它与 vLLM 原生 KV 管理的边界。

2026年7月15日9 minInfra · LLM Serving · KV Cache · vLLM

长上下文、RAG、多轮对话和 agent 系统有一个共同问题：许多请求的开头其实高度重复。system prompt、工具说明、检索到的文档、会话历史都可能已经被别的请求计算过；但如果对应的 KV Cache 已被 GPU 驱逐，传统 serving 路径往往只能重新 prefill。

FlexKV 的定位正是在这里。它不是新的 LLM 推理引擎，也不替代 attention kernel 或 PagedAttention。它是一个面向推理的分布式 KV Store 与多级缓存管理层：把原本主要停留在 GPU 显存里的 KV Cache，扩展到 CPU 内存、本地 SSD 与远端存储，并在需要时将命中的 KV 异步恢复到 GPU。

GPU 只保留当前活跃请求需要的热 KV
CPU / SSD / Remote 保留可复用的历史 KV
新请求先做 prefix match
命中部分回填到 GPU，未命中部分再执行 prefill

这条路线优化的不是单次 attention 的数学计算，而是避免重复 prefill。它最适合长 prompt、高前缀复用、跨实例 serving 的 workload。

1. 从 KV Cache 的生命周期看问题#

decoder-only 模型的每一层都会为历史 token 保留 Key 和 Value。生成下一个 token 时，attention 需要读取所有历史 KV；因此 KV Cache 是 decode 阶段延迟与显存容量的核心约束。

一次请求的理想生命周期可以写成：

prompt
  -> prefill produces KV
  -> decode consumes and appends KV
  -> request ends
  -> KV remains reusable for a later matching prefix

传统本地 prefix cache 已经能复用 GPU 内的相同前缀，但它受 GPU 显存、单个 worker 和局部调度域的限制。FlexKV 延长了最后一步的生命周期：GPU 放不下的 KV 不必立刻失去复用价值，而可以被 offload 到容量更大的层级。

层级	典型位置	作用
L0	GPU HBM	当前请求的热数据，带宽最高、容量最稀缺
L1	CPU memory	较快的本地外部缓存
L2	Local SSD	容量更大的本地持久层
L3	Remote / distributed pool	跨节点复用的可扩展 KV 池

这里的关键取舍很朴素：从 SSD 或远端读取并不免费；只有当读取与搬运成本低于重复 prefill 成本，并且能被异步重叠隐藏时，分层缓存才真正带来端到端收益。

2. 系统分层：控制面决定，数据面搬运#

FlexKV 的 README 将实现拆为 StorageEngine、GlobalCacheEngine 与 TransferEngine。从推理 infra 视角，可以把它理解为控制面和数据面的分工：

vLLM / TensorRT-LLM / Dynamo
            |
            | Connector / Adapter
            v
        KVManager API
 get_async / put_async / prefetch_async / wait
            |
     KVTaskManager / KVTaskEngine
        |                      |
        v                      v
GlobalCacheEngine        TransferManager
match / eviction         registers GPU blocks
placement decision       submits transfer graphs
        |                      |
        v                      v
CPU / SSD / Remote       TransferEngine
                          H2D / D2H / DISK2H / DISK2D / Remote

StorageEngine 负责不同介质上 KV block 的 layout 与物理存储抽象；GlobalCacheEngine 做前缀匹配、空间管理、驱逐及搬运路径规划；TransferEngine 则执行真实的数据拷贝与 I/O。这个边界很重要：缓存命中本身不代表数据已经在 GPU 上，控制面必须生成一份正确的搬运计划，数据面才可以执行它。

3. Block 化与 RadixTree：如何判断哪些 KV 可以复用#

FlexKV 不会按单个 token 管理一份独立的 KV。它按 tokens_per_block 将连续 token 切成 block；一个 block 通常对应这段 token 在所有层的 KV。布局会由 layer 数、KV head 数、head size、block 数与 MLA 等模型属性共同决定。

token ids: [t0, t1, ..., t47]
 
block 0: hash(t0  ... t15)
block 1: hash(t16 ... t31)
block 2: hash(t32 ... t47)

这些 block hash 被插入 RadixTree。新请求到达时，系统对其 block 序列做最长前缀匹配：匹配到的连续 block 可以复用，余下的 suffix 仍需模型计算。

root
  └── B0
       └── B1
            └── B2
 
new request: B0 -> B1 -> Bx
matched KV:   B0, B1
prefill:                  Bx and its suffix

这也是它与 GPU 内部 prefix caching 的根本差异：索引逻辑仍然是 prefix reuse，但命中的数据位置可以是 CPU、SSD 或远端节点，而非仅仅是当前 GPU 的 block table。cache_engine.py 中同时提供 Python 与 C++ 加速的 RadixTree 实现路径。

4. Get：从外部缓存恢复命中的前缀#

KVManager 对上层暴露 get_async、put_async、prefetch_async、launch 与 wait 等异步接口。以 get 为例，控制面首先对 token block 做 prefix match，然后基于每段数据所在层级构建 transfer plan：

CPU hit:  CPU  -> GPU
SSD hit:  SSD  -> CPU -> GPU
GDS path: SSD  -> GPU
remote:   remote -> CPU -> GPU
miss:                    model prefill

GlobalCacheEngine.get() 返回的不只是命中长度，还包括 TransferOpGraph、命中 mask、回调以及用于判断任务结束的 op 信息。也就是说，get 的结果不是同步 memcpy，而是一张带依赖关系的搬运图。

SSD -> CPU -> GPU
       |       |
       +-- callback: release temporary buffer
               |
               +-- task complete

显式图化有两个价值。第一，SSD、CPU、GPU 和远端路径天然可能是多段传输，依赖需要被正确描述。第二，图可被任务层异步提交，允许预取与其他模型计算重叠。数据尚未搬完前，系统不能把它当成可供 attention 使用的 GPU KV。

5. Put：为什么索引状态要区分“存在”和“可读”#

prefill 完成后，GPU 中的 KV 可以被写入外部缓存：

prefill completed
  -> put_async(token_ids, slot_mapping)
  -> allocate CPU / SSD / Remote blocks
  -> insert prefix metadata
  -> D2H / H2DISK / remote transfer
  -> data complete, node becomes ready

这里容易被忽略的并发问题是：索引写入与数据写入是两件事。为了让后续请求尽早看到 prefix，系统可以先插入节点；但在异步 offload 实际完成前，这个节点不能被当作可读取命中。FlexKV 通过 ready 状态与 transfer 完成回调来衔接这段窗口，避免读到尚未完成的数据。

因此它不是“把 KV dump 到磁盘”的脚本，而是把索引、分配、传输、完成事件与回收纳入同一任务生命周期。

6. 驱逐不是物理搬家，而是为 block 重新分配资格#

外部层也不是无限大。GlobalCacheEngine 会维护 CPU、SSD 与 Remote 侧的 mempool，并根据容量阈值和 eviction policy 选择牺牲的 prefix node 或 block。配置可涉及 evict_ratio、evict_start_threshold、hit_reward_seconds 等策略参数，详见 FlexKV 配置文档。

在这个语境中，“逻辑驱逐”通常意味着先移除索引与 block 映射、归还 mempool 的物理 block；它并不等价于把全部数据同步拷贝到另一个层级。这样的设计可避免驱逐本身变成高峰期的额外 I/O 风暴。

策略的目标也不只是传统 LRU。一个昂贵、命中频繁的长前缀，往往比短而冷的前缀更值得保留；是否把 prefill saved、命中率、数据大小和回填代价纳入价值函数，是生产系统仍需按 workload 调整的部分。

7. 分布式复用：索引局部化，数据按需跨节点传输#

单机分层缓存解决的是“GPU 放不下”；分布式复用解决的是“另一个 worker 已经算过”。FlexKV 的文档描述了一条由 Redis、索引快照和 Mooncake Transfer Engine 组成的路径：

Node A computes a long prefix
  -> publishes prefix metadata to Global Meta Store (for example Redis)
  -> Node B refreshes its local global-index snapshot
  -> Node B matches the prefix locally
  -> Node B pulls KV from Node A through Mooncake / RDMA
  -> Node B restores KV locally and then to GPU

关键设计不是让每个请求都同步查询中心元数据，而是让各节点维护 global index 的本地 snapshot，再用 lease 等机制维护跨节点传输期间的数据有效性。这样，频繁的 prefix lookup 留在本地，远端交互主要发生在实际需要的 metadata 更新与 KV 数据传输上。具体机制可参考 distributed reuse 文档。

8. FlexKV 与 vLLM：谁管理 GPU 内的 KV，谁延长 KV 生命周期#

这两个系统的边界最容易混淆。vLLM 原生仍负责模型执行、scheduler、PagedAttention、GPU block table、slot mapping，以及活跃 sequence 的 KV 组织；FlexKV 负责 GPU 外部的 KV 保存、检索与回填。

vLLM:   正在推理的 KV 如何在 GPU 上组织、访问和调度
FlexKV: 已计算但不应一直占用 GPU 的 KV 如何保存、查找和搬回

因此 slot_mapping 是 connector 的关键接口。FlexKV 不决定 attention kernel 如何访问 KV；它使用推理引擎提供的映射，识别目标 GPU slot，再把外部命中的 block 填回正确位置。

按 FlexKV 的 vLLM adapter 文档，支持版本可通过 FlexKVConnectorV1 将其接入 KV transfer connector：

VLLM_USE_V1=1 python -m vllm.entrypoints.cli.main serve Qwen3/Qwen3-32B \
  --tensor-parallel-size 8 \
  --enable-chunked-prefill \
  --enable-prefix-caching \
  --kv-transfer-config \
    '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'

版本兼容性变化很快，部署前应以所用 vLLM release 与 connector 文档为准。上游 PR 中的性能数字也应被当成特定硬件、模型、输入输出长度与并发下的实验结果，而非通用承诺。

9. 什么时候值得使用，什么时候不值得#

FlexKV 的收益依赖命中率与传输成本。下面这些场景通常更有希望：

长 system prompt、固定工具描述或共享 RAG 文档；
多轮对话中反复出现的长历史；
GPU 显存紧张，但本地 CPU/SSD 容量充足；
多个 serving 实例的请求前缀重叠明显；
调度器能够提前发现前缀并发起 prefetch。

反过来，短 prompt、低重叠的随机请求，或 SSD/网络回填慢于重算 prefill 的场景，可能只会增加索引、I/O 与调度开销。实际评估至少应同时观察：prefix hit rate、恢复字节数、回填延迟、被隐藏的传输比例、TTFT、TPOT/ITL，以及因 offload 引入的尾延迟。

10. 对非 CUDA 后端的启示#

FlexKV 的系统思想可以迁移，但当前实现明显带有 CUDA/NVIDIA 假设，例如 CUDA IPC、cudaMemcpyAsync、GDS、CUDA MPS 与 torch.cuda 设备管理。若将类似能力接入 MLU 或其他加速器，不能只改一个 connector 名称；至少要重新实现：

device-memory shared handle / IPC
host <-> device async copy backend
SSD -> device direct-I/O path (if available)
device discovery and stream/event semantics
GPU block registration and slot-mapping handoff

更本质的要求是保持控制面与数据面的分离：RadixTree、mempool、驱逐和 transfer graph 可以尽量复用抽象；真正与硬件绑定的 copy、同步、内存句柄和 direct-storage 路径应被收敛在后端实现中。

11. 总结#

FlexKV 将 KV Cache 从“GPU 显存内、随请求结束而逐渐失去价值的状态”，变成“可索引、可分层、可异步搬运、可跨节点复用的系统资源”。

它的核心路径可以压缩为六步：

1. 用统一 KV block layout 描述不同存储层
2. 用 RadixTree 对 token block 做最长前缀匹配
3. 根据 CPU / SSD / Remote 命中构建 TransferOpGraph
4. 通过异步 task 与 callback 管理 ready、释放与完成事件
5. 通过 connector 将命中 block 回填至推理引擎指定的 GPU slot
6. 用元数据快照与远端传输扩展到跨节点复用

如果把 vLLM 看作“让活跃 KV 在 GPU 上高效运行”的系统，FlexKV 解决的就是另一半问题：让已经计算过、但不该永远占用 GPU 的 KV，仍然能够以合适的代价被再次使用。