vllm学习笔记之 PD分离 kv connector

vllm官方文档中部署PD分离的参数说明：https://docs.vllm.ai/en/latest/features/disagg_prefill.html?h=kv_connector 英文名叫做disaggregated prefilling，即预填充/解码分离。

Disaggregated Prefilling

Prefill 阶段：把用户给的 prompt（输入）喂给模型，一次性处理全上下文，生成一批 KV cache。它在计算上“重”的特点多（因为整个 prompt 都参与 attention、很多操作），但在生成阶段结束后，它的输出（KV cache）会被后续重用。Prefill阶段的计算复杂度是O（N^2）
Decode 阶段：模型基于前面生成／缓存的结果，一次一个 token 地继续生成输出的新内容。由于可以复用prefill阶段预先计算好的kv cache，复杂度降到O（N）。
如果prefill和decode阶段混在一起，Prefill 在占用大量计算资源的时候，可能会“抢”decode 所需的资源（GPU 时间／memory bandwidth），导致 decode 变慢，影响实时性。Decode 又可能在等待 prefill 完成，或者受到 prefill 执行的不确定延迟影响。

什么是Disaggregated Prefilling

将prefill和decode放在两个节点上，可以是两个不同的计算单元，两个不同的进程。 prefill 放在一个专用资源池／机器上，优化它的吞吐（batching、长 prompt 支持等）。 decode 放在另一个池／机器上，优化实时生成、token-by-token延迟。前端系统只需要将 prefill 阶段产生的 KV cache 传输给 decode 端即可，资源利用更高效、两阶段互不干扰、可以按负载独立扩展。

PD分离的两大好处/原因

1. 分别优化首 token 延迟（TTFT）和 token 间延迟（ITL）

PD分离的 prefill 机制让大模型推理的 prefill 阶段和 decode 阶段运行在不同的 vLLM 实例中。
这样你就可以灵活地为它们分配不同的并行策略（例如 tensor parallelism、pipeline parallelism），从而单独调节 TTFT 而不影响 ITL，或者反过来单独优化 ITL 而不影响 TTFT。

2. 控制Tail ITL

在没有使用disaggregated prefill 的情况下，vLLM 可能会在某个请求的decode过程中插入新的 prefill 任务，这会让尾部延迟变高。使用disaggregated prefill 可以避免这种情况，帮助更好地控制tail ITL。虽然“分块 prefill（chunked prefill）”也可以达到类似效果，但在实际中很难精确选择合适的 chunk 大小。因此，disaggregated prefill 是一种更稳定可靠的方式。

注意：对提高throughput无用。

vllm中的–kv-transfer-config参数设置

使用例子
--kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "num_cpu_blocks": 1000}}'

kv_connector

• 含义：选择用哪一种机制来在 prefill 节点和 decode 节点之间传输或者存储 KV 缓存。
• 选择考量：○ 如果部署在同一台机器／同一 GPU 上，可能用 SharedStorageConnector 就足够。○ 如果有跨 GPU／跨节点传输，或者想用高性能传输（比如 NVLink、RDMA），就可能使用 NixlConnector 或 P2pNcclConnector。○ 如果目标是释放 GPU 显存，把 KV 缓存卸载到 CPU 上，也可用 OffloadingConnector。
• 适用场景示例：○ 单机但用两卡分担 prefill + decode → SharedStorageConnector。○ 多机／有高速互联 → NixlConnector。○ 显存受限但 CPU 容量大 → OffloadingConnector。

kv_role

• 含义：标识当前运行实例是“产生 KV”的那一方还是“使用 KV”的那一方。
• 考量：○ Prefill 节点为 kv_producer，它生成 KV 缓存。○ Decode 节点为 kv_consumer，它接收并使用已生成的 KV 缓存。○ 设置为 kv_both，表示该实例可能既生成又使用 KV，适用于单机／无分离部署。
• 选择场景：○ 完全分离部署：一侧只做 prefill，另一侧只做 decode。○ 传统部署（prefill+decode 在同一进程内）：可以用 kv_both。

• 适合的场景：
  多用户／高并发、大批量 prompt，同时生成任务也多。
  Prompt 长度很长、需要做较大的 prefill 批量计算。比如多轮对话场景等。
  生成阶段要求低延迟／高吞吐。
  硬件资源可以分成预填充专用节点 + 生成专用节点。
  网络／存储传输 KV 缓存的成本可控。
