读一点vLLM

我说读一读仓库代码有益身心健康。虽然说是只读了一点和科研要改的部分有关的代码。

vllm的scheduler

我也没有完全看，只能说是cherry pick一点。

`init` 方法

资源限制：

self.max_num_running_reqs = self.scheduler_config.max_num_seqs
self.max_num_scheduled_tokens = self.scheduler_config.max_num_batched_tokens
self.max_model_len = vllm_config.model_config.max_model_len

分别限制最大并发数、每步处理的最大token数、最大上下文长度

队列管理：

# req_id -> Request
self.requests: dict[str, Request] = {}

self.waiting = create_request_queue(self.policy)
self.running: list[Request] = []
self.finished_req_ids: set[str] = set()

还是waiting队列和running队列，其中waiting队列根据policy可以选FCFS或Priority requests做了一个request_id到request的映射 finished_req_ids记录已经完成的request，每次scheduling loop结束的时候，释放相关request的资源，并重置这个记录

各种控制器：

KVCacheManager负责分配和释放PagedAttention的显存块

self.kv_cache_manager = KVCacheManager(
            kv_cache_config=kv_cache_config,
            max_model_len=self.max_model_len,
            enable_caching=bool(self.cache_config.enable_prefix_caching),
            use_eagle=self.use_eagle,
            log_stats=self.log_stats,
            enable_kv_cache_events=self.enable_kv_cache_events,
            dcp_world_size=self.dcp_world_size,
        )

KVConnector用于传递remote KV（比如PD分离场景）

self.connector = KVConnectorFactory.create_connector(
    config=self.vllm_config,
    role=KVConnectorRole.SCHEDULER,
    kv_cache_config=self.kv_cache_config,
)

调度器不仅仅看本地显存，还会通过 connector 询问：“这个请求的 KV Cache 在别的机器上有吗？” （比如PD分离场景需要迁移KV）

EncoderManager 用于管理encoder占的资源（比如多模态场景）

self.encoder_cache_manager = EncoderCacheManager(cache_size=encoder_cache_size)

`schedule` 方法

大量省略号是我没有粘贴过来的代码

不区分prefill和decode阶段，而是采用num_computed_tokens（KVCache里已经有的token）和num_tokens_with_spec（len(prompt_token_ids) + len(output_token_ids) + len(spec_token_ids) ）来描述。

在调度时，首先调度running的request，然后再调度waiting的request。注意running里也可能有prefill request，比如说发生了chunked prefill，这个request已执行了一半的prefill，需要调度上来执行剩下的prefill。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49


while req_index < len(self.running) and token_budget > 0:
    request = self.running[req_index]

#计算该request会产生多少新token
    num_new_tokens = (
        request.num_tokens_with_spec
        + request.num_output_placeholders
        - request.num_computed_tokens
    )
    if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens:
        num_new_tokens = self.scheduler_config.long_prefill_token_threshold
    num_new_tokens = min(num_new_tokens, token_budget)

# 计算encoder部分的cache开销
...

# 分配KVCache空间
    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
# 如果放不下，驱逐running队列里优先级最低的到waiting队列，继续在running队列尝试调度下一个任务
...
        preempted_req = max(
            self.running,
            key=lambda r: (r.priority, r.arrival_time),
         )
         self.running.remove(preempted_req)
...


# 把任务调度上去
     scheduled_running_reqs.append(request)
     req_to_new_blocks[request.request_id] = new_blocks
     num_scheduled_tokens[request.request_id] = num_new_tokens
     token_budget -= num_new_tokens
     req_index += 1

# 投机解码相关
...

# 分配encoder的cache空间
...

如果前面running的任务全都能调度上来（即没有发生驱逐running队列），那就开始调度waiting队列

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20


if not preempted_reqs:
      while self.waiting and token_budget > 0:
           request = self.waiting.peek_request() # 拿一个waiting的任务

# 各种对KV Transfer的讨论，比如说如果远端KV Cache还没到位，那么跳过这个request
...

# 各种load cache
...

# 分配显存，此时遇到无法分配显存的时候就直接退出。
new_blocks = self.kv_cache_manager.allocate_slots(...)
if new_blocks is None:
    # The request cannot be scheduled.
    break
...
request = self.waiting.pop_request()
req_index += 1
self.running.append(request)
...

最后，把调度上的任务整理并发送出去

1
2
3


scheduler_output = SchedulerOutput(...)
...
return scheduler_output

vllm的cuda graph实现

在gpu_model_runner.py的GPUModelRunner内的load_model函数中，对self.model进行了一层wrap，分别有CUDAGraphWrapper和UBatchWrapper两种wrapper。

UBatchWrapper暂不讲，其主要用于Dual Batch Overlap，具体来说是MoE模型开DP+EP的时候对all2all通信做overlap。

CUDAGraphWrapper中，用self.runnable存被包装的模型，用self.concrete_cudagraph_entries来存cudagraph的入口。这个wrapper是包在nn.Module上，要支持__call__函数取代原来的__call__做模型的前向传播，我们就主要看这个__call__。在执行前向传播时，会用带上一个context，存batch_descriptor，此时将其取来，其记录了当前前向传播的token_len等信息，用于决定使用哪个cudagraph，其作为key从self.concrete_cudagraph_entries这个dict中找到对应的entry，即entry = self.concrete_cudagraph_entries[batch_descriptor]。

在entry为None时，执行cudagraph的capture：抽取最重要的代码实际上就三行：

1
2
3


cudagraph = torch.cuda.CUDAGraph()
with torch.cuda.graph(cudagraph, pool=self.graph_pool):
    output = self.runnable(*args, **kwargs)

这里做的优化是对不同token_len的输入，共用了一块graph_pool作为显存地址（因为同一时间只有一个cudagraph会被触发，所以这样是安全的）。

entry非空，可以replay，只需执行entry.cudagraph.replay()，然后return entry.output。

在vllm启动的时候，能看到在capture cudagraph，此时实际上是在跑dummy input，其代码可见GPUModelRunner的_dummy_run函数。

读一点vLLM

vllm的scheduler

__init__ 方法

schedule 方法

vllm的cuda graph实现

`init` 方法

`schedule` 方法