ggaaooppeenngg

走向万卡LLM推理集群之路

发表于 2025-05-25 更新于 2025-05-26 Disqus：

在迈向“万卡集群”的过程中，笔者遇到了许多问题。Kubernetes 标称支持最多 5000 个节点，这个上限通常是指在默认资源配置下的理论规模，实际中包含大量 Pod、CRD 或复杂控制器等对象时，其可用规模会显著缩水。笔者尚未在开源版本中进行系统性压测，尚不清楚极限配置下的真实能力边界。

然而，在使用 AWS 提供的标准 EKS 服务过程中，笔者发现整体性能瓶颈远低于预期：在仅 60 万 QPM（约合 10k QPS）的请求强度下，apiserver 的缓存已无法同步，resourceVersion 出现严重滞后，更新操作频繁失败，集群进入了不可用状态。这一结果显然与多数人对 Kubernetes 扩展能力的认知存在差距。

因此，笔者逐渐意识到，支撑“万卡”级别的集群不仅需要强大的计算与网络资源，更需要在控制面、组件实现上投入大量工程精力。如果不进行多集群拆分，或对 apiserver 本身进行改造优化，仅依赖 EKS 标准配置，在千卡规模就已遭遇明显瓶颈。

需要说明的是，笔者所参与开发的边端组件远未达到 Kubelet 的成熟水平，控制器实现的性能和健壮性也存在差距。这也进一步暴露出：Kubernetes 本身虽具备高扩展性，但真正支撑万卡规模并非“开箱即用”，而是极其依赖控制面的设计和组件性能的协同演进。

以 AWS EKS 默认配置为例，假设集群包含 1000 节点，要想稳定支撑每个节点 20 QPS 的访问压力，apiserver 已接近饱和极限。Kubelet 作为核心组件，或许能够维持这一负载水平，但自研 controller 若性能不佳，极易拖慢整个集群的反应速度。

如果你的万卡集群全部使用的是 Kubernetes 原生组件，且运行稳定，那么恭喜你，确实享受到了 Kubernetes 高稳定性和高性能带来的红利。

在笔者所参与的系统中，千节点规模下，每个节点的访问频率必须严格控制在 20 QPS 以下，才能保证系统正常运行。这是基于 AWS 提供的标准 apiserver 能力测得的，没有采用如阿里云那样深度改造的控制面架构。

根据笔者了解，阿里云通过限制客户端访问模式、优化 apiserver 架构，确实能稳定支撑到 5000 节点规模。但再往上，势必需要对 apiserver 进一步“动刀”。

因此，对于想要迈向万卡规模的团队，笔者认为可以从两个方向入手：

多集群架构：通过拆分集群，缓解单个 apiserver 的压力；
组件限流与优化：在标准 Kubernetes 服务中，尽可能减少对 apiserver 的访问频率，控制各组件 QPS，提升控制面稳定性。

无论 Kubernetes 的理论能力多强，如果不控制客户端行为、不优化组件实现，就很难在实际中获得高扩展性的体验。幸运的是，这其中存在大量可以落地的优化空间。

需要重点注意的几点

1. 控制 API Server 的压力

Kubernetes 的 API Server 实际上远没有我们想象中强大。单个节点超过 20 QPS，在 1000 个节点的场景下就会产生 20,000 QPS 的访问量，对 API Server 是非常大的压力，甚至可能导致服务不可用。因此应尽量将单节点的 QPS 控制在 10 以下。

如何控制单节点 QPS

避免不必要的失败请求：例如由于 resourceVersion 滞后导致的失败，建议使用更严格的 backoff 策略。曾遇到因更新操作使用了 get 而非直接读 cache，导致失败请求级联增加，一次失败引发一次同步读和同步更新。
减少无效更新：Etcd 适合“读多写少”模式，每次调谐应尽量只进行一次更新。确保对象确实发生变更后再发起 update 请求，避免“无脑更新”导致无意义的 QPS 激增。当这也是双刃剑，比如在下载任务中，可能下载时间过长没什么正反馈，作用使用者比较难受。
使用高层框架：尽量采用如 Kubebuilder 提供的高层封装，避免直接使用 controller-runtime 这类底层库。很多开发者对 workqueue 和 informer 等机制理解不深，容易写出高 QPS、低质量的控制器逻辑。
速度并非越快越好：与 API Server 打交道时，“稳”往往比“快”更重要。大规模节点下不必追求极致速度，否则反而扛不住。
CRD 设计需注意：在设计 CRD，尤其是 status 字段时，应避免定义频繁更新的字段，防止出现“有事没事就更新 status”的情况。
Informer的DefaultResyncTime可以优先调成1分钟甚至更大，不要使用秒级别的设置。

2. 提升服务稳定性

在大规模节点扩缩容场景下，服务的稳定性至关重要。

节点优雅退出的挑战

