LLM推理调度器设计与实现：从基础到优化

1. 调度器基础概念与核心挑战

在大规模语言模型推理场景中，调度器（Scheduler）扮演着至关重要的角色。它就像交通指挥中心，需要高效协调来自不同时间、不同长度的用户请求，确保计算资源得到最大化利用。理解调度器的工作原理，是掌握现代LLM推理框架的关键。

1.1 调度器的核心职责

调度器主要解决两个基本问题：

1. 时间维度协调：用户请求的到达时间不可预测，每个请求的生成长度也各不相同。就像高峰期地铁站的人流，有的乘客来得早但行程短，有的来得晚却要坐很多站。
2. 资源维度管理：每个请求都需要占用KV cache资源，这部分内存的大小与请求序列长度成正比。调度器需要像精明的仓库管理员一样，动态分配有限的存储空间。

1.2 现代LLM推理的特殊挑战

与传统任务不同，LLM推理还面临一些独特挑战：

• Continuous batching：需要将不同进度的请求批量处理，就像餐厅同时服务多桌客人，每桌的上菜进度各不相同
• Prefill与decode阶段差异：prompt处理（prefill）和token生成（decode）对计算资源的需求特性完全不同
• KV cache传输：需要高效管理注意力机制中的键值缓存，这对性能影响巨大

在实际工程中，vLLM的调度器代码已经发展到近2000行，包含了各种优化和边缘情况处理。这对初学者来说确实是个不小的挑战。

2. 基础调度器设计与实现

2.1 简化版设计目标

为了帮助理解核心原理，我们设计了一个简化版调度器，保留最关键的几个功能：

• 支持continuous batching机制
• 实现prefill和decode的优先级控制（默认prefill优先）
• 资源不足时的抢占功能
• 剔除prefill和decode混合执行等复杂特性

这个设计就像学车时先用教练车练习，等掌握了基础操作再开复杂车型。

2.2 核心数据结构

2.2.1 Sequence类

每个用户请求被封装为一个Sequence对象，相当于餐厅的"订单"：

python复制class Sequence:
    block_size = 256  # 每个block固定大小
    counter = count()  # 自增ID生成器

    def __init__(self, token_ids: list[int], sampling_params=SamplingParams()):
        self.seq_id = next(Sequence.counter)  # 唯一ID
        self.status = SequenceStatus.WAITING  # 初始状态
        self.token_ids = copy(token_ids)  # token序列
        self.last_token = token_ids[-1]  # 最后一个token
        self.num_tokens = len(self.token_ids)  # 总token数
        self.num_prompt_tokens = len(token_ids)  # prompt长度
        self.num_cached_tokens = 0  # 已缓存token数
        self.block_table = []  # 分配的block列表
        self.temperature = sampling_params.temperature  # 采样参数
        self.max_tokens = sampling_params.max_tokens  # 最大生成长度
        self.ignore_eos = sampling_params.ignore_eos  # 是否忽略结束符

关键字段说明：

• block_table：记录分配的KV cache块，类似餐厅的"座位分配表"
• status：跟踪请求状态（等待/运行/完成）
• 采样参数控制生成多样性，就像调整菜品的咸淡口味

2.2.2 Block管理器

KV cache采用分块管理，就像把大仓库划分成标准货架：

python复制class Block:
    def __init__(self, block_id):
        self.block_id = block_id  # 块唯一ID
        self.ref_count = 0  # 引用计数
        self.hash = -1  # 内容哈希
        self.token_ids = []  # 存储的tokens

    def update(self, hash: int, token_ids: list[int]):
        self.hash = hash
        self.token_ids = token_ids

    def reset(self):
        self.ref_count = 1
        self.hash = -1
        self.token_ids = []

Block管理器的设计要点：

• 固定大小的内存块（如256 tokens/block）
• 引用计数实现安全回收
• 哈希值用于内容校验

2.3 调度器核心逻辑

2.3.1 双队列机制

