Preble
论文地址:https://arxiv.org/pdf/2407.00023 发表于 ICLR 2024
设计
调度分为两个层次:Global and local。所有的请求再经过 tokenizer 后,进入 global scheduler。 global scheduler 管理整个集群,决定请求将被送入某个 local scheduler。每个 local scheduler 管理一个模型实例,负责调度请求在这一个实例上的请求。Local scheduler 与现有的 LLM 推理框架,如 vLLM、SGLang、SpecInfer 等大致处于同一层级。
Global Scheduler
数据结构
基础的数据结构是若干 radix trees,树上每个节点对应一段 prompt。在做前缀匹配时,允许节点分裂为匹配和失配两个节点。树上的节点存储了这三类信息:
- 一段 prompt
- 包含这段 prompt 对应的 kv cache 的 GPU id
- 在过去的一段历史窗口内(3 min.),每个 GPU 使用过该节点的次数 如果某个节点既没有 GPU 存放了对应的 kv cache,过去一段时间内又没有被使用过,该节点就会被清除。
请求调度策略
当前缀的匹配长度大于失配长度时,调度器会直接把请求送入最长前缀匹配的 GPU。如果有多个最长的前缀匹配,调度器将请求送入负载最低的那一个。
如果前缀匹配的长度小于失配的长度,对每个 GPU 估算把请求送给该 GPU 执行的开销,选择最低的那个 GPU 执行本次请求。开销包含三个方面:
- :该 GPU 本身的计算负载。计算方式是过去一段时间窗口(3 min.)内,该 GPU 所有请求的计算耗时求和。每一个请求的耗时计算公式为:。其中两个速率是离线测出来的常数,每种 GPU 型号测算出两个速率数值。decode 长度是时间窗口内所有请求的平均 decode 长度。这样计算耗时可以完全在 global scheduler 里进行,不需要真的记录每一个请求的时间,可以降低负载。使用一段时间区间来计算负载,有两个原因
- 调度发生时的 GPU 负载,和请求真正执行时的 GPU 负载不同
- prefix 位置,相比单次的负载数值,有更长期的影响
- :如果将请求送给该 GPU,进行kv cache 驱逐的代价。先根据驱逐算法,给出如果把请求发送给该 GPU,会有哪些缓存被驱逐。总的驱逐代价是所有被驱逐的缓存节点的代价求和。单个缓存节点的代价公式为:
- :失配 token 的 prefill 耗时:。只考虑了 prefill 而没有 decoding,是因为假设了在所有 GPU 上进行 decoding 速率是一样的。 这三个数值会直接加起来 ,选择使得该数值最小的 GPU。
我的想法
上面的开销公式 ,第一个是针对一个时间区间内所有请求的求和,它衡量的是 GPU 的计算负载高低,时间区间的长度代表着它在 项的权重。但是论文里却说,实验了不同的时间区间长度,发现结果对区间长度不敏感,这比较反常。
论文里说,只在 global scheduler 里记录 token 数量,通过乘以速率来计算时间,相比于真的去记录时间,会降低负载。我觉得记录时间这个操作本身对性能没有那么大的影响。而且就算此处不记录时间,后端的性能监测也要记录时间。我对这个不记时间记 token 的操作表示怀疑。
使用历史区间来计算 GPU 负载的原因提到了调度发生时的 GPU 负载,和请求真正执行时的负载时不同的,使用区间均值可以磨平这个差异。这一点和 Mooncake 的论文是一致的。Mooncake 论文也提到了这一点,但是是用一些预测算法去补偿这个差异。
任务分配后的负载调整
上述负载均衡算法,在请求的 prompt 前缀分布比较稳定的时候表现较好,但是实际的请求会有波动,分布并不一定是稳定的。因此设计了任务分配后的负载调整。负载调整有两种算法:
- 如果当前负载最高的 GPU 和负载最低的 GPU 负载比值超过了某个可配置的阈值,则原本所有的分配给最高负载 GPU 的请求都强制转给最低负载的 GPU,直到最高和最低的负载比值低于阈值。
- 如果在 1. 的平衡之后,某一个前缀对应的 GPU 仍然负载过高(时间窗口内的排队时间翻倍了)。则将 prefix 复制并拆分子树,分给多个 GPU(拆分子树这里论文一句话带过了,没看明白)。
我的问题:为什么到了这里,负载又用排队时间来衡量了呢?论文之前说的不记录时间从而降负载,这里又开始记录时间了。
Prefill-Decoding 平衡
Prefill 和 decoding 的计算-访存比不一样,因此可以在调度过程中区分 prefill/deocding 来做更好的平衡。论文提出,可以将每个请求,按照 prompt 命中缓存的数量,分为以下三种之一:
- 整个 prompt 命中缓存:decoding 请求
- 整个 prompt 都没有命中缓存:prefill 请求
- 部分命中缓存:介于 prefill/decoding 的中间请求 如果一台 GPU 上大部分请求都是 decoding,那说明该 GPU 的计算资源是闲置的,于是调度器就会直接把 prefill 请求发送给该 GPU。直到所有 GPU 上的 prefill-decoding 请求比例基本平衡。这个 prefill-decoding 平衡策略的优先级,要高于前文的基于开销的调度策略。关于部分命中缓存的中间请求在平衡策略中如何处理,论文没有提及。
Local Scheduler
Local scheduler 类似 vLLM,SGLang 等框架,有自己本地的 kv cache,prefix cache tree 等等。新的请求到达 local scheduler,会进行 cache 匹配,并加入等待队列。模型的每一轮执行迭代,都会从等待队列中挑选出一些任务进行执行。当 GPU 显存不够时,会通过 LRU 驱逐一部分树节点。当某个未命中缓存的 token 数太大时,会有 chunked prefill。
Batch 策略
从等待队列挑选请求时,会使用一个兼顾公平性和 cache 复用的策略。将所有请求按照 cache 命中的长度占全部 prompt 长度的比例,将请求分为 组,其中 是一个可配置的参数。例如,设置 ,则请求按照缓存命中 0%~10%,10%~20%, ,90~100% 分为 10 个小组。这 10 个小组的权重比例是 。当需要从等待队列挑选 55 个请求组成一个 batch 时,分别从这 10 个小组挑选 个请求。
消融实验
以 SGLang + round robin 作为基线,逐步加上
- 基于开销的请求调度策略
- 任务分配后的负载平衡
- Prefill-Decoding 平衡
- 等待队列挑选策略
1.,2.,3 都可以让平均延迟和 p99 延迟降低。4. 会让平均延迟略微上涨,但是能降低 p99。