PDBackend (预填充-解码分离后端) 源码分析

< 返回存储架构概览

PDBackend (Prefill-Decode Backend) 是 LMCache 专为 预填充-解码分离 (Prefill-Decode Disaggregation) 架构设计的核心存储后端。在此模式下，计算负载被拆分为两个阶段：

1. Prefill (预填充): 负责处理输入的 Prompt，生成 KV Cache。通常由高算力节点（如 H100）承担。
2. Decode (解码): 负责逐个生成 Token。通常由低显存带宽节点（如 A10/L40）承担。

PDBackend 的核心职责是实现 KV Cache 从 Prefill 实例到 Decode 实例的低延迟、高吞吐、主动推送 (Push-based) 传输。

架构定位: 它不仅是一个存储后端 (Storage Backend)，还取代了 LocalCPUBackend 充当 内存分配器 (Allocator) 的角色，直接管理 Host 内存以实现零拷贝传输。

---

1. 核心架构设计

本章将从宏观视角解析 PDBackend 的设计理念。与传统的 P2PBackend（基于 Pull 模式，即“缺页中断”式拉取）不同，PDBackend 采用 Push 模式：Prefill 节点在生成 KV Cache 后，主动将其推送到 Decode 节点。这种设计消除了 Decode 阶段的首个 Token 生成延迟 (TTFT)，确保推理计算无缝衔接。

1.1 角色分工

PDBackend 根据配置 (pd_role) 分为两种角色，分别对应 PD 分离架构的两端：

• Sender (Prefiller):

  • 行为: 生产者。不保留任何 KV Cache，生成数据后立即推送到 Receiver。
  • 关键逻辑: 远程内存申请 -> 数据推送 -> 通知 Proxy。
  • 存储: 本地 self.data 始终为空，不进行本地存储。

• Receiver (Decoder):

  • 行为: 消费者。被动接收来自 Sender 的数据，并存储在本地内存中供推理引擎读取。
  • 关键逻辑: 监听内存申请请求 -> 分配本地内存 -> 返回地址 -> 等待数据写入。
  • 存储: 维护 self.data，保存接收到的 KV Cache。

1.2 通信协议

PD 模式引入了一套基于 MsgPack 的轻量级控制平面协议，用于协调内存分配。

• AllocRequest: Sender 向 Receiver 申请内存。包含 Key 列表、Shape、Dtype、Format、Last Chunk Tokens 等元数据。
• AllocResponse: Receiver 返回分配结果。包含 remote_indexes (接收端内存地址/索引) 和 already_sent_indexes (已存在的 Key，避免重复传输)。
• ProxyNotif: Sender 通知全局 Proxy 传输完成，Proxy 随后调度 Decode 请求给 Receiver。

1.3 组件交互流程

下图展示了 PD 模式下 KV Cache 的流转过程：

graph LR
    subgraph Sender [Prefill Instance]
        SenderBackend(PDBackend)
        Transfer_S(TransferChannel)
    end

    subgraph Receiver [Decode Instance]
        RecvBackend(PDBackend)
        Transfer_R(TransferChannel)
        AllocThread(Alloc Loop Thread)
    end

    Proxy[Global Proxy]

    %% Flow
    SenderBackend -- 1. AllocRequest (ZMQ) --> AllocThread
    AllocThread -- 2. Alloc Local Mem --> RecvBackend
    AllocThread -- 3. AllocResponse (Remote Addr) --> SenderBackend
    SenderBackend -- 4. Write Data (RDMA/TCP) --> Transfer_S
    Transfer_S == Data Stream ==> Transfer_R
    Transfer_R -.-> RecvBackend
    SenderBackend -- 5. Notify Done --> Proxy

---

2. 核心代码解析

本章深入剖析 PDBackend 的实现细节。作为系统中最复杂的后端组件，PDBackend 不仅实现了标准的存储接口，还集成了 内存分配器 (Allocator)、ZMQ 控制平面 与 数据传输平面。接下来的分析将围绕 Sender (Prefiller) 与 Receiver (Decoder) 的状态流转展开，揭示其如何通过协同工作实现高效的 KV Cache 推送。

代码位于 lmcache/v1/storage_backend/pd_backend.py。

2.1 初始化与配置

PDConfig 类负责解析配置并确定节点角色。

• Sender: 需要配置 proxy_host/port 以通知调度器。
• Receiver: 需要配置 peer_host/port 以绑定监听端口。
• Buffer Device: 自动检测并设置 CUDA 设备，用于 GPU-CPU 内存传输。

