推理工程：编写一个 GPU 调度器 | 第五篇

在上一篇文章中，我们了解了 vLLM 的工作原理：APIServer 和 EngineCore。EngineCore 负责启动 GPU 工作。但在那个循环运行之前，必须有人决定服务器运行在哪台机器上。这就是调度器的工作，在本文中我们来构建一个简单的调度器。

以下是我们将使用的术语：

• 节点（Node）： 拥有空闲 GPU 的机器

• Pod： 在 GPU 上运行的工作负载（例如一个需要 2 块 GPU 的 vLLM 服务器）

• 调度器（Scheduler）： 决定 Pod 将在哪个 Node 上运行。

我们的调度算法将遵循首次适应（first-fit）策略。也就是说，对于输入顺序中的每个 Pod，它会检查哪些节点拥有足够的资源来运行该 GPU 工作负载，并选择第一个满足条件的节点。

规则

• 一个 Pod 必须完全放置在一个节点上。不能从一个节点取 2 块 GPU，再从另一个节点取 1 块。

• GPU 是完整的、相同的单元。

• 首次适应（First-fit）： 按输入顺序处理 Pod。对于每个 Pod，按顺序检查节点，选择第一个拥有足够空闲 GPU 的节点，并从该节点扣减 GPU 数量。

• 如果没有节点能容纳，该 Pod 就是未调度状态，我们继续处理下一个 Pod。

首先我们通过填写表格来模拟一个示例的执行过程：

节点：{a: 4, b: 2} Pod，按顺序：(p1, 3), (p2, 2), (p3, 2), (p4, 1)

Python 代码：

def schedule(nodes, pods):

    assignments = {}
    unscheduled = []
    remaining = dict(nodes)
    
    # 遍历 pod（GPU 工作负载）
    for pod, gpus_needed in pods:
        
        # 遍历节点（拥有空闲 GPU 的机器）
        for node, gpus_free in remaining.items():
            
            if gpus_needed <= gpus_free:
                remaining[node] -= gpus_needed
                assignments[pod] = node
                break
        
        # 如果 for 循环没有触发 break
        else:
            unscheduled.append(pod)
            
    return assignments, unscheduled, remaining

复杂度

在最坏情况下，每个 Pod 检查每个节点：时间复杂度为 O(P × N)，其中 P 是 Pod 数量，N 是节点数量。空间复杂度为 O(N + P)：节点的副本，加上输出结果。

首次适应并非最优

让我们改变 Pod 的顺序： 顺序 A：[3, 2, 2]

顺序 B：[2, 2, 3]

首次适应是贪心地处理每个 Pod，不考虑未来。顺序 B 证明所有 3 个 Pod 都能容纳，但顺序 A 的首次适应只放置了 2 个。

一个常见的改进是在首次适应之前按降序对 Pod 排列（“首次适应递减”）。这就是这里顺序 B 的做法。这通常有帮助，但仍然不能保证是最优的。寻找最佳放置方案是**装箱问题（bin packing）**的一种形式，它是 NP 难的：没有已知方法能在输入规模增长时快速找到保证最优的答案。所以实际的调度器使用快速启发式方法，并接受一些不够完美的放置结果。

如果两个调度器同时运行会怎样？

我们的调度器假设它是唯一的写入者。每个调度器都基于自己对容量的副本工作。两个调度器可能都看到节点 a 有 4 块空闲 GPU，都把一个需要 3 块 GPU 的 Pod 放在那里，然后都报告成功。此时该节点已经被承诺了 6 块 GPU，但实际只有 4 块，而且不会产生任何错误。

解决方案是将容量保存在一个共享存储中，并将每次放置操作变为条件更新：只有在自你读取以来容量未发生变化时才提交；否则重新读取并重试。Kubernetes 通过运行单个活跃调度器、通过一个 API Server 写入、以及在 Pod 到达时由节点重新检查容量来避免这种竞争。

我们可以用 GPU 特有的功能扩展这个算法：

• 拓扑感知（Topology-aware）： 将多 GPU 模型放在通过 NVLink 连接的 GPU 上，或放在同一机架内。

• 集合调度（Gang scheduling）： 需要 8 块 GPU 的模型必须全部获得 8 块，否则一个都不要。Kueue 和 Volcano 等工具支持这一功能。

• 优先级和抢占（Priority and preemption）： 重要工作可以驱逐次要工作。

路由：每个请求分配给哪个副本

当多个副本同时运行时，路由器或负载均衡器为每个请求选择一个副本：

• 轮询（Round-robin）： 简单，但忽略负载。

• 最少负载（Least-load）： 将请求发送给活动请求数最少或队列最短的副本。

• 缓存感知（前缀感知）： 将具有相同提示词前缀（例如相同的系统提示词）的请求发送到其 KV 缓存中已经包含该前缀的副本。这样可以跳过重复的预填充阶段，减少首 Token 延迟（TTFT）。这是 LLM 服务领域的一个重要趋势。

我们学到了什么

调度器写起来简单，但要做到最优却很难。总空闲容量并不能保证能成功放置，请求的顺序也会影响结果。首次适应速度快且可预测，但碎片化可能导致工作负载无法被调度。真实系统还需要一个容量的唯一事实来源，否则两个调度器可能将同一批 GPU 承诺两次。

代码和资源仓库：https://github.com/LogicalGuy77/inference-engineering