对于如 LLM 推理任务，优雅退出困难。例如正在流式输出的对象，可能需要几十秒才能完全释放，需在设计上充分考虑。

错误容忍与重试策略

快速重试优于指数回退：LLM 服务通常对 SLA 要求较高，例如要提供 1 秒内的 TTFT。使用熔断或指数退避策略可能反而破坏 SLA。建议采用快速、少次数（如 3 次）、跨多个对象的重试机制。
容错对象要合理：若当前对象（如某个 vLLM 实例）出现 NCCL 死锁或 runtime error，应立即跳过并选择其他节点，而非对该对象反复重试。
流式处理的异常处理：若输出已经过半再出错，需 proxy 缓存已发送内容并判断是否继续。由于采样是非确定性的，重试结果不一致，建议直接 fail fast，不进行中断续传式的重试。

启动速度优化

模型下载：
- 使用 P2P 技术（如 Dragonfly）从邻近节点拉取模型文件；
- 借助 CDN 提升模型的下载速度。
模型加载：
- 使用内存文件系统效果有限，更直接的做法是预读一遍模型文件至 /dev/null 并用 vmtouch 确保 Page Cache 命中；
- 使用更快的 tensor 格式，比如使用一些优化过的序列化与反序列化库。
预处理加速：
- 如 CUDA Graph 优化：若硬件和模型固定，可预先缓存 CUDA Graph。但考虑到其占整体启动时间比例较小，该优化优先级较低。

3. 服务发现与健康探测

避免依赖 API Server 做服务发现：可以使用配置中心或直接依赖 Kubernetes 的 DNS，如通过 service 的 DNS 解析 pod IP，减少对 API Server 的访问。
健康探针的有效性：
- 简单的端口检查不能有效反映 LLM 的运行状态；
- 推荐使用真实业务接口（如 /v1/chat/completion）做健康探测，以及时发现如 vLLM 死锁等问题。
异常端点的处理：
- 负载均衡策略问题：如使用 LCF（Least Connection First）策略，异常节点会因连接数少反而被更多请求打中。应通过 Redis 等方式缓存失败节点，临时剔除；
- Cache aware 的调度陷阱：缓存命中率高的节点若已宕机，仍可能承接大量请求，也应类似处理。

4. 监控与告警

应关注以下指标：
- SLA：TTFT、TPOT、E2E 时延、吞吐量等；
- Proxy 的调度延迟、负载均衡延迟，这些直接影响 TTFT。

5. PD（Prefill Disaggregation）分离架构

在大规模集群中实现 PD 分离非常具有挑战性，对负载均衡提出更高要求。

即便 PD 单独的计算性质比较良好，若在 P2P 场景下负载均衡不佳，整体系统依然会出现性能瓶颈；
需设计更精细的分发策略，避免热点节点和数据倾斜。

8. 节点资源异构 & NUMA 绑定

问题：大规模集群中，节点配置往往并非完全一致，甚至同一节点上不同 GPU 的性能、带宽、甚至驱动版本都可能略有差异，造成整体推理性能不稳定。
优化建议：
- 在资源调度时考虑 GPU 拓扑结构，避免跨 NUMA 节点调度；
- 使用 cudaVisibleDevices 精准绑定并确保 GPU memory 布局合理；
- 对于 Inference Worker 实现 NUMA-aware 的预绑定机制（如 CPU affinity），可以通过taskset命令绑定CPU。
GPU架构的异同：注意管理Nvidia驱动，Cuda版本，高版本一般对低版本是兼容的，笔者经常碰到的问题就是高版本的Cuda无法跑在低版本的驱动上
优化建议：
- 保证良好的构建系统，对于编译出的vLLM包能够支持不同GPU的com_cap版本，像4090是com_cap8.9的，可能生成的代码无法跑在H20的com_cap9.0上。
- 保证机器上的驱动保持在较新的驱动版本，运维要有比较强的更新节奏和意识。
- 如果可以尽量把Cuda装在容器镜像里面，这样可以和宿主机上的Cuda独立，相对来说可以解决一torch版本和cuda版本的不一致的问题，当然这也会导致镜像体积过大，也是一把双刃剑。

7. 故障演练机制（Chaos Engineering）

问题：在万卡规模下，节点、GPU、网络出现异常是常态。没有预演机制会导致“雪崩式”错误放大。
优化建议：
- 引入如 Chaos Mesh、Litmus 等框架进行 GPU、网络、文件系统等级别的 fault injection；
- 模拟 API Server 限流、模型加载失败、通信超时、session 丢失等场景；
- 在上线前进行 HA 演练，测试自动恢复和降级路径。
- 准备一些fallover的策略，比如保证一个稳定集群，做集群级别的迁移更新，甚至可以用第三方的供应商去做流量迁移，等升级完成之后再从友商切回来。

