nano vllm源码解析-scheduler

1. 核心职责

调度器是推理引擎的“大脑”，负责平衡吞吐量（Throughput）与延迟（Latency）。

• 📋 请求队列管理：维护 waiting、running 和 swapped 队列。
• ⚡ 批处理调度：智能组合多个序列形成 Batch，提升 GPU 利用率。
• 🧠 内存管理：通过 PagedAttention 思想，利用 Block Manager 动态分配 KV Cache。
• 🔄 抢占机制：当显存不足以支撑所有正在生成的序列时，根据优先级或时间戳强制回收资源。

2. 核心概念对比：Prefill vs Decode

3. 调度算法实现 (Code Analysis)

3.1 预填充阶段 (Prefill)

优先处理新进入的请求。如果显存能容纳整个 Prompt 的 KV Cache，则一次性分配空间。

Python

def schedule_prefill(self):    scheduled_seqs = []    num_seqs = 0    num_batched_tokens = 0    while self.waiting and num_seqs < self.max_num_seqs:        seq = self.waiting[0]        # 约束检查：Token 数量限制 或 显存块不足        if (num_batched_tokens + len(seq) > self.max_num_batched_tokens or             not self.block_manager.can_allocate(seq)):            break        num_seqs += 1        self.block_manager.allocate(seq)  # 一次性为 Prompt 分配所有 Block        num_batched_tokens += len(seq) - seq.num_cached_tokens        seq.status = SequenceStatus.RUNNING        self.waiting.popleft()        self.running.append(seq)        scheduled_seqs.append(seq)    if scheduled_seqs:        return scheduled_seqs, True

3.2 解码阶段 (Decode)

如果当前没有新的 Prefill 任务，则继续处理 running 队列。如果显存不足，会触发**抢占 (Preemption)**。

Python

def schedule_decode(self):    scheduled_seqs = []    while self.running and num_seqs < self.max_num_seqs:        seq = self.running.popleft()        # 检查是否能容纳下一个 Token 的 KV Cache        while not self.block_manager.can_append(seq):            if self.running:                # 显存不够，牺牲 running 队列末尾的序列                self.preempt(self.running.pop())            else:                # 连自己都跑不动了，直接抢占自身                self.preempt(seq)                break        else:            num_seqs += 1             self.block_manager.may_append(seq) # 预留 1 个 token 空间            scheduled_seqs.append(seq)    # 将调度的序列放回 running 头部，保持下一轮调度的顺序    self.running.extendleft(reversed(scheduled_seqs))    return scheduled_seqs, False

4. 资源约束与 KV Cache 管理

资源约束矩阵

KV Cache 操作对比

• allocate
  
  : 用于 Prefill，分配足以容纳整个 Prompt 的内存块。若失败，该序列保持 waiting。
• append
  
  : 用于 Decode，每次仅分配 1 个 Token 的空间。若失败，必须通过 preempt 释放其他序列来腾出空间。

5. 辅助方法

• preempt(seq)
  
  : 强制暂停序列。

1. 状态从 RUNNING 变为 WAITING。
2. 释放其占用的所有 KV Cache 物理块。
3. 将其重新推回 waiting 队列头部（通常使用 Recompute 策略重新开始）。

• postprocess(seq)
  
  :
1. 接收模型输出的 Token。
2. 检查是否遇到停止词（Stop Tokens）或达到最大长度。
3. 若结束，调用 block_manager.free(seq) 彻底释放内存。

6. 运行逻辑总结

1. 优先 Prefill
   
   ：只要 waiting 队列有任务且显存够，先做 Prefill。
2. 次选 Decode
   
   ：没有新 Prefill 时，让 running 里的任务继续跑。
3. 动态抢占
   
   ：Decode 过程中显存满了，按照“后进先出”或优先级原则踢掉某些序列，确保前面的序列能顺利完成。