vLLM 分离式 Prefill（Disaggregated Prefill）示例：ExampleConnector 离线 KV Cache 传输的实现与源码剖析

vLLM 分离式 Prefill（Disaggregated Prefill）示例：ExampleConnector 离线 KV Cache 传输的实现与源码剖析

【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm

本文以 examples/disaggregated/example_connector/ 目录下的示例脚本为主体，完整讲解 vLLM 分离式 prefill 的离线（offline）演示：如何用两个 vLLM 实例分别执行 prefill 与 decode，通过ExampleConnector以本地磁盘为共享存储完成 KV cache 的保存与加载。读完本文，你将能够独立运行该示例、理解KVTransferConfig各参数的含义，并掌握ExampleConnector从请求元数据构建、逐层 KV 抽取/注入到基于 prompt 哈希的前缀匹配这一整套底层机制。

1. 什么是分离式 Prefill，为什么需要它

vLLM 官方文档 docs/features/disagg_prefill.md 给出了两点核心动机：

• 独立调优 TTFT 与 ITL：分离式 prefill 把 LLM 推理的 prefill 阶段和 decode 阶段放到不同的 vLLM 实例中，允许为 prefill 实例和 decode 实例分别配置不同的并行策略（如tp、pp），调优首 token 延迟（TTFT）时不影响 token 间延迟（ITL），反之亦然。
• 控制尾部 ITL：不分离时，vLLM 可能在某个请求 decode 过程中插入新的 prefill 任务，导致尾部延迟升高。虽然合适的 chunked prefill 也能达到类似效果，但实践中难以确定正确的 chunk size，分离式方案更可靠。

需要明确的前提：官方文档特别指出，分离式 prefill 并不提升整体吞吐，它是面向延迟控制的架构选择。该功能标记为 experimental，API 可能随版本变化。

在 vLLM 当前的实现体系中，所有分离式 prefill 的实现都位于vllm/distributed/kv_transfer目录下，KV cache 在 prefill 实例与 decode 实例之间通过Connector抽象传递。官方文档列出了 9 类可用连接器，其中ExampleConnector是用于教学与调试的参考实现——它以磁盘作为共享存储，代码路径最短、最易读懂，是理解整个 KV connector 体系的入口。生产环境通常选择 NixlConnector、MooncakeConnector、LMCacheConnectorV1 等第三方连接器（详见 docs/features/disagg_prefill.md）。

2. 示例文件组成与运行流程

README（examples/disaggregated/example_connector/README.md）说明了该目录的三个文件及其职责：

run.sh的完整内容如下，它体现了示例的运行约定：

rm -rf local_storage/ if [ -f "output.txt" ]; then rm output.txt fi # The directory of current script SCRIPT_DIR=$(dirname "$(readlink -f "$0")") VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 "$SCRIPT_DIR/prefill_example.py" VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 "$SCRIPT_DIR/decode_example.py"

两个值得注意的运行细节：

1. 清理共享存储：先删除local_storage/目录和旧的output.txt，保证每次运行都从"冷缓存"开始，prefill 阶段必然发生 KV 落盘，decode 阶段必然命中加载。
2. VLLM_ENABLE_V1_MULTIPROCESSING=0：示例把 V1 引擎的多进程模式关闭，让引擎跑在单进程内，简化调试；CUDA_VISIBLE_DEVICES=0则把两个实例先后约束在同一块 GPU 上，便于最小化演示环境（prefill 实例退出后 decode 实例复用同一张卡）。

整体数据流为：

prefill 实例 磁盘 / output.txt decode 实例 LLM.generate(prompts) ──save──▶ local_storage/<hash>/<layer>.safetensors ──load──▶ LLM.generate(prompts) prompt+1 token ──────────────▶ output.txt

3. Prefill 实例脚本详解

prefill_example.py 的完整核心逻辑：

