case : prefill --pp-size P(P>1),decode --pp-size 1。P个 preill pp rank 各自持有模型的 1 / p 层,各自独立地把自己那一段层的 KV 直接 RDMA 写进 同一个 decode rank 的显存。

0. 背景

当前主流模型部署时存在 prefill 侧开 pp/cp,当 pipe 流水稳定时就会频繁出现 incast,在交换机侧会观察到 pfc (Priority-based Flow Control, PFC),如下方黑色线内的红色虚线时刻。

当前 sglang/vllm 内的实现,如下方图所示,此时就会出现 decode 侧 tp rank 0 拥塞,因为同时被 prefill 的 pp0,pp2,pp3的流量打满。

此时,最 naive 的方案就是,直接给pp0,pp2,pp3此时的每个人要发 kv cache 前进行编排,每一轮一发,下一轮 offset 加一,例如:第一轮每个人同轨先发。但是发现没法保证的就是每个 gpu 的进度都是如此完美,故依旧会出现 incast。

综上,还有两张方案。

1. methods

增加 nvlink 的拷贝,让所有的流量变成同轨,因为 nvlink 够快,把转发藏在 rdma 的时间里。

不论最后是采取哪个方案,在 gpu 自己的视角看,任务只包括三类:

  • 我发 cache 给同轨的 decode tp rank i,或收同轨的 prefill rank i(RDMA)
  • 我要等代理发过来的流量 (nvlink)
  • 我要让别人代理的流量(nvlink)

send-side pxn

优势:可以让发送端所有网卡带宽吃满 缺点:从decode 视角看,接收 c我ache 的上限就是网卡带宽总和。

receive-side pxn