调度器维护两个关键队列：

python复制class Scheduler:
    def __init__(self, config: Config):
        self.waiting = deque()  # 等待队列
        self.running = deque()  # 执行队列
        self.block_manager = BlockManager(...)  # 资源管理器
        # 其他初始化...

• 等待队列：新到达的请求，相当于餐厅的等候区
• 执行队列：正在处理的请求，相当于已入座的顾客

2.3.2 Prefill阶段处理

Prefill就像准备食材的过程，需要较多资源：

python复制def prefill(self):
    scheduled_seqs = []
    while self.waiting and len(scheduled_seqs) < self.max_num_seqs:
        seq = self.waiting[0]
        # 检查资源是否充足
        if (self.num_batched_tokens + len(seq) > self.max_num_batched_tokens 
            or not self.block_manager.can_allocate(seq)):
            break
        # 分配资源并转移队列
        self.block_manager.allocate(seq)
        seq.status = SequenceStatus.RUNNING
        self.waiting.popleft()
        self.running.append(seq)
        scheduled_seqs.append(seq)
    return scheduled_seqs, True

关键判断条件：

1. 当前批次token总数是否超限
2. 是否有足够的KV cache空间
3. 是否达到最大并行请求数

2.3.3 Decode阶段处理

Decode像上菜过程，需要持续服务：

python复制def decode(self):
    scheduled_seqs = []
    while self.running and len(scheduled_seqs) < self.max_num_seqs:
        seq = self.running.popleft()
        # 检查是否能继续追加token
        while not self.block_manager.can_append(seq):
            if self.running:
                self.preempt(self.running.pop())  # 触发抢占
            else:
                self.preempt(seq)
                break
        else:
            self.block_manager.may_append(seq)
            scheduled_seqs.append(seq)
    # 保持原有顺序
    if scheduled_seqs:
        self.running.extendleft(reversed(scheduled_seqs))
    return scheduled_seqs, False

抢占机制特别说明：
当资源不足时，会从执行队列尾部移除请求（类似餐厅让等最久的顾客先离开）

2.3.4 统一调度入口

python复制def schedule(self, prefill_first=True):
    scheduled_seqs = []
    # 清空计数
    self.num_seqs = 0
    self.num_batched_tokens = 0
    
    # 根据优先级决定调用顺序
    first_call, second_call = (self.prefill, self.decode) if prefill_first else (self.decode, self.prefill)
    
    # 尝试第一次调度
    scheduled_seqs, is_prefill = first_call()
    if scheduled_seqs:
        return scheduled_seqs, is_prefill
    
    # 尝试第二次调度
    scheduled_seqs, is_prefill = second_call()
    if scheduled_seqs:
        return scheduled_seqs, is_prefill
    
    raise RuntimeError("No requests scheduled, possible resource starvation")

这个调度策略就像餐厅经理决定是先安排新顾客入座(prefill)，还是继续服务已入座的顾客(decode)。

3. 实战演示与可视化分析

3.1 基础测试流程

典型的测试循环如下：

python复制# 初始化调度器
config = Config(max_num_seqs=3, max_model_len=15)
scheduler = Scheduler(config)

# 添加示例请求
prompts = ["Hello world", "How are you"]
for prompt in prompts:
    token_ids = tokenizer.encode(prompt)
    seq = Sequence(token_ids, SamplingParams())
    scheduler.add(seq)

# 处理循环
while not scheduler.is_finished():
    seqs, is_prefill = scheduler.schedule()
    # 模拟LLM生成
    token_ids = run_fake_model(seqs, config.max_model_len)  
    # 更新状态
    scheduler.postprocess(seqs, token_ids)
    # 打印进度
    for seq in seqs:
        print(f"Seq {seq.seq_id}: {seq.token_ids}")

3.2 可视化调度过程

我们开发了队列可视化工具，可以直观展示调度过程：

场景1：Prefill优先

参数设置：

