ICML'26 | AMPD:PD 分离遇上多轮推理,incremental prefill 到底该在哪跑

ICML’26 | AMPD:PD 分离遇上多轮推理,incremental prefill 到底该在哪跑

原文:Efficient Multi-round LLM Inference over Disaggregated Serving(He et al., ICML 2026)


1. 从单轮问答到多轮 agent,推理负载换了副面孔

想象你在用一个 agent 帮你查资料。你抛给它一个问题,它先在心里盘算一下(生成一段推理文本),决定去调一个搜索工具,拿到结果以后再接着想、再调下一个工具……这样来回好几轮,最后才给你一个答案。从你的视角看,这只是一次对话;但在底下的 GPU 看来,这一次请求被切成了好几段「算一会儿、停下来等外部结果、再接着算」的过程。

要讲清楚为什么这件事对推理系统是个麻烦,得先回到 LLM 推理最基本的两个阶段:prefill 和 decode。

prefill 阶段把你输入的整段 prompt 一次性喂进模型,并行地把每一层的 Key-Value 向量都算出来缓存下来——这份缓存就是 KV cache,后面每生成一个新 token,都要回头看它才能算注意力。因为是把 prompt 里所有 token 一股脑过一遍网络,这一步算力拉满,属于 compute-bound(算力瓶颈)。decode 阶段则相反,它逐个 token 往外蹦,每次只新算一个 token,但每次都要把那份越来越长的 KV cache 从显存里搬出来参与计算,真正卡脖子的是显存带宽,属于 memory-bound(访存瓶颈)。

两种负载的脾气差这么大,如果把它们塞在同一张卡上轮流跑,就会互相抢资源、互相拖后腿,这个现象业界叫 PD interference(prefill-decode 干扰)。为了躲开它,近两年的主流做法是 PD disaggregation(PD 分离):干脆把 prefill 和 decode 拆到不同的 GPU 上,各跑各的,典型代表是 DistServe、Splitwise、Mooncake 这些工作。拆开以后,衡量两个阶段快慢的指标也就自然分开了——prefill 这端看 TTFT(Time-to-First-Token),也就是从请求进来到吐出第一个 token 的延迟,它主要由 prefill 决定,越低越好;decode 这端看 ITL(Inter-Token Latency),也就是生成相邻两个 token 之间的间隔,由每一个 decode 步决定,同样越低越好。

而真正用来给系统打分的,是把这两个延迟指标合起来看的 SLO(Service Level Objective,服务质量目标):给 TTFT 和 ITL 各设一个阈值,一个请求只有当两个指标都压在各自阈值以内,才算这次服务「达标」。把达标请求的占比统计出来,就是 SLO attainment(SLO 达成率),越高说明系统在延迟约束下能服务好的请求越多。这是本文从头到尾的主指标,后面所有的「提升了多少」说的都是它。

PD 分离这套东西,是在单轮推理的设定下被打磨出来的:一个请求进来,prefill 一次,decode 一路到底,结束。可多轮 workflow 把这个干净的假设打破了。下面这张图把单轮和多轮的区别画在了同一条时间轴上,看的时候重点盯住上面那条代表 GPU 执行的色块:橙色是 prefill,蓝色是 decode。

多轮 LLM 推理的执行过程

左边起手是常规的一次 prefill(读 prompt)接一段 decode(生成)。但生成到一半,模型决定调用 Tool 1,于是停下来等工具返回结果;结果回来之后,这段「工具返回的内容」又得作为新的输入被模型读进去——注意,这又是一次 prefill,只不过它发生在 decode 的中途,论文把它叫做 incremental prefill(增量 prefill),与请求刚进来时那次 initial prefill(初始 prefill) 相区别。读完工具结果再接着 decode,又触发 Tool 2,再来一轮 incremental prefill……如此反复。于是从系统的视角看,多轮推理呈现出一种 prefill 和 decode 交织(interleaved) 的负载模式:prefill 不再是开头一锤子买卖,而是零零散散地插在整个 decode 过程里。

这就是矛盾的根源。PD 分离把 prefill 和 decode 拆到两套资源上,前提是「prefill 集中在前、decode 独占在后」;可一旦 prefill 任务开始见缝插针地冒出来,这套分工该怎么继续,就成了一个没人答好的问题。AMPD(Adaptive Multi-round workflows with PD disaggregation)就是冲着这个问题来的——它不换 PD 分离这个大框架,而是在这个框架里重新设计「增量 prefill 该放哪跑、队列该怎么排、两类资源该怎么分」。