基于CPU内存的v1 Transfer Connector

发表于 2025-05-17 更新于 2025-05-19 Disqus：

vLLM最近支持了外部加载Transfer Connector，基于LMCache给出的StorageSharedConnector的例子，我尝试实现了一个基于共享内存的Transfer Connector。

v1的接口是一个可以layer wise的实现。

实际上使用下来其实没有明确的区分Producer和Consumer的角色，在P2P的场景下可能比较明显，实际上谁生产kv cache谁消费kv cache其实没有明确规定。

这个的好处是Prefix Cache和KV Cache Transfer没有明确的区别了，Prefill和Worker之间唯一的区别就变成了max_token=1和max_token为真实值的区别了。

设想几个场景，Worker即是生成者，也可以是消费者，Prefill也可以即是生产者又是消费者：

Worker 生成的对话可以存入一个中心化的缓存当中，在多轮对话的时候Prefill可以直接复用这个缓存，只需要计算新的用户对话。
Prefill 基于新的对话可以生成一个缓存，被Worker使用，也可以被其他的Prefill在新的多轮对话中使用。

有一个比较hack的查看调用栈的方法就是在对应的接口函数抛出异常让程序崩溃，就能在stack trace上看到函数调用的路径了。

当然这个接口也可以实现P2P，毕竟他提供了对应的wait接口，无非是等中心化的缓存是否就绪还是Prefill的直接传输是否就绪区别了，这在提供的接口列表当中可以看到。

实现一个Connector需要关注几个接口。

Worker side

Layer Wise

vllm/vllm/attention/layer.py中可以看到。

def unified_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    layer_name: str,
) -> torch.Tensor:
    wait_for_kv_layer_from_connector(layer_name)

    forward_context: ForwardContext = get_forward_context()
    attn_metadata = forward_context.attn_metadata
    if isinstance(attn_metadata, dict):
        attn_metadata = attn_metadata[layer_name]
    self = forward_context.no_compile_layers[layer_name]
    kv_cache = self.kv_cache[forward_context.virtual_engine]
    output = self.impl.forward(self, query, key, value, kv_cache,
                               attn_metadata)

    maybe_save_kv_layer_to_connector(layer_name, kv_cache)
    return output

每一层会有wait_for_kv_layer_from_connector调用connector.wait_for_layer_load。
在计算结束以后会有maybe_save_kv_layer_to_connector调用connector.save_kv_layer。

他们分别对应了decode和prefill，其中wait是同步的，save是异步的。

Model Wise

在vllm/vllm/attention/layer.py中

self.maybe_setup_kv_connector(scheduler_output)

model_output = self.model(
    input_ids=input_ids,
    positions=positions,
    intermediate_tensors=intermediate_tensors,
    inputs_embeds=inputs_embeds,
)

self.maybe_wait_for_kv_save()
finished_sending, finished_recving = (
    self.get_finished_kv_transfers(scheduler_output))

connector.start_load_kv 来自于 maybe_setup_kv_connector，在decoder的model forward之前调用，用于异步启动kv的load。
connector.wait_for_save 来自于 maybe_wait_for_kv_save，在prefill的model forward之后调用，用于整体的save是同步的。

其他的一些接口包括：

get_finished返回对于给定的request ids对应的已经完成的sending和recving的request ids。
register_kv_caches用于connector提前注册kvcaches，在初始化kvcache的时候调用，这个应该是NIXL需要这样直接读取整个的kvcaches的显存地址做RDMA和注册。
bind_connector_metadata，model forward之前bind metadata，这个数据结构是Metadata的数据结构是实现者自由定义的，在start_load_kv之前调用。
clear_connector_metadata，model forward之后clear。

换成prefiller和decoder的视角来看

Prefiller Connector

每一层调用save_kv_layer，这个可以是异步的，在model foward之后会调用wait_for_save保证kvcache被传输完，不然其中的kvcache可能会被之后的forward所覆盖。
clear_connector_metadata可以帮助清理这次forward相关的metadata。

Decoder Connector

bind_connector_metadata帮助设置forward相关的metadata。
每一层调用start_load_kv，这个可以是异步的，在model forward之前调用，在每一层forward之前wait_for_layer_load，这个是同步的。

Scheduler Side

get_num_new_matched_tokens，基于传入的num_computed_tokens获取可以从外部加载的kvcache，这个是给scheduler用的，表明decoder需要加载的tokens。computed_token代表已经计算过kvcache的token.
调度器要额外分配一个external_computed_tokens的slots给外部加载用并且把这部分也算在computed_token，然后在根据budget_token - computed_token分配new_token。
update_state_after_alloc 在scheduler分配slots以后更新connector内部状态，比如用于告知connector是否要加载kvcache。
build_connector_metadata用于构建connector metadata的相关输出，不能修改输入中的schedulerOutput。
request_finished 在request结束，blocks free之前被调用，可以帮助connector触发相关回调。

