【vllm】(三)vLLM v1 Core — 模块超深度逐行分析之二

第五章: Scheduler.schedule() WAITING 调度逻辑逐行解析

5.1 进入 WAITING 调度的先决条件

python 复制代码
# Next, schedule the WAITING requests.
if not preempted_reqs and self._pause_state == PauseState.UNPAUSED:

逐行解释:

  • 只有在没有请求被抢占 (not preempted_reqs)且调度器处于未暂停状态 (UNPAUSED)时,才会进入 WAITING 调度流程。
  • 设计意图: 如果本轮已经在抢占 RUNNING 请求,说明 KV cache 压力很大,再从 WAITING 调入新请求只会加剧竞争,形成恶性循环。暂停状态(如 PAUSED_ALL)直接跳过一切调度。
python 复制代码
    step_skipped_waiting = create_request_queue(self.policy)
  • 创建本轮的临时跳过队列 step_skipped_waiting,使用与主队列相同的调度策略(FCFS 或 PRIORITY)。
  • 设计意图 : 被异步依赖或约束阻塞的请求暂存于此,本轮结束后会 prepend 回 self.skipped_waiting,保证下次调度时优先处理。

5.2 WAITING 调度主循环

python 复制代码
    while (self.waiting or self.skipped_waiting) and token_budget > 0:
        if len(self.running) == self.max_num_running_reqs:
            break
  • 循环条件 : WAITING 队列或跳过队列非空 且 token 预算仍有剩余。
  • 并发上限检查 : 如果 RUNNING 数已达 max_num_running_reqs(即 max_num_seqs),立即终止------新请求加入必须有空位。
  • 设计意图 : 这是一个多约束退出模式------预算耗尽、并发打满、KV 不足都会终止调度。

5.3 选择调度队列

python 复制代码
        request_queue = self._select_waiting_queue_for_scheduling()
        assert request_queue is not None
  • 调用 _select_waiting_queue_for_scheduling() 决定从 self.waiting 还是 self.skipped_waiting 取请求。
  • 断言确保至少有一个非空队列(循环条件保证)。

_select_waiting_queue_for_scheduling 逻辑(行1576-1586):

  • FCFS 模式 : 优先从 skipped_waiting 取(跳过队列优先,因为这些请求之前因异步依赖被跳过,现在可能已就绪),否则从 waiting 取。
  • PRIORITY 模式 : 当两个队列都非空时,比较队首请求的优先级(< 运算符比较 priority 和 arrival_time),选择优先级更高的。这确保了高优先级请求不会因为被跳过而饿死。

5.4 窥视请求与状态提升

python 复制代码
        request = request_queue.peek_request()
        request_id = request.request_id
  • peek (窥视)而非 pop:先检查请求是否可调度,确认后才移出队列。这是乐观检查 + 延迟提交模式。
python 复制代码
        # try to promote blocked statuses while traversing skipped queue.
        if self._is_blocked_waiting_status(
            request.status
        ) and not self._try_promote_blocked_waiting_request(request):
  • 检查请求是否处于阻塞状态,并尝试提升(promote)到可调度状态。
  • _is_blocked_waiting_status 判断三种阻塞状态:
    • WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR: 等待结构化输出语法编译
    • WAITING_FOR_REMOTE_KVS: 等待远程 KV 异步加载完成
    • WAITING_FOR_STREAMING_REQ: 等待流式输入的下一个 chunk
python 复制代码
            if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
                logger.debug(
                    "%s is still in WAITING_FOR_REMOTE_KVS state.",
                    request_id,
                )
            request_queue.pop_request()
            step_skipped_waiting.prepend_request(request)
            continue
  • 如果提升失败(依赖尚未满足),将请求从当前队列弹出,加入本轮跳过队列,继续处理下一个请求。
  • 设计意图: 阻塞请求不能阻塞整个调度流程,跳过它让其他请求有机会被调度。

5.5 LoRA 约束检查

python 复制代码
        # Check that adding the request still respects the max_loras constraint.
        if (
            self.lora_config
            and request.lora_request
            and (
                len(scheduled_loras) == self.lora_config.max_loras
                and request.lora_request.lora_int_id not in scheduled_loras
            )
        ):
            request_queue.pop_request()
            step_skipped_waiting.prepend_request(request)
            continue
  • LoRA 适配器数量有上限(max_loras),如果当前已调度的 LoRA 数达到上限且该请求需要的 LoRA 不在其中,则跳过。
  • 设计意图: GPU 上同时加载的 LoRA 适配器数量有限制,超出则不能调度。
  • 注意: 这里用 continue 而不是 break,允许后面的请求(可能使用已加载的 LoRA)被调度。

5.6 获取已缓存 Token(前缀缓存 + KV Connector)

python 复制代码
        num_external_computed_tokens = 0
        load_kv_async = False
        connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0
  • 初始化四个变量:
    • num_external_computed_tokens: 远程(外部)已计算的 token 数
    • load_kv_async: 是否需要异步加载远程 KV
    • connector_prefix_cache_queries: 连接器前缀缓存查询总数
    • connector_prefix_cache_hits: 连接器前缀缓存命中数
python 复制代码
        # Get already-cached tokens.
        if request.num_computed_tokens == 0:
  • 条件: 请求尚未计算任何 token(首次调度或抢占后重新调度)。
5.6.1 本地前缀缓存查询
python 复制代码
            new_computed_blocks, num_new_local_computed_tokens = (
                self.kv_cache_manager.get_computed_blocks(request)
            )
  • get_computed_blocks: 在本地 KV cache 中查找请求 prompt 的前缀匹配,返回匹配的 block 列表和匹配的 token 数。
  • 设计意图: 利用前缀缓存(prefix caching)跳过已计算的 prompt 前缀,节省计算。例如,多个请求共享 system prompt 时,只需计算一次。
5.6.2 KV Connector 远程缓存查询
python 复制代码
            if self.connector is not None:
                ext_tokens, load_kv_async = (
                    self.connector.get_num_new_matched_tokens(
                        request, num_new_local_computed_tokens
                    )
                )
  • 如果配置了 KV Connector(用于 P/D 分离或 KV 卸载),查询远程 KV cache 中有多少匹配 token。
  • get_num_new_matched_tokens 返回:
    • ext_tokens: 远程匹配的 token 数(可能为 None 表示暂时无法确定)
    • load_kv_async: 是否需要异步加载
python 复制代码
                if ext_tokens is None:
                    request_queue.pop_request()
                    step_skipped_waiting.prepend_request(request)
                    continue
  • 如果连接器返回 None,说明暂时无法确定远程匹配数(例如正在建立连接),将请求跳过,下轮重试。
python 复制代码
                num_external_computed_tokens = ext_tokens
                connector_prefix_cache_queries = (
                    request.num_tokens - num_new_local_computed_tokens
                )
                connector_prefix_cache_hits = num_external_computed_tokens
  • 记录外部匹配 token 数和前缀缓存统计。
  • queries = 总 token 数 - 本地已缓存数 = 需要从远程查找的 token 数。
  • hits = 远程实际匹配的 token 数。
5.6.3 计算总已缓存 token 数
python 复制代码
            num_computed_tokens = (
                num_new_local_computed_tokens + num_external_computed_tokens
            )
            assert num_computed_tokens <= request.num_tokens
  • 总已缓存 token = 本地 + 外部。
  • 断言确保不超过请求总 token 数。
python 复制代码
            if request.prefill_stats is not None:
                assert num_computed_tokens <= request.num_prompt_tokens
                request.prefill_stats.set(
                    num_prompt_tokens=request.num_prompt_tokens,
                    num_local_cached_tokens=num_new_local_computed_tokens,
                    num_external_cached_tokens=num_external_computed_tokens,
                )
  • 如果请求有 prefill 统计对象(仅首次调度时设置),记录缓存命中信息用于可观测性。
5.6.4 已有计算 token 的情况
python 复制代码
        else:
            new_computed_blocks = self.kv_cache_manager.empty_kv_cache_blocks
            num_new_local_computed_tokens = 0
            num_computed_tokens = request.num_computed_tokens
  • 当 request.num_computed_tokens > 0 时进入此分支。
  • 场景 : KV Transfer 完成后,请求的 num_computed_tokens 已被更新为远程加载的 token 数,回到 WAITING 队列等待调度。
  • 此时不需要重新查前缀缓存,已有 block 信息在 kv_cache_manager 中。

5.7 Encoder 输入调度

python 复制代码
        encoder_inputs_to_schedule = None
        external_load_encoder_input = []
        new_encoder_compute_budget = encoder_compute_budget
  • encoder_inputs_to_schedule: 本步需要计算的 encoder 输入索引列表
  • external_load_encoder_input: 需要从 EC Connector 远程加载的 encoder 输入索引列表
  • new_encoder_compute_budget: 更新后的 encoder 计算预算

5.8 异步 KV 加载分支

python 复制代码
        if load_kv_async:
            assert num_external_computed_tokens > 0
            num_new_tokens = 0
  • 如果 KV Connector 决定异步加载远程 KV,不调度任何新 token。
  • 设计意图: 异步加载时,KV 数据尚未就绪,不能进行计算。只需分配内存块,等待加载完成后在后续步骤中调度。

5.9 同步调度分支:计算新 token 数

python 复制代码
        else:
            num_new_tokens = request.num_tokens - num_computed_tokens
            threshold = self.scheduler_config.long_prefill_token_threshold
            if 0 < threshold < num_new_tokens:
                num_new_tokens = threshold
  • 新 token 数 = 请求总 token 数 - 已计算 token 数。
  • 如果设置了 long_prefill_token_threshold 且新 token 数超过阈值,截断到阈值(实现分块 prefill)。
python 复制代码
            if (
                not self.scheduler_config.enable_chunked_prefill
                and num_new_tokens > token_budget
            ):
                break
  • 关键 : 如果未启用 chunked prefill,且当前 token 预算不够一次性处理所有新 token,直接终止 WAITING 调度(break)。
  • 设计意图: 非分块模式下,请求要么全部 prefill,要么不调度。不允许部分 prefill。
python 复制代码
            num_new_tokens = min(num_new_tokens, token_budget)
            assert num_new_tokens > 0
  • 取新 token 数与剩余预算的较小值。
  • 断言确保至少有一个 token 要调度。

5.10 Encoder 输入调度(WAITING 分支)

python 复制代码
            if request.has_encoder_inputs:
                (
                    encoder_inputs_to_schedule,
                    num_new_tokens,
                    new_encoder_compute_budget,
                    external_load_encoder_input,
                ) = self._try_schedule_encoder_inputs(
                    request,
                    num_computed_tokens,
                    num_new_tokens,
                    encoder_compute_budget,
                    shift_computed_tokens=1 if self.use_eagle else 0,
                )
                if num_new_tokens == 0:
                    break
  • 调用 _try_schedule_encoder_inputs(第八章详述)决定哪些 encoder 输入需要计算。
  • 如果调整后 num_new_tokens == 0,说明 encoder 预算或缓存不足,终止 WAITING 调度。
  • 注意: 这里用 break 而不是 continue------因为 encoder 预算是全局约束,当前请求无法满足,后续请求同样无法满足。
  • shift_computed_tokens=1 if self.use_eagle else 0: EAGLE 模式下,spec token 需要一个额外的位置偏移来正确匹配 encoder 输入范围。

5.11 Mamba Block 对齐切分