2. 把 PD 分离套到多轮推理上,现有系统漏了两个洞

多轮推理加速和 PD 分离,各自都有不少工作,但把两者捏到一起几乎是空白。顺着上一节的交织负载往下想,现有系统在这个场景里有两处明显接不住。

第一个洞:incremental prefill 没有一套自适应的调度。 为了躲开 PD 干扰,直觉上所有 prefill 任务(不管初始还是增量)都该甩给 prefill 实例去跑。但增量 prefill 有个尴尬之处:它要用的历史 KV cache,此刻正躺在负责这个请求 decode 的那张卡的显存里。如果把增量 prefill 送去远程的 prefill worker,就得先把这份历史 KV cache 隔着网络搬过去,算完再把新产生的 KV cache 搬回来。这笔搬运开销到底值不值?论文先用一组测量回答了这个问题。下面这张图的横轴,是「KV cache 传输的估计耗时」除以「这次增量 prefill 的估计计算耗时」的比值(记作 $T_{kv}/T_{pre}$),纵轴是落在各个比值区间里的任务占比和 token 占比;看的时候对比橙色(DuReader)和蓝色(GAIA)两组柱子的分布。

增量 prefill 的 KV 传输开销占比分布

可以看到,绝大多数增量 prefill 任务的传输开销只占计算耗时的 0-1%,比例小得几乎可以忽略——这说明「无脑都走远程」在平均意义上并不亏。但分布有一条不短的尾巴:仍有相当一部分任务(尤其按 token 占比看)落在 5%、10% 甚至 20%-25% 的区间里。对这些重尾任务,一刀切地走远程就要白白付出可观的网络代价。换句话说,走不走远程不该是个固定策略,而要看具体任务、更要看此刻两边资源的忙闲。可现有系统恰恰缺这套自适应机制:它既要决定「这个增量 prefill 是就地在 decode worker 上算,还是路由到某个 prefill worker」,又要决定「prefill worker 上这些排队的任务按什么顺序执行」,两个决策都没人管好。

第二个洞:模型部署对负载特征没有感知。 PD 分离部署时要定两件事——prefill 和 decode 各分几张卡、各自用什么并行策略。这个最优配置和负载特征强相关,但以往的工作定部署时,只看输入 prompt 和输出 response 的长度,而这恰恰是单轮推理的负载画像。多轮推理那种交织的、prefill 负载明显更重的模式,它根本没纳入考量,于是推出来的资源划分和并行策略往往是次优的。

这两个洞对应了 AMPD 的两块设计:一块是在线的自适应调度(解决增量 prefill 往哪跑、怎么排),一块是离线的部署规划(解决两类资源怎么分、并行怎么配)。下面先看整体骨架,再逐块拆。

3. AMPD 全景:离线定部署,在线做调度

AMPD 分成离线和在线两个阶段。下面这张系统总览图信息量不小,建议分两边看:灰底的左上角是离线阶段(Profiler 和 Planner),其余绿底部分是在线阶段,核心是中间的 Coordinator(协调器)、右上的 Prefill Worker 和右下的 Decode Worker。先顺着左下角「Prompt 进来」这条线,看一个请求是怎么在图里流转的。

AMPD 系统总览

离线阶段干两件事。Profiler(性能剖析器)负责测量各类负载的执行时间,为后面的部署规划和在线调度建一个性能模型。这个性能模型用的是分段的 α-β 模型(α-β 模型是经典的通信/计算代价建模思路,把耗时近似成「固定启动开销 + 单位数据量 × 系数」这样的线性关系,分段是指不同规模区间用不同的线性段去拟合),最终形成三个时间代价函数:$T_{pre}(l_{hist}, l_{incr}; \theta)$ 估计在并行策略 $\theta$ 下、带着长度为 $l_{hist}$ 的历史、再 prefill 长度为 $l_{incr}$ 的新输入要多久(初始 prefill 的历史长度就是 0);$T_{dec}(b; \theta)$ 估计 batch size 为 $b$ 时一个 decode 步要多久;$T_{kv}(l_{ctx}; \theta_{src}, \theta_{dst})$ 估计把长度为 $l_{ctx}$ 的 KV cache 从源并行策略 $\theta_{src}$ 搬到目标并行策略 $\theta_{dst}$ 要多久。这三个函数后面会被在线调度反复调用来做「预判」。Planner(部署规划器)则在正式服务前,把资源分配这件事写成一个整数线性规划问题求解,定出最优部署——这块是第 6 节的主角。