上面的接口比如register_kv_caches，bind_connector_metada和clear_connector_metadata不一定要实现，可以把他们理解为一些初始化路径，计算路径上的调用hook，我们希望在相关的hook上处理一些东西就实现这些接口。

CPUMemorySharedConnector

我实现了一个基于共享缓存的实现，Prefill只生成缓存，Worker只消费缓存。
主要是通过用layer和tokens hash做key创建SharedMemory。
完整的项目在这里。

LLM Inference Benchmark

发表于 2025-04-04 Disqus：

参考 Benchmarking Text Generation Inference。

参考 SGLang issue 364。

参考 LLM inference server performances comparison llama.cpp / TGI / vLLM。

测试结果

结果来看，在batch size更大的情况下，TTFT会变得特别长，而TBT也会相应的增加一些但没有TTFT恐怖。
batch size变大以后，TTFT从300s变成了900s，而ITL则从0.2s变成了0.3s。
这和MoonCacke的论文是一致的。

测试一下PD分离的效果，使用vLLM的1P1D。
PD分离以后TTFT可以降低一个数量级，这个效果还是很明显的，直接降了一个数量级。

============ Serving Benchmark Result ============
Backend:                                 vllm
Traffic request rate:                    inf
Max reqeuest concurrency:                not set
Successful requests:                     47
Benchmark duration (s):                  127.03
Total input tokens:                      14545
Total generated tokens:                  2993
Total generated tokens (retokenized):    2992
Request throughput (req/s):              0.37
Input token throughput (tok/s):          114.50
Output token throughput (tok/s):         23.56
Total token throughput (tok/s):          138.06
Concurrency:                             24.49
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   66177.90
Median E2E Latency (ms):                 61336.75
---------------Time to First Token----------------
Mean TTFT (ms):                          39888.70
Median TTFT (ms):                        22421.85
P99 TTFT (ms):                           116090.20
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          491.86
Median TPOT (ms):                        394.97
P99 TPOT (ms):                           1917.39
---------------Inter-token Latency----------------
Mean ITL (ms):                           419.69
Median ITL (ms):                         275.52
P99 ITL (ms):                            1766.40
==================================================

双v100 LLAMA3.2:11b

python -m sglang_router.launch_router --worker-urls http://127.0.0.1:8081 http://127.0.0.1:8082

============ Serving Benchmark Result ============
Backend:                                 vllm
Traffic request rate:                    inf
Max reqeuest concurrency:                not set
Successful requests:                     1000
Benchmark duration (s):                  1247.16
Total input tokens:                      289255
Total generated tokens:                  184429
Total generated tokens (retokenized):    184388
Request throughput (req/s):              0.80
Input token throughput (tok/s):          231.93
Output token throughput (tok/s):         147.88
Total token throughput (tok/s):          379.81
Concurrency:                             470.04
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   586218.50
Median E2E Latency (ms):                 596155.97
---------------Time to First Token----------------
Mean TTFT (ms):                          520113.99
Median TTFT (ms):                        526194.47
P99 TTFT (ms):                           1067230.41
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          363.05
Median TPOT (ms):                        356.14
P99 TPOT (ms):                           736.93
---------------Inter-token Latency----------------
Mean ITL (ms):                           360.61
Median ITL (ms):                         273.54
P99 ITL (ms):                            1525.31
==================================================

双卡的并发的情况下，吞吐可以线性增长，但是相较于1P1D来说，prefill的时间没有改善。

笔者参考dynamo尝试实现了一个基于NCCL版本的P2P的xPyD的PD分离。

基于8卡的L40进行了并发100个prompts的测试。每两卡之间是有一个NVLINK其他的卡之间全部是PCIe。

笔者对于kv 传输的group切分如下。

如果是2P4D的话就是这么划分：

如果是4P4D的话。

从测试结果可以看出来单机多卡的PD分离能够降低TBT（TPOT），一个TP的decode就已经超过8TP的decode了，这里主要是因为没有了prefill的干扰。
但是TTFT相对变大了，这个可能是TTFT多了一次传输的时间，具体原因不知道是不是我的实现方式不对，还是因为4TP的prefill就是要慢一些。
这个可能需要一个更综合性的tuning。

4P4D(4TP+4TP) V0调度器

