vLLM 架构拆解

日本語

简体中文
English
Español
Français
Deutsch
繁體中文

日本語

简体中文
English
Español
Français
Deutsch
繁體中文

Appearance

プリエンプションと prefill throttle:KV cache 逼迫時の自衛 ​

源码版本v0.25.1

役割 ​

KV cache は v1 スケジューラにとって最も希少なリソースです — 一つの GPU 上のメモリブロック数は profiling 段階で確定します。running キューの某 request がさらに多くのトークンを計算する必要がある一方、KVCacheManager.allocate_slots が新しいブロックを返せないとき、スケジューラ (scheduler) はある決定を下す必要があります:running 内の request を一つ選んで追い出し、その KV ブロックを返却させてから再試行する。この仕組みが**プリエンプション (preemption)**で、入口は Scheduler._preempt_request(_preempt_request:1140)です。

もう一つ関連するがより隠れた仕組みが prefill throttle です — DP(データ並列)シナリオでは各 rank の prefill ステップを整列させる必要があり、そうでないと速い rank が遅い rank を待たせてしまいます。スケジューラは throttle_prefills=True を受け取ったとき、in-progress prefill chunk と新規 prefill を整列ステップまで遅延させ、decode でこのステップを埋めます(defer_prefills:436-438)。この二つは「あるステップで自発的に prefill を前に進めるのを控える」という抑制行動なので一緒に説明します。

設計動機 ​

なぜ新規リクエストを拒否するのではなく、プリエンプションを行うのか?

  • 進捗を保つ方がお得:ある request が prefill の大半を終えているとき、そのまま拒否するとそれまで計算した KV がすべて無駄になります。プリエンプション (preemption) は KV ブロックを返却し、ステータスを PREEMPTED に戻し、num_computed_tokens をゼロにしますが(preempt 状態:1152-1153)、request 本体は破棄されず waiting の先頭に戻されます(prepend_request:1161)。次ステップで再 prefill でき(プレフィックスキャッシュ (prefix caching) がヒットした部分は再利用可能)。
  • プリエンプションポリシー切り替え可能:priority モードは (priority, arrival_time) が最大の running を選んで追い出し(priority preempt:547-552)、FCFS モードは直接 self.running.pop()(FCFS preempt:572)。セマンティクスは出隊と一致します。
  • プリエンプションが起きたステップは新規リクエストを受け付けない:schedule() の waiting ループ入口は not preempted_reqs をチェックし(waiting 入口:637)、waiting に戻された request が同じステップで再び取り出され、プリエンプト中の running とリソースを争うのを防ぎます。
  • prefill throttle は DP 公平性:DP rank 間で prefill ステップを整列させる必要があります。そうでないと、ある rank はすでに decode を始めているのに別の rank がまだ prefill という状態になりスループットが低下します。throttle は非整列ステップで decode だけを走らせ、prefill を整列ステップまで遅延させて一緒に行います(defer_prefills:436-438)。
  • capacity-bound で自動上書き:prefill_capacity_bound(prefill_capacity_bound:290)は前回の整列ステップで waiting を空にしたかを記録し、空にしたらもう throttle する必要はありません — throttle を続けても GPU が遊ぶだけです。
  • async KV load の独立予約:_inflight_prefill_reserved_blocks(_inflight_prefill_reserved_blocks:2400)は全 in-flight prefill がさらに必要なブロック数を集計します。非同期ロードの request はプリエンプトされず、十分なブロックを予約しないと入れません。これでデッドロックを回避します。

主要ファイル ​

データフロー ​

プリエンプションのトリガーポイントは running ループ内で、allocate_slots が None を返したとき:

python
# The request cannot be scheduled.
# Preempt the lowest-priority request.
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:
        # 已经在本步分到 token 的牺牲者,把预算还回去
        preempted_req_id = preempted_req.request_id
        scheduled_running_reqs.remove(preempted_req)
        token_budget += num_scheduled_tokens.pop(preempted_req_id)
        req_to_new_blocks.pop(preempted_req_id)
        ...
        req_index -= 1
else:
    preempted_req = self.running.pop()

self._preempt_request(preempted_req, scheduled_timestamp)
preempted_reqs.append(preempted_req)
if preempted_req == request:
    # No more request to preempt. Cannot schedule this request.
    break

この部分は scheduler.py:545-582 にあります。priority モードには一つのディテールがあります:選ばれた被選出者がこのラウンドですでにスケジュールされていた場合(scheduled_running_reqs 内)、トークン予算、新規ブロック、spec token、encoder 予算をすべて返却する必要があり(予算返却:553-570)、running リストが一つ減ったため req_index -= 1 します。preempted_req == request は終了条件です — プリエンプションが自分自身にまで回ってきたら負けを認めてループを break します。