在线阶段由 Coordinator 和一堆 Worker 组成。Coordinator 根据实时负载决定每个任务去哪。每个 prefill / decode worker 各自维护任务队列,并且各自持续记录自己最近一段时间(默认过去 10 秒)的 windowed TTFT / ITL 统计(windowed 指滑动时间窗内的平均值,用它来反映这个 worker 眼下的拥挤程度)。为了让 Coordinator 能随时读到这些信息,队列和窗口统计都用分布式共享内存实现,全局可见。

一个请求的完整流转对着图走一遍:

  • ① Binding(绑定):请求(也叫一个 session)一到,Coordinator 先根据显存占用把它绑定到某个 decode worker 上。这一步定下来之后,这个请求所有的 decode 负载、以及它的 KV cache 管理,都归这张卡负责。绑定的好处是把初始和增量 prefill 的处理统一了起来——不管哪轮的 prefill,历史 KV 的「老家」都明确是这张 decode worker。
  • ② Routing(路由):每当请求需要做 prefill(刚绑定后的初始 prefill,或某轮环境交互后的增量 prefill),Coordinator 用一套自适应路由机制(第 4 节)决定这次 prefill 在哪执行——图里那个「Is Initial Prefill? / Local Exec / Remote Exec」的分叉就是它。Local execution(本地执行) 是直接在负责它的 decode worker 上算,省掉网络传输、也不给 prefill worker 添负担,代价是会打断这张卡正在跑的 decode batch;Remote execution(远程执行) 是路由到某个 prefill worker,躲开了 PD 干扰,但要付出前面说的 KV cache 来回搬运的代价。
  • ③ Prefilling(执行 prefill):被指定的 worker 把这个 prefill 任务塞进自己的 prefill 队列,并用一套 prefill 重排序策略(第 5 节)决定队列里的任务按什么顺序跑。如果指定给了 prefill worker,就按上面说的做 KV 传输。
  • ④ Decoding(解码):第 ③ 步之后,负责的 decode worker 手里就齐了这个请求的全部 KV cache,于是接着往下做自回归解码。

多轮推理里,请求随时可能触发一次和外部环境的交互(比如一次 function call),这通常由 decode worker 的一个后台 CPU 进程处理。交互一结束,就形成一个新的增量 prefill 任务,Coordinator 收到通知,流程回到第 ② 步。这个循环一直转,直到请求推理结束。

看懂了这条主线,就能发现第 ② 步和第 ③ 步正是前面说的两个洞所在——「往哪跑」和「怎么排」。下面两节分别拆开。

4. 在线调度之一:自适应路由决定 prefill 在哪跑

把「往哪跑」这件事再想透一点。如果所有 prefill 都走远程,decode worker 能安安心心做稳态解码,ITL 很漂亮,但代价是额外的 KV 传输、以及 prefill worker 被压得越来越满,TTFT 变差。反过来,让 decode worker 自己顺手把一些 prefill 做了(本地执行),省掉传输、TTFT 受益,但 decode worker 更容易被 PD 干扰波及,ITL 会突然冒尖。所以本地和远程本质是在 TTFT 和 ITL 之间做权衡,没有一边永远更优。

AMPD 的做法是:用前面那套 windowed TTFT / ITL 统计去感知两边此刻的忙闲,哪边有富余就往哪边放,做一个尽量不违反 SLO 的决定。这里引入两个松弛系数 α 和 β(注意,这里的 α、β 是路由里的超参数,和第 3 节性能模型里那个 α-β 通信模型完全是两回事,只是撞了名字):α 用来判断 prefill 这边是否还有余量,β 用来判断 decode 这边是否还有余量。整个路由逻辑如下,看的时候按「先问 prefill 够不够闲 → 再问 decode 够不够闲 → 都不闲就精算」这三层来读:

自适应路由算法