python 复制代码
            if self.need_mamba_block_aligned_split:
                num_new_tokens = self._mamba_block_aligned_split(
                    request,
                    num_new_tokens,
                    num_new_local_computed_tokens,
                    num_external_computed_tokens,
                )
                if num_new_tokens == 0:
                    break
  • 对于含 Mamba 层的混合模型,当 mamba_cache_mode == "align" 时,需要将新 token 数对齐到 block_size 的倍数。
  • 如果对齐后 token 数为 0(不足一个 block),终止调度。
  • 设计意图: Mamba 状态缓存需要 block 对齐,非对齐会导致缓存失效。

5.12 Spec Decode lookahead token 处理

python 复制代码
            effective_lookahead_tokens = (
                0 if request.num_computed_tokens == 0 else self.num_lookahead_tokens
            )
  • 首次 prefill 的请求: lookahead tokens = 0(首次 prefill 不需要 lookahead blocks)。
  • 恢复的请求 (num_computed_tokens > 0): 使用完整的 num_lookahead_tokens。
  • 设计意图: 这是 P/D 分离 + Spec Decode 的边缘情况处理。首次 prefill 时不需要为 speculative tokens 预分配 block,而恢复的 decode 请求可能带有 draft tokens。

5.13 Encoder-Decoder Cross-Attention Block 计算

python 复制代码
            num_encoder_tokens = 0
            if (
                self.is_encoder_decoder
                and request.has_encoder_inputs
                and encoder_inputs_to_schedule
            ):
                num_encoder_tokens = sum(
                    request.get_num_encoder_embeds(i)
                    for i in encoder_inputs_to_schedule
                )
  • 对于 encoder-decoder 模型,计算需要分配的 cross-attention KV block 的 token 数。
  • 设计意图: Encoder-decoder 模型的 decoder 需要额外的 cross-attention KV cache 来存储 encoder 输出,这部分也需要 block 分配。

5.14 全序列保留检查

python 复制代码
            if (
                self.scheduler_reserve_full_isl
                and not self.kv_cache_manager.can_fit_full_sequence(
                    request,
                    num_new_computed_tokens=num_new_local_computed_tokens,
                    new_computed_blocks=new_computed_blocks,
                    num_external_computed_tokens=num_external_computed_tokens,
                    num_encoder_tokens=num_encoder_tokens,
                )
            ):
                if request.has_encoder_inputs:
                    self.encoder_cache_manager.free(request)
                break
  • 当 scheduler_reserve_full_isl=True 时,调度前检查是否有足够 KV cache 空间容纳请求的完整序列。
  • 如果放不下,释放已分配的 encoder cache 并终止 WAITING 调度。
  • 设计意图: 保守策略------确保请求一旦开始就不会因 KV 不足被抢占。适用于低延迟场景,代价是吞吐量降低。

5.15 KV Block 分配

python 复制代码
            new_blocks = self.kv_cache_manager.allocate_slots(
                request,
                num_new_tokens,
                num_new_computed_tokens=num_new_local_computed_tokens,
                new_computed_blocks=new_computed_blocks,
                num_lookahead_tokens=effective_lookahead_tokens,
                num_external_computed_tokens=num_external_computed_tokens,
                delay_cache_blocks=load_kv_async,
                num_encoder_tokens=num_encoder_tokens,
            )
  • 调用 kv_cache_manager.allocate_slots 分配 KV cache block。
  • 参数详解:
    • num_new_tokens: 本步要调度的新 token 数
    • num_new_computed_tokens: 本步新发现的本地已缓存 token 数(前缀匹配)
    • new_computed_blocks: 前缀匹配的 block 列表(复用已有 block,不重新分配)
    • num_lookahead_tokens: 为 speculative decoding 预分配的 block 数
    • num_external_computed_tokens: 远程已缓存的 token 数(需要分配 block 但不立即填充)
    • delay_cache_blocks: 异步加载时延迟缓存 block(避免在数据未到时计算 hash)
    • num_encoder_tokens: cross-attention 需要的额外 block
python 复制代码
            if new_blocks is None:
                if request.has_encoder_inputs:
                    self.encoder_cache_manager.free(request)
                break
  • 分配失败(None)= KV cache 容量不足。
  • 释放已分配的 encoder cache,终止 WAITING 调度。
  • 注意: WAITING 队列中不做抢占------不像 RUNNING 那样抢占低优先级请求来腾空间。

5.16 KV Connector 状态更新

python 复制代码
            if self.connector is not None:
                self.connector.update_state_after_alloc(
                    request,
                    self.kv_cache_manager.get_blocks(request_id),
                    num_external_computed_tokens,
                )
                if (
                    self.connector_prefix_cache_stats is not None
                    and connector_prefix_cache_queries != 0
                ):
                    self.connector_prefix_cache_stats.record(
                        num_tokens=connector_prefix_cache_queries,
                        num_hits=connector_prefix_cache_hits,
                        preempted=request.num_preemptions > 0,
                    )
  • 分配成功后,更新 KV Connector 的内部状态(block 表、外部 token 数等)。
  • Connector 使用这些信息在 _build_kv_connector_meta 阶段决定是否需要发起 KV 加载/存储操作。
  • 记录前缀缓存统计数据(如果启用了统计)。

5.17 从队列移除请求并更新状态

python 复制代码
            request = request_queue.pop_request()
  • 此刻才真正从队列移除 ------之前都是 peek,现在确认可调度后才 pop。
5.17.1 异步加载路径
python 复制代码
            if load_kv_async:
                request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
                step_skipped_waiting.prepend_request(request)
                request.num_computed_tokens = num_computed_tokens
                continue
  • 异步加载时,请求状态变为 WAITING_FOR_REMOTE_KVS,放入跳过队列。
  • 关键 : 设置 num_computed_tokens = num_computed_tokens(本地 + 外部匹配数),尽管 KV 尚未实际加载。
  • 设计意图 : 这个"预设置"使得在 KV 加载完成后,num_computed_tokens 已正确反映实际计算量。如果加载失败,_update_requests_with_invalid_blocks 会修正此值。
  • continue 回到循环顶部,处理下一个请求。
5.17.2 正常调度路径
python 复制代码
            self.running.append(request)
            if self.log_stats:
                request.record_event(
                    EngineCoreEventType.SCHEDULED, scheduled_timestamp
                )
  • 请求移入 RUNNING 队列。
  • 如果启用统计,记录 SCHEDULED 事件时间戳。
python 复制代码
            if request.status == RequestStatus.WAITING:
                scheduled_new_reqs.append(request)
            elif request.status == RequestStatus.PREEMPTED:
                scheduled_resumed_reqs.append(request)
            else:
                raise RuntimeError(f"Invalid request status: {request.status}")
  • 根据先前状态分类:
    • WAITING → 新请求(首次调度)
    • PREEMPTED → 被抢占后恢复的请求
    • 其他状态 → 运行时错误
python 复制代码
            if self.lora_config and request.lora_request:
                scheduled_loras.add(request.lora_request.lora_int_id)
  • 将请求的 LoRA ID 加入已调度 LoRA 集合,用于后续请求的 LoRA 约束检查。
python 复制代码
            req_to_new_blocks[request_id] = self.kv_cache_manager.get_blocks(
                request_id
            )
            num_scheduled_tokens[request_id] = num_new_tokens
            token_budget -= num_new_tokens
  • 记录请求的 block 映射和调度 token 数。
  • 从 token 预算中扣除本请求消耗的 token 数。
python 复制代码
            request.status = RequestStatus.RUNNING
            request.num_computed_tokens = num_computed_tokens
  • 状态更新为 RUNNING。
  • 设置 num_computed_tokens = 本地已缓存 + 外部已缓存(注意: 不包含本步新调度的 token,那些在 _update_after_schedule 中累加)。

5.18 Encoder 缓存分配(WAITING 路径)

python 复制代码
            if encoder_inputs_to_schedule:
                scheduled_encoder_inputs[request_id] = encoder_inputs_to_schedule
                for i in encoder_inputs_to_schedule:
                    self.encoder_cache_manager.allocate(request, i)
                    if self.ec_connector is not None:
                        self.ec_connector.update_state_after_alloc(request, i)
                encoder_compute_budget = new_encoder_compute_budget
            if external_load_encoder_input:
                for i in external_load_encoder_input:
                    self.encoder_cache_manager.allocate(request, i)
                    if self.ec_connector is not None:
                        self.ec_connector.update_state_after_alloc(request, i)
  • 为需要计算的 encoder 输入分配缓存空间。
  • 两类 encoder 输入:
    • encoder_inputs_to_schedule: 本地需要计算的
    • external_load_encoder_input: 从 EC Connector 远程加载的
  • 更新 encoder 计算预算。
  • 通知 EC Connector 状态变更。

5.19 跳过队列回填

python 复制代码
        # re-queue requests skipped in this pass ahead of older skipped items.
        if step_skipped_waiting:
            self.skipped_waiting.prepend_requests(step_skipped_waiting)
  • 本轮结束后,将本轮跳过的请求前插 到 skipped_waiting 队列。
  • 设计意图 : prepend 而非 append 确保这些请求在下一轮调度时优先被检查------它们已经等待了一轮,不应排在更老的跳过请求后面。

第六章: Scheduler.schedule() 后半段 --- 约束检查、Common Prefix、SchedulerOutput 构建

6.1 调度约束断言验证

python 复制代码
        # Check if the scheduling constraints are satisfied.
        total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
        assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens

        assert token_budget >= 0
        assert len(self.running) <= self.max_num_running_reqs
        assert len(scheduled_new_reqs) + len(scheduled_resumed_reqs) + len(
            scheduled_running_reqs
        ) <= len(self.running)
  • 四项开发期断言,生产环境不执行:

    1. 总调度 token 数 ≤ 最大允许调度 token 数
    2. token 预算非负
    3. RUNNING 请求数 ≤ 最大并发数
    4. 被调度的请求总数 ≤ RUNNING 队列长度(有些 RUNNING 请求可能本轮未被调度)
  • 第4项的含义 : 并非所有 RUNNING 请求都会在每个 step 被调度。例如,异步调度模式下已达到 max_tokens 的请求会跳过。

6.2 最长公共前缀计算

python 复制代码
        # Get the longest common prefix among all requests in the running queue.
        # This can be potentially used for cascade attention.
        num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
        with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"):
            if self.running:
                any_request_id = self.running[0].request_id
                num_common_prefix_blocks = (
                    self.kv_cache_manager.get_num_common_prefix_blocks(any_request_id)
                )
  • 为每个 KV cache group 计算 RUNNING 队列中所有请求的最长公共前缀 block 数。
  • 设计意图 : 公共前缀可用于 cascade attention(级联注意力)优化------多个请求共享前缀时,只需对前缀做一次 attention,结果被所有请求复用。
  • 传入 any_request_id: 因为所有 RUNNING 请求的 block 表都在 kv_cache_manager 中,manager 可以遍历所有请求来计算公共前缀。传入任意一个 ID 是实现上的简化。
  • 初始化为全 0 列表(长度 = KV cache group 数),处理空 RUNNING 队列的情况。

6.3 SchedulerOutput 构建

6.3.1 V2 Model Runner 适配
python 复制代码
        if self.use_v2_model_runner:
            scheduled_new_reqs = scheduled_new_reqs + scheduled_resumed_reqs
            scheduled_resumed_reqs = []
            new_reqs_data = [
                NewRequestData.from_request(
                    req,
                    req_to_new_blocks[req.request_id].get_block_ids(),
                    req._all_token_ids,
                )
                for req in scheduled_new_reqs
            ]
  • V2 Model Runner 将新请求和恢复请求合并 为 scheduled_new_reqs,恢复请求列表清空。
  • 为每个请求构建 NewRequestData,包含 prefill_token_ids(req._all_token_ids)------V2 需要完整的 token ID 列表来重建输入。