# 初始化时根据 role 分支处理
if self.pd_config.role == "sender":
    self._init_sender()
    # Sender 需要记录已初始化的 Peer 和对应的内存分配 Socket
    self.initialized_peers: set[str] = set()
    self.mem_alloc_sockets: dict[str, zmq.Socket] = {}
elif self.pd_config.role == "receiver":
    self._init_receiver()

2.2 Sender (Prefiller) 实现

Sender 的核心逻辑在 batched_submit_put_task 中，它重写了标准 StorageBackend 的 put 行为。

2.2.1 远程内存分配

在发送数据前，Sender 必须先知道“往哪里写”。它通过 ZMQ 向 Receiver 发送 AllocRequest。

def batched_submit_put_task(self, keys, memory_objs, transfer_spec):
    # 1. 建立/确认连接 (Lazy Connection)
    # 根据 transfer_spec 中的目标信息建立数据面和控制面连接
    self._ensure_peer_connection(...)

    # 2. 发送分配请求 (AllocRequest)
    # 根据 memory_objs 的 shape/dtype/fmt 构造请求
    alloc_request = self._get_remote_alloc_request(keys, memory_objs)
    # 阻塞等待 Receiver 返回分配结果
    alloc_response = self._remote_allocate(receiver_id, alloc_request)

    # 3. 解析响应
    already_sent_indexes = alloc_response.already_sent_indexes
    remote_indexes = alloc_response.remote_indexes # 目标物理地址列表

2.2.2 数据推送

获得目标地址后，Sender 过滤掉已存在的 Key，利用 TransferChannel 直接将数据写入 Receiver 内存。

    # 4. 过滤无需发送的数据
    mem_objs_to_send = []
    for idx, mem_obj in enumerate(memory_objs):
        if idx in already_sent_indexes:
            mem_obj.ref_count_down() # 释放本地引用
        else:
            mem_objs_to_send.append(mem_obj)

    if mem_objs_to_send:
        # 5. 构造传输规范
        channel_transfer_spec = {
            "receiver_id": receiver_id,
            "remote_indexes": remote_indexes, # 关键：告诉底层往哪写
        }

        # 6. 执行零拷贝传输 (RDMA Write / TCP Send)
        self.transfer_channel.batched_write(
            objects=mem_objs_to_send,
            transfer_spec=channel_transfer_spec,
        )

2.2.3 通知代理

传输完成后，Sender 通知 Proxy，标志着 Prefill 阶段彻底完成，Decode 实例可以开始工作。

    if transfer_spec.is_last_prefill:
        # 发送 ProxyNotif 消息
        notif_msg = ProxyNotif(req_id=transfer_spec.req_id)
        notif_msg_bytes = msgspec.msgpack.encode(notif_msg)
        self.proxy_side_channel.send(notif_msg_bytes)

2.3 Receiver (Decoder) 实现

Receiver 启动一个后台线程 _mem_alloc_loop 专门处理分配请求。

2.3.1 内存分配循环

def _mem_alloc_loop(self):
    while self.running:
        try:
            # 1. 接收请求 (MsgPack 解码)
            alloc_req_bytes = self.alloc_side_channel.recv()
            alloc_req = msgspec.msgpack.decode(alloc_req_bytes, type=PDMsg)

            # 2. 分配内存并注册到 backend
            alloc_resp = self._allocate_and_put(alloc_req)

            # 3. 返回响应 (包含分配到的内存地址)
            self.alloc_side_channel.send(msgspec.msgpack.encode(alloc_resp))
        except Exception as e:
            logger.error(...)

2.3.2 预先注册与忙等待

_allocate_and_put 方法不仅分配内存，还提前将 MemoryObj 放入 self.data 字典中。如果分配失败（显存不足），会进入忙等待 (Busy Loop)。

def _allocate_and_put(self, alloc_request):
    # ...
    # 尝试分配内存
    mem_obj = self.allocate(...)

    # 如果分配失败，进入忙等待 (重试直到成功)
    # TODO: 这里可能会阻塞 Alloc Loop，导致死锁，未来需要优化
    wait_time = 0.01
    while mem_obj is None:
        time.sleep(wait_time)
        mem_obj = self.allocate(...)

    # 获取物理地址 (供 Sender 写入)
    alloc_indexes.append(mem_obj.meta.address)

    # 注册到本地存储 (Key -> MemoryObj)
    # 注意：此时数据尚未到达，但容器已就绪
    self.put(key, mem_obj)

    return AllocResponse(..., remote_indexes=alloc_indexes)