对着伪代码走一遍:

  • 第 1-3 行(prefill 有空就走远程):随机顺序遍历每个 prefill worker $i$,只要它的窗口 TTFT 统计 $\widehat{\text{TTFT}}i \le \alpha \cdot \text{TTFT}{thres}$,说明这个 worker 离 TTFT 的 SLO 阈值还有足够余量(α 默认 0.9,即只要用到阈值的九成以内就算有余量),就直接把任务路由给它,远程执行。随机顺序是为了避免所有任务一股脑涌向同一个「看起来最闲」的 worker。
  • 第 4-5 行(prefill 都满了,看 decode 有没有空):如果所有 prefill worker 都逼近 TTFT 阈值了,就回头看当前这个 decode worker 的窗口 ITL 统计 $\widehat{\text{ITL}} \le \beta \cdot \text{ITL}_{thres}$——decode 这边 ITL 还有余量(β 默认 0.85),就让它本地执行,顺手把 prefill 做了。
  • 第 6-9 行(两边都紧张,就精算一笔账):实在都没有明显余量,就动用性能模型,把本地执行和每个远程执行的代价都估出来,挑最小的那个。

精算用的就是第 3 节建的性能模型。本地执行在 decode worker $d$ 上的代价是:

\[t_{local} = T_{pre}(l_{hist}, l_{incr}; \theta_d) + \sum_{k \in \text{prefill queue}_d} T_{pre}(l_{hist_k}, l_{incr_k}; \theta_d)\]

第一项是这次 prefill 本身要花的时间,第二项是它前面排队的那些任务的时间(也就是排队等待时间)。$\theta_d$ 是 decode worker $d$ 的并行策略。

路由到第 $i$ 个 prefill worker 做远程执行的代价则是三部分之和:

\[t_{remote_i} = t_{pre_i} + t_{kv_i} + t_{queue_i}\]

其中 $t_{pre_i} = T_{pre}(l_{hist}, l_{incr}; \theta_{p_i})$ 是纯计算,$t_{kv_i} = T_{kv}(l_{hist}; \theta_d, \theta_{p_i}) + T_{kv}(l_{incr}; \theta_{p_i}, \theta_d)$ 是 KV cache 一来一回的传输(去的时候把历史 KV 从 decode worker 读过来,回的时候把新算的 KV 送回去),$t_{queue_i}$ 是在这个 prefill worker 上的排队时间。三项加起来,谁最小就路由给谁。

这套机制看着要跨 worker 收集一堆共享元数据,但这些读取是并发做的,而且代价估计极快,所以整体开销很小(第 8 节有实测数字)。

5. 在线调度之二:prefill 重排序榨干 TTFT 的 slack

路由解决了「往哪跑」,但任务到了某个 worker 的 prefill 队列里,还有「按什么顺序跑」的空间可挖。关键观察是:既然有性能模型,就能预判队列里每个任务离违反 TTFT SLO 还有多少余量(slack)。一个离阈值还远得很的任务,稍微晚点跑并不会害它违反 SLO;那就可以把它往后挪,先放那些更紧张的任务,从而让更多任务同时达标。

基于这个想法,AMPD 设计了一个轻量的、TTFT 感知的重排序策略:在队列头部开一个大小为 $w$ 的小窗口(lookahead window,前瞻窗口),只在这几个任务里重排,目标是让窗口内满足 TTFT SLO 的任务数最多。算法如下,看的时候抓住「枚举窗口内的顺序 → 算每种顺序能让几个任务达标 → 取达标最多的那个」这条主干:

prefill 重排序算法

逐行走读:

  • 第 1 行:从队列头部取出至多 $w$ 个元素组成窗口 $W$,当前最优顺序 $\pi^*$ 先置空。
  • 第 2-6 行(枚举顺序,挑达标最多的):枚举 $W$ 上所有可行的排列 $\pi$。对任意一个排列,第 $k$ 个任务的完成时间(相对当前时刻 $T_{now}$)可以预测为前 $k$ 个任务计算耗时的累加:
\[C(\pi(k)) = \sum_{j=1}^{k} t_{pre}^{(\pi(j))}\]

于是这个排列下,窗口内满足 TTFT SLO 的任务数可以预测为:

\[S_\pi = \sum_{k=1}^{|W|} \mathbb{I}\left[ (T_{now} - T_{enq}^{(\pi(k))}) + C(\pi(k)) \le \text{TTFT}_{thres} \right]\]