6.3.2 V1 Model Runner 路径
python 复制代码
        else:
            new_reqs_data = [
                NewRequestData.from_request(
                    req, req_to_new_blocks[req.request_id].get_block_ids()
                )
            for req in scheduled_new_reqs
            ]
  • V1 路径不传 prefill_token_ids,恢复请求单独处理。
6.3.3 构建缓存请求数据
python 复制代码
        with record_function_or_nullcontext("schedule: make_cached_request_data"):
            cached_reqs_data = self._make_cached_request_data(
                scheduled_running_reqs,
                scheduled_resumed_reqs,
                num_scheduled_tokens,
                scheduled_spec_decode_tokens,
                req_to_new_blocks,
            )
  • 调用 _make_cached_request_data(第八章详述)构建已缓存请求的增量数据。
  • 设计意图: Worker 进程已缓存了这些请求的基本信息,只需发送差异(新 block、新 token 等)以减少跨进程通信开销。
6.3.4 记录本轮调度请求 ID
python 复制代码
        self.prev_step_scheduled_req_ids.clear()
        self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
  • 清空并更新"上一步调度的请求 ID 集合"。
  • 用途 : 在 _make_cached_request_data 中判断请求是否在上一步被调度过------如果是,Worker 已有最新 token ID,不需要重复发送 all_token_ids。
6.3.5 新 Block 归零列表
python 复制代码
        new_block_ids_to_zero = (
            (self.kv_cache_manager.take_new_block_ids() or None)
            if self.needs_kv_cache_zeroing
            else None
        )
  • 如果 KV cache 需要归零(needs_kv_cache_zeroing,例如某些 Mamba/SSM 模型),取出本轮新分配的 block ID 列表。
  • Worker 会在使用这些 block 前将 GPU 内存清零,防止脏数据影响计算。
  • take_new_block_ids() 是一次性取出并清空内部列表的操作(消费者模式)。
  • 无新 block 时返回 None。
6.3.6 构建 SchedulerOutput 对象
python 复制代码
        scheduler_output = SchedulerOutput(
            scheduled_new_reqs=new_reqs_data,
            scheduled_cached_reqs=cached_reqs_data,
            num_scheduled_tokens=num_scheduled_tokens,
            total_num_scheduled_tokens=total_num_scheduled_tokens,
            scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
            scheduled_encoder_inputs=scheduled_encoder_inputs,
            num_common_prefix_blocks=num_common_prefix_blocks,
            preempted_req_ids={req.request_id for req in preempted_reqs},
            finished_req_ids=self.finished_req_ids,
            free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
            new_block_ids_to_zero=new_block_ids_to_zero,
        )
  • SchedulerOutput 字段详解 :
    • scheduled_new_reqs: 首次调度请求的完整数据(prompt、block IDs、参数等)
    • scheduled_cached_reqs: 已缓存请求的增量数据
    • num_scheduled_tokens: 每个请求本轮调度的 token 数(dict[str, int])
    • total_num_scheduled_tokens: 总调度 token 数
    • scheduled_spec_decode_tokens: 每个请求的 spec decode token IDs
    • scheduled_encoder_inputs: 每个请求需要计算的 encoder 输入索引
    • num_common_prefix_blocks: 每个 KV cache group 的公共前缀 block 数
    • preempted_req_ids: 被抢占的请求 ID 集合
    • finished_req_ids: 上一步到本步之间完成的请求 ID 集合
    • free_encoder_mm_hashes: 需要从 encoder cache 释放的多模态 hash 列表
    • new_block_ids_to_zero: 需要归零的新 block ID 列表

6.4 KV Connector Metadata 构建

python 复制代码
        if self.connector is not None:
            meta = self._build_kv_connector_meta(self.connector, scheduler_output)
            scheduler_output.kv_connector_metadata = meta
  • 调用 connector.build_connector_meta(scheduler_output) 构建 KV Connector 的元数据。
  • 设计意图 (注释引用):
    1. 规划 KV cache 存储: 决定哪些请求的 KV 需要推送到远程
    2. 包装所有 KV load/save 操作: 生成不透明的元数据对象,Worker 据此执行实际 IO
    3. 清理连接器内部状态: 重置临时跟踪变量

6.5 EC Connector Metadata 构建

python 复制代码
        if self.ec_connector is not None:
            ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta(
                scheduler_output
            )
            scheduler_output.ec_connector_metadata = ec_meta
  • 类似 KV Connector,为 Encoder Cache Connector 构建元数据。
  • EC Connector 处理 encoder 输出的跨节点传输(如视觉编码器输出)。

6.6 调度后状态更新

python 复制代码
        with record_function_or_nullcontext("schedule: update_after_schedule"):
            self._update_after_schedule(scheduler_output)
        return scheduler_output
  • 调用 _update_after_schedule(下一节详述)更新请求的 num_computed_tokens 等状态。
  • 返回构建好的 SchedulerOutput,传递给 Engine Core 的执行层。

第七章: Scheduler.update_from_output() 逐行深度解析

7.1 方法签名与输入

python 复制代码
def update_from_output(
    self,
    scheduler_output: SchedulerOutput,
    model_runner_output: ModelRunnerOutput,
) -> dict[int, EngineCoreOutputs]:
  • 输入 :
    • scheduler_output: 上一步的调度输出(包含哪些请求被调度、调度了多少 token)
    • model_runner_output: Model Runner 执行后的输出(采样 token、logprobs、KV Connector 结果等)
  • 输出 : dict[int, EngineCoreOutputs] --- 按 client_index 分组的引擎核心输出。每个客户端(前端连接)收到自己的请求输出。

7.2 解包 Model Runner 输出

python 复制代码
sampled_token_ids = model_runner_output.sampled_token_ids
logprobs = model_runner_output.logprobs
prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
num_scheduled_tokens = scheduler_output.num_scheduled_tokens
pooler_outputs = model_runner_output.pooler_output
num_nans_in_logits = model_runner_output.num_nans_in_logits
kv_connector_output = model_runner_output.kv_connector_output
cudagraph_stats = model_runner_output.cudagraph_stats
  • 变量含义 :
    • sampled_token_ids: 每个请求的采样 token ID 列表(按 batch 索引排列)
    • logprobs: 采样 token 的 log 概率
    • prompt_logprobs_dict: prompt token 的 log 概率(按请求 ID 索引)
    • num_scheduled_tokens: 来自调度输出,记录每个请求的调度 token 数
    • pooler_outputs: Pooling 模型的输出(如 embedding 模型)
    • num_nans_in_logits: 每个 request 的 NaN logits 数量
    • kv_connector_output: KV Connector 的执行结果(加载完成信号、无效 block 等)
    • cudagraph_stats: CUDA Graph 统计信息

7.3 性能统计

python 复制代码
perf_stats: PerfStats | None = None
if self.perf_metrics and self.perf_metrics.is_enabled():
    perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output)
  • 如果启用了 MFU(Model FLOPs Utilization)指标,计算每 GPU 的性能统计。

7.4 初始化输出容器

python 复制代码
outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
spec_decoding_stats: SpecDecodingStats | None = None
kv_connector_stats: KVConnectorStats | None = (
    kv_connector_output.kv_connector_stats if kv_connector_output else None
)
if kv_connector_stats and self.connector:
    kv_stats = self.connector.get_kv_connector_stats()
    if kv_stats:
        kv_connector_stats = kv_connector_stats.aggregate(kv_stats)
  • outputs: 按客户端索引分组的输出列表,使用 defaultdict(list) 自动创建空列表。
  • spec_decoding_stats: Speculative decoding 统计(跨请求累加)。
  • kv_connector_stats: KV Connector 统计,聚合 Worker 侧和 Scheduler 侧的数据。

7.5 处理 KV 加载失败

python 复制代码
failed_kv_load_req_ids = None
if kv_connector_output and kv_connector_output.invalid_block_ids:
    failed_kv_load_req_ids = self._handle_invalid_blocks(
        kv_connector_output.invalid_block_ids,
        num_scheduled_tokens,
    )
  • 如果 KV Connector 报告有无效 block (加载失败的远程 KV block),调用 _handle_invalid_blocks 处理。
  • 返回受影响的请求 ID 集合,这些请求在主循环中会被跳过。
  • 设计意图: 远程 KV 加载可能部分失败(如网络中断、源节点不可用),需要优雅处理而不是崩溃。

7.6 主循环: 逐请求处理模型输出

python 复制代码
stopped_running_reqs: set[Request] = set()
stopped_preempted_reqs: set[Request] = set()
for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
    assert num_tokens_scheduled > 0
  • 遍历每个被调度的请求。
  • 两个集合分别收集从 RUNNING 和 PREEMPTED 状态停止的请求------后续需要从对应队列移除。
  • 断言: 每个被调度的请求至少调度了一个 token。
7.6.1 跳过 KV 加载失败的请求
python 复制代码
    if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
        continue
  • 如果该请求受 KV 加载失败影响,跳过处理(已在 _handle_invalid_blocks 中调整了 num_computed_tokens)。
7.6.2 跳过已完成请求
python 复制代码
    request = self.requests.get(req_id)
    if request is None or request.is_finished():
        continue
  • 请求可能在本步执行期间被外部中止(如客户端断开连接),此时跳过。
  • PP 和异步调度场景: 请求可能在模型执行过程中被标记为完成。
7.6.3 获取生成 token 和 Spec Decode 处理
python 复制代码
    req_index = model_runner_output.req_id_to_index[req_id]
    generated_token_ids = (
        sampled_token_ids[req_index] if sampled_token_ids else []
    )
  • 通过 req_id_to_index 映射获取请求在 batch 中的索引。
  • 获取该请求的采样 token ID 列表。
python 复制代码
    scheduled_spec_token_ids = (
        scheduler_output.scheduled_spec_decode_tokens.get(req_id)
    )
    if scheduled_spec_token_ids and generated_token_ids:
        num_draft_tokens = len(scheduled_spec_token_ids)
        num_accepted = len(generated_token_ids) - 1
        num_rejected = num_draft_tokens - num_accepted
  • Speculative Decoding 核心逻辑 :
    • num_draft_tokens: draft 模型生成的候选 token 数
    • num_accepted: 被目标模型接受的 token 数 = 总生成数 - 1(第一个是 verify token,不算 draft)
    • num_rejected: 被拒绝的 draft token 数 = draft 数 - 接受数
python 复制代码
        if request.num_computed_tokens > 0:
            request.num_computed_tokens -= num_rejected
        if request.num_output_placeholders > 0:
            request.num_output_placeholders -= num_rejected
  • 关键修正 : _update_after_schedule 乐观地将所有调度 token 计入 num_computed_tokens。现在根据实际拒绝数回退。
  • 同样修正 num_output_placeholders(异步调度场景)。
  • 设计意图 : Spec decode 被拒绝的 token 对应的 KV cache 位置是无效的,需要回退 num_computed_tokens 使得下轮重新计算这些位置。
python 复制代码
        spec_decoding_stats = self.make_spec_decoding_stats(
            spec_decoding_stats,
            num_draft_tokens=num_draft_tokens,
            num_accepted_tokens=num_accepted,
            num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,
            request_id=req_id,
        )
  • 累加 spec decoding 统计数据。
7.6.4 释放 Encoder 输入
python 复制代码
    if request.has_encoder_inputs:
        self._free_encoder_inputs(request)
  • 在模型实际执行后,释放不再需要的 encoder cache 条目。
  • 设计意图: 只有在 step 真正执行后,encoder 输出已被写入 KV cache,才能安全释放 encoder cache。
