@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9a9f48dff7
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[V1] [P/D] Add Support for KV Load Failure Recovery (#19330)

Signed-off-by: David Ben-David <davidb@pliops.com> Co-authored-by: David Ben-David <davidb@pliops.com>

David Ben-David committed Oct 1, 2025 at 00:57 UTC 9a9f48dff7d63f752d1b787499d469f80f1d5f0e
24 files changed +1035 -82
examples/offline_inference/kv_load_failure_recovery/README.md new
+30
@@ -0,0 +1,30 @@
1 +# KV Load Failure Recovery Test
2 +
3 +This example builds upon the `disaggregated-prefill-v1` example in `examples/offline_inference`.
4 +
5 +It demonstrates vLLM's ability to recover from KV load failures in both synchronous and asynchronous loading modes. The goal is to verify that vLLM correctly identifies invalid KV blocks, reschedules the affected requests, and ensures successful and consistent output.
6 +
7 +## Files
8 +
9 +- `prefill_example.py` – performs the prefill stage and saves KV data (same as in `disaggregated-prefill-v1`).
10 +- `decode_example.py` – performs the decode stage. Accepts:
11 + - `--simulate-failure`: simulates KV load failure using a custom connector.
12 + - `--async-load`: enables asynchronous KV loading mode.
13 +- `rogue_shared_storage_connector.py` – defines `RogueSharedStorageConnector`, a subclass of `SharedStorageConnector`, that simulates missing or corrupted external KV blocks by failing to load blocks for the first decode request.
14 +- `run.sh` – orchestrates the test: runs the prefill stage, then three decode stages:
15 + 1. Normal decode (baseline).
16 + 2. Decode with simulated sync KV load failure.
17 + 3. Decode with simulated async KV load failure.
18 +
19 + Finally, it compares the output of the baseline with the recovered outputs to verify correctness.
20 +
21 +## How It Works
22 +
23 +- The test dynamically loads `RogueSharedStorageConnector` via `KVTransferConfig.kv_connector_module_path`, enabling controlled simulation of load failures without modifying the original connector.
24 +- The decode stages that simulate failure are expected to trigger recovery logic in vLLM, resulting in the same output as the baseline decode.
25 +- If recovery fails, the script prints a unified diff of the output mismatch and exits with error.
26 +
27 +## Usage
28 +
29 +```bash
30 +./run.sh
examples/offline_inference/kv_load_failure_recovery/decode_example.py new
+85
@@ -0,0 +1,85 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import argparse
4 +
5 +from vllm import LLM, SamplingParams
6 +from vllm.config import KVTransferConfig
7 +
8 +
9 +def read_prompts():
10 + """Read prompts from prefill_output.txt"""
11 + prompts = []
12 + try:
13 + with open("prefill_output.txt") as f:
14 + for line in f:
15 + prompts.append(line.strip())
16 + print(f"Loaded {len(prompts)} prompts from prefill_output.txt")
17 + return prompts
18 + except FileNotFoundError:
19 + print("Error: prefill_output.txt file not found")
20 + exit(-1)
21 +
22 +
23 +def main():
24 + prompts = read_prompts()
25 + sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
26 +
27 + parser = argparse.ArgumentParser()
28 + parser.add_argument(
29 + "--simulate-failure", action="store_true", help="Simulate KV load failure."
30 + )
31 + parser.add_argument(
32 + "--async-load", action="store_true", help="Simulate async KV load"
33 + )
34 + args = parser.parse_args()
35 +
36 + if args.simulate_failure:
37 + ktc = KVTransferConfig(
38 + kv_connector="RogueSharedStorageConnector",
39 + kv_role="kv_both",
40 + kv_connector_extra_config={
41 + "shared_storage_path": "local_storage",
42 + "async_load": args.async_load,
43 + },
44 + kv_connector_module_path="rogue_shared_storage_connector",
45 + )
46 + out_file = (
47 + "async_decode_recovered_output.txt"
48 + if args.async_load
49 + else "sync_decode_recovered_output.txt"
50 + )
51 + else:
52 + ktc = KVTransferConfig(
53 + kv_connector="SharedStorageConnector",
54 + kv_role="kv_both",
55 + kv_connector_extra_config={
56 + "shared_storage_path": "local_storage",
57 + },
58 + )
59 + out_file = "decode_output.txt"
60 +
61 + llm = LLM(
62 + model="meta-llama/Llama-3.2-1B-Instruct",
63 + enforce_eager=True,
64 + gpu_memory_utilization=0.8,
65 + max_num_batched_tokens=64,
66 + max_num_seqs=16,
67 + kv_transfer_config=ktc,
68 + )
69 +
70 + outputs = llm.generate(prompts, sampling_params)
71 +
72 + sep_str = "-" * 30
73 + with open(out_file, "w", encoding="utf-8") as f:
74 + for output in outputs:
75 + prompt = output.prompt
76 + generated_text = output.outputs[0].text
77 + out_str = f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}"
78 + print(out_str)
79 + print(sep_str)
80 + f.write(out_str)
81 + f.write(sep_str)
82 +
83 +
84 +if __name__ == "__main__":
85 + main()
examples/offline_inference/kv_load_failure_recovery/prefill_example.py new
+58
@@ -0,0 +1,58 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm import LLM, SamplingParams
5 +from vllm.config import KVTransferConfig
6 +
7 +
8 +def read_prompts():
9 + context = "Hi " * 1000
10 + context2 = "Hey " * 500
11 + return [
12 + context + "Hello, my name is",
13 + context + "The capital of France is",
14 + context2 + "Your name is",
15 + context2 + "The capital of China is",
16 + ]
17 +
18 +
19 +def main():
20 + prompts = read_prompts()
21 +
22 + sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
23 +
24 + llm = LLM(
25 + model="meta-llama/Llama-3.2-1B-Instruct",
26 + enforce_eager=True,
27 + gpu_memory_utilization=0.8,
28 + kv_transfer_config=KVTransferConfig(
29 + kv_connector="SharedStorageConnector",
30 + kv_role="kv_both",
31 + kv_connector_extra_config={"shared_storage_path": "local_storage"},
32 + ),
33 + ) # , max_model_len=2048, max_num_batched_tokens=2048)
34 +
35 + # 1ST generation (prefill instance)
36 + outputs = llm.generate(
37 + prompts,
38 + sampling_params,
39 + )
40 +
41 + new_prompts = []
42 + print("-" * 30)
43 + for output in outputs:
44 + prompt = output.prompt
45 + generated_text = output.outputs[0].text
46 + new_prompts.append(prompt + generated_text)
47 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
48 + print("-" * 30)
49 +
50 + # Write new_prompts to prefill_output.txt
51 + with open("prefill_output.txt", "w") as f:
52 + for prompt in new_prompts:
53 + f.write(prompt + "\n")
54 + print(f"Saved {len(new_prompts)} prompts to prefill_output.txt")
55 +
56 +
57 +if __name__ == "__main__":
58 + main()
examples/offline_inference/kv_load_failure_recovery/rogue_shared_storage_connector.py new
+145
@@ -0,0 +1,145 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# ruff: noqa: E501
4 +import logging
5 +from dataclasses import dataclass, field
6 +from typing import TYPE_CHECKING, Optional
7 +
8 +from vllm.config import VllmConfig
9 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
10 + KVConnectorMetadata,
11 + KVConnectorRole,
12 +)
13 +from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import (
14 + SharedStorageConnector,
15 + SharedStorageConnectorMetadata,
16 +)
17 +from vllm.forward_context import ForwardContext
18 +from vllm.v1.core.kv_cache_manager import KVCacheBlocks
19 +from vllm.v1.request import Request
20 +
21 +if TYPE_CHECKING:
22 + from vllm.v1.core.sched.output import SchedulerOutput
23 +
24 +logger = logging.getLogger()
25 +logging.basicConfig(level=logging.INFO)
26 +
27 +
28 +@dataclass
29 +class RogueSharedStorageConnectorMetadata(SharedStorageConnectorMetadata):
30 + req_to_block_ids: dict[str, set[int]] = field(default_factory=dict)
31 +
32 + @classmethod
33 + def from_base(cls, base: SharedStorageConnectorMetadata):
34 + return cls(requests=base.requests)
35 +
36 +
37 +class RogueSharedStorageConnector(SharedStorageConnector):
38 + def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
39 + super().__init__(vllm_config=vllm_config, role=role)
40 + self._async_load = vllm_config.kv_transfer_config.get_from_extra_config(
41 + "async_load", False
42 + )
43 + self._invalid_block_ids: set = None
44 + self._seen_requests: set = set()
45 + self._req_to_block_ids: dict[str, list[int]] = dict()
46 +
47 + def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None:
48 + assert isinstance(connector_metadata, RogueSharedStorageConnectorMetadata)
49 + index, failed_request = next(
50 + (
51 + (i, x)
52 + for i, x in enumerate(connector_metadata.requests)
53 + if not x.is_store
54 + ),
55 + (None, None),
56 + )
57 + if index is not None:
58 + del connector_metadata.requests[index]
59 + self._invalid_block_ids = set(
60 + (
61 + failed_request.slot_mapping[:: self._block_size] // self._block_size
62 + ).tolist()
63 + )
64 + logger.info(
65 + "Simulating failure to load all KV blocks for the "
66 + "first load request. Total blocks: %d",
67 + len(self._invalid_block_ids),
68 + )
69 + super().bind_connector_metadata(connector_metadata)
70 +
71 + def clear_connector_metadata(self) -> None:
72 + self._invalid_block_ids = None
73 + super().clear_connector_metadata()
74 +
75 + def start_load_kv(self, forward_context: ForwardContext, **kwargs) -> None:
76 + if self._async_load and forward_context.attn_metadata is None:
77 + # Bypass sanity check in super().start_load_kv
78 + forward_context.attn_metadata = "None"
79 +
80 + super().start_load_kv(forward_context, **kwargs)
81 +
82 + def get_finished(
83 + self, finished_req_ids: set[str]
84 + ) -> tuple[Optional[set[str]], Optional[set[str]]]:
85 + if self._async_load:
86 + meta = self._get_connector_metadata()
87 + assert isinstance(meta, RogueSharedStorageConnectorMetadata)
88 + if meta.req_to_block_ids:
89 + return None, set(meta.req_to_block_ids)
90 +
91 + return None, None
92 +
93 + def get_block_ids_with_load_errors(self) -> set[int]:
94 + return self._invalid_block_ids
95 +
96 + def get_num_new_matched_tokens(
97 + self,
98 + request: Request,
99 + num_computed_tokens: int,
100 + ) -> tuple[int, bool]:
101 + if request.request_id in self._seen_requests:
102 + return 0, False
103 +
104 + self._seen_requests.add(request.request_id)
105 +
106 + num_tokens, _ = super().get_num_new_matched_tokens(request, num_computed_tokens)
107 + return num_tokens, self._async_load and num_tokens > 0
108 +
109 + def update_state_after_alloc(
110 + self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int
111 + ):
112 + """
113 + Update KVConnector state after block allocation.
114 +
115 + If blocks were allocated, add to _requests_need_load,
116 + such that we load the KVs in the next forward pass.
117 + """
118 + super().update_state_after_alloc(request, blocks, num_external_tokens)
119 +
120 + if num_external_tokens > 0:
121 + self._req_to_block_ids[request.request_id] = blocks.get_block_ids()[0]
122 +
123 + def build_connector_meta(
124 + self,
125 + scheduler_output: "SchedulerOutput",
126 + ) -> KVConnectorMetadata:
127 + if not self._async_load:
128 + base = super().build_connector_meta(scheduler_output)
129 + meta = RogueSharedStorageConnectorMetadata.from_base(base)
130 + else:
131 + meta = RogueSharedStorageConnectorMetadata()
132 + if self._requests_need_load:
133 + for req_id, request in self._requests_need_load.items():
134 + meta.add_request(
135 + token_ids=request.prompt_token_ids,
136 + block_ids=self._req_to_block_ids[req_id],
137 + block_size=self._block_size,
138 + is_store=False,
139 + mm_hashes=[],
140 + )
141 + # Clear state
142 + self._requests_need_load.clear()
143 + meta.req_to_block_ids = self._req_to_block_ids
144 + self._req_to_block_ids = dict()
145 + return meta
examples/offline_inference/kv_load_failure_recovery/run.sh new
+33
@@ -0,0 +1,33 @@
1 +#!/bin/bash
2 +
3 +# Constants
4 +SHARED_STORAGE_DIR="local_storage"
5 +PREFILL_OUTPUT="prefill_output.txt"
6 +DECODE_OUTPUT="decode_output.txt"
7 +SYNC_DECODE_RECOVERED_OUTPUT="sync_decode_recovered_output.txt"
8 +ASYNC_DECODE_RECOVERED_OUTPUT="async_decode_recovered_output.txt"
9 +
10 +# Cleanup
11 +rm -rf "$SHARED_STORAGE_DIR"
12 +rm -f "$PREFILL_OUTPUT" "$DECODE_OUTPUT" "$SYNC_DECODE_RECOVERED_OUTPUT" "$ASYNC_DECODE_RECOVERED_OUTPUT"
13 +
14 +# Run inference examples
15 +VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 prefill_example.py
16 +VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 decode_example.py
17 +VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 decode_example.py --simulate-failure
18 +VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 decode_example.py --simulate-failure --async-load
19 +
20 +# Compare outputs
21 +if ! cmp -s "$DECODE_OUTPUT" "$SYNC_DECODE_RECOVERED_OUTPUT"; then
22 + echo "❌ Outputs differ: sync recovery failed."
23 + diff -u "$DECODE_OUTPUT" "$SYNC_DECODE_RECOVERED_OUTPUT"
24 + exit 1
25 +fi
26 +
27 +if ! cmp -s "$DECODE_OUTPUT" "$ASYNC_DECODE_RECOVERED_OUTPUT"; then
28 + echo "❌ Outputs differ: async recovery failed."
29 + diff -u "$DECODE_OUTPUT" "$ASYNC_DECODE_RECOVERED_OUTPUT"
30 + exit 1
31 +fi
32 +
33 +echo "✅ Outputs match: recovery successful."
tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py new
+341
@@ -0,0 +1,341 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Callable
5 +from unittest.mock import Mock
6 +
7 +import pytest
8 +
9 +from vllm.v1.core.sched.scheduler import Scheduler
10 +from vllm.v1.request import Request, RequestStatus
11 +
12 +from .utils import (create_model_runner_output, create_request,
13 + create_scheduler, create_vllm_config)
14 +
15 +
16 +def _make_get_num_new_matched_tokens(
17 + req_num_new_matched_tokens: dict[str, int],
18 + async_load,
19 +) -> Callable[[Request, int], tuple[int, bool]]:
20 +
21 + def get_num_new_matched_tokens(request: Request,
22 + _: int) -> tuple[int, bool]:
23 + value = req_num_new_matched_tokens.get(request.request_id, 0)
24 + return value, async_load
25 +
26 + return get_num_new_matched_tokens
27 +
28 +
29 +@pytest.fixture
30 +def scheduler():
31 + vllm_config = create_vllm_config()
32 + return create_scheduler(vllm_config)
33 +
34 +
35 +@pytest.mark.parametrize(
36 + "num_prompt_blocks,"
37 + "num_external_computed_blocks,"
38 + "invalid_block_idxs",
39 + [
40 + (100, 99, {0, 98}),
41 + (100, 99, {50, 98}),
42 + (100, 99, {98}),
43 + ],
44 +)
45 +def test_async_load_failure(
46 + scheduler: Scheduler,
47 + num_prompt_blocks: int,
48 + num_external_computed_blocks: int,
49 + invalid_block_idxs: set[int],
50 +):
51 + assert num_prompt_blocks >= num_external_computed_blocks
52 +
53 + num_prompt_tokens = num_prompt_blocks * scheduler.block_size
54 + num_external_computed_tokens = (num_external_computed_blocks *
55 + scheduler.block_size)
56 +
57 + request1 = create_request(num_tokens=num_prompt_tokens)
58 + scheduler.add_request(request=request1)
59 + request2 = create_request(num_tokens=num_prompt_tokens)
60 + scheduler.add_request(request=request2)
61 + request3 = create_request(num_tokens=num_prompt_tokens)
62 + scheduler.add_request(request=request3)
63 +
64 + # Mock KV connector method.
65 + # req_id -> num_external_computed_tokens
66 + req_num_new_matched_tokens = {
67 + request1.request_id: num_external_computed_tokens,
68 + request2.request_id: num_external_computed_tokens,
69 + request3.request_id: num_external_computed_tokens,
70 + }
71 +
72 + scheduler.connector = Mock()
73 + scheduler.connector.get_num_new_matched_tokens.side_effect = (
74 + _make_get_num_new_matched_tokens(req_num_new_matched_tokens,
75 + async_load=True))
76 + scheduler.connector.take_events.return_value = ()
77 +
78 + scheduler_output = scheduler.schedule()
79 +
80 + assert len(scheduler.waiting) == 3
81 + for request in scheduler.waiting:
82 + assert request.num_computed_tokens == 0
83 + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
84 + assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
85 +
86 + # Simulate a failure in loading some of request2 blocks.
87 + (req2_block_ids, ) = scheduler.kv_cache_manager.get_block_ids(
88 + request2.request_id)
89 + invalid_block_ids = {req2_block_ids[i] for i in invalid_block_idxs}
90 + model_runner_output = create_model_runner_output(
91 + reqs=[],
92 + finished_recving={request1.request_id, request3.request_id},
93 + invalid_block_ids=invalid_block_ids,
94 + use_eos=True)
95 +
96 + scheduler.update_from_output(scheduler_output, model_runner_output)
97 +
98 + min_invalid_block_idx = min(invalid_block_idxs)
99 +
100 + assert len(scheduler.waiting) == 3
101 + for request in scheduler.waiting:
102 + if request.request_id == request2.request_id:
103 + assert request.num_computed_tokens == (min_invalid_block_idx *
104 + scheduler.block_size)
105 + else:
106 + assert request.num_computed_tokens == 0
107 + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
108 + assert scheduler.failed_recving_kv_req_ids == {request2.request_id}
109 + assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
110 +
111 +
112 +@pytest.mark.parametrize(
113 + "num_prompt_blocks,"
114 + "num_external_computed_blocks,"
115 + "invalid_block_idxs",
116 + [
117 + (100, 99, {0, 98}),
118 + (100, 99, {50, 98}),
119 + (100, 99, {98}),
120 + ],
121 +)
122 +def test_sync_load_failure(
123 + scheduler: Scheduler,
124 + num_prompt_blocks: int,
125 + num_external_computed_blocks: int,
126 + invalid_block_idxs: set[int],
127 +):
128 + assert num_prompt_blocks >= num_external_computed_blocks
129 +
130 + num_prompt_tokens = num_prompt_blocks * scheduler.block_size
131 + num_external_computed_tokens = (num_external_computed_blocks *
132 + scheduler.block_size)
133 +
134 + request1 = create_request(num_tokens=num_prompt_tokens)
135 + scheduler.add_request(request=request1)
136 + request2 = create_request(num_tokens=num_prompt_tokens)
137 + scheduler.add_request(request=request2)
138 + request3 = create_request(num_tokens=num_prompt_tokens)
139 + scheduler.add_request(request=request3)
140 +
141 + # Mock KV connector method.
142 + # req_id -> num_external_computed_tokens
143 + req_num_new_matched_tokens = {
144 + request1.request_id: num_external_computed_tokens,
145 + request2.request_id: num_external_computed_tokens,
146 + request3.request_id: num_external_computed_tokens,
147 + }
148 +
149 + scheduler.connector = Mock()
150 + scheduler.connector.get_num_new_matched_tokens.side_effect = (
151 + _make_get_num_new_matched_tokens(req_num_new_matched_tokens,
152 + async_load=False))
153 + scheduler.connector.request_finished.return_value = (False, None)
154 + scheduler.connector.take_events.return_value = ()
155 +
156 + scheduler_output = scheduler.schedule()
157 +
158 + # req_id -> num_computed_tokens
159 + expected_computed_tokens = {
160 + request1.request_id: num_external_computed_tokens,
161 + request2.request_id: num_external_computed_tokens,
162 + request3.request_id: num_external_computed_tokens,
163 + }
164 +
165 + assert len(scheduler.running) == 3
166 + assert len(scheduler_output.scheduled_new_reqs) == 3
167 + for request in scheduler_output.scheduled_new_reqs:
168 + assert request.num_computed_tokens == expected_computed_tokens[
169 + request.req_id]
170 + assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
171 +
172 + # Simulate a failure in loading some of request2 blocks.
173 + req2_block_ids = scheduler_output.scheduled_new_reqs[1].block_ids[0]
174 + invalid_block_ids = {req2_block_ids[i] for i in invalid_block_idxs}
175 + model_runner_output = create_model_runner_output(
176 + [request1, request2, request3],
177 + invalid_block_ids=invalid_block_ids,
178 + use_eos=True)
179 +
180 + scheduler.update_from_output(scheduler_output, model_runner_output)
181 +
182 + assert len(scheduler.running) == 1
183 + assert scheduler.running[0].request_id == request2.request_id
184 + assert scheduler.running[0].num_computed_tokens == (
185 + min(invalid_block_idxs) * scheduler.block_size)
186 + assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
187 + assert scheduler.connector.request_finished.call_count == 2
188 +
189 +
190 +@pytest.mark.parametrize(
191 + "num_prompt_blocks,"
192 + "num_external_computed_blocks,"
193 + "num_common_prefix_blocks,"
194 + "invalid_block_idxs",
195 + [
196 + (100, 99, 50, {0, 49}),
197 + (100, 99, 50, {25, 49}),
198 + (100, 99, 50, {49}),
199 + ],
200 +)
201 +def test_sync_load_failure_with_shared_blocks(
202 + scheduler: Scheduler,
203 + num_prompt_blocks: int,
204 + num_external_computed_blocks: int,
205 + num_common_prefix_blocks: int,
206 + invalid_block_idxs: set[int],
207 +):
208 + assert (num_prompt_blocks >= num_external_computed_blocks >=
209 + num_common_prefix_blocks)
210 +
211 + num_prompt_tokens = num_prompt_blocks * scheduler.block_size
212 + num_external_computed_tokens = (num_external_computed_blocks *
213 + scheduler.block_size)
214 + common_prefix_len = num_common_prefix_blocks * scheduler.block_size
215 +
216 + request1 = create_request(num_tokens=num_prompt_tokens,
217 + common_prefix_len=common_prefix_len)
218 + scheduler.add_request(request=request1)
219 + request2 = create_request(num_tokens=num_prompt_tokens,
220 + common_prefix_len=common_prefix_len)
221 + scheduler.add_request(request=request2)
222 +
223 + # Mock KV connector method.
224 + # req_id -> num_external_computed_tokens
225 + req_num_new_matched_tokens = {
226 + request1.request_id: num_external_computed_tokens,
227 + }
228 +
229 + scheduler.connector = Mock()
230 + scheduler.connector.get_num_new_matched_tokens.side_effect = (
231 + _make_get_num_new_matched_tokens(req_num_new_matched_tokens,
232 + async_load=False))
233 + scheduler.connector.take_events.return_value = ()
234 +
235 + scheduler_output = scheduler.schedule()
236 +
237 + # req_id -> num_computed_tokens
238 + expected_computed_tokens = {
239 + request1.request_id: num_external_computed_tokens,
240 + request2.request_id: common_prefix_len,
241 + }
242 +
243 + assert len(scheduler.running) == 2
244 + assert len(scheduler_output.scheduled_new_reqs) == 2
245 + for request in scheduler_output.scheduled_new_reqs:
246 + assert request.num_computed_tokens == expected_computed_tokens[
247 + request.req_id]
248 + assert scheduler.connector.get_num_new_matched_tokens.call_count == 2
249 +
250 + # Simulate a failure in loading some of the shared blocks.
251 + req1_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0]
252 + invalid_block_ids = {req1_block_ids[i] for i in invalid_block_idxs}
253 + model_runner_output = create_model_runner_output(
254 + [request1, request2],
255 + invalid_block_ids=invalid_block_ids,
256 + use_eos=True)
257 +
258 + scheduler.update_from_output(scheduler_output, model_runner_output)
259 +
260 + # req_id -> num_computed_tokens
261 + # all the common prefix blocks will be computed by request1
262 + expected_computed_tokens = {
263 + request1.request_id: min(invalid_block_idxs) * scheduler.block_size,
264 + request2.request_id: common_prefix_len,
265 + }
266 +
267 + assert len(scheduler.running) == 2
268 + for request in scheduler.running:
269 + assert request.num_computed_tokens == expected_computed_tokens[
270 + request.request_id]
271 + assert scheduler.connector.get_num_new_matched_tokens.call_count == 2
272 +
273 +
274 +@pytest.mark.parametrize(
275 + "num_prompt_blocks,"
276 + "num_external_computed_blocks,"
277 + "invalid_block_idxs",
278 + [
279 + (100, 99, {0, 50, 98}),
280 + (100, 99, {98, 50, 0}),
281 + ],
282 +)
283 +def test_async_progressive_load_failure(
284 + scheduler: Scheduler,
285 + num_prompt_blocks: int,
286 + num_external_computed_blocks: int,
287 + invalid_block_idxs: set[int],
288 +):
289 + assert num_prompt_blocks >= num_external_computed_blocks
290 +
291 + num_prompt_tokens = num_prompt_blocks * scheduler.block_size
292 + num_external_computed_tokens = (num_external_computed_blocks *
293 + scheduler.block_size)
294 +
295 + request = create_request(num_tokens=num_prompt_tokens)
296 + scheduler.add_request(request=request)
297 +
298 + # Mock KV connector method.
299 + # req_id -> num_external_computed_tokens
300 + req_num_new_matched_tokens = {
301 + request.request_id: num_external_computed_tokens,
302 + }
303 +
304 + scheduler.connector = Mock()
305 + scheduler.connector.get_num_new_matched_tokens.side_effect = (
306 + _make_get_num_new_matched_tokens(req_num_new_matched_tokens,
307 + async_load=True))
308 + scheduler.connector.take_events.return_value = ()
309 +
310 + scheduler_output = scheduler.schedule()
311 +
312 + assert len(scheduler.waiting) == 1
313 + assert scheduler.waiting.peek_request().request_id == request.request_id
314 + assert request.num_computed_tokens == 0
315 + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
316 + assert scheduler.connector.get_num_new_matched_tokens.call_count == 1
317 +
318 + min_invalid_block_idx = max(invalid_block_idxs) + 1
319 + # Simulate failures when progressively loading request blocks.
320 + for invalid_block_idx in invalid_block_idxs:
321 + (req_block_ids, ) = scheduler.kv_cache_manager.get_block_ids(
322 + request.request_id)
323 + invalid_block_ids = {req_block_ids[invalid_block_idx]}
324 + model_runner_output = create_model_runner_output(
325 + reqs=[],
326 + finished_recving=set(),
327 + invalid_block_ids=invalid_block_ids,
328 + use_eos=True)
329 +
330 + scheduler.update_from_output(scheduler_output, model_runner_output)
331 +
332 + min_invalid_block_idx = min(min_invalid_block_idx, invalid_block_idx)
333 +
334 + assert len(scheduler.waiting) == 1
335 + assert scheduler.waiting.peek_request(
336 + ).request_id == request.request_id
337 + assert request.num_computed_tokens == (min_invalid_block_idx *
338 + scheduler.block_size)
339 + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
340 + assert scheduler.failed_recving_kv_req_ids == {request.request_id}
341 + assert scheduler.connector.get_num_new_matched_tokens.call_count == 1
tests/v1/kv_connector/unit/test_offloading_connector.py
+2 -2
@@ -281,8 +281,8 @@ class RequestRunner:
281
282 model_runner_output = create_model_runner_output(
283 reqs=self.scheduler.running,
284 - finished_sending=list(finished_sending),
285 - finished_recving=list(finished_recving),
284 + finished_sending=finished_sending,
285 + finished_recving=finished_recving,
286 token_id=token_id)
287
288 if self.scheduler.running:
tests/v1/kv_connector/unit/test_output_aggreagator.py
+51 -23
@@ -15,26 +15,38 @@ class DummyModelRunnerOutput(ModelRunnerOutput):
15
16 def __init__(self,
17 finished_sending: Optional[set[str]] = None,
18 - finished_recving: Optional[set[str]] = None):
18 + finished_recving: Optional[set[str]] = None,
19 + invalid_block_ids: Optional[set[int]] = None):
20 self.kv_connector_output = KVConnectorOutput(
21 finished_sending=finished_sending,
22 finished_recving=finished_recving,
22 - )
23 + invalid_block_ids=invalid_block_ids or set())
24
25 def __repr__(self):
26 return (
27 f"DummyModelRunnerOutput("
28 f"finished_sending={self.kv_connector_output.finished_sending},"
28 - f"finished_recving={self.kv_connector_output.finished_recving})")
29 + f"finished_recving={self.kv_connector_output.finished_recving})"
30 + f"invalid_block_ids={self.kv_connector_output.invalid_block_ids})")
31
32
33 def test_aggregate_workers_output():
34 aggregator = KVOutputAggregator(world_size=2)
35
36 + output1 = DummyModelRunnerOutput()
37 + output2 = DummyModelRunnerOutput()
38 +
39 + aggregated = aggregator.aggregate([output1, output2])
40 +
41 + assert aggregated is output1
42 + aggregated = aggregated.kv_connector_output
43 + assert aggregated.finished_sending is None
44 + assert aggregated.finished_recving is None
45 + assert not aggregated.invalid_block_ids
46 +
47 output1 = DummyModelRunnerOutput(finished_sending={'req1'},
48 finished_recving={'req2'})
36 - output2 = DummyModelRunnerOutput(finished_sending=None,
37 - finished_recving=None)
49 + output2 = DummyModelRunnerOutput(invalid_block_ids={1})
50
51 aggregated = aggregator.aggregate([output1, output2])
52
@@ -42,11 +54,10 @@ def test_aggregate_workers_output():
54 aggregated = aggregated.kv_connector_output
55 assert aggregated.finished_sending is None
56 assert aggregated.finished_recving is None
57 + assert aggregated.invalid_block_ids == {1}
58
46 - output1 = DummyModelRunnerOutput(finished_sending=None,
47 - finished_recving=None)
48 - output2 = DummyModelRunnerOutput(finished_sending={'req1'},
49 - finished_recving=None)
59 + output1 = DummyModelRunnerOutput(invalid_block_ids={2})
60 + output2 = DummyModelRunnerOutput(finished_sending={'req1'})
61
62 aggregated = aggregator.aggregate([output1, output2])
63
@@ -54,11 +65,11 @@ def test_aggregate_workers_output():
65 aggregated = aggregated.kv_connector_output
66 assert aggregated.finished_sending == {'req1'}
67 assert aggregated.finished_recving is None
68 + assert aggregated.invalid_block_ids == {2}
69
58 - output1 = DummyModelRunnerOutput(finished_sending=None,
59 - finished_recving=None)
60 - output2 = DummyModelRunnerOutput(finished_sending={'req1'},
61 - finished_recving={'req2'})
70 + output1 = DummyModelRunnerOutput(invalid_block_ids={3, 4})
71 + output2 = DummyModelRunnerOutput(finished_recving={'req2'},
72 + invalid_block_ids={4, 5})
73
74 aggregated = aggregator.aggregate([output1, output2])
75
@@ -66,6 +77,7 @@ def test_aggregate_workers_output():
77 aggregated = aggregated.kv_connector_output
78 assert aggregated.finished_sending is None
79 assert aggregated.finished_recving == {'req2'}
80 + assert aggregated.invalid_block_ids == {3, 4, 5}
81
82
83 def test_async_aggregate_workers_output():
@@ -75,10 +87,26 @@ def test_async_aggregate_workers_output():
87 future2: Future[DummyModelRunnerOutput] = Future()
88 result_future = aggregator.async_aggregate([future1, future2])
89
90 + output1 = DummyModelRunnerOutput()
91 + output2 = DummyModelRunnerOutput()
92 + future1.set_result(output1)
93 + future2.set_result(output2)
94 +
95 + assert result_future.done()
96 + aggregated = result_future.result()
97 + assert aggregated is output1
98 + aggregated = aggregated.kv_connector_output
99 + assert aggregated.finished_sending is None
100 + assert aggregated.finished_recving is None
101 + assert not aggregated.invalid_block_ids
102 +
103 + future1 = Future()
104 + future2 = Future()
105 + result_future = aggregator.async_aggregate([future1, future2])
106 +
107 output1 = DummyModelRunnerOutput(finished_sending={'req1'},
108 finished_recving={'req2'})
80 - output2 = DummyModelRunnerOutput(finished_sending=None,
81 - finished_recving=None)
109 + output2 = DummyModelRunnerOutput(invalid_block_ids={1})
110 future1.set_result(output1)
111 future2.set_result(output2)
112
@@ -88,15 +116,14 @@ def test_async_aggregate_workers_output():
116 aggregated = aggregated.kv_connector_output
117 assert aggregated.finished_sending is None
118 assert aggregated.finished_recving is None
119 + assert aggregated.invalid_block_ids == {1}
120
121 future1 = Future()
122 future2 = Future()
123 result_future = aggregator.async_aggregate([future1, future2])
124
96 - output1 = DummyModelRunnerOutput(finished_sending=None,
97 - finished_recving=None)
98 - output2 = DummyModelRunnerOutput(finished_sending={'req1'},
99 - finished_recving=None)
125 + output1 = DummyModelRunnerOutput(invalid_block_ids={2})
126 + output2 = DummyModelRunnerOutput(finished_sending={'req1'})
127 future1.set_result(output1)
128 future2.set_result(output2)
129
@@ -106,15 +133,15 @@ def test_async_aggregate_workers_output():
133 aggregated = aggregated.kv_connector_output
134 assert aggregated.finished_sending == {'req1'}
135 assert aggregated.finished_recving is None
136 + assert aggregated.invalid_block_ids == {2}
137
138 future1 = Future()
139 future2 = Future()
140 result_future = aggregator.async_aggregate([future1, future2])
141
114 - output1 = DummyModelRunnerOutput(finished_sending=None,
115 - finished_recving=None)
116 - output2 = DummyModelRunnerOutput(finished_sending={'req1'},
117 - finished_recving={'req2'})
142 + output1 = DummyModelRunnerOutput(invalid_block_ids={3, 4})
143 + output2 = DummyModelRunnerOutput(finished_recving={'req2'},
144 + invalid_block_ids={4, 5})
145 future1.set_result(output1)
146 future2.set_result(output2)
147
@@ -124,3 +151,4 @@ def test_async_aggregate_workers_output():
151 aggregated = aggregated.kv_connector_output
152 assert aggregated.finished_sending is None
153 assert aggregated.finished_recving == {'req2'}
154 + assert aggregated.invalid_block_ids == {3, 4, 5}
tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py
+3 -3
@@ -92,7 +92,7 @@ def test_basic_lifecycle():
92 # (3b): execute_model()
93 model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
94 model_runner_output.kv_connector_output = KVConnectorOutput(
95 - finished_sending=[request_id])
95 + finished_sending={request_id})
96
97 # (3c): update_from_output()
98 scheduler.update_from_output(scheduler_output, model_runner_output)
@@ -139,7 +139,7 @@ def test_short_prompt_lifecycle():
139 scheduler_output = scheduler.schedule()
140 # Use create_model_runner_output to pass kv_connector_output along
141 model_runner_output = create_model_runner_output(
142 - reqs=[request], finished_sending=[request.request_id])
142 + reqs=[request], finished_sending={request.request_id})
143 scheduler.update_from_output(scheduler_output, model_runner_output)
144 assert_scheduler_empty(scheduler)
145
@@ -195,6 +195,6 @@ def test_prefix_cache_lifecycle():
195 scheduler_output = scheduler.schedule()
196 model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
197 model_runner_output.kv_connector_output = KVConnectorOutput(
198 - finished_sending=[request_remote.request_id])
198 + finished_sending={request_remote.request_id})
199 scheduler.update_from_output(scheduler_output, model_runner_output)
200 assert_scheduler_empty(scheduler)
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
+7 -7
@@ -78,7 +78,7 @@ def test_basic_lifecycle():
78 # (2b): forward(): request finishes recv.
79 model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
80 model_runner_output.kv_connector_output = KVConnectorOutput(
81 - finished_recving=[request_id])
81 + finished_recving={request_id})
82
83 # (2c): update_from_output():
84 engine_core_outputs = scheduler.update_from_output(scheduler_output,
@@ -197,7 +197,7 @@ def test_interleaved_lifecycle():
197
198 model_runner_output = create_model_runner_output(
199 [request_local_a, request_local_b],
200 - finished_recving=[request_remote.request_id])
200 + finished_recving={request_remote.request_id})
201 scheduler.update_from_output(scheduler_output, model_runner_output)
202
203 # STEP 5: RECVed KVs are sent to ModelRunner.
@@ -246,16 +246,16 @@ def test_no_spurious_prefix_caching():
246 request_id=1,
247 block_size=BLOCK_SIZE,
248 num_tokens=NUM_TOKENS,
249 + common_prefix_len=NUM_TOKENS,
250 do_remote_prefill=True,
250 - use_all_1s_for_prompt_tokens=True,
251 )
252
253 request_local = create_request(
254 request_id=2,
255 block_size=BLOCK_SIZE,
256 num_tokens=NUM_TOKENS,
257 + common_prefix_len=NUM_TOKENS,
258 do_remote_prefill=False,
258 - use_all_1s_for_prompt_tokens=True,
259 )
260
261 # Schedule the remote prefill request. This should not
@@ -322,7 +322,7 @@ def test_full_block_prompt():
322 scheduler_output = scheduler.schedule()
323 model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
324 model_runner_output.kv_connector_output = KVConnectorOutput(
325 - finished_recving=[request_id])
325 + finished_recving={request_id})
326 scheduler.update_from_output(scheduler_output, model_runner_output)
327 assert len(scheduler.waiting) == 1
328 assert (request_id in scheduler.finished_recving_kv_req_ids)
@@ -402,7 +402,7 @@ def test_cannot_schedule_after_recv():
402 # Step 3: finish recving (5 blocks in use)
403 scheduler_output = scheduler.schedule()
404 model_runner_output = create_model_runner_output(
405 - reqs=[request_normal], finished_recving=[request_remote.request_id])
405 + reqs=[request_normal], finished_recving={request_remote.request_id})
406 scheduler.update_from_output(scheduler_output, model_runner_output)
407 assert len(scheduler.running) == 1
408 assert len(scheduler.waiting) == 1
@@ -516,7 +516,7 @@ def test_cannot_recv():
516 # Step 5: finish recving (5 blocks in use)
517 scheduler_output = scheduler.schedule()
518 model_runner_output = create_model_runner_output(
519 - reqs=[], finished_recving=[request_remote.request_id])
519 + reqs=[], finished_recving={request_remote.request_id})
520 scheduler.update_from_output(scheduler_output, model_runner_output)
521 assert len(scheduler.running) == 0
522 assert len(scheduler.waiting) == 1
tests/v1/kv_connector/unit/utils.py
+31 -18
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import tempfile
4 from collections import defaultdict
5 +from itertools import count
6 from typing import Any, Callable, Optional
7
8 import torch
@@ -61,12 +62,15 @@ def create_vllm_config(
62 max_num_seqs: int = 16,
63 max_num_batched_tokens: int = 64,
64 block_size: int = 16,
65 + max_model_len: int = 10000,
66 + enable_chunked_prefill: bool = True,
67 ) -> VllmConfig:
68 """Initialize VllmConfig For Testing."""
69 scheduler_config = SchedulerConfig(
70 max_num_seqs=max_num_seqs,
71 max_num_batched_tokens=max_num_batched_tokens,
69 - max_model_len=max_num_batched_tokens,
72 + max_model_len=max_model_len,
73 + enable_chunked_prefill=enable_chunked_prefill,
74 )
75 model_config = ModelConfig(
76 model=model,
@@ -117,19 +121,27 @@ def create_scheduler(
121 )
122
123
124 +_request_count = count(1)
125 _none_hash_initialized = False
126
127
123 -def create_request(request_id: int,
124 - num_tokens: int = 10,
125 - max_tokens: int = 16,
126 - do_remote_decode: bool = False,
127 - do_remote_prefill: bool = False,
128 - use_all_1s_for_prompt_tokens: bool = False,
129 - num_remote_blocks: int = 3,
130 - block_size: int = 16,
131 - hash_fn: Callable = sha256) -> Request:
128 +def create_request(
129 + request_id: Optional[int] = None,
130 + num_tokens: int = 10,
131 + common_prefix_len=0,
132 + max_tokens: int = 16,
133 + do_remote_decode: bool = False,
134 + do_remote_prefill: bool = False,
135 + num_remote_blocks: int = 3,
136 + block_size: int = 16,
137 + hash_fn: Callable = sha256,
138 +) -> Request:
139 """Make dummy request for testing."""
140 + assert num_tokens >= common_prefix_len >= 0
141 +
142 + if request_id is None:
143 + request_id = next(_request_count)
144 +
145 global _none_hash_initialized
146 if not _none_hash_initialized:
147 init_none_hash(hash_fn)
@@ -153,10 +165,9 @@ def create_request(request_id: int,
165 max_tokens = 1 if do_remote_decode else max_tokens
166 sampling_params = SamplingParams(max_tokens=max_tokens)
167
156 - if use_all_1s_for_prompt_tokens:
157 - prompt_token_ids = [1] * num_tokens
158 - else:
159 - prompt_token_ids = [i * request_id for i in range(num_tokens)]
168 + common_prefix = [1] * common_prefix_len if common_prefix_len > 0 else []
169 + suffix = [i * request_id for i in range(num_tokens - common_prefix_len)]
170 + prompt_token_ids = common_prefix + suffix
171
172 req = Request(
173 request_id=f"id-{request_id}",
@@ -173,8 +184,9 @@ def create_request(request_id: int,
184
185 def create_model_runner_output(
186 reqs: list[Request],
176 - finished_sending: Optional[list[str]] = None,
177 - finished_recving: Optional[list[str]] = None,
187 + finished_sending: Optional[set[str]] = None,
188 + finished_recving: Optional[set[str]] = None,
189 + invalid_block_ids: Optional[set[int]] = None,
190 use_eos: bool = False,
191 token_id: int = 0,
192 ) -> ModelRunnerOutput:
@@ -189,10 +201,11 @@ def create_model_runner_output(
201 sampled_token_ids = [[sampled_token] for _ in req_ids]
202
203 kv_connector_output = None if (
192 - finished_sending is None
193 - and finished_recving is None) else KVConnectorOutput(
204 + finished_sending is None and finished_recving is None
205 + and invalid_block_ids is None) else KVConnectorOutput(
206 finished_sending=finished_sending,
207 finished_recving=finished_recving,
208 + invalid_block_ids=invalid_block_ids or set(),
209 )
210
211 # Make output data structure.
tests/v1/worker/test_gpu_model_runner.py
+1
@@ -250,6 +250,7 @@ def test_update_states_request_resumed(model_runner, dist_init):
250 new_token_ids=[[]],
251 new_block_ids=([[0]], ),
252 num_computed_tokens=[0],
253 + num_output_tokens=[0],
254 )
255
256 scheduler_output = SchedulerOutput(
vllm/distributed/kv_transfer/kv_connector/utils.py
+5 -1
@@ -117,7 +117,7 @@ def get_kv_connector_cache_layout():
117
118
119 class KVOutputAggregator:
120 - """Utility class to aggregate the output of all workers into a single
120 + """Utility class to aggregate the output of all workers into a single
121 output corresponding to Rank 0 for scheduler."""
122
123 def __init__(self, world_size: int):
@@ -143,6 +143,7 @@ class KVOutputAggregator:
143 finished_sending = set[str]()
144 finished_recving = set[str]()
145 aggregated_kv_connector_stats = None
146 + invalid_block_ids = set[int]()
147 for model_runner_output in outputs:
148 output = model_runner_output.kv_connector_output
149 if not output:
@@ -165,6 +166,8 @@ class KVOutputAggregator:
166 aggregated_kv_connector_stats = \
167 aggregated_kv_connector_stats.aggregate(kv_connector_stats)
168
169 + invalid_block_ids |= output.invalid_block_ids
170 +
171 # select output of the worker specified by output_rank
172 output = outputs[output_rank]
173
@@ -172,6 +175,7 @@ class KVOutputAggregator:
175 finished_sending=finished_sending or None,
176 finished_recving=finished_recving or None,
177 kv_connector_stats=aggregated_kv_connector_stats or None,
178 + invalid_block_ids=invalid_block_ids,
179 )
180
181 return output
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+29 -2
@@ -229,6 +229,26 @@ class KVConnectorBase_V1(ABC):
229 """
230 return None, None
231
232 + def get_block_ids_with_load_errors(self) -> set[int]:
233 + """
234 + Get the set of block IDs that failed to load.
235 +
236 + Returns:
237 + Set of block IDs that encountered load errors.
238 + Empty set if no load errors occurred.
239 +
240 + Notes:
241 + - Applies to both sync- and async-loading requests.
242 + - Async loading: failed blocks may be reported in any forward pass
243 + up to and including the pass where the request ID is returned by
244 + `get_finished()`. Even if failures occur, the request must still
245 + be reported via `get_finished()`, and the failed block IDs must
246 + appear here no later than that same pass.
247 + - Sync loading: failed blocks should be reported in the forward
248 + pass in which they are detected.
249 + """
250 + return set()
251 +
252 def shutdown(self):
253 """
254 Shutdown the connector. This is called when the worker process
@@ -264,14 +284,21 @@ class KVConnectorBase_V1(ABC):
284
285 Returns:
286 A tuple with the following elements:
267 - - An optional number of tokens that can be loaded from the
268 - external KV cache beyond what is already computed.
287 + - An optional number of tokens that can be loaded from the
288 + external KV cache beyond what is already computed.
289 If None, it means that the connector needs more time to
290 determine the number of matched tokens, and the scheduler
291 should query for this request again later.
292 - `True` if external KV cache tokens will be loaded
293 asynchronously (between scheduler steps). Must be
294 'False' if the first element is 0.
295 +
296 + Notes:
297 + The connector should only consider the largest prefix of prompt-
298 + tokens for which KV cache is actually available at the time of the
299 + call. If the cache cannot be loaded for some tokens (e.g., due to
300 + connectivity issues or eviction), those tokens must not be taken
301 + into account.
302 """
303 pass
304
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+6
@@ -189,6 +189,12 @@ class MultiConnector(KVConnectorBase_V1):
189
190 return finished_sending or None, finished_recving or None
191
192 + def get_block_ids_with_load_errors(self) -> set[int]:
193 + agg_block_ids: set[int] = set()
194 + for c in self._connectors:
195 + agg_block_ids |= c.get_block_ids_with_load_errors()
196 + return agg_block_ids
197 +
198 # ==============================
199 # Scheduler-side methods
200 # ==============================
vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py
+2 -5
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import hashlib
4 import os
5 -from dataclasses import dataclass
5 +from dataclasses import dataclass, field
6 from typing import TYPE_CHECKING, Any, Optional
7
8 import safetensors
@@ -55,10 +55,7 @@ class ReqMeta:
55
56 @dataclass
57 class SharedStorageConnectorMetadata(KVConnectorMetadata):
58 - requests: list[ReqMeta]
59 -
60 - def __init__(self):
61 - self.requests = []
58 + requests: list[ReqMeta] = field(default_factory=list)
59
60 def add_request(
61 self,
vllm/v1/core/block_pool.py
+1 -1
@@ -211,7 +211,7 @@ class BlockPool:
211 block_size: Number of tokens in each block.
212 kv_cache_group_id: The id of the KV cache group.
213 """
214 - if num_cached_blocks == num_full_blocks:
214 + if num_cached_blocks >= num_full_blocks:
215 return
216 new_full_blocks = blocks[num_cached_blocks:num_full_blocks]
217 assert len(request.block_hashes) >= num_full_blocks
vllm/v1/core/sched/output.py
+2
@@ -101,6 +101,7 @@ class CachedRequestData:
101 new_token_ids: list[list[int]]
102 new_block_ids: list[Optional[tuple[list[int], ...]]]
103 num_computed_tokens: list[int]
104 + num_output_tokens: list[int]
105
106 @property
107 def num_reqs(self) -> int:
@@ -114,6 +115,7 @@ class CachedRequestData:
115 new_token_ids=[],
116 new_block_ids=[],
117 num_computed_tokens=[],
118 + num_output_tokens=[],
119 )
120
121
vllm/v1/core/sched/scheduler.py
+174 -14
@@ -133,6 +133,7 @@ class Scheduler(SchedulerInterface):
133
134 # KV Connector: requests in process of async KV loading or recving
135 self.finished_recving_kv_req_ids: set[str] = set()
136 + self.failed_recving_kv_req_ids: set[str] = set()
137
138 # Encoder-related.
139 # Calculate encoder cache size if applicable
@@ -671,6 +672,7 @@ class Scheduler(SchedulerInterface):
672 new_token_ids: list[list[int]] = []
673 new_block_ids: list[Optional[tuple[list[int], ...]]] = []
674 num_computed_tokens: list[int] = []
675 + num_output_tokens: list[int] = []
676
677 use_connector = self.connector is not None
678 for req in itertools.chain(running_reqs, resumed_reqs):
@@ -695,6 +697,7 @@ class Scheduler(SchedulerInterface):
697 new_block_ids.append(
698 req_to_new_blocks[req_id].get_block_ids(allow_none=True))
699 num_computed_tokens.append(req.num_computed_tokens)
700 + num_output_tokens.append(len(req.output_token_ids))
701 # Because resumed_reqs is usually empty, it is more efficient to do
702 # in-place appending so that we don't need to allocate a new list.
703 resumed_from_preemption = [False] * len(running_reqs)
@@ -706,6 +709,7 @@ class Scheduler(SchedulerInterface):
709 new_token_ids=new_token_ids,
710 new_block_ids=new_block_ids,
711 num_computed_tokens=num_computed_tokens,
712 + num_output_tokens=num_output_tokens,
713 )
714
715 def _try_schedule_encoder_inputs(
@@ -878,6 +882,14 @@ class Scheduler(SchedulerInterface):
882 kv_connector_stats = (kv_connector_output.kv_connector_stats
883 if kv_connector_output else None)
884
885 + failed_kv_load_req_ids = None
886 + if kv_connector_output and kv_connector_output.invalid_block_ids:
887 + # These blocks contain externally computed tokens that failed to
888 + # load. Identify affected requests and adjust their computed token
889 + # count to trigger recomputation of the invalid blocks.
890 + failed_kv_load_req_ids = self._handle_invalid_blocks(
891 + kv_connector_output.invalid_block_ids)
892 +
893 # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
894 # the below loop can be a performance bottleneck. We should do our best
895 # to avoid expensive operations inside the loop.
@@ -885,6 +897,9 @@ class Scheduler(SchedulerInterface):
897 stopped_preempted_reqs: set[Request] = set()
898 for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
899 assert num_tokens_scheduled > 0
900 + if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
901 + # Skip requests that were recovered from KV load failure
902 + continue
903 request = self.requests.get(req_id)
904 if request is None:
905 # The request is already finished. This can happen if the
@@ -988,9 +1003,8 @@ class Scheduler(SchedulerInterface):
1003 self.waiting.remove_requests(stopped_preempted_reqs)
1004
1005 # KV Connector: update state for finished KV Transfers.
991 - if model_runner_output.kv_connector_output:
992 - self._update_from_kv_xfer_finished(
993 - model_runner_output.kv_connector_output)
1006 + if kv_connector_output:
1007 + self._update_from_kv_xfer_finished(kv_connector_output)
1008
1009 # Create EngineCoreOutputs for all clients that have requests with
1010 # outputs in this step.
@@ -1252,18 +1266,33 @@ class Scheduler(SchedulerInterface):
1266 if request.request_id not in self.finished_recving_kv_req_ids:
1267 return False
1268
1255 - # Now that the blocks are ready, actually cache them.
1256 - (block_ids, ) = self.kv_cache_manager.get_block_ids(request.request_id)
1257 - num_computed_tokens = len(block_ids) * self.block_size
1258 - # Handle the case where num request tokens less than one block.
1259 - num_computed_tokens = min(num_computed_tokens, request.num_tokens)
1260 - if num_computed_tokens == request.num_tokens:
1261 - num_computed_tokens -= 1
1262 - # This will cache the blocks iff caching is enabled.
1263 - self.kv_cache_manager.cache_blocks(request, num_computed_tokens)
1269 + if request.request_id in self.failed_recving_kv_req_ids:
1270 + # Request had KV load failures; num_computed_tokens was already
1271 + # updated in _update_requests_with_invalid_blocks
1272 + if request.num_computed_tokens:
1273 + # Cache any valid computed tokens.
1274 + self.kv_cache_manager.cache_blocks(request,
1275 + request.num_computed_tokens)
1276 + else:
1277 + # No valid computed tokens, release allocated blocks.
1278 + # There may be a local cache hit on retry.
1279 + self.kv_cache_manager.free(request)
1280
1265 - # Update the request state for scheduling.
1266 - request.num_computed_tokens = num_computed_tokens
1281 + self.failed_recving_kv_req_ids.remove(request.request_id)
1282 + else:
1283 + # Now that the blocks are ready, actually cache them.
1284 + (block_ids, ) = self.kv_cache_manager.get_block_ids(
1285 + request.request_id)
1286 + num_computed_tokens = len(block_ids) * self.block_size
1287 + # Handle the case where num request tokens less than one block.
1288 + num_computed_tokens = min(num_computed_tokens, request.num_tokens)
1289 + if num_computed_tokens == request.num_tokens:
1290 + num_computed_tokens -= 1
1291 + # This will cache the blocks iff caching is enabled.
1292 + self.kv_cache_manager.cache_blocks(request, num_computed_tokens)
1293 +
1294 + # Update the request state for scheduling.
1295 + request.num_computed_tokens = num_computed_tokens
1296
1297 # Return that we are ready.
1298 self.finished_recving_kv_req_ids.remove(request.request_id)
@@ -1296,3 +1325,134 @@ class Scheduler(SchedulerInterface):
1325 "but the request is already freed.", req_id)
1326 else:
1327 self._free_blocks(self.requests[req_id])
1328 +
1329 + def _update_requests_with_invalid_blocks(
1330 + self, requests: Iterable[Request],
1331 + invalid_block_ids: set[int]) -> tuple[set[str], int]:
1332 + """
1333 + Identify and update requests affected by invalid KV cache blocks.
1334 +
1335 + This method scans the given requests, detects those with invalid blocks
1336 + and adjusts their `num_computed_tokens` to the longest valid prefix.
1337 + For observability, it also accumulates the total number of tokens that
1338 + will need to be recomputed across all affected requests.
1339 +
1340 + Args:
1341 + requests: The set of requests to scan for invalid blocks.
1342 + invalid_block_ids: IDs of invalid blocks.
1343 +
1344 + Returns:
1345 + tuple:
1346 + - affected_req_ids (set[str]): IDs of requests impacted by
1347 + invalid blocks.
1348 + - total_affected_tokens (int): Total number of tokens that must
1349 + be recomputed across all affected requests (for observability).
1350 + """
1351 + affected_req_ids: set[str] = set()
1352 + total_affected_tokens = 0
1353 + # If a block is invalid and shared by multiple requests in the batch,
1354 + # these requests must be rescheduled, but only the first will recompute
1355 + # it. This set tracks blocks already marked for recomputation.
1356 + marked_invalid_block_ids: set[int] = set()
1357 + for request in requests:
1358 + is_affected = False
1359 + marked_invalid_block = False
1360 + req_id = request.request_id
1361 + # TODO (davidb): add support for hybrid memory allocator
1362 + (req_block_ids, ) = self.kv_cache_manager.get_block_ids(req_id)
1363 + # We iterate only over blocks that may contain externally computed
1364 + # tokens
1365 + if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
1366 + # Async loading. If num_computed_tokens is set it implies we
1367 + # already processed some block failures for it in a prior step
1368 + req_num_computed_tokens = (
1369 + request.num_computed_tokens if req_id
1370 + in self.failed_recving_kv_req_ids else len(req_block_ids) *
1371 + self.block_size)
1372 + else:
1373 + # Sync loading. num_computed_tokens includes new tokens
1374 + req_num_computed_tokens = request.num_cached_tokens
1375 +
1376 + req_num_computed_blocks = (req_num_computed_tokens +
1377 + self.block_size - 1) // self.block_size
1378 + for idx, block_id in zip(range(req_num_computed_blocks),
1379 + req_block_ids):
1380 +
1381 + if block_id not in invalid_block_ids:
1382 + continue
1383 +
1384 + is_affected = True
1385 +
1386 + if block_id in marked_invalid_block_ids:
1387 + # This invalid block is shared with a previous request
1388 + # and was already marked for recomputation.
1389 + # This means this request can still consider this block
1390 + # as computed when rescheduled.
1391 + # Currently this only applies to sync loading; Async
1392 + # loading does not yet support block sharing
1393 + continue
1394 +
1395 + marked_invalid_block_ids.add(block_id)
1396 +
1397 + if marked_invalid_block:
1398 + # This request has already marked an invalid block for
1399 + # recomputation and updated its num_computed_tokens.
1400 + continue
1401 +
1402 + marked_invalid_block = True
1403 + # Truncate the computed tokens at the first failed block
1404 + request.num_computed_tokens = idx * self.block_size
1405 + total_affected_tokens += (req_num_computed_tokens -
1406 + request.num_computed_tokens)
1407 +
1408 + if is_affected:
1409 + if not marked_invalid_block:
1410 + # All invalid blocks of this request are shared with
1411 + # previous requests and will be recomputed by them.
1412 + # Revert to considering only cached tokens as computed.
1413 + # Currently this only applies to sync loading; Async
1414 + # loading does not yet support block sharing
1415 + total_affected_tokens += (request.num_computed_tokens -
1416 + request.num_cached_tokens)
1417 + request.num_computed_tokens = request.num_cached_tokens
1418 +
1419 + affected_req_ids.add(request.request_id)
1420 +
1421 + return (affected_req_ids, total_affected_tokens)
1422 +
1423 + def _handle_invalid_blocks(self, invalid_block_ids: set[int]) -> set[str]:
1424 + total_requests_to_reschedule = 0
1425 + total_tokens_to_reschedule = 0
1426 +
1427 + # --- Handle async KV loads (WAITING_FOR_REMOTE_KVS) ---
1428 + async_load_reqs = (
1429 + req for req in self.waiting
1430 + if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS)
1431 + async_affected_req_ids, num_tokens_to_reschedule = (
1432 + self._update_requests_with_invalid_blocks(async_load_reqs,
1433 + invalid_block_ids))
1434 +
1435 + total_requests_to_reschedule += len(async_affected_req_ids)
1436 + total_tokens_to_reschedule += num_tokens_to_reschedule
1437 +
1438 + # Mark requests with async KV load failures; they will be rescheduled
1439 + # once loading completes
1440 + self.failed_recving_kv_req_ids |= async_affected_req_ids
1441 +
1442 + # --- Handle sync KV loads (running requests) ---
1443 + sync_affected_req_ids, num_tokens_to_reschedule = (
1444 + self._update_requests_with_invalid_blocks(self.running,
1445 + invalid_block_ids))
1446 +
1447 + total_requests_to_reschedule += len(sync_affected_req_ids)
1448 + total_tokens_to_reschedule += num_tokens_to_reschedule
1449 +
1450 + if total_requests_to_reschedule:
1451 + logger.warning(
1452 + "Recovered from KV load failure: "
1453 + "%d request(s) rescheduled (%d tokens affected).",
1454 + total_requests_to_reschedule, total_tokens_to_reschedule)
1455 +
1456 + # Return the IDs of affected running requests to skip in
1457 + # update_from_output.
1458 + return sync_affected_req_ids
vllm/v1/core/single_type_kv_cache_manager.py
+3
@@ -142,6 +142,9 @@ class SingleTypeKVCacheManager(ABC):
142 num_cached_blocks = self.num_cached_block[request.request_id]
143 num_full_blocks = num_tokens // self.block_size
144
145 + if num_cached_blocks >= num_full_blocks:
146 + return
147 +
148 self.block_pool.cache_full_blocks(
149 request=request,
150 blocks=self.req_to_blocks[request.request_id],
vllm/v1/outputs.py
+5 -2
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from abc import ABC, abstractmethod
5 -from dataclasses import dataclass
5 +from dataclasses import dataclass, field
6 from typing import TYPE_CHECKING, NamedTuple, Optional, Union
7
8 import torch
@@ -87,10 +87,13 @@ class KVConnectorOutput:
87 finished_sending: Optional[set[str]] = None
88 finished_recving: Optional[set[str]] = None
89 kv_connector_stats: Optional["KVConnectorStats"] = None
90 + # IDs of externally computed KV blocks that failed to load.
91 + # Requests referencing these blocks should be rescheduled to recompute them.
92 + invalid_block_ids: set[int] = field(default_factory=set)
93
94 def is_empty(self):
95 return (not self.finished_sending and not self.finished_recving
93 - and not self.kv_connector_stats)
96 + and not self.kv_connector_stats and not self.invalid_block_ids)
97
98
99 # ModelRunnerOutput is serialized and sent to the scheduler process.
vllm/v1/worker/gpu_model_runner.py
+17
@@ -634,8 +634,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
634 num_computed_tokens = req_data.num_computed_tokens[i]
635 new_block_ids = req_data.new_block_ids[i]
636 resumed_from_preemption = req_data.resumed_from_preemption[i]
637 + num_output_tokens = req_data.num_output_tokens[i]
638
639 # Update the cached states.
640 +
641 req_state.num_computed_tokens = num_computed_tokens
642
643 if not is_last_rank:
@@ -653,6 +655,21 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
655 elif num_new_tokens > 0:
656 req_state.output_token_ids.extend(
657 new_token_ids[-num_new_tokens:])
658 + elif num_output_tokens < len(req_state.output_token_ids):
659 + # Some output tokens were discarded due to a sync-KV-load
660 + # failure. Align the cached state.
661 + del req_state.output_token_ids[num_output_tokens:]
662 +
663 + req_index = self.input_batch.req_id_to_index.get(req_id)
664 + if req_index is not None:
665 + old_end_idx = self.input_batch.num_tokens_no_spec[
666 + req_index]
667 + end_idx = self.input_batch.num_prompt_tokens[
668 + req_index] + num_output_tokens
669 + self.input_batch.num_tokens[req_index] = end_idx
670 + self.input_batch.num_tokens_no_spec[req_index] = end_idx
671 + self.input_batch.is_token_ids[req_index,
672 + end_idx:old_end_idx] = False
673
674 # Update the block IDs.
675 if not resumed_from_preemption:
vllm/v1/worker/gpu_worker.py
+1 -2
@@ -464,8 +464,7 @@ class Worker(WorkerBase):
464
465 # In case of PP with kv transfer, we need to pass through the
466 # kv_connector_output
467 - if (not kv_connector_output.finished_sending
468 - and not kv_connector_output.finished_recving):
467 + if kv_connector_output.is_empty():
468 return EMPTY_MODEL_RUNNER_OUTPUT
469
470 output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
vllm/v1/worker/kv_connector_model_runner_mixin.py
+3 -2
@@ -75,8 +75,7 @@ class KVConnectorModelRunnerMixin:
75 scheduler_output, wait_for_save=False) as kv_connector_output:
76 pass
77
78 - if (not kv_connector_output.finished_sending
79 - and not kv_connector_output.finished_recving):
78 + if kv_connector_output.is_empty():
79 return EMPTY_MODEL_RUNNER_OUTPUT
80
81 output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
@@ -120,6 +119,8 @@ class KVConnectorModelRunnerMixin:
119
120 output.finished_sending, output.finished_recving = (
121 kv_connector.get_finished(scheduler_output.finished_req_ids))
122 + output.invalid_block_ids = (
123 + kv_connector.get_block_ids_with_load_errors())
124
125 output.kv_connector_stats = KVConnectorModelRunnerMixin.\
126 get_kv_connector_stats()