from vllm import LLM, SamplingParams from vllm.config import KVTransferConfig def read_prompts(): context = "Hi " * 1000 context2 = "Hey " * 500 return [ context + "Hello, my name is", context + "The capital of France is", context2 + "Your name is", context2 + "The capital of China is", ] def main(): prompts = read_prompts() sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1) llm = LLM( model="meta-llama/Llama-3.2-1B-Instruct", enforce_eager=True, gpu_memory_utilization=0.8, kv_transfer_config=KVTransferConfig( kv_connector="ExampleConnector", kv_role="kv_both", kv_connector_extra_config={"shared_storage_path": "local_storage"}, ), ) # 1ST generation (prefill instance) outputs = llm.generate(prompts, sampling_params) new_prompts = [] for output in outputs: prompt = output.prompt generated_text = output.outputs[0].text new_prompts.append(prompt + generated_text) # Write new_prompts to output.txt with open("output.txt", "w") as f: for prompt in new_prompts: f.write(prompt + "\n")

几个关键设计点：

• max_tokens=1：prefill 实例只额外生成 1 个 token。这样 prefill 阶段计算的 KV 基本覆盖整个 prompt，随后这 1 个 token 被拼接到 prompt 后面写入output.txt，成为 decode 实例的输入。
• kv_connector="ExampleConnector"：通过连接器名称从工厂（vllm/distributed/kv_transfer/kv_connector/factory.py）动态实例化对应连接器类。
• kv_role="kv_both"：该实例既可作为 KV 生产者（store）也可作为消费者（load）。示例中 prefill 脚本实际上只做 store、decode 脚本实际上只做 load，但两边都用kv_both以便复用同一套配置。
• shared_storage_path: "local_storage"：传给ExampleConnector的额外配置，指定 KV 落盘目录；不配置时源码默认回退到/tmp（见 example_connector.py）。
• enforce_eager=True与gpu_memory_utilization=0.8：关闭 CUDA graph 以保证连接器逐层钩子（store/load 发生在 attention 层前后）行为确定、易调试。

4. Decode 实例脚本详解

decode_example.py 与 prefill 脚本结构几乎一致，差异集中在三点：

def read_prompts(): """Read prompts from output.txt""" prompts = [] try: with open("output.txt") as f: for line in f: prompts.append(line.strip()) except FileNotFoundError: print("Error: output.txt file not found") exit(-1) ... sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) llm = LLM( model="meta-llama/Llama-3.2-1B-Instruct", enforce_eager=True, gpu_memory_utilization=0.8, max_num_batched_tokens=64, max_num_seqs=16, kv_transfer_config=KVTransferConfig( kv_connector="ExampleConnector", kv_role="kv_both", kv_connector_extra_config={"shared_storage_path": "local_storage"}, ), ) outputs = llm.generate(prompts, sampling_params)

• 输入来自output.txt：即 prefill 阶段产出的prompt + 1 token序列。decode 实例按这些 token 计算出与磁盘目录一致的哈希，从而"命中"外部缓存。
• max_tokens=10：decode 实例真正负责续写，生成 10 个 token。
• max_num_batched_tokens=64/max_num_seqs=16：把 decode 实例调度成典型的小 batch、纯 decode 形态——这正是分离式部署中 decode 实例的常见调参思路（小 batch、低 ITL 敏感），与 prefill 实例的配置解耦。

decode 阶段命中外部缓存后，调度器把对应 KV 的加载任务交给连接器，模型无需重算 prompt 部分的 attention KV，直接进入逐 token 生成。

5. ExampleConnector 源码实现剖析

实现位于 vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py。源码注释开宗明义："This is Simple debug implementation of the KV connector. It save / load the KV cache to / from the disk."，即这是一个以磁盘为后端的调试版连接器，继承自KVConnectorBase_V1（example_connector.py#L85）。

5.1 请求元数据：ReqMeta与build_connector_meta

每个需要 store 或 load KV 的请求被封装成ReqMeta（example_connector.py#L32-L65）：

