vLLM主要模块Scheduler详解
在 vLLM 中有许多的模块，而在这篇文章中，我们主要来介绍 vLLM 中如调度管理prompt的。
本文章是按照vLLM版本：v0.11.0
官方代码入口
python
# 以下是官方调用代码
from vllm import LLM, SamplingParams
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
llm = LLM(model="facebook/opt-125m")
outputs = llm.generate(prompts, sampling_params)我们从这里为入口了解vLLM是如何调用scheduler的。
bash
scheduler是vLLM最重要模块之一，它负责对prompts，KV cache的调度和管理vllm/entrypoints/llm.py
我们从源码得知llm.generate 是调用的，vllm/entrypoints/llm.py的generate方法，
python
class LLM:
def gennerate(...) -> list[RequestOutput]:
... # 省略其他代码
self._validate_and_add_requests(
prompts=prompts,
params=sampling_params,
use_tqdm=use_tqdm,
lora_request=lora_request,
priority=priority,
)
outputs = self._run_engine(use_tqdm=use_tqdm)
return self.engine_class.validate_outputs(outputs, RequestOutput)vllm/v1/engine/llm_engine.py
再按照self._validate_and_add_requests ，可以最终看到，vllm/v1/engine/llm_engine.py的add_request,我们要关注request的流向。
python
class LLMEngine:
def add_request(...) -> None:
... # 省略其他代码
n = params.n if isinstance(params, SamplingParams) else 1
if n == 1:
# Make a new RequestState and queue.
self.output_processor.add_request(request, prompt_text, None, 0)
# Add the request to EngineCore.
self.engine_core.add_request(request)
return
...vllm/v1/engine/core.py
我们需要查看到self.engine_core.add_request(request) ，
最终调用了vllm/v1/engine/core.py中EngineCore中的add_request方法并在其中调用了self.scheduler.add_request(request)
python
# vllm/v1/core/sched/scheduler.py
class Scheduler(SchedulerInterFace):
def add_request(self, request: Request) -> None:
self.waiting.add_request(request) # 重要代码
self.requests[request.request_id] = request
if self.log_stats:
request.record_event(EngineCoreEventType.QUEUED)当获得prompt后，执行这些prompt是self._run_engine，最终会调用到vllm/v1/engine/core.py的step()
python
class EngineCore:
def step(...):
...
with record_function_or_nullcontext("core step: schedule"):
scheduler_output = self.scheduler.schedule() # 重要代码
...最终选择哪个prompt，执行仍然是调用到self.scheduler.schedule()
总结
从以上的代码可以知道以下重要知识点：
- 上层llm.generate是最终调用到底层的scheduler中。
- 上层的prompt会被包装成request，在V1中首先存入到waiting中。
- waiting是scheduler中重要的队列，会根据Policy是FCFS还是PRIORITY，来决定是双端队列还是普通队列
- waiting是来存储准备要做推理的prompt，prefill阶段都是来自于此
Scheduler.schedule()重要代码详解
重要变量概念
在解释这个算法前，需要了解几个变量，num_token_with_spec, num_computed_tokens, token_budget,
整个算法都是围绕这些变量变化做的,并且这个算法都是将prompts进行token化，对prompts进行管理.
python
num_token_with_spec:len(prompt_token_ids) + len(output_token_ids) + len(spec_token_ids)
token_budget=max_num_batched_tokens,单次迭代处理的最大token数量
num_computed_tokens:被计算过的 KV cache的token数算法整体思路
首先说说schedule的意图和整体思路：先后看running和waiting中所有request，为他们分配合理的显存blocks,后续会说如何分配block。
running的调度
接下来我们来看看源码中是如何处理request的。
python
def schedule():
...
# 首先执行running
req_index = 0
while req_index < len(self.running) and token_budget > 0:
request = self.running[req_index]
# 获取所需新token数
num_new_tokens = (
request.num_tokens_with_spec
+ request.num_output_placeholders
- request.num_computed_tokens
)
...
with record_function_or_nullcontext("schedule: allocate_slots"):
while True:
# 分配显存
new_blocks = self.kv_cache_manager.allocate_slots(
request,
num_new_tokens,
num_lookahead_tokens=self.num_lookahead_tokens,
)
# 如果分配成功，退出此循环
if new_blocks is not None:
# The request can be scheduled.
break
# 分配不成功，按照策略选择一个request，放入waiting队列中
if self.policy == SchedulingPolicy.PRIORITY:
preempted_req = max(
self.running,
key=lambda r: (r.priority, r.arrival_time),
)
self.running.remove(preempted_req)
if preempted_req in scheduled_running_reqs:
scheduled_running_reqs.remove(preempted_req)
token_budget += num_scheduled_tokens[
preempted_req.request_id
]
req_to_new_blocks.pop(preempted_req.request_id)
num_scheduled_tokens.pop(preempted_req.request_id)
req_index -= 1
else:
preempted_req = self.running.pop()
self.kv_cache_manager.free(preempted_req)
self.encoder_cache_manager.free(preempted_req)
preempted_req.status = RequestStatus.PREEMPTED
preempted_req.num_computed_tokens = 0
preempted_req.num_preemptions += 1
# 存入waiting和被抢占队列中
self.waiting.prepend_request(preempted_req)
preempted_reqs.append(preempted_req)
...
token_budget -= num_new_tokens
req_index += 1
...算法总结
running队列中从第一个request开始，目的是为它分配显存block，
如果blocks不满足条件，从按照一定的策略（FCFS或PRIORITY），从running总挑选一个request，通过释放它的显存，再将被挑出来的request放入waiting队列第一个位置，同时也存放到preempted_reqs中，从而满足一开始的request显存要求，直到满足分配完所有running队列并且token_budget>0
waiting的调度
python
# 接着是waiting
if not preempted_reqs:
while self.waiting and token_budget > 0:
if len(self.running) == self.max_num_running_reqs:
break
# 取出第一个但不移除的request
request = self.waiting.peek_request()
...
# 如果new_blocks为空提前结束对waiting的处理
if new_blocks is None:
# The request cannot be scheduled.
break
...
request = self.waiting.pop_request()
# 存入running中
self.running.append(request)
...算法总结
在新抢占队列preempted_reqs为空时，waiting和token_budget都有满足条件时，取出第一个但不移除的request，为其分配显存block。
这里request有两类，一是在prefill阶段的token，二是在decode阶段，所以也可看出并没有明显区分这两个阶段，由于running不满足显存条件从而被抢占后被存入waiting中的token，如果当前request分配当显存不够时，就会提前结束对waiting的处理，
反之则会继续，将当前request从waiting转移到running中。
num_new_tokens的含义
以上算法中，对于到底需要分配多少block，有一个共同的公式。
python
num_new_tokens = (
request.num_tokens_with_spec
+ request.num_output_placeholders
- request.num_computed_tokens
)
request.num_tokens_with_spec:len(prompt_token_ids) + len(output_token_ids) + len(spec_token_ids)
prompt_token_ids:当前prompt转成token数。
output_token_ids:已经生成的回答token数。从prefill开始产生第一个token开始计算，直到decode结束，所以的动态增加的，output_token_ids.append(token_ids)
spec_token_ids:这个与推测解码相关。
request.num_output_placeholders:为输出预先分配的占位符。
request.num_computed_tokens:计算过的（被分配过显存了）token数。num_new_tokens代表是需要被分配显存block数。
以上内容基本就是vLLM的调度策略。接下来会对kv cache manager进行学习。