在当今复杂多变的技术领域,任何足够复杂的推理业务,都必然要紧密结合推理引擎行为,精心设计出高效的调度系统。原因主要有以下三点:
其一,复杂推理服务本质上是分布式的;
其二,分布式系统处理请求时,调度不可或缺,若缺乏精细调度,各节点工作负载将不均衡;
其三,要满足调度需求,必须深入了解推理请求在引擎中的处理方式,并预判最优调度方式。

目前,大模型推理基本采用 PD 分离模式,这使得推理调度涵盖 prefill 阶段多个推理引擎间的负载均衡以及 decode 阶段多个推理引擎间的负载均衡。再考虑诸如模型多版本灰度、LORA 调度等因素,整个推理请求调度过程极为复杂。不过,本文仅对推理调度进行简单探讨,聚焦于最基础的 prefill 阶段负载均衡,不考虑 LORA 之外的其他因素。

Prefill 阶段具有两大特点:
一是计算密集型操作;
二是计算量取决于 input token 未命中 prefix cache 的数量。

接下来的讨论将围绕这两个特点展开。

鉴于调度系统需结合推理引擎行为,实时掌握引擎信息至关重要。由此衍生出两个关键问题:获取哪些信息以及如何实时获取信息。

针对第一个问题,通过分析引擎侧可获取数据,运用排除法,可筛选出以下关键信息:

  • 过去一段时间内处理的 input token 数量:这是调度的结果,我们倾向于使用领先指标来预判下一步的行为
  • 当前 GPU/kvcache 的使用量:这种指标不一定能直接反映计算量。举个例子,假设有两个引擎 A、B,A 上有当前请求对应的 kvcache,而 B 没有。因为有了这个 kvcache,A 的内存用量比 B 更高。如果是优先调度到内存用量小的节点,那么应该调度到 B,尽管调度到 A 能命中缓存。
  • GPU 利用率:对于 PD 不分离的场景,有可能因为内存带宽限制,导致计算上不去。调度多个请求过去,增大 batch size 是一个好的策略。但是在 PD 分离的场景,作为计算密集型的操作,prefill 阶段往往能打满 GPU,基于 GPU 利用率来做调度意义不大。
  • cache 分布:好的指标。它决定了 input token 在换算成计算量时需要打多少折扣(注意这里不是线性的关系)。
  • request queue size 或当前处理的 input token 数量:好的指标。它可以换算成计算量。
  • waiting queue size:如果调度时已经考虑到了每个节点的能力,而且容量预估正确,那么 waiting queue 只是 request queue 的暂存处。只看 waiting queue size 的调度算法意义不大 —— 因为大部分时候 waiting queue size 都很小甚至为 0,没办法充分反映状况。

经筛选,确定关键信息为:cache 分布、request queue size 或当前处理的 input token 数量、waiting queue size。

对于如何采集这些信息,一个初步想法是定期采集引擎的 Prometheus 格式 metrics endpoint,但此方式存在两大缺陷:

一是无法获取 cache 分布,因为 metrics 无法体现哪些 prompt 被缓存、哪些被淘汰

二是采集效率低下,以 1000 台推理引擎的集群为例,若采集间隔过长,数据实时性无保障,若间隔过短,将产生大量无变化的 metrics,造成资源浪费。例如,假设集群平均 TPOT 为 50ms,batch size 为 50,每个请求平均 output token 数量为 1000 个,则 QPS 为 (1000ms / 50ms 50) / 1000 token 1000 = 1000,为保证不漏掉每个事件,采集间隔设为与 TPOT 相同,即每台推理引擎每秒采集 20 次,对应 QPS 为 20000,是集群推理 QPS 的 20 倍,若由 20 台网关采集 1000 台推理引擎,则总量达 400K QPS。

为解决上述问题,可采用改进方案:让引擎上报数据而非频繁轮询。技术上难度不大,因引擎的请求解析、kv cache block 管理及每次推理完更新指标均用 Python 代码实现,在其中添加上报数据代码的开发量和难度可控。

推理引擎上报数据,通常便是上报到中心化的 broker 集群。这种服务需要通过多实例来做高可用,比如采用 redis sentinel,master 节点故障后提拔一个 replica 作为 master。

我们可以简单算算这种 broker 集群需要满足的容量。再次以前面提到的 1000 台推理引擎的集群为例。

数据上报的频率为 收到的任务数 x 5(进入 waiting queue 一条、转换到 request queue 一条,处理完成一条,新增 kvcache 一条,汰换 kvcache 一条),不考虑诸如 pipeline 这样合并优化方式,每秒 5000 请求。对于多数支持频繁写入的集群,此写入频率无大碍。

上报的数据里,metrics 总量取决于引擎数量,只有 1000,对比于 cache 分布来说,可以忽略。cache 分布的数据才是大头。假设 1k token 对应的 kvcache block 大小为 100 MB,每台推理引擎留出 500 GB 存储 kvcache,则 kvcache token 总量约为 5g(500GB / 100MB 1k 1000),按每个 token 8 字节计算,需约 40GB 存储,但因推理引擎开启 prefix cache,相同前缀 prompt 共享 kvcache block,若 broker 不按前缀树存储 prompt,内存占用会更大。当然以上全是估算值,实际情况可能和这个有数量级的差异。如单机 8 卡 H20 提供 768GB 显存,部署全规格 deepseek r1 后,仅剩几十 GB 存储 kvcache。建议感兴趣的读者根据自己的实际情况核算。

在调度时,调度器需订阅 broker 数据,重建前缀树计算请求在各节点的缓存复用情况,并结合各节点当前计算量,以实现计算量均衡。虽然推理引擎以 block 粒度判断缓存复用,调度器不做 tokenize 时无法感知 block 数量,只能以 byte 粒度进行前缀匹配。但单个 block 只有几十个 tokens,与用户请求输入的上千甚至更多 tokens 相比,误差影响不大。

依前文计算,前缀树可能占约 40gb,基于此有以下两种选择:
一是在采用大实例(如 32c64g)的网关上完成调度;
二是采用小实例(如 4c4g)的网关扮演传统无状态角色,将 model 和 prompt 发给 32c64g 的调度节点,由调度节点确定最佳节点并告知网关。

前者实现简便,后者更契合微服务按资源使用拆分的思路。现阶段,相较于节省 GPU 花费,减少 CPU 和内存花费的优先级较低。


spacewander
5.6k 声望1.5k 粉丝

make building blocks that people can understand and use easily, and people will work together to solve the very largest problems.