============ Serving Benchmark Result ============    
Backend:                                 vllm        
Traffic request rate:                    inf         
Max reqeuest concurrency:                not set     
Successful requests:                     100         
Benchmark duration (s):                  37.95       
Total input tokens:                      34965       
Total generated tokens:                  20654       
Total generated tokens (retokenized):    20654       
Request throughput (req/s):              2.64        
Input token throughput (tok/s):          921.42      
Output token throughput (tok/s):         544.29      
Total token throughput (tok/s):          1465.71     
Concurrency:                             46.49       
----------------End-to-End Latency----------------    
Mean E2E Latency (ms):                   17642.84    
Median E2E Latency (ms):                 14298.02    
---------------Time to First Token----------------    
Mean TTFT (ms):                          8147.33     
Median TTFT (ms):                        8393.70     
P99 TTFT (ms):                           8834.48     
-----Time per Output Token (excl. 1st token)------    
Mean TPOT (ms):                          58.38       
Median TPOT (ms):                        50.79       
P99 TPOT (ms):                           162.16      
---------------Inter-token Latency----------------    
Mean ITL (ms):                           46.27       
Median ITL (ms):                         42.21       
P99 ITL (ms):                            56.21       
==================================================

8TP V0调度器

============ Serving Benchmark Result ============    
Backend:                                 vllm        
Traffic request rate:                    inf         
Max reqeuest concurrency:                not set     
Successful requests:                     100         
Benchmark duration (s):                  27.86       
Total input tokens:                      29552       
Total generated tokens:                  24879       
Total generated tokens (retokenized):    24875       
Request throughput (req/s):              3.59        
Input token throughput (tok/s):          1060.84     
Output token throughput (tok/s):         893.10      
Total token throughput (tok/s):          1953.94     
Concurrency:                             54.24       
----------------End-to-End Latency----------------    
Mean E2E Latency (ms):                   15109.72    
Median E2E Latency (ms):                 15397.55    
---------------Time to First Token----------------    
Mean TTFT (ms):                          5398.04     
Median TTFT (ms):                        6015.93     
P99 TTFT (ms):                           7251.63     
-----Time per Output Token (excl. 1st token)------    
Mean TPOT (ms):                          71.76       
Median TPOT (ms):                        42.68       
P99 TPOT (ms):                           307.25      
---------------Inter-token Latency----------------    
Mean ITL (ms):                           39.23       
Median ITL (ms):                         32.54       
P99 ITL (ms):                            43.65       
==================================================

8TP V1 调度器

============ Serving Benchmark Result ============    
Backend:                                 vllm        
Traffic request rate:                    inf         
Max reqeuest concurrency:                not set     
Successful requests:                     100         
Benchmark duration (s):                  24.07       
Total input tokens:                      21404       
Total generated tokens:                  20379       
Total generated tokens (retokenized):    20377       
Request throughput (req/s):              4.15        
Input token throughput (tok/s):          889.06      
Output token throughput (tok/s):         846.49      
Total token throughput (tok/s):          1735.55     
Concurrency:                             41.30       
----------------End-to-End Latency----------------    
Mean E2E Latency (ms):                   9943.29     
Median E2E Latency (ms):                 9369.62     
---------------Time to First Token----------------    
Mean TTFT (ms):                          2798.48     
Median TTFT (ms):                        2731.82     
P99 TTFT (ms):                           4323.56     
-----Time per Output Token (excl. 1st token)------    
Mean TPOT (ms):                          55.11       
Median TPOT (ms):                        39.79       
P99 TPOT (ms):                           307.50      
---------------Inter-token Latency----------------    
Mean ITL (ms):                           36.45       
Median ITL (ms):                         31.66       
P99 ITL (ms):                            341.00      
==================================================

多机器配置

DeepSeek R1 8xH20 x2 台机器，每台机器RDMA配置16个 MT2910 Family [ConnectX-7] 做8个bond。

8TP x 2PP 的部署方案，如果后面EP支持的话可能会有更好的效果。

============ Serving Benchmark Result ============
Backend:                                 vllm
Traffic request rate:                    inf
Max reqeuest concurrency:                not set
Successful requests:                     1000
Benchmark duration (s):                  234.47
Total input tokens:                      303481
Total generated tokens:                  187870
Total generated tokens (retokenized):    186116
Request throughput (req/s):              4.26
Input token throughput (tok/s):          1294.33
Output token throughput (tok/s):         801.26
Total token throughput (tok/s):          2095.59
Concurrency:                             363.04
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   85122.29
Median E2E Latency (ms):                 82826.18
---------------Time to First Token----------------
Mean TTFT (ms):                          31789.26
Median TTFT (ms):                        17669.77
P99 TTFT (ms):                           100110.92
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          770.73
Median TPOT (ms):                        341.77
P99 TPOT (ms):                           9445.55
---------------Inter-token Latency----------------
Mean ITL (ms):                           284.74
Median ITL (ms):                         214.68
P99 ITL (ms):                            745.14
==================================================

sglang tp 16的配置，sglang不支持pp，sglang明显要快一些，主要原因应该是sglang支持了MTP，vLLM目前还没有。