追い出された request は _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._free_request_blocks(request)
    self.encoder_cache_manager.free(request)
    self._inflight_prefills.discard(request)
    request.status = RequestStatus.PREEMPTED
    request.num_computed_tokens = 0
    if request.spec_token_ids:
        request.spec_token_ids = []
    request.num_preemptions += 1
    ...
    # Put the request back to the waiting queue.
    self.waiting.prepend_request(request)
    self.reset_preempted_req_ids.add(request.request_id)

この部分は _preempt_request:1140-1162 にあります。三つのアクションに注意:① _free_request_blocks が KV ブロックを block pool に返却(defer_block_free が有効なときは deferred_frees キューに入り fence を待つ)(_free_request_blocks:2130-2143);② encoder_cache_manager.free が同期的に encoder キャッシュを解放;③ num_computed_tokens = 0 で進捗をゼロに — ただし prefix cache ヒットのブロックは _free_request_blocks 内部で保持されます(ヒットブロックは共有で、この request の専有ではないため)。次回の再 prefill で必ずしも再計算するとは限りません。

prefill throttle の経路は独立しています。defer_prefills は schedule() の冒頭で一度計算されます(defer_prefills:436-438):throttle_prefills and not self.prefill_capacity_bound and any(not r.is_prefill_chunk for r in self.running)。三つの条件がすべて満たされたときだけ遅延します — throttle フラグは DP engine core 由来、prefill_capacity_bound は前回 waiting を空にしたかのフラグ、最後の条件は running 内に decode があるときだけ prefill を遅延させる(decode は in-progress prefill chunk より優先度が低いため)保証です。その後二箇所で効きます:running 内の in-progress prefill chunk はスキップ(running 遅延:467-471)、waiting 内の新規 prefill は直接 break(waiting 遅延:801-804)。非 throttle ステップ終了時に prefill_capacity_bound = bool(self.waiting) を更新し(capacity_bound 更新:1019-1020)、次ラウンドの throttle ステップで waiting がすでに空なら throttle を続けません。

境界と失敗 ​

  • preempt できるのは RUNNING のみ:_preempt_request の冒頭で request.status == RequestStatus.RUNNING をアサート(assert:1146-1148)。waiting 内の request は「再プリエンプト」できません。KV ブロックを占有していないためです。
  • assert: request は running から削除済み:docstring に「popped from the running queue outside of this method」と明記(NOTE:1143-1144)。_preempt_request 内部は状態の変更と waiting への戻すことだけを担当し、self.running リストには触れません。呼び出し側が先に pop しなければなりません。
  • spec_token_ids をクリア:プリエンプト時に request.spec_token_ids = [](spec クリア:1154-1155)。次回再整列時に spec decoder が期限切れの draft token を返すのを防ぎます。
  • num_preemptions をインクリメント:毎回プリエンプトで request.num_preemptions += 1(num_preemptions:1156)。このカウントは prefix cache stats で preempted=True としてマークされ(preempted マーク:721)、可観測性に使われます。
  • deferred free で並行書き込みを防止:_free_request_blocks は defer_block_free=True かつ last_sched_seq > processed_step_seq のとき、即座に free せず deferred_frees キューに push して fence を待ちます(_free_request_blocks:2130-2143)。async scheduling 下で別の in-flight step がまだこれらのブロックに書き込んでいる可能性があるためです。
  • reset_preempted_req_ids で worker に通知:プリエンプトされた request id は reset_preempted_req_ids に入り(reset_preempted:1162)、SchedulerOutput.preempted_req_ids に詰め込まれます(SchedulerOutput.preempted_req_ids:1100)。worker は対応する KV cache 状態をクリアする必要があることを知ります。
  • throttle は decode に影響しない:defer_prefills は is_prefill_chunk の request だけをスキップ(running 遅延条件:467)。純 decode の running request は通常通りスケジュールされ、throttle ステップでも GPU が遊がないことを保証します。
  • async KV load はプリエンプト不可:load_kv_async を使う request は WAITING_FOR_REMOTE_KVS 状態に入り(WAITING_FOR_REMOTE_KVS:950)、running には入りません。_inflight_prefill_reserved_blocks(_inflight_prefill_reserved_blocks:2400)でブロックを予約し、running に入った後にまたプリエンプトされて KV 転送状態が混乱するのを防ぎます。

まとめ ​

プリエンプションは KV cache が逼迫したときの自衛です:running から低優先度のものを一つ選んで追い出し、ブロックを解放して高優先度 request に譲り、追い出されたものは waiting の先頭に戻して次ステップで再整列を待ちます。prefill throttle は DP シナリオでの別の抑制です:非整列ステップでは decode だけを走らせ、prefill を整列ステップまで遅延させて一緒に行います。両者ともスケジューラが「もう一つ prefill を詰め込む」のを自発的に諦めてシステムの安定性を取る行動です。プリエンプションがスケジューリング主ループのどこで起きるかは Scheduler.schedule を、追い出された request がどう再整列されるかは RequestQueue を参照してください。

公式資料:vLLM ドキュメント · README