• token_ids：参与传输的 token 序列，长度对齐到 block 边界（align_to_block_size按(n-1) // block_size * block_size向下取整，见 example_connector.py#L441-L443）；
• slot_mapping：由block_ids与block_size推导出的全局 slot 索引（slot = block_id * block_size + offset），与 token 一一对应；
• is_store：标记该请求本步是保存还是加载；
• mm_hashes：多模态输入的特征标识列表，参与缓存键计算。

build_connector_meta（example_connector.py#L299-L373）在每步调度后运行，决定本步各请求的角色：

1. 遍历scheduler_output.scheduled_new_reqs：若该请求此前被get_num_new_matched_tokens判定命中并登记到_requests_need_load，则以is_store=False加入元数据（本步需要加载）；否则，若磁盘上还没有该 prompt 的缓存目录，则以is_store=True加入（本步需要保存）。
2. 对从抢占中恢复（resumed from preemption）的请求，同样登记为 load，并在结尾用assert total_need_load == len(self._requests_need_load)保证所有待加载请求都在一步内被排上，随后清空登记表。

注释中也说明了这个简化实现的取舍：store 与 load 被设为互斥处理（"a single request can have both store and load"，但调试版只缓存原始 prompt token）。

5.2 保存路径：save_kv_layer逐层抽取 KV 并写盘

save_kv_layer（example_connector.py#L201-L245）是 worker 侧钩子，在模型 forward 过程中被逐层调用。其内层函数extract_kv_from_layer从每层的 paged KV buffer 中按slot_mapping抽取本请求的 KV：

def extract_kv_from_layer(layer: torch.Tensor, slot_mapping: torch.Tensor) -> torch.Tensor: """The layer is a standardized [B, H, N, C] per-layer view (H == 1 for MLA).""" slot_mapping = slot_mapping.to(layer.device, non_blocking=True) if isinstance(attn_metadata, MLACommonMetadata): # [B, 1, N, C] -> [B * N, C]; slot_mapping indexes B * N slots. return layer.reshape(-1, layer.shape[-1])[slot_mapping, ...] block_idxs = slot_mapping // self._block_size offsets = slot_mapping % self._block_size return layer[block_idxs, :, offsets]

普通 attention 布局下，层视图为[num_blocks, 2*H, block_size, head_dim]，通过slot // block_size与slot % block_size定位到块内偏移；MLA 布局则 flatten 成[B*N, C]后直接用 slot 索引。抽取结果以{"kv_cache": tensor}形式同步拷回 CPU 并用safetensors落盘到：

<shared_storage_path>/<prompt_hash>/<layer_name>.safetensors

即一个请求一个哈希目录，一层一个 safetensors 文件。这也解释了示例中local_storage/目录的内容组织方式。

5.3 加载路径：start_load_kv读盘并注入 paged buffer

start_load_kv（example_connector.py#L110-L188）在 forward 开始前运行，遍历元数据中is_store=False的请求，对forward_context.no_compile_layers中的每一层：跳过没有kv_cache属性的非 attention 层（如 MoE/MLP），从对应 safetensors 文件读回kv_cache，搬到 KV 所在设备，再由inject_kv_into_layer按slot_mapping写回 paged KV buffer——即save_kv_layer的逆操作（dst[block_idxs, :, offsets] = src_kv_cache，MLA 分支同样按 flatten 后的 slot 索引写入）。

wait_for_layer_load/wait_for_save均为空实现：因为磁盘读写在本实现中是同步完成的，没有需要等待的异步传输；官方文档中提到的逐层流水线（layer-by-layer pipelining）接口在此仅为占位。

5.4 前缀匹配与哈希命名：get_num_new_matched_tokens

decode 实例如何知道磁盘里有它要的 KV？答案是以 prompt token 序列的哈希作为缓存键。

get_num_new_matched_tokens（example_connector.py#L250-L285）在调度阶段被调用，注释给出了该调试版的假设：prompt 形如cached_prompt + newly_generated_single_token，因此用prompt_token_ids[:-1]（去掉最后 1 个新生成 token）来定位目录：