============ Serving Benchmark Result ============
Backend:                                 sglang
Traffic request rate:                    inf
Max reqeuest concurrency:                not set
Successful requests:                     1000
Benchmark duration (s):                  190.92
Total input tokens:                      306113
Total generated tokens:                  197108
Total generated tokens (retokenized):    195033
Request throughput (req/s):              5.24
Input token throughput (tok/s):          1603.38
Output token throughput (tok/s):         1032.43
Total token throughput (tok/s):          2635.81
Concurrency:                             488.50
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   93263.23
Median E2E Latency (ms):                 86230.17
---------------Time to First Token----------------
Mean TTFT (ms):                          39722.57
Median TTFT (ms):                        43590.80
P99 TTFT (ms):                           60010.86
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          1529.43
Median TPOT (ms):                        270.69
P99 TPOT (ms):                           37619.47
---------------Inter-token Latency----------------
Mean ITL (ms):                           276.88
Median ITL (ms):                         158.45
P99 ITL (ms):                            945.60
==================================================

Dynamo的xPyD

发表于 2025-03-31 更新于 2025-05-19 Disqus：

dynamo对于vLLM的给懂是一个patch文件，vLLM给了一个PR方便在线对比。

PD分离的背景介绍

在大模型推理中，Prefill 和 Decode 的计算特性存在显著差异：

Prefill：计算密集型，计算比例更高。
Decode：访存密集型，访存比例更高。

Decode 的计算依赖于 Prefill 生成的 KV Cache。在没有 PD 分离的情况下，较长的 Prefill 通常会优先占用计算资源，导致长 Prompt 的 Prefill 时间过长，从而增加 Decode 的延迟。

Chunked Prefill 的局限性

为了解决上述问题，Chunked Prefill 是一种常见的优化方案，但它也存在以下挑战：

适合超长 Context Length：Chunked Prefill 能有效降低中间显存的占用，但仅适用于超长上下文场景。
大 Chunked Prefill 的影响：当大 Chunk Prefill 和 Decode 同时出现在一个 Batch 中时，会显著拖慢 Decode 的速度。
小 Chunked Prefill 的退化：小 Chunk Prefill 容易退化为 Memory Bound，导致计算单元的利用率（MFU）下降。

因此，PD 分离的优势在于可以分别优化 Prefill 和 Decode 的性能。

Dynamo 的 PD 分离策略

Dynamo 使用了一种条件 PD 分离策略：

仅在满足特定条件时，Prefill 才会远程计算；否则，仍然在本地进行 PD 混合计算。

PD 分离的实现主要包括两部分：

模型的切分与传输：通过合理的切分策略和高效的传输机制实现计算资源的分离。
高效的异步传输或存储引擎：这是性能优化的关键，尤其是 KV Cache 的传输或存储。

KV Cache 传输的挑战

以 A100 为例，在 8B LLAMA 模型下，Prefill 的计算速度可达 1 万 tokens/s，这会产生约 3GB 的 KV Cache 数据，给传输带宽带来极大压力。
对于计算速度更快的 H100，这种传输需求会进一步增加，对带宽提出了更高的要求。

xPyD 的主要设计概要

以下是 xPyD 的主要设计概要，不同框架可能在 PD 负载均衡方式、KV 传输和存储上有所不同。

KV Cache 的切分

TP 条件下的切分

在 Tensor Parallel (TP) 条件下，KV Cache 按照 head 进行切分。例如，对于 Qwen 小模型，其 KV head 数为 8，Q head 数为 40，hidden size 为 5120。
一个 token 的大小计算如下：
8 * 5120 / 40 = 1024

如果 P/D TP比例为 2，则 P 会沿着 head 切分为两部分。传输引擎会将切分后的 tensor 发送给每个 D。
因此，一个 token 在 D 上的大小为：
8 / 2 * 5120 / 40 = 512

需要注意，上述计算未包含数据宽度。如果数据类型为 FP8，则宽度为 1；如果为 BF16，则宽度为 2。

PP 条件下的切分

在 Pipeline Parallel (PP) 条件下，切分相对简单，直接沿着层进行切分即可。例如，如果 P/D PP比例为 2，且模型有 64 层，则：

层 0-31 分配给 D0
层 32-63 分配给 D1

这种切分方式无需对 tensor 进行额外处理。

DP 和 EP 条件下的切分

EP：EP 主要用于 FFN，与注意力机制无关，因此无需考虑 EP 条件下的切分。
DP：DP 条件下无需切分。如果 P/D 比例为 2，直接将 P 的副本同时发送给两个 D 即可。

传输方式

P2P 传输

P2P 传输采用点对点方式：

P 向 D 建立 RDMA 连接，并申请 RDMA 的 VRAM。
P 直接将数据发送给 D。

KV Cache Store

KV Cache Store 属于 Pooling 模式，支持以下功能：

中间存储：P 将数据存储到 KV Store，D 从 KV Store 中领取数据。
缓存优化：Store 也可以基于 P2P 实现，支持显存缓存、内存或 SSD 上的 KV Cache Swap。
Prefix Cache 共享：通过中间存储，可以实现跨请求的 Prefix Cache 共享。