7.6.5 处理生成 token 和停止检查
python 复制代码
    stopped = False
    new_logprobs = None
    new_token_ids = generated_token_ids
    pooler_output = pooler_outputs[req_index] if pooler_outputs else None
    kv_transfer_params = None
    status_before_stop = request.status
  • 初始化本轮处理的变量:
    • stopped: 请求是否已停止
    • new_logprobs: 采样 logprobs
    • new_token_ids: 生成的 token IDs
    • pooler_output: pooling 模型输出
    • kv_transfer_params: KV 传输参数(完成后设置)
    • status_before_stop: 记录停止前的状态(用于分类停止请求)
python 复制代码
    if new_token_ids:
        new_token_ids, stopped = self._update_request_with_output(
            request, new_token_ids
        )
    elif request.pooling_params and pooler_output is not None:
        request.status = RequestStatus.FINISHED_STOPPED
        stopped = True
  • 有生成 token : 调用 _update_request_with_output 追加 token 并检查停止条件。
  • Pooling 模型: 只要有输出就标记为完成(pooling 模型只做前向传播,不生成 token 序列)。
7.6.6 结构化输出验证
python 复制代码
    if new_token_ids and self.structured_output_manager.should_advance(request):
        struct_output_request = request.structured_output_request
        assert struct_output_request is not None
        assert struct_output_request.grammar is not None
        if not struct_output_request.grammar.accept_tokens(
            req_id, new_token_ids
        ):
            logger.error(...)
            request.status = RequestStatus.FINISHED_ERROR
            request.resumable = False
            stopped = True
  • 对使用结构化输出(JSON schema、regex 等)的请求,验证生成的 token 是否符合语法。
  • should_advance: 只在 decode 阶段(非 prefill chunk)检查。
  • 如果语法拒绝 token,将请求标记为错误完成。
  • 设计意图: 结构化输出的 grammar 约束是硬性的------违反语法的输出不能返回给用户。
7.6.7 停止请求处理
python 复制代码
    routed_experts = None
    finish_reason = None
    if stopped:
        routed_experts = self._get_routed_experts(request)

        finish_reason = request.get_finished_reason()
        finished = self._handle_stopped_request(request)
        if finished:
            kv_transfer_params = self._free_request(request)

        if status_before_stop == RequestStatus.RUNNING:
            stopped_running_reqs.add(request)
        else:
            stopped_preempted_reqs.add(request)
  • routed_experts : 如果请求完成且启用了 enable_return_routed_experts,读取 MoE 路由专家信息。
  • finish_reason : 在 _handle_stopped_request 之前捕获------该方法可能修改状态(如流式请求继续等待)。
  • _handle_stopped_request: 处理停止的请求,返回是否真正完成(流式请求可能继续)。
  • _free_request: 如果请求真正完成,释放所有资源,返回 KV 传输参数。
  • 分类收集: 根据停止前的状态(RUNNING vs PREEMPTED),加入对应集合,后续从对应队列移除。
7.6.8 Logprobs 提取
python 复制代码
    if (
        request.sampling_params is not None
        and request.sampling_params.logprobs is not None
        and logprobs
    ):
        new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
  • 如果请求要求返回 logprobs 且模型输出了 logprobs,按请求索引和 token 数切片提取。
  • slice_request 避免返回多余数据。
7.6.9 NaN Logits 统计
python 复制代码
    if num_nans_in_logits is not None and req_id in num_nans_in_logits:
        request.num_nans_in_logits = num_nans_in_logits[req_id]
  • 记录请求的 NaN logits 数量,用于诊断和可观测性。
7.6.10 Prompt Logprobs 提取
python 复制代码
    prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
  • 从模型输出中获取该请求的 prompt logprobs。
7.6.11 构建 EngineCoreOutput
python 复制代码
    if (
        new_token_ids
        or pooler_output is not None
        or kv_transfer_params
        or stopped
    ):
        outputs[request.client_index].append(
            EngineCoreOutput(
                request_id=req_id,
                new_token_ids=new_token_ids,
                finish_reason=finish_reason,
                new_logprobs=new_logprobs,
                new_prompt_logprobs_tensors=prompt_logprobs_tensors,
                pooling_output=pooler_output,
                stop_reason=request.stop_reason,
                events=request.take_events(),
                prefill_stats=request.take_prefill_stats(),
                kv_transfer_params=kv_transfer_params,
                trace_headers=request.trace_headers,
                routed_experts=routed_experts,
                num_nans_in_logits=request.num_nans_in_logits,
            )
        )
    else:
        assert not prompt_logprobs_tensors
  • 有内容可输出时 : 创建 EngineCoreOutput 并加入对应客户端的列表。
  • 无内容时: 不创建输出(部分 prefill 不产生输出),断言确保没有 prompt logprobs(因为部分 prefill 不应返回 logprobs)。
  • take_events() / take_prefill_stats(): 一次性取出并清空请求的事件和 prefill 统计。

7.7 从队列移除已停止请求

python 复制代码
if stopped_running_reqs:
    self.running = remove_all(self.running, stopped_running_reqs)
if stopped_preempted_reqs:
    self.waiting.remove_requests(stopped_preempted_reqs)
  • remove_all: 高效地从列表中移除集合中的所有元素。
  • RUNNING 中停止的请求从 self.running 移除。
  • PREEMPTED 中停止的请求从 self.waiting 移除(罕见情况,通常是请求被中止时)。

7.8 KV 加载失败处理(fail 策略)

python 复制代码
if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
    requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
    self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
    for request in requests:
        outputs[request.client_index].append(
            EngineCoreOutput(
                request_id=request.request_id,
                new_token_ids=[],
                finish_reason=request.get_finished_reason(),
                events=request.take_events(),
                trace_headers=request.trace_headers,
            )
        )
  • 当 kv_load_failure_policy == "fail"(而非 "recompute")时,受影响的请求直接标记为错误完成。
  • 为每个失败请求创建包含 finish_reason 的输出,通知客户端。

7.9 KV Connector 传输完成处理

python 复制代码
if kv_connector_output:
    self._update_from_kv_xfer_finished(kv_connector_output)
  • 处理 KV Connector 报告的传输完成事件(详见 8.16 节)。

7.10 KV Cache 事件发布

python 复制代码
events = self.kv_cache_manager.take_events()

if self.connector is not None:
    connector_events = self.connector.take_events()
    if connector_events:
        if events is None:
            events = list(connector_events)
        else:
            events.extend(connector_events)

if events:
    batch = KVEventBatch(ts=time.time(), events=events)
    self.kv_event_publisher.publish(batch)
  • 从 KV cache manager 和 KV Connector 收集事件(block 分配、释放、共享等)。
  • 合并后打包为 KVEventBatch,通过事件发布器发布。
  • 设计意图: KV 事件用于分布式 KV cache 的监控和调试,如 P/D 分离场景下的缓存一致性跟踪。

7.11 构建最终 EngineCoreOutputs

python 复制代码
engine_core_outputs = {
    client_index: EngineCoreOutputs(outputs=outs)
    for client_index, outs in outputs.items()
}
  • 将每个客户端的输出列表包装为 EngineCoreOutputs 对象。

7.12 已完成请求 ID 集合(多引擎场景)

python 复制代码
finished_req_ids = self.finished_req_ids_dict
if finished_req_ids:
    for client_index, finished_set in finished_req_ids.items():
        if (eco := engine_core_outputs.get(client_index)) is not None:
            eco.finished_requests = finished_set
        else:
            engine_core_outputs[client_index] = EngineCoreOutputs(
                finished_requests=finished_set
            )
    finished_req_ids.clear()
  • 在多引擎场景(include_finished_set=True)中,将自上次输出以来完成的请求 ID 集合附加到对应客户端的输出。
  • 设计意图: 多引擎场景需要高效跟踪请求生命周期,避免重复扫描。

7.13 附加统计信息

python 复制代码
if (
    stats := self.make_stats(
        spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats
    )
) is not None:
    if (eco := next(iter(engine_core_outputs.values()), None)) is None:
        engine_core_outputs[0] = eco = EngineCoreOutputs()
    eco.scheduler_stats = stats
  • 如果启用了统计(log_stats=True),将调度统计附加到第一个客户端的输出。
  • 如果没有请求输出(空步),创建一个空的 EngineCoreOutputs 来承载统计。
  • 设计意图: 统计信息只需发送给一个前端,避免重复。
python 复制代码
return engine_core_outputs
  • 返回按客户端分组的引擎核心输出。

第八章: Scheduler 辅助方法逐行解析

8.1 _preempt_request --- 请求抢占

python 复制代码
def _preempt_request(self, request: Request, timestamp: float) -> None:
    """Preempt a request and put it back to the waiting queue.

    NOTE: The request should be popped from the running queue outside of this method.
    """
    assert request.status == RequestStatus.RUNNING, (
        "Only running requests can be preempted"
    )
    self.kv_cache_manager.free(request)
    self.encoder_cache_manager.free(request)
    request.status = RequestStatus.PREEMPTED
    request.num_computed_tokens = 0
    if request.spec_token_ids:
        request.spec_token_ids = []
    request.num_preemptions += 1
    if self.log_stats:
        request.record_event(EngineCoreEventType.PREEMPTED, timestamp)
    self.waiting.prepend_request(request)

逐行解析:

  1. 断言 : 只有 RUNNING 状态的请求才能被抢占。调用方需先将请求从 self.running 中移除。
  2. kv_cache_manager.free(request): 释放请求占用的所有 KV cache block。这是抢占的核心代价------下次调度时需要重新 prefill。
  3. encoder_cache_manager.free(request): 释放 encoder cache。
  4. 状态 → PREEMPTED: 标记为被抢占。
  5. num_computed_tokens = 0 : 归零!抢占意味着放弃所有已计算的 KV cache,下次从头 prefill。
  6. 清空 spec_token_ids: 被抢占时,draft tokens 作废。
  7. num_preemptions += 1: 累计抢占次数,用于统计和优先级调整。
  8. 日志事件: 记录 PREEMPTED 事件时间戳。
  9. waiting.prepend_request : 前插到等待队列------被抢占的请求应优先恢复(已经等待了更久)。

设计要点 : 抢占是"全量回退"策略,不像某些系统支持部分保留 KV cache。这简化了实现但代价较高。prepend 确保公平性。

8.2 _update_after_schedule --- 调度后状态更新

python 复制代码
def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
    num_scheduled_tokens = scheduler_output.num_scheduled_tokens
    for req_id, num_scheduled_token in num_scheduled_tokens.items():
        request = self.requests[req_id]
        request.num_computed_tokens += num_scheduled_token
        request.is_prefill_chunk = request.num_computed_tokens < (
            request.num_tokens + request.num_output_placeholders
        )
        scheduler_output.has_structured_output_requests |= (
            request.use_structured_output and not request.is_prefill_chunk
        )

    self.finished_req_ids = set()

逐行解析:

  1. 遍历所有被调度的请求 ,乐观地推进 num_computed_tokens。
  2. num_computed_tokens += num_scheduled_token : 加上本轮调度的 token 数。这是乐观更新 ------如果 spec decode 有拒绝,后续 update_from_output 会修正。
  3. is_prefill_chunk : 判断请求是否仍在 prefill 阶段。条件: num_computed_tokens < num_tokens + num_output_placeholders。如果还有 token 没计算完,就是 prefill chunk。
  4. has_structured_output_requests: 如果任何请求使用结构化输出且不在 prefill chunk 阶段,标记为 True。这控制是否需要计算 grammar bitmask。
  5. self.finished_req_ids = set() : 用新 set 替换(而非 .clear()),因为 scheduler_output.finished_req_ids 引用的是旧 set。

设计意图: 注释解释了三个关键原因:

  • SchedulerOutput 需要包含原始调度 token 数来确定输入 IDs
  • 提前推进 num_computed_tokens 允许下步立即重新调度 prefill 请求
  • 后续 update_from_output 修正 spec token 拒绝导致的偏差

8.3 _update_request_as_session --- 流式请求更新

python 复制代码
def _update_request_as_session(
    self, session: Request, update: StreamingUpdate
) -> None:
    num_computed_tokens = session.num_computed_tokens
    kept_output_tokens = session._all_token_ids[
        session.num_prompt_tokens : num_computed_tokens
    ]
    del session._all_token_ids[num_computed_tokens:]
    session._output_token_ids.clear()

逐行解析:

  1. kept_output_tokens: 提取已计算的输出 token(从 prompt 结束到计算位置)。
  2. 截断 _all_token_ids: 删除计算位置之后的所有 token(未验证的 spec tokens 等)。
  3. 清空 _output_token_ids: 输出 token 将被重新并入 prompt。
python 复制代码
    assert session.prompt_token_ids is not None
    session.prompt_token_ids.extend(kept_output_tokens)
  • 将已计算的输出 token 追加到 prompt: 流式请求的每个 chunk 将之前的输出 token 视为新 prompt 的一部分。
python 复制代码
    if update.mm_features:
        base = session.num_tokens
        for mm_feature in update.mm_features:
            mm_feature.mm_position = replace(
                mm_feature.mm_position, offset=mm_feature.mm_position.offset + base
            )
        session.mm_features.extend(update.mm_features)
  • 如果更新包含多模态特征,调整其位置偏移(基于当前 token 数),然后追加到请求。
python 复制代码
    session._all_token_ids.extend(update.prompt_token_ids or ())
    session.prompt_token_ids.extend(update.prompt_token_ids or ())
    session.update_block_hashes()
    session.num_prompt_tokens = len(session.prompt_token_ids)
    session.arrival_time = update.arrival_time
    session.sampling_params = update.sampling_params
    if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
        self.num_waiting_for_streaming_input -= 1
    session.status = RequestStatus.WAITING

    if self.log_stats:
        session.record_event(EngineCoreEventType.QUEUED)
  • 追加新 prompt token,更新 block hash(用于前缀缓存),更新 prompt token 数、到达时间、采样参数。
  • 如果之前在等待流式输入,减少计数器。
  • 状态变为 WAITING,重新进入调度。
  • 设计意图: 流式请求(如多轮对话)在每个 chunk 到达时将请求"重启"------之前的输出变成新 prompt,然后等待下一个 chunk。

8.4 _make_cached_request_data --- 构建缓存请求数据

python 复制代码
def _make_cached_request_data(
    self,
    running_reqs: list[Request],
    resumed_reqs: list[Request],
    num_scheduled_tokens: dict[str, int],
    spec_decode_tokens: dict[str, list[int]],
    req_to_new_blocks: dict[str, KVCacheBlocks],
) -> CachedRequestData:

逐行解析:

python 复制代码
    req_ids: list[str] = []
    new_token_ids: list[list[int]] = []
    new_block_ids: list[tuple[list[int], ...] | None] = []
    all_token_ids: dict[str, list[int]] = {}
    num_computed_tokens: list[int] = []
    num_output_tokens: list[int] = []
    resumed_req_ids = set()
  • 初始化七个数据结构,对应 CachedRequestData 的各字段。
python 复制代码
    num_running_reqs = len(running_reqs)
    for idx, req in enumerate(itertools.chain(running_reqs, resumed_reqs)):
  • 用 itertools.chain 合并遍历 RUNNING 和 RESUMED 请求。
  • idx 用于区分两者(idx >= num_running_reqs 的是 RESUMED)。
python 复制代码
        req_id = req.request_id
        req_ids.append(req_id)
  • 收集请求 ID。
python 复制代码
        if self.use_pp and not self.scheduler_config.async_scheduling:
            num_tokens = num_scheduled_tokens[req_id] - len(
                spec_decode_tokens.get(req_id, ())
            )
            token_ids = req.all_token_ids[
                req.num_computed_tokens : req.num_computed_tokens + num_tokens
            ]
            new_token_ids.append(token_ids)
  • Pipeline Parallelism (PP) 非异步模式: 需要发送采样 token 回 Worker,因为 PP 中首阶段和末阶段没有直接通信通道。
  • 切片提取本轮新计算的 token IDs(减去 spec decode tokens)。
python 复制代码
        scheduled_in_prev_step = req_id in self.prev_step_scheduled_req_ids
        if idx >= num_running_reqs:
            assert not scheduled_in_prev_step
            resumed_req_ids.add(req_id)
        if not scheduled_in_prev_step:
            all_token_ids[req_id] = req.all_token_ids.copy()
  • scheduled_in_prev_step: 如果请求在上一步也被调度过,Worker 已有最新 token ID,不需要重复发送。
  • RESUMED 请求: 断言不在上步调度中(刚从 WAITING 恢复)。
  • 非上步调度的请求: 发送完整 token ID 列表(Worker 可能过期)。
python 复制代码
        new_block_ids.append(
            req_to_new_blocks[req_id].get_block_ids(allow_none=True)
        )
        num_computed_tokens.append(req.num_computed_tokens)
        num_output_tokens.append(
            req.num_output_tokens + req.num_output_placeholders
        )
  • 收集新 block IDs、已计算 token 数、输出 token 数(含占位符)。
python 复制代码
    return CachedRequestData(
        req_ids=req_ids,
        resumed_req_ids=resumed_req_ids,
        new_token_ids=new_token_ids,
        all_token_ids=all_token_ids,
        new_block_ids=new_block_ids,
        num_computed_tokens=num_computed_tokens,
        num_output_tokens=num_output_tokens,
    )
  • 构建并返回 CachedRequestData。
  • resumed_req_ids 的用途 : Worker 对 RESUMED 请求使用 new_block_ids 替换 block 表(而非追加),因为抢占后 block 表已完全重建。

8.5 _try_schedule_encoder_inputs --- Encoder 输入调度

python 复制代码
def _try_schedule_encoder_inputs(
    self,
    request: Request,
    num_computed_tokens: int,
    num_new_tokens: int,
    encoder_compute_budget: int,
    shift_computed_tokens: int = 0,
) -> tuple[list[int], int, int, list[int]]:

方法目的 : 确定哪些 encoder 输入需要在本步调度,并相应调整 num_new_tokens 和 encoder 预算。

返回值 : (encoder_inputs_to_schedule, num_new_tokens, encoder_compute_budget, external_load_encoder_input)

python 复制代码
    if num_new_tokens == 0 or not request.has_encoder_inputs:
        return [], num_new_tokens, encoder_compute_budget, []
  • 快速路径: 无新 token 或无 encoder 输入。
python 复制代码
    encoder_inputs_to_schedule: list[int] = []
    mm_features = request.mm_features
    assert mm_features is not None
    assert len(mm_features) > 0
    external_load_encoder_input = []
  • encoder_inputs_to_schedule: 本地需计算的 encoder 输入索引
  • external_load_encoder_input: 从 EC Connector 远程加载的索引
python 复制代码
    mm_hashes_to_schedule = set()
    num_embeds_to_schedule = 0
    for i, mm_feature in enumerate(mm_features):
        start_pos = mm_feature.mm_position.offset
        num_encoder_tokens = mm_feature.mm_position.length
        num_encoder_embeds = mm_feature.mm_position.get_num_embeds()
        item_identifier = mm_feature.identifier
  • 遍历每个多模态特征:
    • start_pos: encoder 输出在 decoder 输入中的起始位置
    • num_encoder_tokens: encoder 输出对应的 placeholder token 数
    • num_encoder_embeds: encoder 输出的 embedding 数(可能与 token 数不同,如视觉 patch)
    • item_identifier: 用于缓存去重的 hash/ID
python 复制代码
        if (
            start_pos
            >= num_computed_tokens + num_new_tokens + shift_computed_tokens
        ):
            break
  • 如果 encoder 输入在本步计算范围之外,终止遍历 (后续特征的 start_pos 只会更大)。
  • shift_computed_tokens: EAGLE 模式下的偏移补偿。
python 复制代码
        if self.is_encoder_decoder and num_computed_tokens > 0:
            assert start_pos == 0
            continue
  • Encoder-decoder 模型: encoder 输入总是在序列开头(start_pos=0)。一旦 decoder 已计算了一些 token,encoder 输入必然已完成。
python 复制代码
        elif start_pos + num_encoder_tokens <= num_computed_tokens:
            continue
  • 如果 encoder 输出已完全在已计算范围内,跳过(KV cache 中已有)。
python 复制代码
        if not self.is_encoder_decoder:
            if item_identifier in mm_hashes_to_schedule:
                continue
            if self.encoder_cache_manager.check_and_update_cache(request, i):
                continue
  • 去重: 同一 hash 的 encoder 输入只调度一次。
  • 缓存命中: 如果 encoder cache 中已有该输入的计算结果,跳过。
python 复制代码
        if (
            self.scheduler_config.disable_chunked_mm_input
            and num_computed_tokens < start_pos
            and (num_computed_tokens + num_new_tokens)
            < (start_pos + num_encoder_tokens)
        ):
            num_new_tokens = max(
                0, start_pos - (num_computed_tokens + shift_computed_tokens)
            )
            break
  • 禁止分块多模态输入: 如果调度范围无法覆盖整个 mm 输入,回退到 mm 输入之前的位置。
  • 设计意图: 某些 encoder(如视觉编码器)使用双向注意力,必须一次性处理完整输入。部分处理会导致结果不正确。
python 复制代码
        if not self.encoder_cache_manager.can_allocate(
            request, i, encoder_compute_budget, num_embeds_to_schedule
        ):
            if num_computed_tokens + shift_computed_tokens < start_pos:
                num_new_tokens = start_pos - (
                    num_computed_tokens + shift_computed_tokens
                )
            else:
                num_new_tokens = 0
            break
  • 缓存/预算不足 : 如果 encoder cache 满或计算预算不足:
    • 如果 mm 输入尚未开始,截断 num_new_tokens 到 mm 输入之前
    • 如果已开始(前缀缓存导致 num_computed_tokens >= start_pos),只能设为 0
python 复制代码
        start_idx_rel = max(0, num_computed_tokens - start_pos)
        end_idx_rel = min(
            num_encoder_tokens, num_computed_tokens + num_new_tokens - start_pos
        )
        curr_embeds_start, curr_embeds_end = (
            mm_feature.mm_position.get_embeds_indices_in_range(
                start_idx_rel, end_idx_rel
            )
        )
        if curr_embeds_end - curr_embeds_start == 0:
            continue
  • 计算本步调度范围内该 encoder 输入对应的 embedding 索引范围。
  • 如果范围内无 embedding(如 placeholder token 不对应 embedding),跳过。
python 复制代码
        if self.ec_connector is not None and self.ec_connector.has_cache_item(
            item_identifier
        ):
            mm_hashes_to_schedule.add(item_identifier)
            external_load_encoder_input.append(i)
            num_embeds_to_schedule += num_encoder_embeds
            continue
  • EC Connector 远程加载: 如果远程缓存中有该 encoder 输入,标记为远程加载(不消耗本地计算预算,但需要分配缓存空间)。
python 复制代码
        num_embeds_to_schedule += num_encoder_embeds
        encoder_compute_budget -= num_encoder_embeds
        mm_hashes_to_schedule.add(item_identifier)
        encoder_inputs_to_schedule.append(i)
  • 本地计算: 扣除 encoder 计算预算,加入调度列表。

8.6 _handle_stopped_request --- 处理已停止请求

python 复制代码
def _handle_stopped_request(self, request: Request) -> bool:
    """Return True if finished (can be False for resumable requests)."""
    if not request.resumable:
        return True

    if request.streaming_queue:
        update = request.streaming_queue.popleft()
        if update is None:
            return True
        self._update_request_as_session(request, update)
    else:
        request.status = RequestStatus.WAITING_FOR_STREAMING_REQ
        self.num_waiting_for_streaming_input += 1

    self._enqueue_waiting_request(request)
    return False

逐行解析:

  1. 非可恢复请求 : 直接返回 True(已完成)。
  2. 流式队列有更新 : 取出下一个 chunk,调用 _update_request_as_session 更新请求。返回 False(未完成,继续调度)。
  3. update is None : 流式结束哨兵,返回 True。
  4. 流式队列为空 : 请求进入 WAITING_FOR_STREAMING_REQ 状态,等待下一个 chunk。增加等待计数。
  5. _enqueue_waiting_request : 根据 status 决定放入 waiting 还是 skipped_waiting(WAITING_FOR_STREAMING_REQ 是阻塞状态,放入 skipped_waiting)。
  6. 返回 False: 请求未完成,只是暂停等待。

设计意图: 流式请求(如聊天场景的逐条消息)在一条消息完成后不立即结束,而是等待下一条消息。这使得同一个 request ID 可以跨多轮对话复用。

8.7 _update_request_with_output --- 追加输出 Token 并检查停止

python 复制代码
def _update_request_with_output(
    self, request: Request, new_token_ids: list[int]
) -> tuple[list[int], bool]:
    stopped = False
    for num_new, output_token_id in enumerate(new_token_ids, 1):
        request.append_output_token_ids(output_token_id)
        stopped = check_stop(request, self.max_model_len)
        if stopped:
            del new_token_ids[num_new:]
            break
    return new_token_ids, stopped

逐行解析:

  1. 逐个追加输出 token 到请求。
  2. 每追加一个就调用 check_stop 检查是否应该停止(EOS token、stop words、max_tokens 等)。
  3. 如果停止,截断 new_token_ids(只保留停止前的 token),跳出循环。
  4. 返回截断后的 token 列表和停止标志。
  • 设计意图 : 逐 token 检查确保不会返回超出停止条件后的 token。enumerate(..., 1) 使 num_new 从 1 开始,方便切片。

8.8 _free_encoder_inputs --- 释放已处理的 Encoder 输入

python 复制代码
def _free_encoder_inputs(self, request: Request) -> None:
    cached_encoder_input_ids = self.encoder_cache_manager.get_cached_input_ids(
        request
    )
    if not cached_encoder_input_ids:
        return

    for input_id in list(cached_encoder_input_ids):
        mm_feature = request.mm_features[input_id]
        start_pos = mm_feature.mm_position.offset
        num_tokens = mm_feature.mm_position.length
        if self.is_encoder_decoder and request.num_computed_tokens > 0:
            self.encoder_cache_manager.free_encoder_input(request, input_id)
        elif start_pos + num_tokens <= request.num_computed_tokens:
            self.encoder_cache_manager.free_encoder_input(request, input_id)

逐行解析:

  1. 获取请求在 encoder cache 中的所有输入 ID。
  2. 空 set 直接返回。
  3. list(cached_encoder_input_ids): 复制为列表再遍历,避免在迭代时修改 set。
  4. Encoder-decoder 模型: 一旦生成了任何 decode token,encoder 输出已被写入 cross-attention KV cache,可安全释放。
  5. Decoder-only 模型 : 当 encoder 输出的所有 token 都已被计算(start_pos + num_tokens <= num_computed_tokens),说明 encoder 输出已完全融入 decoder 的 KV cache,可释放。
  • 设计意图: 及时释放 encoder cache 空间给其他请求使用。Encoder cache 大小有限,不及时释放会导致新请求无法调度。

8.9 update_draft_token_ids --- 更新 Draft Token IDs

python 复制代码
def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
    for req_id, spec_token_ids in zip(
        draft_token_ids.req_ids,
        draft_token_ids.draft_token_ids,
    ):
        request = self.requests.get(req_id)
        if request is None or request.is_finished():
            continue

        if request.is_prefill_chunk:
            if request.spec_token_ids:
                request.spec_token_ids = []
            continue

        if self.structured_output_manager.should_advance(request):
            metadata = request.structured_output_request
            spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids)
        request.spec_token_ids = spec_token_ids

逐行解析:

  1. 遍历 draft model 生成的 token IDs。
  2. 跳过已完成或不存在的请求。
  3. Prefill chunk 阶段: 忽略 draft tokens(prefill 阶段不使用 speculative decoding)。清空已有的 spec tokens。
  4. 结构化输出: 通过 grammar 验证 draft tokens,过滤不符合语法的 token。
  5. 设置请求的 spec_token_ids,下步调度时这些 token 会被一起调度并验证。

设计意图: Speculative decoding 的核心机制------draft model 快速生成候选 token,目标模型批量验证。Grammar 验证确保 spec tokens 不会违反结构化输出约束。

8.10 update_draft_token_ids_in_output --- 在输出中更新 Draft Token IDs

python 复制代码
def update_draft_token_ids_in_output(
    self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput
) -> None:
    num_invalid_spec_tokens: dict[str, int] = {}

    sched_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
    for req_id, spec_token_ids in zip(
        draft_token_ids.req_ids,
        draft_token_ids.draft_token_ids,
    ):
  • 与 update_draft_token_ids 类似,但直接修改 scheduler_output 中已调度的 spec tokens。
  • 使用场景: 在调度完成后、模型执行前,draft model 生成了新的 spec tokens,需要更新到调度输出中。
python 复制代码
        request = self.requests.get(req_id)
        if request is None or request.is_finished():
            continue

        placeholder_spec_tokens = sched_spec_tokens.get(req_id)
        if not placeholder_spec_tokens:
            continue
  • 只处理已有调度 spec tokens 的请求。
python 复制代码
        orig_num_spec_tokens = len(placeholder_spec_tokens)
        del spec_token_ids[orig_num_spec_tokens:]
  • 截断 draft tokens 到调度时的数量(chunked prefill 可能只调度了部分 spec tokens)。
python 复制代码
        if self.structured_output_manager.should_advance(request):
            metadata = request.structured_output_request
            assert metadata is not None and metadata.grammar is not None
            spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids)
  • Grammar 验证。
python 复制代码
        num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids)
        if num_invalid_tokens:
            spec_token_ids.extend([-1] * num_invalid_tokens)
            num_invalid_spec_tokens[req_id] = num_invalid_tokens

        sched_spec_tokens[req_id] = spec_token_ids
  • Padding : 被grammar拒绝的token位置用 -1 填充,保持长度不变。Worker 需要固定数量的 slot 来处理。
  • 记录每个请求的无效 spec token 数。
python 复制代码
    scheduler_output.num_invalid_spec_tokens = num_invalid_spec_tokens
  • 设置到 scheduler_output,供 update_from_output 中 spec decode 统计使用。

8.11 finish_requests --- 批量完成请求

python 复制代码
def finish_requests(
    self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus
) -> list[tuple[str, int]]:

方法目的: 从调度器外部(如 API server)触发请求完成,用于中止、超时等场景。

返回值 : (req_id, client_index) 列表,表示实际被中止的请求。

python 复制代码
    assert RequestStatus.is_finished(finished_status)
    if isinstance(request_ids, str):
        request_ids = (request_ids,)
    elif request_ids is not None:
        request_ids = set(request_ids)
    else:
        request_ids = self.requests.keys()
  • 参数规范化: 支持 string、可迭代对象、None(全部完成)三种输入。
python 复制代码
    running_requests_to_remove = set()
    waiting_requests_to_remove = []
    valid_requests = []

    for req_id in request_ids:
        request = self.requests.get(req_id)
        if request is None or request.is_finished():
            continue

        valid_requests.append(request)
        if request.status == RequestStatus.RUNNING:
            running_requests_to_remove.add(request)
        else:
            if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
                self.num_waiting_for_streaming_input -= 1
            waiting_requests_to_remove.append(request)
  • 两遍扫描的第一遍: 收集需要从队列移除的请求。
  • 分类为 RUNNING(批量移除效率高)和非 RUNNING。
  • 调整流式等待计数器。
python 复制代码
    if running_requests_to_remove:
        self.running = remove_all(self.running, running_requests_to_remove)
    if waiting_requests_to_remove:
        self.waiting.remove_requests(waiting_requests_to_remove)
        self.skipped_waiting.remove_requests(waiting_requests_to_remove)
  • 批量从队列移除(比逐个移除高效)。
  • 非运行请求可能在 waiting 或 skipped_waiting 中,两个都清理。
python 复制代码
    for request in valid_requests:
        delay_free_blocks = False
        if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
            delay_free_blocks = (
                request.request_id not in self.finished_recving_kv_req_ids
            )
            self.finished_recving_kv_req_ids.discard(request.request_id)
            self.failed_recving_kv_req_ids.discard(request.request_id)

        request.status = finished_status
        self._free_request(request, delay_free_blocks=delay_free_blocks)
  • 第二遍: 设置状态并释放资源。
  • KV Connector 延迟释放: 如果请求正在等待远程 KV 加载且尚未完成,延迟释放 block(等加载完成后再释放,避免内存泄漏)。
  • 清理 KV 加载状态跟踪。

8.12 _free_request --- 释放请求资源

python 复制代码
def _free_request(
    self, request: Request, delay_free_blocks: bool = False
) -> dict[str, Any] | None:
    assert request.is_finished()

    connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
    self.encoder_cache_manager.free(request)
    request_id = request.request_id
    self.finished_req_ids.add(request_id)
    if self.finished_req_ids_dict is not None:
        self.finished_req_ids_dict[request.client_index].add(request_id)

    delay_free_blocks |= connector_delay_free_blocks
    if not delay_free_blocks:
        self._free_blocks(request)

    return kv_xfer_params

逐行解析:

  1. 断言: 请求必须处于完成状态。
  2. _connector_finished: 通知 KV Connector 请求完成,获取是否延迟释放 block 和 KV 传输参数。
  3. 释放 encoder cache。
  4. 记录 finished_req_ids: 两处------全局 set 和按客户端的 dict(多引擎场景)。
  5. 合并延迟释放标志: 外部指定或 Connector 要求的延迟。
  6. 非延迟时释放 block 并删除请求记录 : _free_blocks 释放 KV cache 并从 self.requests 中删除。
  7. 返回 KV 传输参数: 用于通知客户端 KV 存储位置(P/D 分离场景)。

设计意图 : delay_free_blocks 处理异步 KV 传输场景------请求虽已完成,但 KV 数据可能还在传输中,过早释放会导致数据丢失。

8.13 _free_blocks --- 释放 KV Block

python 复制代码
def _free_blocks(self, request: Request):
    assert request.is_finished()
    self.kv_cache_manager.free(request)
    del self.requests[request.request_id]
  • 释放 KV cache block,从请求字典中删除请求对象。
  • 注意: 这是最终清理步骤,调用后请求对象不再可访问。

8.14 reset_prefix_cache --- 重置前缀缓存

python 复制代码
def reset_prefix_cache(
    self, reset_running_requests: bool = False, reset_connector: bool = False
) -> bool:

参数:

  • reset_running_requests: 是否抢占所有 RUNNING 请求以强制释放 KV cache
  • reset_connector: 是否同时重置 KV Connector 缓存
python 复制代码
    if reset_running_requests:
        timestamp = time.monotonic()
        while self.running:
            request = self.running.pop()
            self._preempt_request(request, timestamp)
            request.num_output_placeholders = 0
            request.discard_latest_async_tokens = True

        self.prev_step_scheduled_req_ids.clear()
  • 抢占所有 RUNNING 请求(逆序弹出,使得重入时顺序正确)。
  • num_output_placeholders = 0: 清除异步调度的占位符。
  • discard_latest_async_tokens = True: 标记丢弃最新的异步输出 token,避免重调度时重复输出。
  • 清除 prev_step_scheduled_req_ids: 因为所有请求都将被重新调度,上步调度记录无效。
python 复制代码
    reset_successful = self.kv_cache_manager.reset_prefix_cache()
    if reset_running_requests and not reset_successful:
        raise RuntimeError(...)
  • 实际重置 KV cache manager 的前缀缓存。
  • 如果强制清空了 RUNNING 请求但仍失败,抛出异常(可能因为有 WAITING_FOR_REMOTE_KVS 的请求持有 block)。
python 复制代码
    if reset_connector:
        reset_successful = self.reset_connector_cache() and reset_successful

    return reset_successful
  • 可选重置 Connector 缓存。
  • 返回是否全部成功。

8.15 reset_connector_cache / reset_encoder_cache

python 复制代码
def reset_connector_cache(self) -> bool:
    if self.connector is None:
        logger.warning("reset_connector called but no KV connector is configured.")
        return False
    if self.connector.reset_cache() is False:
        return False
    if self.log_stats:
        assert self.connector_prefix_cache_stats is not None
        self.connector_prefix_cache_stats.reset = True
    return True
  • 重置 KV Connector 的远程缓存。
  • 记录统计标记。
python 复制代码
def reset_encoder_cache(self) -> None:
    self.encoder_cache_manager.reset()
  • 重置 encoder cache,使所有缓存的 encoder 输出失效。
  • 使用场景: 模型权重更新后,旧的视觉 embedding 不再有效。

8.16 _update_from_kv_xfer_finished --- KV 传输完成更新

python 复制代码
def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):
    if self.connector is not None:
        self.connector.update_connector_output(kv_connector_output)

    for req_id in kv_connector_output.finished_recving or ():
        logger.debug("Finished recving KV transfer for request %s", req_id)
        assert req_id in self.requests
        req = self.requests[req_id]
        if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
            self.finished_recving_kv_req_ids.add(req_id)
        else:
            assert RequestStatus.is_finished(req.status)
            self._free_blocks(self.requests[req_id])
    for req_id in kv_connector_output.finished_sending or ():
        logger.debug("Finished sending KV transfer for request %s", req_id)
        assert req_id in self.requests
        self._free_blocks(self.requests[req_id])

逐行解析:

  1. 更新 Connector 状态: 将 Worker 侧的输出传给 Connector,更新内部跟踪。
  2. finished_recving : 远程 KV 加载完成的请求:
    • 如果在 WAITING_FOR_REMOTE_KVS 状态: 加入 finished_recving_kv_req_ids,下步调度时 _try_promote_blocked_waiting_request 会将其提升回 WAITING。
    • 如果已完成: 直接释放 block(请求已完成,KV 加载也完成了,可以安全释放)。
  3. finished_sending : KV 推送完成的请求:
    • 推送完成后释放 block(P/D 分离中,Prefill 节点推送 KV 给 Decode 节点,推送完成后可释放)。

8.17 _update_waiting_for_remote_kv --- 远程 KV 等待完成更新

python 复制代码
def _update_waiting_for_remote_kv(self, request: Request) -> None:
    assert self.connector is not None

    if request.request_id in self.failed_recving_kv_req_ids:
        if request.num_computed_tokens:
            self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens)
        else:
            self.kv_cache_manager.free(request)
        self.failed_recving_kv_req_ids.remove(request.request_id)
    else:
        self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens)
        if request.num_computed_tokens == request.num_tokens:
            request.num_computed_tokens = request.num_tokens - 1
    self.finished_recving_kv_req_ids.remove(request.request_id)

逐行解析:

  1. 有加载失败 :
    • 如果有部分有效 token: 缓存有效的 block
    • 如果全部无效: 释放所有 block(可能有本地缓存命中可重试)
    • 从失败集合中移除
  2. 加载成功 : 缓存 block,并检查全量命中的边缘情况。
  3. 全量命中修正 : 如果 num_computed_tokens == num_tokens(所有 token 都从远程获取),需要回退 1 个 token(num_tokens - 1),因为最后一个 token 需要重新计算以采样下一个 token。
  4. 从完成集合中移除。

设计意图: 全量命中时,请求的所有 token 都在远程 KV cache 中,但 LLM 的采样需要基于最后一个 token 的 logits,所以必须重新计算最后一个 token。

8.18 _try_promote_blocked_waiting_request --- 尝试提升阻塞请求

python 复制代码
def _try_promote_blocked_waiting_request(self, request: Request) -> bool:
    if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
        if request.request_id not in self.finished_recving_kv_req_ids:
            return False
        self._update_waiting_for_remote_kv(request)
        if request.num_preemptions:
            request.status = RequestStatus.PREEMPTED
        else:
            request.status = RequestStatus.WAITING
        return True
  • WAITING_FOR_REMOTE_KVS: 检查 KV 是否已加载完成。如果是,更新请求状态并缓存 block。
  • 状态恢复: 有抢占记录 → PREEMPTED(需要完整 prefill),否则 → WAITING(可直接调度剩余 token)。
python 复制代码
    if request.status == RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR:
        structured_output_req = request.structured_output_request
        if not (structured_output_req and structured_output_req.grammar):
            return False
        request.status = RequestStatus.WAITING
        return True
  • WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR: 检查 grammar 是否编译完成。如果是,恢复为 WAITING。
python 复制代码
    if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
        assert not request.streaming_queue
        return False
  • WAITING_FOR_STREAMING_REQ: 流式队列空时无法提升(等待外部输入)。断言确保队列为空。
python 复制代码
    raise AssertionError(
        "Unexpected blocked waiting status in promotion: "
        f"{request.status.name} for request {request.request_id}"
    )
  • 任何其他阻塞状态都是编程错误。

8.19 _connector_finished --- Connector 完成处理

python 复制代码
def _connector_finished(
    self, request: Request
) -> tuple[bool, dict[str, Any] | None]:
    if self.connector is None:
        return False, None

    self.kv_cache_manager.remove_skipped_blocks(
        request_id=request.request_id,
        total_computed_tokens=request.num_tokens,
    )

    block_ids = self.kv_cache_manager.get_block_ids(request.request_id)

    if not isinstance(self.connector, SupportsHMA):
        assert len(self.kv_cache_config.kv_cache_groups) == 1
        return self.connector.request_finished(request, block_ids[0])

    return self.connector.request_finished_all_groups(request, block_ids)

逐行解析:

  1. 无 Connector: 直接返回。
  2. remove_skipped_blocks: 在将 block 表传给 Connector 前,移除超出窗口的前缀 block(滑动窗口注意力场景)。
  3. 获取 block IDs: 传给 Connector 用于 KV 存储规划。
  4. HMA(Hybrid Memory Allocator)分支 : 如果 Connector 支持 HMA,调用 request_finished_all_groups 处理多个 KV cache group。否则,断言只有一个 group,调用 request_finished。
  5. 返回值 : (是否延迟释放block, KV传输参数)。

8.20 _update_requests_with_invalid_blocks --- 无效 Block 更新

python 复制代码
def _update_requests_with_invalid_blocks(
    self,
    requests: Iterable[Request],
    invalid_block_ids: set[int],
    num_scheduled_tokens: dict[str, int],
    evict_blocks: bool = True,
) -> tuple[set[str], int, set[int]]:

方法目的 : 扫描请求,检测受无效 block 影响的请求,调整 num_computed_tokens 到最长有效前缀。

返回值 : (受影响请求ID集合, 需要重计算的token总数, 需要驱逐的block集合)

python 复制代码
    affected_req_ids: set[str] = set()
    total_affected_tokens = 0
    blocks_to_evict: set[int] = set()
    marked_invalid_block_ids: set[int] = set()
  • marked_invalid_block_ids: 已标记为需要重计算的 block------共享 block 只需一次重计算。
python 复制代码
    for request in requests:
        is_affected = False
        marked_invalid_block = False
        req_id = request.request_id
        (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id)
  • 遍历请求,获取 block ID 列表。
python 复制代码
        req_num_computed_tokens = (
            request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0)
        )
  • 调整: 当前步的调度 token 还未实际计算(正在执行中),减去得到执行前的已计算 token 数。
python 复制代码
        req_num_computed_blocks = (
            req_num_computed_tokens + self.block_size - 1
        ) // self.block_size
  • 计算已计算 token 对应的 block 数(向上取整)。
python 复制代码
        for idx, block_id in zip(range(req_num_computed_blocks), req_block_ids):
            if block_id not in invalid_block_ids:
                continue

            is_affected = True

            if block_id in marked_invalid_block_ids:
                continue

            marked_invalid_block_ids.add(block_id)

            if marked_invalid_block:
                continue

            marked_invalid_block = True
            request.num_computed_tokens = idx * self.block_size
            num_affected_tokens = (
                req_num_computed_tokens - request.num_computed_tokens
            )
            total_affected_tokens += num_affected_tokens

            if evict_blocks:
                blocks_to_evict.update(req_block_ids[idx:])

逐行解析:

  1. 遍历请求的已计算 block,检查是否在无效 block 集合中。
  2. 如果受影响但 block 已被其他请求标记重计算,跳过(共享 block 只需计算一次)。
  3. 首次遇到无效 block : 截断 num_computed_tokens 到无效 block 之前的位置(idx * block_size)。
  4. 计算受影响 token 数(需要重计算的量)。
  5. 收集无效 block 及其下游所有 block 用于驱逐(req_block_ids[idx:])。
python 复制代码
        if is_affected:
            if not marked_invalid_block:
                total_affected_tokens += (
                    request.num_computed_tokens - req_num_computed_tokens
                )
                request.num_computed_tokens = req_num_computed_tokens

            affected_req_ids.add(request.request_id)
  • 所有无效 block 都被其他请求标记的情况: 回退到执行前的已计算 token 数(因为那些"有效"的 block 依赖于已被标记重计算的上游 block)。

8.21 _handle_invalid_blocks --- 处理无效 Block 总入口

python 复制代码
def _handle_invalid_blocks(
    self, invalid_block_ids: set[int], num_scheduled_tokens: dict[str, int]
) -> set[str]:
    should_fail = not self.recompute_kv_load_failures
  • 根据 kv_load_failure_policy 决定策略:
    • "recompute": 保留请求,重计算无效 block
    • "fail": 直接失败请求
python 复制代码
    async_load_reqs = (
        req
        for req in self.skipped_waiting
        if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
    )
    async_failed_req_ids, num_failed_tokens, _ = (
        self._update_requests_with_invalid_blocks(
            async_load_reqs,
            invalid_block_ids,
            num_scheduled_tokens,
            evict_blocks=False,
        )
    )
  • 异步加载请求 : 在 skipped_waiting 中等待远程 KV 的请求。
  • evict_blocks=False: 异步请求的 block 尚未缓存,不需要驱逐。