if not self._found_match_for_request(request): return 0, False logger.info("External Cache Hit!") token_ids = request.prompt_token_ids or [] num_tokens_to_check = align_to_block_size(len(token_ids) - 1, self._block_size) return num_tokens_to_check - num_computed_tokens, False

命中时返回可外部加载的 token 数（按 block 粒度对齐，因为 v1 调度器要求返回值与num_computed_tokens对齐到 block 边界），未命中返回0, False。_found_match_for_request最终归结为一次os.path.exists(foldername)检查（example_connector.py#L379-L402）——这是磁盘版连接器"查缓存"的全部逻辑。

目录名由_generate_foldername_debug（example_connector.py#L404-L424）生成：将 token 张量的字节序列转 bytes，若存在多模态输入再把mm_hashes拼接进被哈希的字节串（注释说明这样做既防止路径穿越，也让多模态输入形成规范化缓存键），经safe_hash得到十六进制摘要，与shared_storage_path拼接成目录。

6.KVTransferConfig参数说明

示例中用到的KVTransferConfig定义在 vllm/config/kv_transfer.py。除示例涉及的三个字段外，完整字段一览：

配置生效的判定见 kv_transfer.py#L108-L118：is_kv_transfer_instance要求kv_connector非空且kv_role合法，引擎据此决定是否把本实例接入 KV 传输体系。ExampleConnector 自身只消费kv_connector_extra_config里的shared_storage_path一项（example_connector.py#L104-L108），其余字段是为 NixlConnector、TorchDistributedConnector 等其他连接器服务的。

7. 测试用例与哈希语义验证

单元测试 tests/v1/kv_connector/unit/test_example_connector.py 用多模态模型（Qwen2.5-VL）系统验证了上述哈希语义：以tmp_path作为shared_storage_path，断言存储目录下生成的哈希目录数量随输入变化而精确递增——

• 同文本、不同图片（同尺寸）→ 新增一个哈希目录；
• 同一张图再次输入 → 目录数不变（命中缓存）；
• 两张图的顺序互换 → 新增目录（顺序参与缓存键）;
• 纯文本输入作为对照 → 同样按内容去重。

测试同时参数化了 attention backend（CUDA 下FLASH_ATTN/TRITON_ATTN），确认extract_kv_from_layer的按 slot 抽取逻辑在不同后端布局下都成立。这组用例也佐证了 5.4 节所述的结论：ExampleConnector的缓存键 = prompt token 字节哈希 ⊕ 多模态特征哈希，与磁盘目录一一对应。

8. 适用边界与后续方向

从源码注释和文档定位看，ExampleConnector明确是教学/调试实现，存在几处刻意留下的简化：

• 磁盘读写同步执行、无异步传输与逐层流水线，wait_for_save/wait_for_layer_load为空操作；
• store 与 load 按步互斥，只缓存原始 prompt token（不缓存生成过程）；
• 源码注释承认它"会做多余工作、覆盖 GPU 上已有的 prefix cache"，需要额外 mask 才能消除该开销；
• 缓存命中依赖 prompt 前缀精确匹配（且要求 decode 侧 prompt = prefill prompt + 1 个 token 的特定形态），不适用于一般性的前缀共享场景。

因此生产级分离式部署应选用官方文档列出的其他连接器（NixlConnector、MooncakeConnector、LMCacheConnectorV1、OffloadingConnector、FlexKVConnectorV1 等），配置方式相同——都是KVTransferConfig中改kv_connector名称并补充对应的kv_connector_extra_config。若要从这个最小示例出发扩展自己的连接器，参考 docs/features/disagg_prefill.md 给出的三条实现路径（完全自定义 connector、数据库式 LookupBuffer、分布式 P2P Pipe），并对照 tests/v1/kv_connector/unit/test_example_connector.py 的断言方式编写验证用例，即可在 vLLM 的 KV connector 框架内安全迭代。

【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm