nanoPD:一个 LLM P/D 分离推理引擎的实现笔记

这个项目的起因是读vllm, DistServe 和 Mooncake 的时候有些地方没有完全想清楚，觉得与其在论文层面反复打转，不如自己动手实现一遍，于是花了一段时间和claude大人一起从零写了一个支持 Prefill/Decode 分离调度的推理引擎，覆盖了从 CUDA 内核到自适应路由器的完整栈，代码大约 2000 行 Python 加 400行 CUDA C++，每个模块有单独的文档，也顺手做了中英双语版。(菜菜勿喷呜呜）

背景和动机#

LLM 推理的两个阶段在计算特性上的差异非常显著，prefill 是一次性处理整个 prompt 的过程，本质上是大规模的GEMM，算力密集，GPU 的计算单元在这个阶段是满载的；而 decode 每步只生成一个token，每次前向传播只有一个（或很少几个）新 token 需要计算，但需要从显存里读取所有历史 token 的 KV cache，是典型的memory bound操作，GPU 的计算单元大部分时间在等数据，实际上 decode 阶段的算术强度非常低，主要的开销是 HBM 的读取带宽。

这两种操作对 GPU 资源的需求模式截然不同，把它们放在同一张卡上并发运行时会产生 SM 资源的竞争，prefill 的矩阵乘法会占用大量 SM，导致同时在跑的 decode 请求的 attention 计算被推迟，表现出来就是 decode的延迟在有 prefill 并发时会显著上升，这种干扰在高并发场景下尤其明显，也是 vLLM 等系统在负载较高时尾延迟劣化的一个重要原因。

P/D 分离（Disaggregated Serving）的思路是把 prefill 和 decode 分配到专用的 GPU 上，让两个阶段互不干扰，这个方向在工业界和学术界都有比较多的工作，DistServe 比较系统地分析了分离的收益，Mooncake 则是 月之暗面在生产系统中实践分离调度的工程经验，两篇文章读起来都很有意思，但读完之后我对一些具体的设计决策仍然有疑问，比如路由策略在不同硬件上的表现差异有多大，代价模型里的参数对结论的敏感性如何，这些问题通过实现一遍会有更直接的感受。

实现栈，从底到上#

CUDA 内核部分手写了 paged attention kernel 和 KV store ops，主要是想理解非连续内存块上的 attention 是怎么做的，传统的 attention 假设KV cache 存储在连续内存里，但 paged KV cache 把显存切成固定大小的物理块，序列的 KV 数据分散在这些块中，attention 计算需要根据 blocktable 做间接寻址，实现上需要在 CUDA kernel 里根据 token 的位置先找到对应的物理块再读取数据，这个 gather 操作相比连续 KV cache 有额外的开销，但换来的是显存利用率的显著提升，因为不再需要为每条序列预先分配最大长度的连续显存，这个 tradeoff 在长序列和高并发场景下非常值得。内核的实现参考了 vLLM 的设计，但为了保持简单没有做 Flash Attention 那样的 tiling 优化，所以在长序列上性能差很多，这部分如果要真正优化的话工作量还是比较大的。（next work预定了说是）

块管理器以 block 为粒度管理显存的分配和释放，逻辑上类似操作系统的虚拟内存管理，每个物理块有引用计数，引用计数降为零时才真正释放，支持 Copy-on-Write fork，这个特性在 beam search 或者需要复制序列状态的场景下有用，fork 的时候共享物理块，只有在某条路径需要写入新token 时才触发实际的物理块复制，避免了不必要的显存拷贝。

推理引擎实现了 chunked prefill，长 prompt 会被切成固定大小的 chunk和当前正在 decode 的请求交错执行，而不是一次性把整个 prompt 打进去阻塞所有 decode 请求，这个设计的好处是降低了 prefill 对 decode 的干扰，代价是单个请求的 prefill 总时间会因为被切分而略有增加，调度器维护 waiting、prefilling、running 三个队列，每一步决定哪些请求进入 prefill、哪些进行 decode、chunk 大小是多少，这些调度决策对系统的整体延迟和吞吐影响很大，实现上采用了比较简单的启发式策略，没有做复杂的动态调整。

Worker 层分三类，

• CollocatedWorker 在单卡上同时做 prefill 和 decode，内部复用了推理引擎的调度器；
• PrefillWorker 专门处理 prefill，完成后把生成的 KV cache 从 GPU 显存提取到 pinned memory buffer（为什么我的服务器连PCIe都没有！！！），准备传输；
• DecodeWorker 接收传输过来的 KV cache，加载到自己的显存后把对应的序列加入 decode 队列。 三类 Worker 可以在不同 GPU 上并发运行，由上层的 CentralScheduler 协调，CentralScheduler 把 collocated 和disaggregated 两条 pipeline 跑在独立的线程上，每个 step 并发执行，然后汇总结果。

KV 传输部分用了独立的 transfer stream，目标是和 compute stream 尽量overlap，减少等待时间，代码里用 torch.cuda.Event 来同步两个 stream之间的依赖，实现上比较简单，没有做精细的 pipeline 重叠，实际测下来overlap 的效果取决于传输数据量和计算量的比例，在小 batch 或短序列的时候效果不太明显。同时会自动检测当前硬件是否支持 P2P 直传，通过torch 的 _check_p2p 接口查询，不支持的话 fallback 到 pinned memory relay，也就是先把数据从 GPU 拷到 CPU 的 pinned memory，再从 pinned memory 拷到目标 GPU，这个路径的带宽会受限于 PCIe，在没有 NVLink 的多卡环境下（比如 8 张 4090 通过 PCIe 互联）这个开销是比较显著的，实测大约 12.9 GB/s，而有 NVLink 的 H20 可以达到约 392 GB/s，差了将近 30 倍，这个差距直接决定了路由判断在两种硬件上的结论会有多大的不同。

代价模型和路由#

为了决定每个请求走合并路径还是分离路径，我实现了一个解析式代价模型，思路是先在真实设备上跑 micro-benchmark 测四个参数，prefill 速度α（ms/token）、decode 单步延迟 β（ms）、prefill 对 decode 的干扰系数 γ（ms/token）、以及 GPU 间传输带宽 bw（GB/s），然后用这四个参数建立两条路径的延迟估算公式，每个请求来了之后实时计算预估延迟并选更小的那条，所有参数都来自实测，没有用论文里的理论值，因为不同的硬件和软件环境下这些数字差异很大，直接用实测值比从理论推导更可靠。

profiling 的过程大概是：prefill latency 用不同长度的 prompt 各跑若干次取中位数，用线性回归拟合得到 α；decode latency 在不同的 KV 长度和 batch size 下各测一遍，用来拟合 β 和 batch_thresh（decode 从内存带宽瓶颈切换到算力瓶颈的拐点）；interference 通过对比纯 decode和混合 prefill+decode 下 decode 延迟的差值来测量，用线性回归拟合得到 γ；P2P 带宽直接测一次大块数据传输的时间来估算。

路由判断最后可以化简成比较两个都正比于 prompt 长度 L 的量，分离路径相比合并路径的额外代价是传输 KV cache 的时间，大约是 transfer_rate× L，合并路径相比分离路径的额外代价是 prefill 对并发 decode 的干扰，大约是 γ × L × (system_load / batch_thresh)，分离更划算的条件可以化简为 γ / transfer_rate > batch_thresh / system_load，其中γ / transfer_rate 是一个只取决于硬件的比值，反映了”干扰有多贵”相对于”传输有多贵”的比例，而 batch_thresh / system_load 则反映了当前系统的负载程度。

在 RTX 4090 上实测，γ / transfer_rate ≈ 7.6，代入 batch_thresh = 16可以得到 system_load ≥ 约 2.1 时分离就开始划算，也就是只要有两到三个并发请求在跑，分离路径在延迟上就已经比合并路径更好了；而在H20 上这个比值达到 346，几乎在任何非零负载下分离都是更优的选择，两种硬件上的结论差距这么大主要是因为传输带宽差了将近 30 倍，γ 的差异相对没那么大（0.087 vs 0.130 ms/token），所以比值主要是被transfer_rate 拉开的。这个结论确实感觉很不合理，直接原因就导致路由决策非常单一，理论上能充分展示路由决策的机器本人无钱无时间找到（谁来帮我考算分期中！），于是懒得改了。

此外还实现了一个输出长度预测器，因为路由需要估算 decode 阶段的代价，而 decode 代价正比于输出长度，但输出长度在请求来的时候是未知的，用了一个在线的贝叶斯预测器，按 prompt 长度分桶，每桶内维护历史输出长度的统计，新请求来了用对应桶的均值作为预测，桶内样本不足时fallback 到全局均值，这部分比较粗糙，输出长度预测本身就是一个很难的问题，实际上即使是 vLLM 这样的成熟系统目前也还没有特别好的解法，这里的实现只是为了让路由可以跑起来，准确性有限。

测试和结果#

写了一个 Poisson 到达过程的benchmark，模拟真实服务里请求按泊松过程到达的场景，固定到达率 λ，跑固定时长，然后统计完成的请求数、端到端延迟的分布、以及 drop 的请求数，这比简单的串行测试更接近实际的服务场景，因为串行测试本质上是测单个请求的延迟，没有并发，没有排队，和真实负载差距比较大。

结果上，在 RTX 4090 × 8 的环境下，adaptive 路由在中等到达率下的吞吐和延迟比 collocated 有一定改善，在 H20 上分离路径的优势更明显，和代价模型的预测基本一致，但实际的性能数字和成熟框架比是没有可比性的，因为缺少了 CUDA Graph、Flash Attention、算子融合等优化，这里就不列具体数字了，主要是看各策略之间的相对趋势是否符合理论预期，从这个角度来看结果是比较合理的。 性能出现巨大问题的原因还有disagg路径没有像collocated做那么多小优化，such as cotinious batching等等，在项目的文档中我做了具体说明。

另外一个有意思的观察是，在 H20 上 adaptive 的峰值吞吐反而低于RTX 4090，原因是 H20 的 decode 步更快（β = 33ms vs 51ms），更多请求在 collocated GPU 上就快速完成了，disaggregated 路径的利用率相对 更低，多卡的优势没有完全发挥出来，这有点反直觉，但想清楚之后也是合理的，更快的 decode 反而减少了分离的必要性，加之我较为naive的paged attention实现（肉眼可见有一处就可以改为reduce求和），使得max_request_num调大了就会影响TTFT等乱七八糟的事，本人又买不起autodl的H20 * 4呜呜呜，就这样吧。

一些感受#

实现上没有做 CUDA Graph、Flash Attention 或量化这些会显著影响性能的优化（实际上懒了写不了了），主要是想保持每一层的逻辑足够清晰，方便对照论文理解设计决策，也方便文档解释，所以性能上和成熟的推理框架没有可比性，仅仅是作为理解系统设计的实践项目。

写完之后最大的收获可能不是哪个具体的技术细节，而是对整个系统各层之间的依赖关系有了更清晰的感受，代价模型的准确性依赖 micro-benchmark 的质量，micro-benchmark 又依赖引擎本身的实现是否足够接近真实推理路径，路由的效果依赖输出长度预测的准确性，调度器的策略影响 profiling 时测出来的 interference 数值，这些依赖关系在读论文的时候是模糊的，实现一遍之后会更具体地感受到每一层的假设在什么条件下成立，什么条件下会失效，以及各层之间的耦合程度，这种感受很难单纯通过读代码或读论文获得。

已经放到了github上，文档里对每层的设计决策有比较详细的说明。