python 复制代码
    sync_failed_req_ids, num_failed_tokens, sync_blocks_to_evict = (
        self._update_requests_with_invalid_blocks(
            self.running, invalid_block_ids, num_scheduled_tokens, evict_blocks=True
        )
    )
  • 同步加载请求: 在 RUNNING 队列中的请求,block 已缓存,需要驱逐无效 block。
python 复制代码
    if not total_failed_requests:
        return set()
  • 无受影响请求,返回空集合。
python 复制代码
    if sync_blocks_to_evict and not self.recompute_kv_load_failures:
        self.kv_cache_manager.evict_blocks(sync_blocks_to_evict)
  • fail 策略下: 驱逐无效 block 及下游 block(这些请求即将被终止,block 不再需要)。
  • recompute 策略下: 不驱逐------重计算后 block 会被复用。
python 复制代码
    if should_fail:
        all_failed_req_ids = async_failed_req_ids | sync_failed_req_ids
        logger.error(...)
        return all_failed_req_ids
  • fail 策略 : 返回所有受影响请求 ID,主循环会跳过这些请求,然后 finish_requests 会将它们标记为错误。
python 复制代码
    logger.warning("Recovered from KV load failure: ...")
    self.failed_recving_kv_req_ids |= async_failed_req_ids
    return sync_failed_req_ids
  • recompute 策略 :
    • 异步请求加入 failed_recving_kv_req_ids,在 _update_waiting_for_remote_kv 中会缓存有效 token。
    • 返回同步受影响请求 ID,主循环跳过这些请求(它们的 num_computed_tokens 已被调整,下步会重计算)。

8.22 get_grammar_bitmask --- 获取 Grammar Bitmask

python 复制代码
def get_grammar_bitmask(
    self, scheduler_output: SchedulerOutput
) -> GrammarOutput | None:
    if not scheduler_output.has_structured_output_requests:
        return None

    structured_output_request_ids = [
        req_id
        for req_id in scheduler_output.num_scheduled_tokens
        if (req := self.requests.get(req_id))
        and (req.use_structured_output and not req.is_prefill_chunk)
    ]
    if not structured_output_request_ids:
        return None

    bitmask = self.structured_output_manager.grammar_bitmask(
        self.requests,
        structured_output_request_ids,
        scheduler_output.scheduled_spec_decode_tokens,
    )
    return GrammarOutput(structured_output_request_ids, bitmask)

逐行解析:

  1. 如果没有结构化输出请求,返回 None。
  2. 收集使用结构化输出且非 prefill chunk 的请求 ID(只有 decode 阶段才需要 grammar bitmask)。
  3. 调用 structured_output_manager 计算 bitmask(每个请求一行,每个 token 一列,标记哪些 token 被语法允许)。
  4. 返回 GrammarOutput,传给 Worker 用于约束采样。
  • 设计意图: Grammar bitmask 是结构化输出的核心机制------在采样时强制只选择符合语法的 token,确保输出始终有效。

8.23 其他辅助方法简述

get_request_counts
python 复制代码
def get_request_counts(self) -> tuple[int, int]:
    return len(self.running), len(self.waiting) + len(self.skipped_waiting)
  • 返回 (RUNNING 数, WAITING 数)。用于监控和限流。
add_request
python 复制代码
def add_request(self, request: Request) -> None:
    existing = self.requests.get(request.request_id)
    if existing is not None:
        update = StreamingUpdate.from_request(request)
        if existing.status != RequestStatus.WAITING_FOR_STREAMING_REQ:
            existing.streaming_queue.append(update)
        elif update is not None:
            self._update_request_as_session(existing, update)
        else:
            self.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED)
    else:
        if request.resumable:
            request.streaming_queue = deque()
        self._enqueue_waiting_request(request)
        self.requests[request.request_id] = request
        if self.log_stats:
            request.record_event(EngineCoreEventType.QUEUED)
  • 重复 ID 处理: 流式请求通过相同 ID 发送新 chunk,检测到已有请求时更新而非创建新的。
  • 新请求 : 根据状态选择队列(_enqueue_waiting_request),注册到请求字典。
pause_state / set_pause_state
python 复制代码
@property
def pause_state(self) -> PauseState:
    return self._pause_state

def set_pause_state(self, pause_state: PauseState) -> None:
    self._pause_state = pause_state
  • 暂停状态管理,支持 UNPAUSED、PAUSED_NEW(不调度新请求)、PAUSED_ALL(完全暂停)。
get_num_unfinished_requests
python 复制代码
def get_num_unfinished_requests(self) -> int:
    if self._pause_state == PauseState.PAUSED_ALL:
        return 0
    if self._pause_state == PauseState.PAUSED_NEW:
        return len(self.running)
    num_waiting = (
        len(self.waiting)
        + len(self.skipped_waiting)
        - self.num_waiting_for_streaming_input
    )
    return num_waiting + len(self.running)
  • 计算未完成请求数,排除流式等待中的请求(它们不算"真正"在等待)。
make_stats / make_spec_decoding_stats
python 复制代码
def make_stats(...) -> SchedulerStats | None:
    if not self.log_stats:
        return None
    ...
    return SchedulerStats(
        num_running_reqs=...,
        num_waiting_reqs=...,
        kv_cache_usage=...,
        prefix_cache_stats=...,
        spec_decoding_stats=...,
        kv_connector_stats=...,
        ...
    )
  • 构建调度统计对象,包含所有运行时指标。
shutdown
python 复制代码
def shutdown(self) -> None:
    if self.kv_event_publisher:
        self.kv_event_publisher.shutdown()
    if self.connector is not None:
        self.connector.shutdown()
  • 优雅关闭事件发布器和 KV Connector。
_is_blocked_waiting_request / _enqueue_waiting_request
python 复制代码
@staticmethod
def _is_blocked_waiting_status(status: RequestStatus) -> bool:
    return status in (
        RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR,
        RequestStatus.WAITING_FOR_REMOTE_KVS,
        RequestStatus.WAITING_FOR_STREAMING_REQ,
    )

def _enqueue_waiting_request(self, request: Request) -> None:
    if self._is_blocked_waiting_status(request.status):
        self.skipped_waiting.add_request(request)
    else:
        self.waiting.add_request(request)
  • 阻塞状态的请求进入 skipped_waiting,正常请求进入 waiting。
_get_routed_experts --- MoE 路由专家读取
python 复制代码
def _get_routed_experts(self, request: Request) -> np.ndarray | None:
    if not self.vllm_config.model_config.enable_return_routed_experts:
        return None
    kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)
    block_ids = kv_blocks.get_block_ids()[self.routed_experts_attn_gid]
    num_tokens = request.num_tokens - 1
    block_ids_array = np.array(block_ids, dtype=np.int32)
    ...
    slot_mapping = (
        block_offsets.reshape((1, block_size))
        + block_ids_array.reshape((num_blocks, 1)) * block_size
    ).flatten()[:num_tokens]
    return self.routed_experts_reader.get_routed_experts(indices=slot_mapping)
  • 计算请求的 slot mapping(block_id * block_size + offset),从 GPU buffer 读取 MoE 路由专家信息。
  • 设计意图: 用于 MoE 模型的可观测性和分析,了解哪些专家被激活。

总结: 核心调度流程全景图

复制代码
schedule()
├── 初始化调度变量与预算
├── kv_cache_manager.new_step_starts()
├── RUNNING 调度循环
│   ├── 计算新 token 数
│   ├── Encoder 输入调度 (_try_schedule_encoder_inputs)
│   ├── Mamba block 对齐
│   ├── KV block 分配 (allocate_slots)
│   │   └── 失败时抢占 (_preempt_request)
│   └── 更新调度状态
├── WAITING 调度循环
│   ├── 队列选择 (_select_waiting_queue_for_scheduling)
│   ├── 阻塞状态提升 (_try_promote_blocked_waiting_request)
│   ├── LoRA 约束检查
│   ├── 前缀缓存查询 (get_computed_blocks)
│   ├── KV Connector 远程查询
│   ├── 异步加载分支 → WAITING_FOR_REMOTE_KVS
│   ├── 同步调度分支
│   │   ├── 新 token 数计算
│   │   ├── Encoder 输入调度
│   │   ├── Mamba 对齐
│   │   ├── 全序列保留检查
│   │   ├── KV block 分配
│   │   └── 请求状态 → RUNNING
│   └── 跳过队列回填
├── 约束断言验证
├── 公共前缀计算
├── SchedulerOutput 构建
│   ├── NewRequestData / CachedRequestData
│   ├── KV Connector metadata
│   └── EC Connector metadata
└── _update_after_schedule (推进 num_computed_tokens)

update_from_output()
├── 解包 ModelRunnerOutput
├── KV 加载失败处理 (_handle_invalid_blocks)
├── 逐请求主循环
│   ├── Spec decode 拒绝修正
│   ├── Encoder 输入释放
│   ├── 输出 token 追加与停止检查
│   ├── 结构化输出验证
│   ├── 停止请求处理
│   └── EngineCoreOutput 构建
├── 队列清理 (remove stopped)
├── KV 传输完成处理 (_update_from_kv_xfer_finished)
├── KV 事件发布
└── 统计信息附加

核心设计哲学:

  1. 乐观调度 + 悲观修正 : _update_after_schedule 乐观推进 num_computed_tokens,update_from_output 根据实际结果(spec token 拒绝、KV 加载失败)修正。这避免了调度和执行之间的复杂同步。

  2. 两阶段队列管理 : waiting + skipped_waiting 的双队列设计,将正常请求和阻塞请求分离,避免阻塞请求干扰调度流程,同时保证它们在依赖满足后被优先处理。

  3. 延迟释放模式 : KV Connector 的异步传输场景下,block 释放可能延迟到传输完成,避免数据丢失。这通过 delay_free_blocks 标志贯穿 _free_request → _free_blocks 调用链。

  4. 全量回退抢占 : 抢占时完全释放 KV cache 并归零 num_computed_tokens,实现简单但代价高。未来可能支持部分保留。

  5. 多约束协同调度: Token 预算、并发上限、KV cache 容量、encoder 预算、LoRA 上限等多个约束同时生效,任何约束不足都终止对应方向的调度,但不影响其他方向。

本文档逐行解析 vLLM v1 Core 模块中 KV Cache 工具体系、Block Pool、KV Cache Manager 体系及辅助模块的完整实现。

相关推荐
源流之道2 小时前
联想Y700 五代(TB323FU)刷 ColorOS 16 教程指南
网络·ai·刷机指南·刷机心得
无忧智库4 小时前
AI分镜提示词怎么选
人工智能·ai
小凡geo4 小时前
本地商家 GEO:用脚本一键生成 FAQPage 结构化数据,让 AI 切片更稳
开发语言·人工智能·python·microsoft·搜索引擎·ai
天远Date Lab5 小时前
零信任架构实战:基于天远全能消金报告(标准版)构建自动化骑手资信评估网关
运维·人工智能·架构·自动化
Shulex5 小时前
从问答到执行:跨境电商场景下抑制大模型幻觉与实现端到端任务闭环的工程架构实践
架构
努力努力再努力wz7 小时前
【CUDA 入门系列】:Orin、GPU 并行执行模型与 CUDA 核心机制
架构
wp123_17 小时前
TLVR 电感在 AI 服务器电源中的应用与市场前景分析
服务器·人工智能·科技·ai·硬件工程
萧瑟余晖7 小时前
Dubbo 综合实战与性能调优详解
架构·dubbo
秃了也弱了。8 小时前
自适应类架构详解:自适应轮询、自适应重试、自适应采样
架构