PD 顺序

Dynamo 的策略：Dynamo 采用先 Decode 后 Prefill 的条件PD策略。请求首先发送到 Decode 实例。如果是短的 Prefill，Decode 实例会直接计算，无需触发远端 Prefill；如果需要远端 Prefill，则会触发远端计算。
其他框架的策略：大多数框架采用先 Prefill 后 Decode 的策略。请求先由 Prefill 计算出第一个bonus token，然后转交给 Decode 继续计算。

Dynamo 实现分析

本文仅分析 Dynamo 对 vLLM 本身的一些改动，不涉及 Dynamo 在上层的工作，例如全局基于消息队列的 PrefillQueue 是在上层实现的。
vLLM 在被请求时会告知自己是否是 Prefill 请求，或者是否需要远程 Prefill 的 Decode 请求，这一层逻辑由 Dynamo 在上层完成。

Dynamo 基于 vLLM V0 的调度器实现。V0 的调度器主要将 Prefill 和 Decode 明确分开，而 V1 的调度器则考虑了 Chunked Prefill，不再在调度器内部区分 Prefill 和 Decode 两种 sequence。

vLLM V0 调度器回顾

在 vLLM V0 的实现中，engine 的 step 方法会调用调度器，调度器负责给出需要执行的 sequence group request，然后调用模型执行器（model executor）执行这些请求。执行完成后，调度器会处理结果并更新被调度请求的状态。

LLMEngine

LLMEngine 是 vLLM 的核心组件之一，负责协调调度器和模型执行器的工作。Dynamo 在此基础上进行了扩展，以支持 PD 分离的功能。

增加_finished_transfers和_finished_prefills用于保存prefill的传输结束的request和decode接收传输的request，这两个变量会传给调度器。

remote_prefill_requests保存在remote prefill中的requests。

调度结束以后会拆分出running request和remote prefill requests。

对于running request逻辑不变，但是remote prefill requests会在model execution执行前先发出去。方法是给seq的remote_prefill_request_callback添加remote_prefill_request。
这个callback是dynamo上层的worker设置的，他对应的是向全局的消息队列PrefillQueue发送Prefill消息。

通过比较computed_block_nums是否等于block_table标记完成并且放入到本地的_finished_prefills当中。
这是decode视角：把调度器调度的prefill的request发送出去不在本地计算。

到prefill视角来看，会在memory_transfer_reqs中加入已经计算好的computed block和requestid等信息构造的MemoryTransferRequest。

excute_model_req会增加一个需要传输的requests。

execute_model_req.memory_transfer_requests = memory_transfer_reqs

然后开始执行model_excutor的excute_model。

根据执行结果返回的request_notif_counter和request_done_counter更新对应的_finished_prefills和_finished_transfers。

上层需要初始化LLMEngine的NIXL Agent。

Schedule

调度器的改动其实比较简单，对于prefill角度，prefill结束的释放掉request，decode角度把remote prefill结束的变成running走原本的decode流程。

除了running之外额外增加了一个remote_prefilling的queue用于管理在远端prefill的请求，他和running queue的关系是有相似性的，
比如判断有无未完成的请求时会同时看running和remote_prefilling，但他们的区别在于是不是在本地running。
也增加了prefill_sending用户标记正在传输的prefill。

调度主体会接收finished_prefills和finished_transfers用于D标记的远端prefill结束（已经传到了本地）和P标记已经完成传输的requests。

当remote_prefilling的中的request在finished_prefills中时代表prefill结束，会把状态设置为running并且开始decode调度。
当prefill_sending的中的request在finished_transfers中时代表prefill传输结束，会free掉这个request。

prefill_sending和finished_transfers是一对，是对于prefill instance来说的。
remote_prefilling和finished_prefills是一对，是对于decode instance来说的。

seq_group中会添加is_remote_decode这个用于标记这个请求只在自己这里prefill，decode要在decode instance上做，这个标记是上层的worker设置的，不在vLLM层。

每个sequence group会添加一个标记，seq_group.remote_prefill_params.is_remote_prefill，标记了就加入到remote_prefilling队列中
不然就走老的流程。这也是上层决定。

EventManager

只和 router 负载均衡有关。

worker 上有 KVPublisher 负责发送 kvcache 的创建和删除事件，同时 KvMetricsPublisher 用于发送监控指标（如队列排队长度等）。
router 上则包含 KVIndexer，用于收集 kvcache 的事件并建立前缀树，同时 KvMetricsAggregator 用于收集监控指标。

路由策略基于 KV match rate - Load 的最大值，旨在平衡负载与 kvcache 的匹配度。

PrefixCachingBlockAllocator加入了event_manager。KVCacheEventManager实现就不细说了，就是一个事件收发器。