• 不太适合或需谨慎的场景：
  Prompt 很短、或者几乎都是 “小 prompt +快速生成” 的轻量任务，此时 prefill/ decode 混在一起可能更简单更快。
  无法保证两端之间传输 KV cache 有足够带宽或低延迟。传输成本高反而抵消分离带来的收益。
  硬件资源受限，无法独立地为 prefill & decode 分配节点。
  系统对复杂部署（多节点、缓存传输、同步机制）不熟悉、维护成本较高。

KVConnectorBase_V1 class模板解析

KVConnector 是 “prefill → decode” 之间传输 KV 缓存的桥梁，除了使用vllm已经提供的几种kv connector外，我们也可以自定义结合硬件的kv connector， 参考vLLM v1 中实现 KV Connector 的模板接口（KVConnectorBase_V1）。主要要求实现以下内容：
Scheduler 端（调度器）要怎么跟 worker 协作
Worker 端（执行推理的 GPU 节点）要怎么加载/保存 KV
以及两边之间如何通信、同步状态

Scheduler vs Worker

class KVConnectorRole(enum.Enum):SCHEDULER = 0WORKER = 1

Scheduler 端：掌握全局状态，知道哪些请求要分配、哪些已经完成
Worker 端：真正运行模型 forward() 的那台机器（GPU worker）
每个 connector 都在这两边各有一个实例，互相协作。

Scheduler side method（调度器逻辑）

方法名	功能	作用阶段
get_num_new_matched_tokens()	询问当前请求还能从外部 KV 加载多少个 token	Scheduler 在分配显存前调用
update_state_after_alloc()	通知 connector：显存块已经分配好，可以加载了	调度器分配 buffer 后
build_connector_meta()	生成 metadata（元信息），发给 Worker，用来指导加载	Prefill → Decode 之间
request_finished()	通知 connector 请求完成，是否需要异步保存 KV	生成完毕后清理

Worker Side Method 实际的执行者

方法名	功能	说明
bind_connector_metadata()	绑定 scheduler 发来的 metadata	每次执行前
start_load_kv()	异步加载 KV 缓存（可以边加载边算）	Prefill → Decode 过渡
wait_for_layer_load()	等待某一层 KV 加载完	解码过程中，确保同步
save_kv_layer()	异步保存某层 KV（Prefill 结束时）	反方向传输
wait_for_save()	等待所有层保存完成	避免覆盖数据
get_finished()	通知哪些请求的 KV 传输已完成	异步确认状态

Metadata:
存kvcache的tokenid blockid定位，存一些存储状态。目的是让worker找到并加载kv cache。
以下是vllm V1中SharedStorageConnectorMetadata的实现。

class ReqMeta:# Request tokenstoken_ids: torch.Tensor# Slot mappings, should have the same length as token_idsslot_mapping: torch.Tensor# Is store or loadis_store: bool@staticmethoddef make_meta(token_ids: list[int], block_ids: list[int], block_size: int,is_store: bool) -> "ReqMeta":valid_num_tokens = align_to_block_size(len(token_ids), block_size)token_ids_tensor = torch.tensor(token_ids)[:valid_num_tokens]block_ids_tensor = torch.tensor(block_ids)num_blocks = block_ids_tensor.shape[0]block_offsets = torch.arange(0, block_size)slot_mapping = block_offsets.reshape((1, block_size)) + \block_ids_tensor.reshape((num_blocks, 1)) * block_sizeslot_mapping = slot_mapping.flatten()[:valid_num_tokens]return ReqMeta(token_ids=token_ids_tensor,slot_mapping=slot_mapping,is_store=is_store,)@dataclass
class SharedStorageConnectorMetadata(KVConnectorMetadata):requests: list[ReqMeta]def __init__(self):self.requests = []def add_request(self,token_ids: list[int],block_ids: list[int],block_size: int,is_store: bool,) -> None:self.requests.append(ReqMeta.make_meta(token_ids, block_ids, block_size, is_store))

学习资料参考

https://devpress.csdn.net/aibjcy/68d01bfa8867235e1386e0dd.html?spm=1001.2101.3001.6650.1&utm_medium=distribute.pc_relevant.none-task-blog-2%7Edefault%7EBlogOpenSearchComplete%7Eactivity-1-151937207-blog-154493882.235%5Ev43%5Epc_blog_bottom_relevance_base9&depth_1-utm_source=distribute.pc_relevant.none-task-blog-2%7Edefault%7EBlogOpenSearchComplete%7Eactivity-1-151937207-blog-154493882.235%5Ev43%5Epc_blog_bottom_relevance_base9&utm_relevant_index=1