这里 $T_{enq}$ 是任务入队的时刻,$(T_{now} - T_{enq})$ 是它已经等了多久,加上预测的完成时间,只要不超过 TTFT 阈值,这个任务就算达标,示性函数 $\mathbb{I}[\cdot]$ 记 1。每枚举一个排列就更新一次最优 $\pi^*$,留下 $S_\pi$ 最大的那个。

  • 第 3-4 行、第 7 行(防饿死):光追求达标数最多,可能让某个「反正也快违约了」的任务被无限往后踢,饿死在队列里。为此每个任务配一个 postponement counter(延后计数器),每被往后挪一次就加一;一旦有任务的计数器到了 $w$,就跳过所有还会继续推迟它的排列(第 3-4 行的 continue)。换句话说,任何任务最多被延后 $w$ 次,到点必须让它上。
  • 第 8-9 行:按选定的最优顺序重排队列头部,然后出队第一个,交给执行。

因为窗口 $w$ 很小(实践中通常小于 5),枚举排列和算达标数都极快,重排序的开销可以忽略。有一个细节论文自己点了出来:前一节远程代价里的排队时间 $t_{queue_i}$ 并没有把这个重排序考虑进去,但因为窗口相对整个 prefill 队列很小,这个近似不影响路由的有效性。

6. 离线部署规划:把资源分配写成一个整数线性规划

在线调度解决的是「任务来了怎么分」,但有个更前置的问题:服务开始之前,手里这一堆 GPU 到底该怎么切给 prefill 和 decode?这就是第 2 节说的第二个洞。先看清楚部署这件事里有哪两个自由度。

prefill 与 decode worker 可以用不同的并行配置

这张图把分离式部署的两个维度画了出来,左边是 prefill workers,右边是 decode workers,中间箭头是两者之间的 KV cache 流动(增量 KV 往左送、历史 KV 往右读)。看每一侧内部:一个维度是用几个独立副本去扛并发请求,图里 R1/R2/R3 就是副本,这对应 data parallelism(DP,数据并行);另一个维度是把单个副本切到几张 GPU 上以满足显存和算力需求,这对应 model parallelism(模型并行,实现上通常就是 tensor parallelism / TP,张量并行,把每一层的权重和计算切到多张卡上并行算)。所以给 prefill 和 decode 分别定部署,就是在固定的 GPU 总数约束下,为两类 worker 各自定一组 DP(复制几份)和 MP/TP(每份切几张卡)的配置。

AMPD 把这件事形式化成一个资源受限的优化问题,并进一步落成整数线性规划(Integer Linear Programming, ILP)。先看决策变量。由于模型并行度通常是 2 的幂,可以用「每种并行度下部署了几个副本」来表示整个部署方案:给定一组支持的模型并行度集合 $T$(比如 ${1,2,4,8}$),决策变量是两个整数向量 $x, y \in \mathbb{Z}_{\ge 0}^{ T }$,其中 $x^{(n)}$ 表示用并行度 $n$(即每个副本切在 $n$ 张卡上)部署的 prefill worker 数量,$y^{(n)}$ 是 decode 这边对应的数量。举个例子,$x = [x^{(1)}, x^{(2)}, x^{(4)}, x^{(8)}]$ 这个向量就把「prefill 侧 TP=1 的有几个、TP=2 的有几个……」全描述清楚了。值得一提的是,这种写法允许 prefill 和 decode 用不同、甚至不规则的 TP 大小,比 prefill/decode 必须对称的方案灵活。

优化目标上,论文引入一个辅助变量 $Z$,代表所有 worker 里最大的 P95 延迟(P95 延迟指把延迟从小到大排、第 95 百分位的那个值,用它刻画「绝大多数请求能保证的延迟上界」,比平均值更贴近 SLO 这种尾部敏感的目标)。整个问题写成:

\[\arg\min_{x, y \in \mathbb{Z}_{\ge 0}^{|T|}} Z\] \[\text{s.t.} \quad (C1)\ Z \ge \tau_{pre}(n),\ \forall n \in T\ \text{where}\ x^{(n)} \ge 1\] \[(C2)\ Z \ge \tau_{dec}(n),\ \forall n \in T\ \text{where}\ y^{(n)} \ge 1\] \[(C3)\ \sum_{n \in T}(x^{(n)} \cdot n) + \sum_{n \in T}(y^{(n)} \cdot n) \le N\]