在evict block的时候，发送删除事件event_manager.enqueue_removed_event(content_hash_to_evict)
当block被填满变成 immutable block 的时候，发送分配事件event_manager.enqueue_stored_event(block.prev_block, block)。

NIXL Transfer

数据传输的切分

假设：

PS(prefill parallel size)=1
DS(decode parallel size)=1
PTP(prefill tensor parallel size)=2
DTP(decode tensor parallel size)=4

存在两个 kvgroup：P0,D0,D1 和 P1,D2,D3。

计算可得：
TPM(tensor parallel multiplier) = DTP/PTP = 2

由此可以推导出以下关系：

kv_rank：表示每个 P 或者 D 实例的序号，在本示例中分别为 0 和 1。
p_kv_group_rank = kv_rank
d_kv_group_rank = PS + (kv_rank - PS) * TPM + rank % TPM
kv_world_size = PS + DS * TPM
p_kv_global_rank = kv_rank * PTP + rank
d_kv_global_rank = PS * PTP + (kv_rank - PS) * DTP + rank

关于集合通信的端口号的 base：

p_port_offset_base = 2 * rank + 1
d_port_offset_base = 2 * rank//TPM + 1

以下是具体的参数值表格：

role	rank	kv_rank	kv_group_rank	kv_world_size	kv_global_rank	port_offset_base
P	0	0	0	3	0	1
P	1	0	0	3	1	3
D	0	1	1	3	2	1
D	1	1	2	3	3	1
D	2	1	1	3	4	3
D	3	1	2	3	5	3

初始化 rank 0 收集all_gather所有的parallel_config

每有一个kv_role = kv_producer则kv_producers_parallel_size就+1。

第一步当然要支持xPyD的配置，比如下面的配置。

P的并行规模

1	kv_producers_parallel_size: Optional[int] = None

P的TP

1	kv_producers_tensor_parallel_size: Optional[int] = None

P的PP

1	kv_producers_pipeline_parallel_size: Optional[int] = None

D的TP

1	kv_consumers_tensor_parallel_size: Optional[int] = None

D的PP

1	kv_consumers_pipeline_parallel_size: Optional[int] = None

属性函数
其中tensor_parallel_multiplier是 p的tp // d的tp

然后总结一下各种pp,tp,rank的关系。

TP的倍率由D的TP地板除P的TP：tensor_parallel_multiplier = self.kv_consumers_tensor_parallel_size // self.kv_producers_tensor_parallel_size
D的并行规模是整体并行规模减去P的并行规模： kv_consumers_parallel_size = self.kv_parallel_size - self.kv_producers_parallel_size
kv_world_size = self.kv_producers_parallel_size + self.kv_consumers_parallel_size * self.tensor_parallel_multiplier

发送的入口是send_kv_caches_and_hidden_states

这个函数的主要工作是根据TP的大小，从Prefill worker上切出Decode worker上需要的tensor
并且发送给Decode worker。

在这个函数中根据自己的rank计算对应的D的rank。
PP的话比较容易，直接按层的range就行了，TP的话需要在给定层的range下做TP的切分。从代码来看dynamo只支持PP=1。

笔者写了一个简易版的示例带代码方便查看对应的shape

import torch
# kv_cache 是一个包含张量的列表
# head 数量是 8, hidden_state per head 也是8
# key_cache 是一个连续的空间，有slot_mapping做索引
key_cache = torch.randn(10, 8, 8)
print("原始 key_cache shape")
print(key_cache.shape)
# output: torch.Size([10, 8, 8])
# 这次request使用的是第0和5条的key cache
current_slot_mapping = [0,5]
# prefill tp=2
p_tp = 2
# decode tp=4
d_tp = 4
tp_multiplier = d_tp // p_tp
# 考虑decode worker 的 rank 0 的情况
target_rank = 0
# num_heads_per_rank = 1 也就是每个decode rank分到一个head
num_heads_per_rank = 8 // p_tp // d_tp
# 计算head的range
head_start = target_rank * num_heads_per_rank
head_end = head_start + num_heads_per_rank
# 按 p_tp reshape
key_cache = key_cache.reshape(-1,8//p_tp,8)
print("按 prefill tp 切分")
print(key_cache.shape)
# output: torch.Size([20, 4, 8])
print("decode 选择器（在第0维度）", current_slot_mapping, ",head range（在第1维度）", str(head_start)+":"+str(head_end))
# output: decode 选择器（在第0维度） [0, 5] ,head range（在第1维度） 0:1
d_key_cache = key_cache[current_slot_mapping, head_start:head_end]
print("获取 d key cache")
print(d_key_cache.shape)
# output: torch.Size([2, 1, 8])

相对应的接收函数的入口recv_kv_caches_and_hidden_states 没啥特殊处理，直接已经切好，收到以后直接cache住就行。