注意: 这里存在一个隐式的同步假设。Receiver 提前 put 了空的 MemoryObj，如果 Decode 引擎立即读取，可能会读到空数据。因此，全局 Proxy 的调度至关重要——它必须确保 Sender 发送完 ProxyNotif 后，才向 Receiver 发送推理请求。

---

3. 内存分配器实现细节

与 LocalCPUBackend 类似，PDBackend 也实现了 AllocatorBackendInterface，这意味着它不仅仅是 KV Cache 的存储者，更是物理内存的管理者。然而，由于 PD 场景的特殊性（高吞吐、单向流动），其内存管理策略与 L1 CPU 缓存有显著差异。

3.1 Allocator 的核心职责

在 PD 架构中，Allocator 负责管理 数据传输缓冲区。

• 对于 Sender: 分配的内存作为发送缓冲区。数据从计算流生成后，暂存于此，随后被 RDMA/TCP 引擎读取并发送。
• 对于 Receiver: 分配的内存作为接收缓冲区。这块内存的物理地址会被发送给 Sender，以便远程节点直接进行 RDMA Write 操作。

3.2 实现逻辑

PDBackend 复用了 PagedCpuGpuMemoryAllocator，通过指定 allocator_type="gpu" 来直接管理显存（或 Pinned Memory），以支持 GPUDirect RDMA 技术。

def allocate(self, shapes, dtypes, ...):
    """
    PDBackend 的内存分配逻辑
    """
    # 直接调用底层的 PagedCpuGpuMemoryAllocator
    # allocator_type="gpu" 确保分配的是适合 RDMA 传输的显存或 Pinned 内存
    return self.memory_allocator.allocate(
        shapes,
        dtypes,
        fmt,
        allocator_type="gpu"
    )

3.3 与 LocalCPUBackend 的区别

虽然两者都扮演 Allocator 的角色，但设计目标截然不同：

特性	LocalCPUBackend	PDBackend
管理介质	主要管理 Host Memory (CPU RAM)	主要管理 Device Memory (GPU VRAM) 或 Pinned Host Memory
逐出策略 (Eviction)	支持。内存不足时，通过 LRU 策略逐出旧数据 (Allocate-Evict Loop)。	不支持。PD 模式下的数据通常是一次性的流式数据，不涉及长期缓存，因此不进行复杂的逐出。
使用场景	通用 L1 缓存，服务于反复访问的热点数据。	专用的流式传输通道，服务于一次性的 Prefill-to-Decode 传输。

注意: PD 模式下通常不进行 Eviction。如果显存不足，分配请求会直接失败或阻塞，这依赖于上层调度器（Proxy）精确控制并发量，避免“爆显存”。

3.4 关于共存性的说明

Q: 启用 PD 模式后，是否还需要 LocalCPUBackend？

A: 不需要，两者在架构设计上是互斥的。

这不仅仅是功能重叠的问题，更核心的原因在于 KV Cache 在 PD 场景下的角色发生了根本转变：

1. 流式传输 vs 静态缓存:

   • 在 PD 架构中，KV Cache 是连接 Prefill 和 Decode 阶段的 一次性数据流 (Transient Stream)。Sender 生成数据后立即推送，Receiver 接收后立即使用，整个过程不涉及“缓存复用”。

2. 分配器 (Allocator) 的排他性:

   • LMCache 系统要求内存分配器是唯一的。PD 模式下，KV Cache 的生命周期（申请->传输->释放）由 PDBackend 全权管理。LocalCPUBackend 的 LRU 逐出策略并不适用于这种流式场景。

---

4. 总结

PDBackend 作为 LMCache 对分离式推理架构的深度定制，体现了“存储服务于计算”的设计哲学：

1. 消灭首字延迟 (TTFT): 采用 Push-based 主动推送机制，将 KV Cache 的传输隐藏在调度间隙中。相比于 P2P 的“按需拉取 (Pull)”，它彻底消除了首次解码时的网络 I/O 等待。
2. 极致传输效率: 通过 控制面 (ZMQ) 与数据面 (RDMA/TCP) 分离，配合地址预协商机制，实现了 Sender 显存到 Receiver 显存的 零拷贝 (Zero-Copy) 直达，最大化利用物理带宽。
3. 深度调度协同: 突破了传统存储后端的被动边界，通过 ProxyNotif 机制与全局调度器紧密耦合，确保了“数据到位即计算开始”的流水线并行能力。