拆开看这三条约束:$\tau_{pre}(n)$、$\tau_{dec}(n)$ 是用并行度 $n$ 配置的单个 prefill / decode 副本的估计 P95 延迟。(C1)(C2) 要求 $Z$ 不小于任何一个被实例化的 worker 副本的 P95 延迟——于是最小化 $Z$,等价于最小化所有副本里最坏的那个 P95 延迟,也就是在优化瓶颈 worker。(C3) 是 GPU 总量约束:每种配置吃掉的 GPU 数等于「副本数 × 并行度」,所有配置加起来不能超过集群可用的 $N$ 张卡。

这个形式其实是多维无界背包问题(unbounded knapsack) 的一个变体,求解器要在合法的并行度组合里挑出最优资源分配。论文用标准的混合整数线性规划(MILP)求解器来解,$\tau_{pre}(n)$、$\tau_{dec}(n)$ 作为常量系数喂进去,求解器在整数向量 $x, y$ 定义的可行域里搜索。好处是:对典型集群规模,这个求解计算上很便宜,而且能保证找到全局最优。实测里,离线规划在 8 卡上只要 0.54 秒,一路到 256 卡也不过 61.67 秒(中间 16/32/64/128 卡分别是 1.93s / 5.06s / 14.50s / 28.88s),一分钟出头就能搞定一个相当常见的服务集群规模。

这里有个诚实的细节值得点出来:优化目标用的是 P95 延迟,而系统真正想最大化的是 SLO 达成率,两者并不完全是一回事。为什么不直接优化 SLO 达成率?因为 SLO 达成是个二值判断(达标/不达标),这种离散指标没法当成连续的目标函数塞进线性规划里高效求解,所以只能退而求其次用连续的 P95 延迟做代理目标。那这个代理靠不靠谱?下面这张表给了经验证据。表里对每个「模型 × trace」组合,列出了规划器推演出的 top-3 部署配置,和真实上机服务测出来的 top-3,格式是 P:<TP=?, DP=?>, D:<TP=?, DP=?>(P 是 prefill 侧、D 是 decode 侧)。

规划器推演的 top-3 部署配置与真实服务对照

对着表看一眼就明白了:Planner 列和 Real-system Serving 列,从第 1 名到第 3 名,一字不差地对齐。也就是说,尽管优化目标和真实目标之间隔着一层 P95 代理,规划器照样稳稳地挑中了真实服务里表现最好的那个配置。论文也补了反面情形的讨论:当 TTFT/ITL 阈值被设得极端严或极端松时,绝大多数部署方案的 SLO 达成率会一起塌到接近 0% 或冲到接近 100%,此时二值指标占主导,优化连续的 P95 就可能和真实最优脱节;但现实服务里阈值一般设在合理区间,这个代理因此在经验上是可靠的。

7. 工程实现:Dynamo + NIXL + Redis 把这套跑起来

这套设计落到代码上,AMPD 建在 NVIDIA Dynamo(一个数据中心规模的分布式推理服务框架)之上。全局共享的队列和 windowed TTFT/ITL 统计用 Redis 实现,前面那个 ILP 用 SCIP 库求解。

KV cache 的跨 worker 传输是性能关键,这里用 NVIDIA 的 NIXL(Inference Xfer Library)做点对点的 RDMA 网络通信(RDMA 让一台机器直接读写另一台的内存、绕开 CPU,延迟和吞吐都比走常规网络栈好)。还有两个减负的工程细节:一是历史 KV cache 用 lazy read(惰性读取)——路由阶段只把元数据发给 prefill worker,只有当这个增量 prefill 真的被调度执行时(第 5 节出队那一刻),prefill worker 才去 decode worker 那儿读历史 KV,避免过早搬运;二是 KV cache 的传输和前一个任务的计算是重叠进行的,用计算把网络延迟藏起来。另外因为每个 decode worker 本地维护着 prefix cache,收到的增量 KV 可以和本地已存的历史 KV 直接拼起来,所以只需要传增量那一段 KV,把网络成本压到最低。

8. 效果:SLO 达成率提升与延迟拆解

实验环境是 4 台服务器、每台 8 张 NVIDIA H20(96GB),一共 32 张卡;机内 GPU 走 NVLink(900GB/s),机间走 InfiniBand(200GB/s)。模型选了三个有代表性的:Qwen3-32B、Llama3.1-70B,以及一个 MoE 模型 Mixtral-8x7B。负载 trace 来自四个公开数据集,它们的平均轮数、prefill 长度、decode 长度差异很大,正好覆盖不同的多轮形态:

四个 trace 的平均轮数与长度

对着表读:ToolBench 和 GAIA 代表 agent 式 workflow——其中 GAIA 又长又多轮(平均 11.32 轮、prefill 长达 6161 token、decode 528 token),是最吃资源的一类;ToolBench 则短平快(3.96 轮、prefill 703 token)。HotpotQA 和 DuReader 代表迭代式 RAG(各跑 3 轮检索),DuReader 的 prefill 偏长(3081 token)。按规模匹配资源:HotpotQA 和 ToolBench 跑在单台服务器(8 卡),DuReader 两台(16 卡),GAIA 四台(32 卡)。

对比的基线有三个:Dynamo(原版 PD 分离、不带本文的技术)、vLLM(最知名的 PD 共置系统之一)、以及 vLLM-Continuum(基于 vLLM、专门为多轮 workflow 定制的调度)。为了公平,AMPD 用自己的离线规划器定部署,而所有基线都手工调过部署、报告各自最好的结果。默认超参 α=0.9、β=0.85、w=3,请求到达按泊松过程生成,主指标是 SLO 达成率。

先看最关键的端到端对比。下面这张大图是 3 个模型 × 4 个 trace 共 12 个子图,每个子图横轴是请求到达率(request/s,越往右负载越重),纵轴是 SLO 达成率,四根柱子是四个系统;柱子上方标的百分比是 AMPD 相对各基线的提升幅度。看的时候抓两点:蓝色柱(AMPD)是不是基本都最高,以及负载越重时差距是不是越明显。

端到端 SLO 达成率对比

结论很清楚:AMPD 在所有实验里都拿到了最高的 SLO 达成率。 汇总成一句话——相比分离式基线 Dynamo,平均提升 67.29%、最高处提升 967.54%;相比共置式基线 vLLM,平均提升 339.74%、最高处提升 3435.1%。差距在重负载区间尤其大,比如 Qwen3-32B 的 ToolBench、HotpotQA、DuReader 这几栏,到达率一上去基线就断崖式下跌,而 AMPD 还能稳住。

光看达成率还不够过瘾,提升到底从哪来?下面这张延迟拆解图把 Llama3.1-70B 在 2 reqs/s 下的延迟拆成了 TTFT(上排)和 ITL(下排),四个 trace 各一列。读的时候对比同一列里四根柱子的高低:谁的 TTFT 低、谁的 ITL 低。

延迟拆解:TTFT 与 ITL

拆解图把两类基线各自的「病」暴露得很直白。共置式(vLLM) 为了连续批处理,倾向于优先/抢占着把 prefill 早点做掉,于是 TTFT 压得很低(上排里 vLLM 的柱子矮),但这打断了 decode 的节奏,ITL 被拖高(下排里 vLLM 的柱子最高)。分离式(Dynamo) 正相反,它把 prefill 和 decode 解耦、decode 这端干扰小,ITL 通常更低,但 TTFT 常常成为瓶颈(上排里 Dynamo 在 DuReader、GAIA 上的 TTFT 柱子冲得老高)——因为分出去的 prefill 资源变少,又多了 KV 传输开销。于是这俩基线总是被 TTFT 或 ITL 其中一个指标拖住,达成率上不去。

而 AMPD 既吃到了 PD 分离的红利(ITL 不像 vLLM 那样被拖垮),又靠自适应调度把 TTFT 压了下来(相比 Dynamo 明显更低)。说白了,AMPD 是在 TTFT 和 ITL 之间找到了一个更稳的平衡点,而不是像基线那样顾此失彼。至于 vLLM-Continuum,它沿用共置架构、靠优先调度同一 session 内可复用 KV 的任务来缩短排队、降低 TTFT,但很多设定下这点 TTFT 下降没能转化成 SLO 达成率的实质提升,所以它经常和 vLLM 打平。

那这套在线调度本身贵不贵? 实测里自适应路由平均耗时 0.39ms(其中靠分布式共享内存读元数据只占 0.15ms),prefill 重排序平均 0.45ms。考虑到 LLM 的一次 prefill 计算或 decode 步动辄几十到几百毫秒,这点开销可以忽略,而且它们跑在 CPU 上、和 GPU 计算重叠,不占用算力。