• max_num_seqs=3
• init_reqs=5
• total_reqs=8
• 第15步加入新请求

关键观察点：

1. 新请求到达时（第19帧）立即获得执行权
2. 正在decode的请求会被暂时搁置
3. 体现了"尽快开始新请求"的策略

场景2：资源抢占

参数设置：

• num_kvcache_blocks=3
• kvcache_block_size=10

关键帧分析（第9帧）：

1. 总资源限制为30 tokens
2. 当运行中请求总长度超过限制时
3. 执行队列尾部的请求被移回等待队列
4. 体现了"公平性"与"资源保障"的平衡

场景3：Decode优先

行为特点：

1. 一旦请求进入decode阶段就会持续执行
2. 新请求必须等待当前decode完成
3. 适合对延迟不敏感但对吞吐敏感的场景

4. 进阶思考与优化方向

4.1 当前实现的局限性

虽然我们的基础调度器已经能演示核心原理，但与生产级系统相比还有明显差距：

1. 缺乏混合执行：无法同时处理prefill和decode请求，就像餐厅不能边准备食材边上菜
2. 简单抢占策略：总是从尾部抢占，可能不是最优选择
3. 无优先级控制：所有请求平等对待，无法实现VIP通道
4. 固定块大小：实际系统可能需要动态块大小

4.2 常见问题排查指南

在实际使用中可能会遇到以下问题：

4.3 性能优化技巧

根据实际经验，分享几个有效的优化手段：

1. 动态批次大小：根据当前负载自动调整max_num_seqs

python复制# 简单实现示例
def auto_adjust_batch_size():
    avg_latency = monitor.get_avg_latency()
    if avg_latency < 50ms:
        self.max_num_seqs += 1
    elif avg_latency > 200ms:
        self.max_num_seqs = max(1, self.max_num_seqs - 1)

2. 智能抢占策略：考虑多个因素决定抢占谁

python复制def select_preempt_victim():
    # 综合考虑：生成进度、等待时间、优先级等
    return min(self.running, key=lambda x: (
        x.num_completion_tokens / x.max_tokens,
        -x.wait_time,
        -x.priority
    ))

3. Block预分配：提前分配一批block减少碎片

python复制class BlockManager:
    def __init__(self):
        self.pool = deque([Block() for _ in range(10)])  # 预热
        
    def allocate(self, seq):
        if not self.pool:
            self.pool.extend(Block() for _ in range(5))  # 动态扩容
        block = self.pool.popleft()
        # ...其余分配逻辑

4.4 扩展功能建议

当熟悉基础调度器后，可以考虑添加这些进阶功能：

1. Chunked prefill：将长prompt分块处理，缓解内存压力
2. 异步执行：解耦调度与执行，提高并行度
3. 优先级队列：实现多级服务质量(QoS)
4. 弹性KV cache：支持动态调整block大小
5. 推测解码：尝试多个生成路径提升吞吐

这些优化就像给基础引擎添加涡轮增压，可以显著提升系统性能，但也会增加复杂度。建议循序渐进地实现，并充分测试每个改动。

5. 从基础到生产级的跨越

理解这个基础调度器后，再看vLLM的实际实现会轻松很多。生产级调度器主要在这些方面做了增强：

1. 精细化的状态管理：增加了更多中间状态和转换条件
2. 混合执行支持：精心协调prefill和decode的资源分配
3. 高级调度策略：支持多种优先级和公平性算法
4. 异常处理：完善的错误恢复和资源回收机制
5. 性能监控：丰富的metrics收集和暴露接口

建议的学习路径：

1. 充分理解本文的基础实现
2. 阅读vLLM的Scheduler V0版本代码
3. 逐步研究V1版本的各个优化点
4. 尝试自己实现某个进阶功能

调度器设计就像下棋，规则简单但变化无穷。希望这个基础实现能成为你探索LLM推理系统的好起点。当你在实际项目中遇到调度问题时，不妨回想这个简化版的设计，往往能找到解决问题的灵感。