两个技术各自贡献多少? 下面这张消融图,左、中两组是在 DuReader、GAIA 上逐步加回技术(w/o AR+PR 是都不要、w/o PR 是只加自适应路由、AMPD 是全加),右边一组是本地/远程执行的比例分布。

消融实验

走读一遍:自适应路由(AR)把 13.9%-31.7% 的 prefill 任务分流到 decode worker 本地执行(右图里 DuReader 本地 31.7%、GAIA 31.3%、HotpotQA 15.5%、ToolBench 13.9%),缓解了 prefill worker 的负担,单这一项就把 SLO 达成率拉高 27.37%-350%;在此之上再加 prefill 重排序(PR),进一步贡献 13.42%-14.81%(靠削掉 TTFT 的长尾)。两者合起来带来 44.47%-402% 的提升。有意思的是本地执行比例的规律:prefill 越短的 trace(ToolBench、HotpotQA)本地执行占比越低,因为它们的 prefill worker 本来就没那么挤,没必要往本地分流。

超参敏不敏感? 下面这张图扫了窗口 $w$、以及 α、β 三个超参,上下两排分别是 Llama3.1-70B/DuReader 和 Qwen3-32B/HotpotQA。

超参敏感性

$w$ 从 2 到 5,达成率的波动在 3% 以内——说明一个很短的前瞻窗口就够捕捉有用的负载信号了,不用开大。α、β 则呈现出明显的「中间最好」:α 太小会过早判定 prefill 拥挤、导致 prefill worker 利用不足,β 太小会过早判定 decode 拥挤;两个取适中值(默认 α=0.9、β=0.85)能在 prefill 和 decode 之间平衡得最好。

9. 总结

AMPD 要解决的问题可以浓缩成一句话:当 LLM 推理从单轮问答走向多轮 agent / 迭代 RAG,prefill 不再集中在开头,而是以 incremental prefill 的形式零散地插在整个 decode 过程里,这打破了 PD 分离赖以成立的「prefill 在前、decode 在后」的假设。AMPD 没有推翻 PD 分离,而是在这个框架里补上两块:在线用自适应路由决定每个增量 prefill 是在本地 decode worker 上就地算、还是路由到远程 prefill worker,用prefill 重排序在队列头部的小窗口里重排以让更多任务卡进 TTFT 阈值;离线把资源分配和并行策略写成一个整数线性规划,用 P95 延迟做代理目标求全局最优。实验上,相比共置和分离两类基线,SLO 达成率平均提升 67.29%-339.74%,而且这套在线调度本身开销只有亚毫秒级。

平心而论,这是一篇扎实但偏增量的 serving 调度工作:PD 分离、多轮推理加速都不是它发明的,它的价值在于第一个把这两条线认真缝到了一起,并且找准了「incremental prefill 的归属」这个被忽略的关键点——而且诚实地用实测(KV 传输开销分布、消融、规划器 vs 真实服务的 top-3 对照)把每个设计的收益和适用边界都说清楚了。离线规划那块的 ILP 形式化尤其干净,P95 代理和 SLO 目标之间的 gap 也没藏着掖着,而是正面讨论了它什么时候会失效。

论文自己也点了几个待补的方向:把 chunked prefill(把长 prefill 切成小块)接进来,在本地执行时减少对 decode 的阻塞;离线规划器目前吃的是固定负载和硬件假设,一旦负载漂移或硬件可用性变化(弹性伸缩、故障恢复)就可能次优,可以引入「监控-重规划」并在后台 CPU 线程上做滚动迁移(逐副本调整),避免停服。这些方向补齐之后,这套东西离真正的生产可用会更近一步。


扯一句题外话。AMPD 这篇是纯粹的推理服务调度,属于「模型本身不动,优化怎么把它部署和调度得更好」这条系统工程的路子,和我们之前整理的《动手学 AutoML:从 NAS 到大语言模型优化实战》方向其实隔得有点远——那本书讲的是 NAS、超参搜索,以及 LLM 的剪枝、量化、模型融合这类「把模型本身变小变快」的手段。角度不一样,但两条线想达到的终点是一致的:让大模型跑得又快又省。如果你对「从模型侧做自动化优化」这条线感兴趣,可以翻翻那本书。

动手学AutoML书籍封面

Flag Counter