@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a9e4106f28
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 11, 17:30 UTC

[P/D] KV Load Failure Recovery/Abort Configuration (#26813)

Signed-off-by: Will Eaton <weaton@redhat.com> Signed-off-by: Will Eaton <me@wseaton.com> Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com> Co-authored-by: chaunceyjiang <chaunceyjiang@gmail.com>

Will Eaton committed Dec 10, 2025 at 14:00 UTC a9e4106f28834315de4bfb1cb1186c9a2dc95856
16 files changed +1553 -49
tests/entrypoints/openai/test_chat_error.py new
+228
@@ -0,0 +1,228 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from dataclasses import dataclass, field
5 +from http import HTTPStatus
6 +from typing import Any
7 +from unittest.mock import AsyncMock, MagicMock
8 +
9 +import pytest
10 +
11 +from vllm.config.multimodal import MultiModalConfig
12 +from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ErrorResponse
13 +from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
14 +from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels
15 +from vllm.outputs import CompletionOutput, RequestOutput
16 +from vllm.transformers_utils.tokenizer import get_tokenizer
17 +from vllm.v1.engine.async_llm import AsyncLLM
18 +
19 +MODEL_NAME = "openai-community/gpt2"
20 +MODEL_NAME_SHORT = "gpt2"
21 +BASE_MODEL_PATHS = [
22 + BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME),
23 + BaseModelPath(name=MODEL_NAME_SHORT, model_path=MODEL_NAME_SHORT),
24 +]
25 +
26 +
27 +@dataclass
28 +class MockHFConfig:
29 + model_type: str = "any"
30 +
31 +
32 +@dataclass
33 +class MockModelConfig:
34 + task = "generate"
35 + runner_type = "generate"
36 + tokenizer = MODEL_NAME
37 + trust_remote_code = False
38 + tokenizer_mode = "auto"
39 + max_model_len = 100
40 + tokenizer_revision = None
41 + multimodal_config = MultiModalConfig()
42 + hf_config = MockHFConfig()
43 + logits_processor_pattern = None
44 + logits_processors: list[str] | None = None
45 + diff_sampling_param: dict | None = None
46 + allowed_local_media_path: str = ""
47 + allowed_media_domains: list[str] | None = None
48 + encoder_config = None
49 + generation_config: str = "auto"
50 + media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
51 + skip_tokenizer_init = False
52 +
53 + def get_diff_sampling_param(self):
54 + return self.diff_sampling_param or {}
55 +
56 +
57 +def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat:
58 + models = OpenAIServingModels(
59 + engine_client=engine,
60 + base_model_paths=BASE_MODEL_PATHS,
61 + )
62 + serving_chat = OpenAIServingChat(
63 + engine,
64 + models,
65 + response_role="assistant",
66 + request_logger=None,
67 + chat_template=None,
68 + chat_template_content_format="auto",
69 + )
70 +
71 + async def _fake_process_inputs(
72 + request_id,
73 + engine_prompt,
74 + sampling_params,
75 + *,
76 + lora_request,
77 + trace_headers,
78 + priority,
79 + ):
80 + return dict(engine_prompt), {}
81 +
82 + async def _fake_preprocess_chat(*args, **kwargs):
83 + # return conversation, request_prompts, engine_prompts
84 + return (
85 + [{"role": "user", "content": "Test"}],
86 + [[1, 2, 3]],
87 + [{"prompt_token_ids": [1, 2, 3]}],
88 + )
89 +
90 + serving_chat._process_inputs = AsyncMock(side_effect=_fake_process_inputs)
91 + serving_chat._preprocess_chat = AsyncMock(side_effect=_fake_preprocess_chat)
92 + return serving_chat
93 +
94 +
95 +@pytest.mark.asyncio
96 +async def test_chat_error_non_stream():
97 + """test finish_reason='error' returns 500 InternalServerError (non-streaming)"""
98 + mock_engine = MagicMock(spec=AsyncLLM)
99 + mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
100 + mock_engine.errored = False
101 + mock_engine.model_config = MockModelConfig()
102 + mock_engine.input_processor = MagicMock()
103 + mock_engine.io_processor = MagicMock()
104 +
105 + serving_chat = _build_serving_chat(mock_engine)
106 +
107 + completion_output = CompletionOutput(
108 + index=0,
109 + text="",
110 + token_ids=[],
111 + cumulative_logprob=None,
112 + logprobs=None,
113 + finish_reason="error",
114 + )
115 +
116 + request_output = RequestOutput(
117 + request_id="test-id",
118 + prompt="Test prompt",
119 + prompt_token_ids=[1, 2, 3],
120 + prompt_logprobs=None,
121 + outputs=[completion_output],
122 + finished=True,
123 + metrics=None,
124 + lora_request=None,
125 + encoder_prompt=None,
126 + encoder_prompt_token_ids=None,
127 + )
128 +
129 + async def mock_generate(*args, **kwargs):
130 + yield request_output
131 +
132 + mock_engine.generate = MagicMock(side_effect=mock_generate)
133 +
134 + request = ChatCompletionRequest(
135 + model=MODEL_NAME,
136 + messages=[{"role": "user", "content": "Test prompt"}],
137 + max_tokens=10,
138 + stream=False,
139 + )
140 +
141 + response = await serving_chat.create_chat_completion(request)
142 +
143 + assert isinstance(response, ErrorResponse)
144 + assert response.error.type == "InternalServerError"
145 + assert response.error.message == "Internal server error"
146 + assert response.error.code == HTTPStatus.INTERNAL_SERVER_ERROR
147 +
148 +
149 +@pytest.mark.asyncio
150 +async def test_chat_error_stream():
151 + """test finish_reason='error' returns 500 InternalServerError (streaming)"""
152 + mock_engine = MagicMock(spec=AsyncLLM)
153 + mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
154 + mock_engine.errored = False
155 + mock_engine.model_config = MockModelConfig()
156 + mock_engine.input_processor = MagicMock()
157 + mock_engine.io_processor = MagicMock()
158 +
159 + serving_chat = _build_serving_chat(mock_engine)
160 +
161 + completion_output_1 = CompletionOutput(
162 + index=0,
163 + text="Hello",
164 + token_ids=[100],
165 + cumulative_logprob=None,
166 + logprobs=None,
167 + finish_reason=None,
168 + )
169 +
170 + request_output_1 = RequestOutput(
171 + request_id="test-id",
172 + prompt="Test prompt",
173 + prompt_token_ids=[1, 2, 3],
174 + prompt_logprobs=None,
175 + outputs=[completion_output_1],
176 + finished=False,
177 + metrics=None,
178 + lora_request=None,
179 + encoder_prompt=None,
180 + encoder_prompt_token_ids=None,
181 + )
182 +
183 + completion_output_2 = CompletionOutput(
184 + index=0,
185 + text="Hello",
186 + token_ids=[100],
187 + cumulative_logprob=None,
188 + logprobs=None,
189 + finish_reason="error",
190 + )
191 +
192 + request_output_2 = RequestOutput(
193 + request_id="test-id",
194 + prompt="Test prompt",
195 + prompt_token_ids=[1, 2, 3],
196 + prompt_logprobs=None,
197 + outputs=[completion_output_2],
198 + finished=True,
199 + metrics=None,
200 + lora_request=None,
201 + encoder_prompt=None,
202 + encoder_prompt_token_ids=None,
203 + )
204 +
205 + async def mock_generate(*args, **kwargs):
206 + yield request_output_1
207 + yield request_output_2
208 +
209 + mock_engine.generate = MagicMock(side_effect=mock_generate)
210 +
211 + request = ChatCompletionRequest(
212 + model=MODEL_NAME,
213 + messages=[{"role": "user", "content": "Test prompt"}],
214 + max_tokens=10,
215 + stream=True,
216 + )
217 +
218 + response = await serving_chat.create_chat_completion(request)
219 +
220 + chunks = []
221 + async for chunk in response:
222 + chunks.append(chunk)
223 +
224 + assert len(chunks) >= 2
225 + assert any("Internal server error" in chunk for chunk in chunks), (
226 + f"Expected error message in chunks: {chunks}"
227 + )
228 + assert chunks[-1] == "data: [DONE]\n\n"
tests/entrypoints/openai/test_completion_error.py new
+216
@@ -0,0 +1,216 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from dataclasses import dataclass, field
5 +from http import HTTPStatus
6 +from typing import Any
7 +from unittest.mock import AsyncMock, MagicMock
8 +
9 +import pytest
10 +
11 +from vllm.config.multimodal import MultiModalConfig
12 +from vllm.entrypoints.openai.protocol import CompletionRequest, ErrorResponse
13 +from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
14 +from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels
15 +from vllm.outputs import CompletionOutput, RequestOutput
16 +from vllm.transformers_utils.tokenizer import get_tokenizer
17 +from vllm.v1.engine.async_llm import AsyncLLM
18 +
19 +MODEL_NAME = "openai-community/gpt2"
20 +MODEL_NAME_SHORT = "gpt2"
21 +BASE_MODEL_PATHS = [
22 + BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME),
23 + BaseModelPath(name=MODEL_NAME_SHORT, model_path=MODEL_NAME_SHORT),
24 +]
25 +
26 +
27 +@dataclass
28 +class MockHFConfig:
29 + model_type: str = "any"
30 +
31 +
32 +@dataclass
33 +class MockModelConfig:
34 + task = "generate"
35 + runner_type = "generate"
36 + tokenizer = MODEL_NAME
37 + trust_remote_code = False
38 + tokenizer_mode = "auto"
39 + max_model_len = 100
40 + tokenizer_revision = None
41 + multimodal_config = MultiModalConfig()
42 + hf_config = MockHFConfig()
43 + logits_processor_pattern = None
44 + logits_processors: list[str] | None = None
45 + diff_sampling_param: dict | None = None
46 + allowed_local_media_path: str = ""
47 + allowed_media_domains: list[str] | None = None
48 + encoder_config = None
49 + generation_config: str = "auto"
50 + media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
51 + skip_tokenizer_init = False
52 +
53 + def get_diff_sampling_param(self):
54 + return self.diff_sampling_param or {}
55 +
56 +
57 +def _build_serving_completion(engine: AsyncLLM) -> OpenAIServingCompletion:
58 + models = OpenAIServingModels(
59 + engine_client=engine,
60 + base_model_paths=BASE_MODEL_PATHS,
61 + )
62 + serving_completion = OpenAIServingCompletion(
63 + engine,
64 + models,
65 + request_logger=None,
66 + )
67 +
68 + async def _fake_process_inputs(
69 + request_id,
70 + engine_prompt,
71 + sampling_params,
72 + *,
73 + lora_request,
74 + trace_headers,
75 + priority,
76 + ):
77 + return dict(engine_prompt), {}
78 +
79 + serving_completion._process_inputs = AsyncMock(side_effect=_fake_process_inputs)
80 + return serving_completion
81 +
82 +
83 +@pytest.mark.asyncio
84 +async def test_completion_error_non_stream():
85 + """test finish_reason='error' returns 500 InternalServerError (non-streaming)"""
86 + mock_engine = MagicMock(spec=AsyncLLM)
87 + mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
88 + mock_engine.errored = False
89 + mock_engine.model_config = MockModelConfig()
90 + mock_engine.input_processor = MagicMock()
91 + mock_engine.io_processor = MagicMock()
92 +
93 + serving_completion = _build_serving_completion(mock_engine)
94 +
95 + completion_output = CompletionOutput(
96 + index=0,
97 + text="",
98 + token_ids=[],
99 + cumulative_logprob=None,
100 + logprobs=None,
101 + finish_reason="error",
102 + )
103 +
104 + request_output = RequestOutput(
105 + request_id="test-id",
106 + prompt="Test prompt",
107 + prompt_token_ids=[1, 2, 3],
108 + prompt_logprobs=None,
109 + outputs=[completion_output],
110 + finished=True,
111 + metrics=None,
112 + lora_request=None,
113 + encoder_prompt=None,
114 + encoder_prompt_token_ids=None,
115 + )
116 +
117 + async def mock_generate(*args, **kwargs):
118 + yield request_output
119 +
120 + mock_engine.generate = MagicMock(side_effect=mock_generate)
121 +
122 + request = CompletionRequest(
123 + model=MODEL_NAME,
124 + prompt="Test prompt",
125 + max_tokens=10,
126 + stream=False,
127 + )
128 +
129 + response = await serving_completion.create_completion(request)
130 +
131 + assert isinstance(response, ErrorResponse)
132 + assert response.error.type == "InternalServerError"
133 + assert response.error.message == "Internal server error"
134 + assert response.error.code == HTTPStatus.INTERNAL_SERVER_ERROR
135 +
136 +
137 +@pytest.mark.asyncio
138 +async def test_completion_error_stream():
139 + """test finish_reason='error' returns 500 InternalServerError (streaming)"""
140 + mock_engine = MagicMock(spec=AsyncLLM)
141 + mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)
142 + mock_engine.errored = False
143 + mock_engine.model_config = MockModelConfig()
144 + mock_engine.input_processor = MagicMock()
145 + mock_engine.io_processor = MagicMock()
146 +
147 + serving_completion = _build_serving_completion(mock_engine)
148 +
149 + completion_output_1 = CompletionOutput(
150 + index=0,
151 + text="Hello",
152 + token_ids=[100],
153 + cumulative_logprob=None,
154 + logprobs=None,
155 + finish_reason=None,
156 + )
157 +
158 + request_output_1 = RequestOutput(
159 + request_id="test-id",
160 + prompt="Test prompt",
161 + prompt_token_ids=[1, 2, 3],
162 + prompt_logprobs=None,
163 + outputs=[completion_output_1],
164 + finished=False,
165 + metrics=None,
166 + lora_request=None,
167 + encoder_prompt=None,
168 + encoder_prompt_token_ids=None,
169 + )
170 +
171 + completion_output_2 = CompletionOutput(
172 + index=0,
173 + text="Hello",
174 + token_ids=[100],
175 + cumulative_logprob=None,
176 + logprobs=None,
177 + finish_reason="error",
178 + )
179 +
180 + request_output_2 = RequestOutput(
181 + request_id="test-id",
182 + prompt="Test prompt",
183 + prompt_token_ids=[1, 2, 3],
184 + prompt_logprobs=None,
185 + outputs=[completion_output_2],
186 + finished=True,
187 + metrics=None,
188 + lora_request=None,
189 + encoder_prompt=None,
190 + encoder_prompt_token_ids=None,
191 + )
192 +
193 + async def mock_generate(*args, **kwargs):
194 + yield request_output_1
195 + yield request_output_2
196 +
197 + mock_engine.generate = MagicMock(side_effect=mock_generate)
198 +
199 + request = CompletionRequest(
200 + model=MODEL_NAME,
201 + prompt="Test prompt",
202 + max_tokens=10,
203 + stream=True,
204 + )
205 +
206 + response = await serving_completion.create_completion(request)
207 +
208 + chunks = []
209 + async for chunk in response:
210 + chunks.append(chunk)
211 +
212 + assert len(chunks) >= 2
213 + assert any("Internal server error" in chunk for chunk in chunks), (
214 + f"Expected error message in chunks: {chunks}"
215 + )
216 + assert chunks[-1] == "data: [DONE]\n\n"
tests/entrypoints/openai/test_responses_error.py new
+89
@@ -0,0 +1,89 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from http import HTTPStatus
5 +from unittest.mock import MagicMock
6 +
7 +import pytest
8 +
9 +from vllm.entrypoints.openai.protocol import ErrorResponse
10 +from vllm.entrypoints.openai.serving_engine import GenerationError, OpenAIServing
11 +
12 +
13 +@pytest.mark.asyncio
14 +async def test_raise_if_error_raises_generation_error():
15 + """test _raise_if_error raises GenerationError"""
16 + # create a minimal OpenAIServing instance
17 + mock_engine = MagicMock()
18 + mock_engine.model_config = MagicMock()
19 + mock_engine.model_config.max_model_len = 100
20 + mock_models = MagicMock()
21 +
22 + serving = OpenAIServing(
23 + engine_client=mock_engine,
24 + models=mock_models,
25 + request_logger=None,
26 + )
27 +
28 + # test that error finish_reason raises GenerationError
29 + with pytest.raises(GenerationError) as exc_info:
30 + serving._raise_if_error("error", "test-request-id")
31 +
32 + assert str(exc_info.value) == "Internal server error"
33 + assert exc_info.value.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
34 +
35 + # test that other finish_reasons don't raise
36 + serving._raise_if_error("stop", "test-request-id") # should not raise
37 + serving._raise_if_error("length", "test-request-id") # should not raise
38 + serving._raise_if_error(None, "test-request-id") # should not raise
39 +
40 +
41 +@pytest.mark.asyncio
42 +async def test_convert_generation_error_to_response():
43 + """test _convert_generation_error_to_response creates proper ErrorResponse"""
44 + mock_engine = MagicMock()
45 + mock_engine.model_config = MagicMock()
46 + mock_engine.model_config.max_model_len = 100
47 + mock_models = MagicMock()
48 +
49 + serving = OpenAIServing(
50 + engine_client=mock_engine,
51 + models=mock_models,
52 + request_logger=None,
53 + )
54 +
55 + # create a GenerationError
56 + gen_error = GenerationError("Internal server error")
57 +
58 + # convert to ErrorResponse
59 + error_response = serving._convert_generation_error_to_response(gen_error)
60 +
61 + assert isinstance(error_response, ErrorResponse)
62 + assert error_response.error.type == "InternalServerError"
63 + assert error_response.error.message == "Internal server error"
64 + assert error_response.error.code == HTTPStatus.INTERNAL_SERVER_ERROR
65 +
66 +
67 +@pytest.mark.asyncio
68 +async def test_convert_generation_error_to_streaming_response():
69 + """test _convert_generation_error_to_streaming_response output"""
70 + mock_engine = MagicMock()
71 + mock_engine.model_config = MagicMock()
72 + mock_engine.model_config.max_model_len = 100
73 + mock_models = MagicMock()
74 +
75 + serving = OpenAIServing(
76 + engine_client=mock_engine,
77 + models=mock_models,
78 + request_logger=None,
79 + )
80 +
81 + # create a GenerationError
82 + gen_error = GenerationError("Internal server error")
83 +
84 + # convert to streaming error response
85 + error_json = serving._convert_generation_error_to_streaming_response(gen_error)
86 +
87 + assert isinstance(error_json, str)
88 + assert "Internal server error" in error_json
89 + assert "InternalServerError" in error_json
tests/v1/kv_connector/unit/test_cache_pollution_prevention.py new
+163
@@ -0,0 +1,163 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +test that invalid blocks are evicted from prefix cache to prevent pollution.
6 +
7 +verifies that when sync-loading fails, invalid blocks are removed from the
8 +prefix cache hash table so future requests cannot match and reuse corrupted data.
9 +"""
10 +
11 +from collections.abc import Callable
12 +from unittest.mock import Mock
13 +
14 +import pytest
15 +
16 +from vllm.v1.core.sched.scheduler import Scheduler
17 +from vllm.v1.request import Request, RequestStatus
18 +
19 +from .utils import (
20 + create_model_runner_output,
21 + create_request,
22 + create_scheduler,
23 + create_vllm_config,
24 +)
25 +
26 +pytestmark = pytest.mark.cpu_test
27 +
28 +
29 +def _make_get_num_new_matched_tokens(
30 + req_num_new_matched_tokens: dict[str, int],
31 + async_load: bool,
32 +) -> Callable[[Request, int], tuple[int, bool]]:
33 + def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]:
34 + value = req_num_new_matched_tokens.get(request.request_id, 0)
35 + return value, async_load
36 +
37 + return get_num_new_matched_tokens
38 +
39 +
40 +@pytest.fixture
41 +def fail_scheduler():
42 + """scheduler with kv_load_failure_policy='fail'"""
43 + vllm_config = create_vllm_config()
44 + vllm_config.kv_transfer_config.kv_load_failure_policy = "fail"
45 + return create_scheduler(vllm_config)
46 +
47 +
48 +def test_invalid_blocks_evicted_prevents_cache_pollution(
49 + fail_scheduler: Scheduler,
50 +):
51 + """
52 + verify invalid blocks are evicted to prevent future cache hits.
53 +
54 + scenario:
55 + 1. request 1 loads externally-computed blocks (sync mode)
56 + 2. some blocks fail to load and are marked invalid
57 + 3. with fail policy, invalid blocks should be evicted from prefix cache
58 + 4. request is marked as FINISHED_ERROR
59 + """
60 + num_prompt_blocks = 100
61 + num_external_computed_blocks = 99
62 + invalid_block_idx = 50
63 +
64 + num_prompt_tokens = num_prompt_blocks * fail_scheduler.block_size
65 + num_external_computed_tokens = (
66 + num_external_computed_blocks * fail_scheduler.block_size
67 + )
68 +
69 + # request 1: will have invalid blocks
70 + request1 = create_request(num_tokens=num_prompt_tokens, request_id=1)
71 + fail_scheduler.add_request(request=request1)
72 +
73 + req_num_new_matched_tokens = {
74 + request1.request_id: num_external_computed_tokens,
75 + }
76 +
77 + # mock connector indicating sync load
78 + fail_scheduler.connector = Mock()
79 + fail_scheduler.connector.get_num_new_matched_tokens.side_effect = (
80 + _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False)
81 + )
82 + fail_scheduler.connector.request_finished.return_value = (False, None)
83 + fail_scheduler.connector.take_events.return_value = ()
84 +
85 + scheduler_output = fail_scheduler.schedule()
86 +
87 + # request should be running with sync KV load
88 + assert len(fail_scheduler.running) == 1
89 + assert request1.status == RequestStatus.RUNNING
90 +
91 + # get allocated block IDs
92 + req_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0]
93 + invalid_block_id = req_block_ids[invalid_block_idx]
94 + invalid_block_ids = {invalid_block_id}
95 +
96 + # get the block object to verify eviction later
97 + block = fail_scheduler.kv_cache_manager.block_pool.blocks[invalid_block_id]
98 +
99 + # cache the blocks to simulate they've been computed and cached
100 + # (in real scenario blocks would be cached after compute)
101 + fail_scheduler.kv_cache_manager.cache_blocks(request1, num_external_computed_tokens)
102 +
103 + # verify block has a hash (is cached) before reporting invalid blocks
104 + assert block.block_hash is not None, (
105 + f"block {invalid_block_id} should be cached (have a hash) before "
106 + f"eviction test, but hash is None"
107 + )
108 +
109 + # report invalid blocks
110 + model_runner_output = create_model_runner_output(
111 + [request1],
112 + invalid_block_ids=invalid_block_ids,
113 + use_eos=False,
114 + )
115 +
116 + fail_scheduler.update_from_output(scheduler_output, model_runner_output)
117 +
118 + # verify request finished with error (fail policy)
119 + assert request1.status == RequestStatus.FINISHED_ERROR
120 +
121 + # critical assertion: invalid block and all subsequent blocks should be evicted
122 + # all blocks from invalid_block_idx onwards become invalid since they were
123 + # computed based on the failed block
124 + for idx in range(invalid_block_idx, len(req_block_ids)):
125 + block_id = req_block_ids[idx]
126 + block_obj = fail_scheduler.kv_cache_manager.block_pool.blocks[block_id]
127 + assert block_obj.block_hash is None, (
128 + f"block {block_id} at index {idx} should have been evicted "
129 + f"(hash reset to None), but hash is {block_obj.block_hash}. "
130 + f"All blocks from index {invalid_block_idx} onwards should be evicted "
131 + f"since they depend on the invalid block at index {invalid_block_idx}."
132 + )
133 +
134 + # verify cache contains exactly the valid blocks (before first affected block)
135 + # and none of the invalid blocks (from first affected block onwards)
136 +
137 + # valid blocks: all blocks before invalid_block_idx should be cached
138 + for idx in range(invalid_block_idx):
139 + block_id = req_block_ids[idx]
140 + block_obj = fail_scheduler.kv_cache_manager.block_pool.blocks[block_id]
141 + assert block_obj.block_hash is not None, (
142 + f"valid block {block_id} at index {idx} should still be cached "
143 + f"(have a hash), but hash is None. Only blocks from index "
144 + f"{invalid_block_idx} onwards should be evicted."
145 + )
146 +
147 + # invalid blocks: verify they're not in the cached_block_hash_to_block map
148 + cached_blocks = (
149 + fail_scheduler.kv_cache_manager.block_pool.cached_block_hash_to_block
150 + )
151 + cached_block_ids = {
152 + b.block_id
153 + for blocks_val in cached_blocks._cache.values()
154 + for b in (
155 + [blocks_val] if not isinstance(blocks_val, dict) else blocks_val.values()
156 + )
157 + }
158 +
159 + for idx in range(invalid_block_idx, len(req_block_ids)):
160 + block_id = req_block_ids[idx]
161 + assert block_id not in cached_block_ids, (
162 + f"invalid block {block_id} at index {idx} should not be in cache hash table"
163 + )
tests/v1/kv_connector/unit/test_error_propagation.py new
+147
@@ -0,0 +1,147 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc 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 FinishReason, Request, RequestStatus
11 +
12 +from .utils import (
13 + create_model_runner_output,
14 + create_request,
15 + create_scheduler,
16 + create_vllm_config,
17 +)
18 +
19 +pytestmark = pytest.mark.cpu_test
20 +
21 +
22 +def _make_get_num_new_matched_tokens(
23 + req_num_new_matched_tokens: dict[str, int],
24 + async_load: bool,
25 +) -> Callable[[Request, int], tuple[int, bool]]:
26 + def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]:
27 + value = req_num_new_matched_tokens.get(request.request_id, 0)
28 + return value, async_load
29 +
30 + return get_num_new_matched_tokens
31 +
32 +
33 +@pytest.fixture
34 +def fail_scheduler():
35 + """scheduler with kv_load_failure_policy='fail'"""
36 + vllm_config = create_vllm_config()
37 + vllm_config.kv_transfer_config.kv_load_failure_policy = "fail"
38 + return create_scheduler(vllm_config)
39 +
40 +
41 +def test_error_propagation_sync_load(fail_scheduler: Scheduler):
42 + """test invalid_block_ids with fail policy -> FINISHED_ERROR (sync load)"""
43 + num_prompt_blocks = 100
44 + num_external_computed_blocks = 99
45 + invalid_block_idx = 50
46 +
47 + num_prompt_tokens = num_prompt_blocks * fail_scheduler.block_size
48 + num_external_computed_tokens = (
49 + num_external_computed_blocks * fail_scheduler.block_size
50 + )
51 +
52 + request = create_request(num_tokens=num_prompt_tokens)
53 + fail_scheduler.add_request(request=request)
54 +
55 + req_num_new_matched_tokens = {
56 + request.request_id: num_external_computed_tokens,
57 + }
58 +
59 + fail_scheduler.connector = Mock()
60 + fail_scheduler.connector.get_num_new_matched_tokens.side_effect = (
61 + _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False)
62 + )
63 + fail_scheduler.connector.request_finished.return_value = (False, None)
64 + fail_scheduler.connector.take_events.return_value = ()
65 +
66 + scheduler_output = fail_scheduler.schedule()
67 +
68 + assert len(fail_scheduler.running) == 1
69 + assert len(scheduler_output.scheduled_new_reqs) == 1
70 + assert fail_scheduler.connector.get_num_new_matched_tokens.call_count == 1
71 +
72 + req_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0]
73 + invalid_block_ids = {req_block_ids[invalid_block_idx]}
74 + model_runner_output = create_model_runner_output(
75 + [request],
76 + invalid_block_ids=invalid_block_ids,
77 + use_eos=True,
78 + )
79 +
80 + outputs = fail_scheduler.update_from_output(scheduler_output, model_runner_output)
81 +
82 + assert request.status == RequestStatus.FINISHED_ERROR
83 + assert request.get_finished_reason() == FinishReason.ERROR
84 +
85 + assert len(outputs) == 1
86 + engine_outputs = next(iter(outputs.values()))
87 + assert len(engine_outputs.outputs) == 1
88 + output = engine_outputs.outputs[0]
89 + assert output.request_id == request.request_id
90 + assert output.finish_reason == FinishReason.ERROR
91 +
92 + assert len(fail_scheduler.running) == 0
93 +
94 +
95 +def test_error_propagation_async_load(fail_scheduler: Scheduler):
96 + """test invalid_block_ids with fail policy -> FINISHED_ERROR (async load)"""
97 + num_prompt_blocks = 100
98 + num_external_computed_blocks = 99
99 + invalid_block_idx = 50
100 +
101 + num_prompt_tokens = num_prompt_blocks * fail_scheduler.block_size
102 + num_external_computed_tokens = (
103 + num_external_computed_blocks * fail_scheduler.block_size
104 + )
105 +
106 + request = create_request(num_tokens=num_prompt_tokens)
107 + fail_scheduler.add_request(request=request)
108 +
109 + req_num_new_matched_tokens = {
110 + request.request_id: num_external_computed_tokens,
111 + }
112 +
113 + fail_scheduler.connector = Mock()
114 + fail_scheduler.connector.get_num_new_matched_tokens.side_effect = (
115 + _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True)
116 + )
117 + fail_scheduler.connector.request_finished.return_value = (False, None)
118 + fail_scheduler.connector.take_events.return_value = ()
119 +
120 + scheduler_output = fail_scheduler.schedule()
121 +
122 + assert len(fail_scheduler.waiting) == 1
123 + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
124 + assert request.num_computed_tokens == 0
125 +
126 + (req_block_ids,) = fail_scheduler.kv_cache_manager.get_block_ids(request.request_id)
127 + invalid_block_ids = {req_block_ids[invalid_block_idx]}
128 + model_runner_output = create_model_runner_output(
129 + reqs=[],
130 + finished_recving=set(),
131 + invalid_block_ids=invalid_block_ids,
132 + use_eos=True,
133 + )
134 +
135 + outputs = fail_scheduler.update_from_output(scheduler_output, model_runner_output)
136 +
137 + assert request.status == RequestStatus.FINISHED_ERROR
138 + assert request.get_finished_reason() == FinishReason.ERROR
139 +
140 + assert len(outputs) == 1
141 + engine_outputs = next(iter(outputs.values()))
142 + assert len(engine_outputs.outputs) == 1
143 + output = engine_outputs.outputs[0]
144 + assert output.request_id == request.request_id
145 + assert output.finish_reason == FinishReason.ERROR
146 +
147 + assert len(fail_scheduler.waiting) == 0
tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py new
+454
@@ -0,0 +1,454 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +Tests for correctness in invalid block handling.
6 +
7 +These tests verify correct behavior in three scenarios:
8 +1. Sync recompute case: Blocks should not be freed for running requests
9 + that need to recompute invalid blocks
10 +2. Sync fail case: Invalid blocks must be evicted from cache when request fails
11 +3. Async recompute case: Invalid blocks should not be cached after transfer
12 +"""
13 +
14 +from collections.abc import Callable
15 +from unittest.mock import Mock
16 +
17 +import pytest
18 +
19 +from vllm.v1.core.sched.scheduler import Scheduler
20 +from vllm.v1.request import FinishReason, Request, RequestStatus
21 +
22 +from .utils import (
23 + create_model_runner_output,
24 + create_request,
25 + create_scheduler,
26 + create_vllm_config,
27 +)
28 +
29 +pytestmark = pytest.mark.cpu_test
30 +
31 +
32 +def _make_get_num_new_matched_tokens(
33 + req_num_new_matched_tokens: dict[str, int],
34 + async_load: bool,
35 +) -> Callable[[Request, int], tuple[int, bool]]:
36 + def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]:
37 + value = req_num_new_matched_tokens.get(request.request_id, 0)
38 + return value, async_load
39 +
40 + return get_num_new_matched_tokens
41 +
42 +
43 +@pytest.fixture
44 +def fail_scheduler():
45 + """scheduler with kv_load_failure_policy='fail'"""
46 + vllm_config = create_vllm_config()
47 + vllm_config.kv_transfer_config.kv_load_failure_policy = "fail"
48 + return create_scheduler(vllm_config)
49 +
50 +
51 +@pytest.fixture
52 +def recompute_scheduler():
53 + """scheduler with kv_load_failure_policy='recompute'"""
54 + vllm_config = create_vllm_config()
55 + vllm_config.kv_transfer_config.kv_load_failure_policy = "recompute"
56 + return create_scheduler(vllm_config)
57 +
58 +
59 +def test_sync_recompute_blocks_not_freed_for_running_requests(
60 + recompute_scheduler: Scheduler,
61 +):
62 + """
63 + Test sync recompute case - blocks must not be freed for running requests.
64 +
65 + When a running request has invalid blocks and retry_policy is 'recompute':
66 + 1. Request should remain in RUNNING state
67 + 2. num_computed_tokens should be truncated to invalid block boundary
68 + 3. Blocks should NOT be freed (request still needs them for recomputation)
69 + 4. Request should remain in scheduler.requests and scheduler.running
70 + """
71 + num_prompt_blocks = 100
72 + num_external_computed_blocks = 99
73 + invalid_block_idx = 50
74 +
75 + num_prompt_tokens = num_prompt_blocks * recompute_scheduler.block_size
76 + num_external_computed_tokens = (
77 + num_external_computed_blocks * recompute_scheduler.block_size
78 + )
79 +
80 + request = create_request(num_tokens=num_prompt_tokens)
81 + recompute_scheduler.add_request(request=request)
82 +
83 + req_num_new_matched_tokens = {
84 + request.request_id: num_external_computed_tokens,
85 + }
86 +
87 + # mock connector indicating sync load
88 + recompute_scheduler.connector = Mock()
89 + recompute_scheduler.connector.get_num_new_matched_tokens.side_effect = (
90 + _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False)
91 + )
92 + recompute_scheduler.connector.request_finished.return_value = (False, None)
93 + recompute_scheduler.connector.take_events.return_value = ()
94 +
95 + scheduler_output = recompute_scheduler.schedule()
96 +
97 + # request should be running with sync KV load
98 + assert len(recompute_scheduler.running) == 1
99 + assert len(scheduler_output.scheduled_new_reqs) == 1
100 + assert request.status == RequestStatus.RUNNING
101 +
102 + # get the allocated block IDs before invalid blocks are reported
103 + req_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0]
104 + invalid_block_ids = {req_block_ids[invalid_block_idx]}
105 +
106 + # store original num_computed_tokens for comparison
107 + original_num_computed_tokens = request.num_computed_tokens
108 +
109 + model_runner_output = create_model_runner_output(
110 + [request],
111 + invalid_block_ids=invalid_block_ids,
112 + use_eos=False, # not finished - should continue running
113 + )
114 +
115 + outputs = recompute_scheduler.update_from_output(
116 + scheduler_output, model_runner_output
117 + )
118 +
119 + # critical assertions for recompute case:
120 +
121 + # 1. request should still be RUNNING (not finished, not aborted)
122 + assert request.status == RequestStatus.RUNNING, (
123 + f"Request should remain RUNNING for recompute, got {request.status}"
124 + )
125 +
126 + # 2. num_computed_tokens should be truncated to first invalid block
127 + expected_truncated_tokens = invalid_block_idx * recompute_scheduler.block_size
128 + assert request.num_computed_tokens == expected_truncated_tokens, (
129 + f"num_computed_tokens should be truncated to {expected_truncated_tokens}, "
130 + f"got {request.num_computed_tokens}"
131 + )
132 + assert request.num_computed_tokens < original_num_computed_tokens, (
133 + "num_computed_tokens should be reduced after invalid block detection"
134 + )
135 +
136 + # 3. no output should be generated (request is still running)
137 + # the request should be skipped in the output loop
138 + assert len(outputs) == 0 or request.request_id not in [
139 + out.request_id for outs in outputs.values() for out in outs.outputs
140 + ], "No output should be generated for recompute requests"
141 +
142 + # 4. request should still be in running queue
143 + assert request in recompute_scheduler.running, (
144 + "Request should remain in running queue for recomputation"
145 + )
146 +
147 + # 5. request should still be in scheduler.requests (not deleted)
148 + assert request.request_id in recompute_scheduler.requests, (
149 + "Request should not be deleted from scheduler.requests"
150 + )
151 +
152 + # 6. blocks should NOT be freed - verify blocks are still allocated
153 + try:
154 + allocated_blocks = recompute_scheduler.kv_cache_manager.get_block_ids(
155 + request.request_id
156 + )
157 + assert allocated_blocks is not None
158 + assert len(allocated_blocks[0]) > 0, (
159 + "Blocks should still be allocated for recomputation"
160 + )
161 + except KeyError:
162 + pytest.fail(
163 + "Blocks were freed incorrectly! Running requests need their blocks "
164 + "to recompute invalid portions."
165 + )
166 +
167 + # 7. verify request can be rescheduled in next step
168 + scheduler_output_2 = recompute_scheduler.schedule()
169 +
170 + # request should appear in the new schedule to recompute invalid blocks
171 + scheduled_req_ids = [
172 + req.request_id for req in scheduler_output_2.scheduled_new_reqs
173 + ]
174 + if scheduler_output_2.num_scheduled_tokens:
175 + scheduled_req_ids.extend(scheduler_output_2.num_scheduled_tokens.keys())
176 +
177 + assert (
178 + request.request_id in scheduled_req_ids or len(recompute_scheduler.running) > 0
179 + ), "Request should be reschedulable for recomputation"
180 +
181 +
182 +def test_sync_fail_invalid_blocks_evicted(fail_scheduler: Scheduler):
183 + """
184 + Test sync fail case - invalid blocks must be evicted from cache.
185 +
186 + When a request fails with policy='fail' and has invalid blocks from sync loading:
187 + 1. Request should be finished with FINISHED_ERROR
188 + 2. Invalid blocks should be evicted from the KV cache
189 + 3. Valid blocks (if shared) should remain in cache
190 + 4. Future requests should not reuse the invalid blocks
191 +
192 + This test verifies that invalid blocks are properly evicted to prevent
193 + cache corruption and reuse of invalid data.
194 + """
195 + num_prompt_blocks = 100
196 + num_external_computed_blocks = 99
197 + invalid_block_idx = 50
198 +
199 + num_prompt_tokens = num_prompt_blocks * fail_scheduler.block_size
200 + num_external_computed_tokens = (
201 + num_external_computed_blocks * fail_scheduler.block_size
202 + )
203 +
204 + request = create_request(num_tokens=num_prompt_tokens)
205 + fail_scheduler.add_request(request=request)
206 +
207 + req_num_new_matched_tokens = {
208 + request.request_id: num_external_computed_tokens,
209 + }
210 +
211 + # mock connector indicating sync load
212 + fail_scheduler.connector = Mock()
213 + fail_scheduler.connector.get_num_new_matched_tokens.side_effect = (
214 + _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False)
215 + )
216 + fail_scheduler.connector.request_finished.return_value = (False, None)
217 + fail_scheduler.connector.take_events.return_value = ()
218 +
219 + scheduler_output = fail_scheduler.schedule()
220 +
221 + # request should be running with sync KV load
222 + assert len(fail_scheduler.running) == 1
223 + assert request.status == RequestStatus.RUNNING
224 +
225 + # get allocated block IDs
226 + req_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0]
227 + invalid_block_id = req_block_ids[invalid_block_idx]
228 + invalid_block_ids = {invalid_block_id}
229 +
230 + # verify the block is in the block pool before we report it as invalid
231 + block = fail_scheduler.kv_cache_manager.block_pool.blocks[invalid_block_id]
232 + assert block is not None
233 +
234 + # report invalid blocks - request should fail
235 + model_runner_output = create_model_runner_output(
236 + [request],
237 + invalid_block_ids=invalid_block_ids,
238 + use_eos=True,
239 + )
240 +
241 + outputs = fail_scheduler.update_from_output(scheduler_output, model_runner_output)
242 +
243 + # verify request is finished with error
244 + assert request.status == RequestStatus.FINISHED_ERROR
245 + assert request.get_finished_reason() == FinishReason.ERROR
246 +
247 + # verify output is generated
248 + assert len(outputs) == 1
249 + engine_outputs = next(iter(outputs.values()))
250 + assert len(engine_outputs.outputs) == 1
251 + output = engine_outputs.outputs[0]
252 + assert output.request_id == request.request_id
253 + assert output.finish_reason == FinishReason.ERROR
254 +
255 + # verify the request was removed from scheduler
256 + assert request.request_id not in fail_scheduler.requests
257 + assert len(fail_scheduler.running) == 0
258 +
259 + # critical: verify invalid block was actually freed from cache
260 + # this is the key assertion - the invalid block should no longer be
261 + # tracked by the KV cache manager for this request
262 + # if it's still there, a future request could reuse the invalid data
263 + try:
264 + block_ids = fail_scheduler.kv_cache_manager.get_block_ids(request.request_id)
265 + # if we get here, check if blocks were actually freed
266 + if block_ids is not None and len(block_ids[0]) > 0:
267 + pytest.fail(
268 + f"Invalid blocks still tracked for finished request! "
269 + f"Request {request.request_id} should have been freed but "
270 + f"still has {len(block_ids[0])} blocks allocated."
271 + )
272 + # blocks list exists but is empty - this is fine, they were freed
273 + except KeyError:
274 + # expected - request completely removed from tracking
275 + pass
276 +
277 + # critical: verify invalid block was evicted from prefix cache
278 + # the block should no longer have a hash (hash is reset on eviction)
279 + assert block.block_hash is None, (
280 + f"Invalid block {invalid_block_id} should have been evicted from cache "
281 + f"(hash should be None), but hash is still {block.block_hash}"
282 + )
283 +
284 +
285 +def test_async_recompute_blocks_not_cached_when_invalid(
286 + recompute_scheduler: Scheduler,
287 +):
288 + """
289 + Test async recompute case - invalid blocks not cached after transfer.
290 +
291 + When async KV loading has invalid blocks and retry_policy is 'recompute':
292 + 1. Blocks are allocated but not cached yet
293 + 2. When async transfer completes, only valid blocks should be cached
294 + 3. Invalid blocks should never enter the prefix cache
295 +
296 + This test verifies correctness, the failed_recving_kv_req_ids protection
297 + ensures only valid blocks are cached when the transfer completes, and we
298 + only evict blocks from cache that are already hashed in the block table.
299 + """
300 + from unittest.mock import patch
301 +
302 + num_prompt_blocks = 100
303 + num_external_computed_blocks = 99
304 + invalid_block_idx = 50
305 +
306 + num_prompt_tokens = num_prompt_blocks * recompute_scheduler.block_size
307 + num_external_computed_tokens = (
308 + num_external_computed_blocks * recompute_scheduler.block_size
309 + )
310 +
311 + request = create_request(num_tokens=num_prompt_tokens)
312 + recompute_scheduler.add_request(request=request)
313 +
314 + req_num_new_matched_tokens = {
315 + request.request_id: num_external_computed_tokens,
316 + }
317 +
318 + # mock connector indicating async load
319 + recompute_scheduler.connector = Mock()
320 + recompute_scheduler.connector.get_num_new_matched_tokens.side_effect = (
321 + _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True)
322 + )
323 + recompute_scheduler.connector.request_finished.return_value = (False, None)
324 + recompute_scheduler.connector.take_events.return_value = ()
325 +
326 + scheduler_output = recompute_scheduler.schedule()
327 +
328 + # request should be waiting for remote KVs
329 + assert len(recompute_scheduler.waiting) == 1
330 + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
331 + assert request.num_computed_tokens == 0
332 +
333 + # get the allocated block IDs
334 + (req_block_ids,) = recompute_scheduler.kv_cache_manager.get_block_ids(
335 + request.request_id
336 + )
337 + invalid_block_id = req_block_ids[invalid_block_idx]
338 + invalid_block_ids = {invalid_block_id}
339 +
340 + # get the block object to verify it's not cached yet and stays uncached
341 + block = recompute_scheduler.kv_cache_manager.block_pool.blocks[invalid_block_id]
342 +
343 + # verify block has no hash before invalid blocks are reported
344 + assert block.block_hash is None, (
345 + "Async loading blocks should not be cached yet (no hash)"
346 + )
347 +
348 + # report invalid blocks (transfer not finished yet)
349 + model_runner_output = create_model_runner_output(
350 + reqs=[],
351 + finished_recving=None, # transfer NOT finished
352 + invalid_block_ids=invalid_block_ids,
353 + use_eos=False,
354 + )
355 +
356 + # critical: spy on evict_blocks to verify it's NOT called for async blocks
357 + original_evict_blocks = recompute_scheduler.kv_cache_manager.evict_blocks
358 + evict_blocks_calls = []
359 +
360 + def evict_blocks_spy(block_ids):
361 + evict_blocks_calls.append(set(block_ids))
362 + return original_evict_blocks(block_ids)
363 +
364 + with patch.object(
365 + recompute_scheduler.kv_cache_manager, "evict_blocks", evict_blocks_spy
366 + ):
367 + recompute_scheduler.update_from_output(scheduler_output, model_runner_output)
368 +
369 + # verify evict_blocks was NOT called (async blocks excluded from eviction)
370 + assert len(evict_blocks_calls) == 0, (
371 + f"evict_blocks should not be called for async-only invalid blocks, "
372 + f"but was called {len(evict_blocks_calls)} time(s) with {evict_blocks_calls}"
373 + )
374 +
375 + # request should still be waiting (not finished with error due to recompute policy)
376 + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
377 + assert request.request_id in recompute_scheduler.failed_recving_kv_req_ids
378 +
379 + # verify num_computed_tokens was truncated to before invalid block
380 + expected_valid_tokens = invalid_block_idx * recompute_scheduler.block_size
381 + assert request.num_computed_tokens == expected_valid_tokens
382 +
383 + # verify invalid block still has no hash (was not evicted)
384 + assert block.block_hash is None, (
385 + f"Async loading blocks shouldn't be cached or evicted. "
386 + f"Block {invalid_block_id} hash should be None but is {block.block_hash}"
387 + )
388 +
389 + # now simulate async transfer completing
390 + model_runner_output_2 = create_model_runner_output(
391 + reqs=[],
392 + finished_recving={request.request_id},
393 + invalid_block_ids=None,
394 + use_eos=False,
395 + )
396 +
397 + recompute_scheduler.update_from_output(scheduler_output, model_runner_output_2)
398 +
399 + # verify request is now marked as finished receiving and ready to be processed
400 + assert request.request_id in recompute_scheduler.finished_recving_kv_req_ids
401 + assert request.request_id in recompute_scheduler.failed_recving_kv_req_ids
402 +
403 + # critical: verify invalid block still has no hash before recompute
404 + # the async transfer invalid data was never cached
405 + assert block.block_hash is None, (
406 + f"Invalid block {invalid_block_id} should not be cached before recompute "
407 + f"(hash should be None), but hash is {block.block_hash}"
408 + )
409 +
410 + # critical end-to-end test: spy on cache_blocks to verify it's called with
411 + # the truncated num_computed_tokens value
412 + original_cache_blocks = recompute_scheduler.kv_cache_manager.cache_blocks
413 + cache_blocks_calls = []
414 +
415 + def cache_blocks_spy(req, num_tokens):
416 + cache_blocks_calls.append((req.request_id, num_tokens))
417 + return original_cache_blocks(req, num_tokens)
418 +
419 + with patch.object(
420 + recompute_scheduler.kv_cache_manager, "cache_blocks", cache_blocks_spy
421 + ):
422 + # call schedule() again - this triggers _update_waiting_for_remote_kv()
423 + # which should call cache_blocks with the truncated value
424 + recompute_scheduler.schedule()
425 +
426 + # verify cache_blocks was called with the truncated value
427 + assert len(cache_blocks_calls) == 1, (
428 + f"cache_blocks should be called exactly once, "
429 + f"got {len(cache_blocks_calls)} calls"
430 + )
431 + cached_req_id, cached_num_tokens = cache_blocks_calls[0]
432 + assert cached_req_id == request.request_id
433 + assert cached_num_tokens == expected_valid_tokens, (
434 + f"cache_blocks should be called with truncated value {expected_valid_tokens}, "
435 + f"but was called with {cached_num_tokens}"
436 + )
437 +
438 + # request should now be RUNNING (scheduled immediately after transfer completes)
439 + # the flow is: WAITING_FOR_REMOTE_KVS -> WAITING -> RUNNING in same schedule() call
440 + assert request.status == RequestStatus.RUNNING
441 +
442 + # num_computed_tokens should be >= expected_valid_tokens because the scheduler
443 + # will schedule additional new tokens (up to max_num_batched_tokens) for the request
444 + assert request.num_computed_tokens >= expected_valid_tokens, (
445 + f"num_computed_tokens should be at least {expected_valid_tokens}, "
446 + f"got {request.num_computed_tokens}"
447 + )
448 +
449 + # request should no longer be in the failed/finished receiving sets
450 + assert request.request_id not in recompute_scheduler.failed_recving_kv_req_ids
451 + assert request.request_id not in recompute_scheduler.finished_recving_kv_req_ids
452 +
453 + # request should be in the running queue
454 + assert request in recompute_scheduler.running
vllm/config/kv_transfer.py
+5
@@ -64,6 +64,11 @@ class KVTransferConfig:
64 enable_permute_local_kv: bool = False
65 """Experiment feature flag to enable HND to NHD KV Transfer"""
66
67 + kv_load_failure_policy: Literal["recompute", "fail"] = "recompute"
68 + """Policy for handling KV cache load failures.
69 + 'recompute': reschedule the request to recompute failed blocks (default)
70 + 'fail': immediately fail the request with an error finish reason"""
71 +
72 def compute_hash(self) -> str:
73 """
74 WARNING: Whenever a new field is added to this config,
vllm/entrypoints/openai/serving_chat.py
+16 -1
@@ -51,7 +51,11 @@ from vllm.entrypoints.openai.protocol import (
51 ToolCall,
52 UsageInfo,
53 )
54 -from vllm.entrypoints.openai.serving_engine import OpenAIServing, clamp_prompt_logprobs
54 +from vllm.entrypoints.openai.serving_engine import (
55 + GenerationError,
56 + OpenAIServing,
57 + clamp_prompt_logprobs,
58 +)
59 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
60 from vllm.entrypoints.openai.tool_parsers import ToolParser
61 from vllm.entrypoints.openai.tool_parsers.mistral_tool_parser import MistralToolCall
@@ -380,6 +384,8 @@ class OpenAIServingChat(OpenAIServing):
384 tokenizer,
385 request_metadata,
386 )
387 + except GenerationError as e:
388 + return self._convert_generation_error_to_response(e)
389 except ValueError as e:
390 # TODO: Use a vllm-specific Validation Error
391 return self.create_error_response(str(e))
@@ -1120,6 +1126,10 @@ class OpenAIServingChat(OpenAIServing):
1126
1127 # if the model is finished generating
1128 else:
1129 + # check for error finish reason and abort streaming
1130 + # finish_reason='error' indicates a retryable error
1131 + self._raise_if_error(output.finish_reason, request_id)
1132 +
1133 # check to make sure we haven't "forgotten" to stream
1134 # any tokens that were generated but previously
1135 # matched by partial json parsing
@@ -1287,6 +1297,8 @@ class OpenAIServingChat(OpenAIServing):
1297 delta=False,
1298 )
1299
1300 + except GenerationError as e:
1301 + yield f"data: {self._convert_generation_error_to_streaming_response(e)}\n\n"
1302 except Exception as e:
1303 # TODO: Use a vllm-specific Validation Error
1304 logger.exception("Error in chat completion stream generator.")
@@ -1327,6 +1339,9 @@ class OpenAIServingChat(OpenAIServing):
1339
1340 role = self.get_chat_request_role(request)
1341 for output in final_res.outputs:
1342 + # check for error finish reason and raise GenerationError
1343 + # finish_reason='error' indicates a retryable request-level internal error
1344 + self._raise_if_error(output.finish_reason, request_id)
1345 token_ids = output.token_ids
1346 out_logprobs = output.logprobs
1347 tool_call_info = None
vllm/entrypoints/openai/serving_completion.py
+14 -1
@@ -24,7 +24,11 @@ from vllm.entrypoints.openai.protocol import (
24 RequestResponseMetadata,
25 UsageInfo,
26 )
27 -from vllm.entrypoints.openai.serving_engine import OpenAIServing, clamp_prompt_logprobs
27 +from vllm.entrypoints.openai.serving_engine import (
28 + GenerationError,
29 + OpenAIServing,
30 + clamp_prompt_logprobs,
31 +)
32 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
33 from vllm.entrypoints.renderer import RenderConfig
34 from vllm.entrypoints.utils import get_max_tokens, should_include_usage
@@ -300,6 +304,8 @@ class OpenAIServingCompletion(OpenAIServing):
304 )
305 except asyncio.CancelledError:
306 return self.create_error_response("Client disconnected")
307 + except GenerationError as e:
308 + return self._convert_generation_error_to_response(e)
309 except ValueError as e:
310 # TODO: Use a vllm-specific Validation Error
311 return self.create_error_response(str(e))
@@ -437,6 +443,8 @@ class OpenAIServingCompletion(OpenAIServing):
443 finish_reason = output.finish_reason
444 stop_reason = output.stop_reason
445
446 + self._raise_if_error(finish_reason, request_id)
447 +
448 chunk = CompletionStreamResponse(
449 id=request_id,
450 created=created_time,
@@ -498,8 +506,11 @@ class OpenAIServingCompletion(OpenAIServing):
506 # report to FastAPI middleware aggregate usage across all choices
507 request_metadata.final_usage_info = final_usage_info
508
509 + except GenerationError as e:
510 + yield f"data: {self._convert_generation_error_to_streaming_response(e)}\n\n"
511 except Exception as e:
512 # TODO: Use a vllm-specific Validation Error
513 + logger.exception("Error in completion stream generator.")
514 data = self.create_streaming_error_response(str(e))
515 yield f"data: {data}\n\n"
516 yield "data: [DONE]\n\n"
@@ -530,6 +541,8 @@ class OpenAIServingCompletion(OpenAIServing):
541 out_logprobs: GenericSequence[dict[int, Logprob] | None] | None
542
543 for output in final_res.outputs:
544 + self._raise_if_error(output.finish_reason, request_id)
545 +
546 assert request.max_tokens is not None
547 if request.echo:
548 if request.return_token_ids:
vllm/entrypoints/openai/serving_engine.py
+61
@@ -133,6 +133,15 @@ from vllm.utils.async_utils import (
133 from vllm.utils.collection_utils import is_list_of
134 from vllm.v1.engine import EngineCoreRequest
135
136 +
137 +class GenerationError(Exception):
138 + """raised when finish_reason indicates internal server error (500)"""
139 +
140 + def __init__(self, message: str = "Internal server error"):
141 + super().__init__(message)
142 + self.status_code = HTTPStatus.INTERNAL_SERVER_ERROR
143 +
144 +
145 logger = init_logger(__name__)
146
147 CompletionLikeRequest: TypeAlias = (
@@ -456,6 +465,29 @@ class OpenAIServing:
465 # Iterate through all beam inference results
466 for i, result in enumerate(output):
467 current_beam = all_beams[i]
468 +
469 + # check for error finish reason and abort beam search
470 + if result.outputs[0].finish_reason == "error":
471 + # yield error output and terminate beam search
472 + yield RequestOutput(
473 + request_id=request_id,
474 + prompt=prompt_text,
475 + outputs=[
476 + CompletionOutput(
477 + index=0,
478 + text="",
479 + token_ids=[],
480 + cumulative_logprob=None,
481 + logprobs=None,
482 + finish_reason="error",
483 + )
484 + ],
485 + finished=True,
486 + prompt_token_ids=prompt_token_ids,
487 + prompt_logprobs=None,
488 + )
489 + return
490 +
491 if result.outputs[0].logprobs is not None:
492 logprobs = result.outputs[0].logprobs[0]
493 all_beams_token_id.extend(list(logprobs.keys()))
@@ -780,6 +812,35 @@ class OpenAIServing:
812 )
813 return json_str
814
815 + def _raise_if_error(self, finish_reason: str | None, request_id: str) -> None:
816 + """Raise GenerationError if finish_reason indicates an error."""
817 + if finish_reason == "error":
818 + logger.error(
819 + "Request %s failed with an internal error during generation",
820 + request_id,
821 + )
822 + raise GenerationError("Internal server error")
823 +
824 + def _convert_generation_error_to_response(
825 + self, e: GenerationError
826 + ) -> ErrorResponse:
827 + """Convert GenerationError to ErrorResponse."""
828 + return self.create_error_response(
829 + str(e),
830 + err_type="InternalServerError",
831 + status_code=e.status_code,
832 + )
833 +
834 + def _convert_generation_error_to_streaming_response(
835 + self, e: GenerationError
836 + ) -> str:
837 + """Convert GenerationError to streaming error response."""
838 + return self.create_streaming_error_response(
839 + str(e),
840 + err_type="InternalServerError",
841 + status_code=e.status_code,
842 + )
843 +
844 async def _check_model(
845 self,
846 request: AnyRequest,
vllm/entrypoints/openai/serving_responses.py
+40 -13
@@ -50,6 +50,7 @@ from openai.types.responses.response_reasoning_item import (
50 )
51 from openai.types.responses.tool import Mcp, Tool
52 from openai_harmony import Message as OpenAIHarmonyMessage
53 +from pydantic import TypeAdapter
54
55 from vllm import envs
56 from vllm.engine.protocol import EngineClient
@@ -94,7 +95,10 @@ from vllm.entrypoints.openai.protocol import (
95 ResponseUsage,
96 StreamingResponsesResponse,
97 )
97 -from vllm.entrypoints.openai.serving_engine import OpenAIServing
98 +from vllm.entrypoints.openai.serving_engine import (
99 + GenerationError,
100 + OpenAIServing,
101 +)
102 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
103 from vllm.entrypoints.responses_utils import (
104 construct_input_messages,
@@ -541,6 +545,8 @@ class OpenAIServingResponses(OpenAIServing):
545 tokenizer,
546 request_metadata,
547 )
548 + except GenerationError as e:
549 + return self._convert_generation_error_to_response(e)
550 except Exception as e:
551 return self.create_error_response(str(e))
552
@@ -648,6 +654,8 @@ class OpenAIServingResponses(OpenAIServing):
654 status = "incomplete"
655 elif context.finish_reason == "abort":
656 status = "cancelled"
657 + else:
658 + self._raise_if_error(context.finish_reason, request.request_id)
659 else:
660 status = "incomplete"
661 elif isinstance(context, ParsableContext):
@@ -673,6 +681,9 @@ class OpenAIServingResponses(OpenAIServing):
681 assert len(final_res.outputs) == 1
682 final_output = final_res.outputs[0]
683
684 + # finish_reason='error' indicates retryable internal error
685 + self._raise_if_error(final_output.finish_reason, request.request_id)
686 +
687 output = self._make_response_output_items(request, final_output, tokenizer)
688
689 if request.enable_response_messages:
@@ -1066,6 +1077,8 @@ class OpenAIServingResponses(OpenAIServing):
1077 async for event in generator:
1078 event_deque.append(event)
1079 new_event_signal.set() # Signal new event available
1080 + except GenerationError as e:
1081 + response = self._convert_generation_error_to_response(e)
1082 except Exception as e:
1083 logger.exception("Background request failed for %s", request.request_id)
1084 response = self.create_error_response(str(e))
@@ -1089,6 +1102,8 @@ class OpenAIServingResponses(OpenAIServing):
1102 ):
1103 try:
1104 response = await self.responses_full_generator(request, *args, **kwargs)
1105 + except GenerationError as e:
1106 + response = self._convert_generation_error_to_response(e)
1107 except Exception as e:
1108 logger.exception("Background request failed for %s", request.request_id)
1109 response = self.create_error_response(str(e))
@@ -1227,6 +1242,8 @@ class OpenAIServingResponses(OpenAIServing):
1242 continue
1243 if ctx.last_output.outputs:
1244 output = ctx.last_output.outputs[0]
1245 + # finish_reason='error' indicates a retryable error
1246 + self._raise_if_error(output.finish_reason, request.request_id)
1247 if reasoning_parser:
1248 delta_message = reasoning_parser.extract_reasoning_streaming(
1249 previous_text=previous_text,
@@ -1522,6 +1539,9 @@ class OpenAIServingResponses(OpenAIServing):
1539 async for ctx in result_generator:
1540 assert isinstance(ctx, StreamingHarmonyContext)
1541
1542 + # finish_reason='error' indicates a retryable error
1543 + self._raise_if_error(ctx.finish_reason, request.request_id)
1544 +
1545 if ctx.is_expecting_start():
1546 current_output_index += 1
1547 sent_output_item_added = False
@@ -2016,18 +2036,25 @@ class OpenAIServingResponses(OpenAIServing):
2036 )
2037 )
2038
2019 - async for event_data in processer(
2020 - request,
2021 - sampling_params,
2022 - result_generator,
2023 - context,
2024 - model_name,
2025 - tokenizer,
2026 - request_metadata,
2027 - created_time,
2028 - _increment_sequence_number_and_return,
2029 - ):
2030 - yield event_data
2039 + try:
2040 + async for event_data in processer(
2041 + request,
2042 + sampling_params,
2043 + result_generator,
2044 + context,
2045 + model_name,
2046 + tokenizer,
2047 + request_metadata,
2048 + created_time,
2049 + _increment_sequence_number_and_return,
2050 + ):
2051 + yield event_data
2052 + except GenerationError as e:
2053 + error_json = self._convert_generation_error_to_streaming_response(e)
2054 + yield _increment_sequence_number_and_return(
2055 + TypeAdapter(StreamingResponsesResponse).validate_json(error_json)
2056 + )
2057 + return
2058
2059 async def empty_async_generator():
2060 # A hack to trick Python to think this is a generator but
vllm/v1/core/block_pool.py
+19
@@ -397,6 +397,25 @@ class BlockPool:
397 [block for block in blocks_list if block.ref_cnt == 0 and not block.is_null]
398 )
399
400 + def evict_blocks(self, block_ids: set[int]) -> None:
401 + """evict blocks from the prefix cache by their block IDs.
402 +
403 + only evicts blocks that are currently cached (have a hash). blocks
404 + with ref_cnt > 0 are not freed from the block pool, only evicted
405 + from the prefix cache hash table.
406 +
407 + Args:
408 + block_ids: Set of block IDs to evict from cache.
409 + """
410 + for block_id in block_ids:
411 + assert block_id < len(self.blocks), (
412 + f"Invalid block_id {block_id} >= {len(self.blocks)}. "
413 + f"This indicates a bug in the KV connector - workers should "
414 + f"only report block IDs that were allocated by the scheduler."
415 + )
416 + block = self.blocks[block_id]
417 + self._maybe_evict_cached_block(block)
418 +
419 def reset_prefix_cache(self) -> bool:
420 """Reset prefix cache. This function may be used in RLHF
421 flows to invalid prefix caching after the weights are updated,
vllm/v1/core/kv_cache_manager.py
+8
@@ -333,6 +333,14 @@ class KVCacheManager:
333 """
334 self.coordinator.free(request.request_id)
335
336 + def evict_blocks(self, block_ids: set[int]) -> None:
337 + """evict blocks from the prefix cache by their block IDs.
338 +
339 + Args:
340 + block_ids: Set of block IDs to evict from cache.
341 + """
342 + self.block_pool.evict_blocks(block_ids)
343 +
344 def reset_prefix_cache(self) -> bool:
345 """Reset prefix cache. This function may be used in RLHF
346 flows to invalidate prefix caching after the weights are updated,
vllm/v1/core/sched/scheduler.py
+85 -31
@@ -106,6 +106,7 @@ class Scheduler(SchedulerInterface):
106 # KV Connector pushes/pull of remote KVs for P/D and offloading.
107 self.connector = None
108 self.connector_prefix_cache_stats: PrefixCacheStats | None = None
109 + self.recompute_kv_load_failures = True
110 if self.vllm_config.kv_transfer_config is not None:
111 assert not self.is_encoder_decoder, (
112 "Encoder-decoder models are not currently supported with KV connectors"
@@ -117,6 +118,10 @@ class Scheduler(SchedulerInterface):
118 )
119 if self.log_stats:
120 self.connector_prefix_cache_stats = PrefixCacheStats()
121 + kv_load_failure_policy = (
122 + self.vllm_config.kv_transfer_config.kv_load_failure_policy
123 + )
124 + self.recompute_kv_load_failures = kv_load_failure_policy == "recompute"
125
126 self.kv_event_publisher = EventPublisherFactory.create(
127 self.kv_events_config,
@@ -1066,7 +1071,7 @@ class Scheduler(SchedulerInterface):
1071 for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
1072 assert num_tokens_scheduled > 0
1073 if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids:
1069 - # Skip requests that were recovered from KV load failure
1074 + # skip failed or rescheduled requests from KV load failure
1075 continue
1076 request = self.requests.get(req_id)
1077 if request is None:
@@ -1177,6 +1182,21 @@ class Scheduler(SchedulerInterface):
1182 # This is a rare case and unlikely to impact performance.
1183 self.waiting.remove_requests(stopped_preempted_reqs)
1184
1185 + if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
1186 + requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
1187 + self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
1188 + for request in requests:
1189 + outputs[request.client_index].append(
1190 + EngineCoreOutput(
1191 + request_id=request.request_id,
1192 + new_token_ids=[],
1193 + finish_reason=request.get_finished_reason(),
1194 + events=request.take_events(),
1195 + trace_headers=request.trace_headers,
1196 + num_cached_tokens=request.num_cached_tokens,
1197 + )
1198 + )
1199 +
1200 # KV Connector: update state for finished KV Transfers.
1201 if kv_connector_output:
1202 self._update_from_kv_xfer_finished(kv_connector_output)
@@ -1610,8 +1630,11 @@ class Scheduler(SchedulerInterface):
1630 self._free_blocks(self.requests[req_id])
1631
1632 def _update_requests_with_invalid_blocks(
1613 - self, requests: Iterable[Request], invalid_block_ids: set[int]
1614 - ) -> tuple[set[str], int]:
1633 + self,
1634 + requests: Iterable[Request],
1635 + invalid_block_ids: set[int],
1636 + evict_blocks: bool = True,
1637 + ) -> tuple[set[str], int, set[int]]:
1638 """
1639 Identify and update requests affected by invalid KV cache blocks.
1640
@@ -1623,16 +1646,21 @@ class Scheduler(SchedulerInterface):
1646 Args:
1647 requests: The set of requests to scan for invalid blocks.
1648 invalid_block_ids: IDs of invalid blocks.
1649 + evict_blocks: Whether to collect blocks for eviction (False for
1650 + async requests which aren't cached yet).
1651
1652 Returns:
1653 tuple:
1654 - affected_req_ids (set[str]): IDs of requests impacted by
1655 invalid blocks.
1656 - total_affected_tokens (int): Total number of tokens that must
1632 - be recomputed across all affected requests (for observability).
1657 + be recomputed across all affected requests.
1658 + - blocks_to_evict (set[int]): Block IDs to evict from cache,
1659 + including invalid blocks and downstream dependent blocks.
1660 """
1661 affected_req_ids: set[str] = set()
1662 total_affected_tokens = 0
1663 + blocks_to_evict: set[int] = set()
1664 # If a block is invalid and shared by multiple requests in the batch,
1665 # these requests must be rescheduled, but only the first will recompute
1666 # it. This set tracks blocks already marked for recomputation.
@@ -1690,6 +1718,9 @@ class Scheduler(SchedulerInterface):
1718 )
1719 total_affected_tokens += num_affected_tokens
1720 request.num_external_computed_tokens -= num_affected_tokens
1721 + # collect invalid block and all downstream dependent blocks
1722 + if evict_blocks:
1723 + blocks_to_evict.update(req_block_ids[idx:])
1724
1725 if is_affected:
1726 if not marked_invalid_block:
@@ -1705,47 +1736,70 @@ class Scheduler(SchedulerInterface):
1736
1737 affected_req_ids.add(request.request_id)
1738
1708 - return affected_req_ids, total_affected_tokens
1739 + return affected_req_ids, total_affected_tokens, blocks_to_evict
1740
1741 def _handle_invalid_blocks(self, invalid_block_ids: set[int]) -> set[str]:
1711 - total_requests_to_reschedule = 0
1712 - total_tokens_to_reschedule = 0
1742 + """
1743 + Handle requests affected by invalid KV cache blocks.
1744
1714 - # --- Handle async KV loads (WAITING_FOR_REMOTE_KVS) ---
1745 + Returns:
1746 + Set of affected request IDs to skip in update_from_output main loop.
1747 + """
1748 + should_fail = not self.recompute_kv_load_failures
1749 +
1750 + # handle async KV loads (not cached yet, evict_blocks=False)
1751 async_load_reqs = (
1752 req
1753 for req in self.waiting
1754 if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
1755 )
1720 - async_affected_req_ids, num_tokens_to_reschedule = (
1756 + async_failed_req_ids, num_failed_tokens, _ = (
1757 self._update_requests_with_invalid_blocks(
1722 - async_load_reqs, invalid_block_ids
1758 + async_load_reqs, invalid_block_ids, evict_blocks=False
1759 )
1760 )
1761
1726 - total_requests_to_reschedule += len(async_affected_req_ids)
1727 - total_tokens_to_reschedule += num_tokens_to_reschedule
1728 -
1729 - # Mark requests with async KV load failures; they will be rescheduled
1730 - # once loading completes.
1731 - self.failed_recving_kv_req_ids |= async_affected_req_ids
1762 + total_failed_requests = len(async_failed_req_ids)
1763 + total_failed_tokens = num_failed_tokens
1764
1733 - # --- Handle sync KV loads (running requests) ---
1734 - sync_affected_req_ids, num_tokens_to_reschedule = (
1735 - self._update_requests_with_invalid_blocks(self.running, invalid_block_ids)
1765 + # handle sync loads (may be cached, collect blocks for eviction)
1766 + sync_failed_req_ids, num_failed_tokens, sync_blocks_to_evict = (
1767 + self._update_requests_with_invalid_blocks(
1768 + self.running, invalid_block_ids, evict_blocks=True
1769 + )
1770 )
1771
1738 - total_requests_to_reschedule += len(sync_affected_req_ids)
1739 - total_tokens_to_reschedule += num_tokens_to_reschedule
1740 -
1741 - if total_requests_to_reschedule:
1742 - logger.warning(
1743 - "Recovered from KV load failure: "
1744 - "%d request(s) rescheduled (%d tokens affected).",
1745 - total_requests_to_reschedule,
1746 - total_tokens_to_reschedule,
1772 + total_failed_requests += len(sync_failed_req_ids)
1773 + total_failed_tokens += num_failed_tokens
1774 +
1775 + if not total_failed_requests:
1776 + return set()
1777 +
1778 + # evict invalid blocks and downstream dependent blocks from cache
1779 + # only when not using recompute policy (where blocks will be recomputed
1780 + # and reused by other requests sharing them)
1781 + if sync_blocks_to_evict and not self.recompute_kv_load_failures:
1782 + self.kv_cache_manager.evict_blocks(sync_blocks_to_evict)
1783 +
1784 + if should_fail:
1785 + all_failed_req_ids = async_failed_req_ids | sync_failed_req_ids
1786 + logger.error(
1787 + "Failing %d request(s) due to KV load failure "
1788 + "(failure_policy=fail, %d tokens affected). Request IDs: %s",
1789 + total_failed_requests,
1790 + total_failed_tokens,
1791 + all_failed_req_ids,
1792 )
1793 + return all_failed_req_ids
1794 +
1795 + logger.warning(
1796 + "Recovered from KV load failure: "
1797 + "%d request(s) rescheduled (%d tokens affected).",
1798 + total_failed_requests,
1799 + total_failed_tokens,
1800 + )
1801
1749 - # Return the IDs of affected running requests to skip in
1750 - # update_from_output.
1751 - return sync_affected_req_ids
1802 + # Mark async requests with KV load failures for retry once loading completes
1803 + self.failed_recving_kv_req_ids |= async_failed_req_ids
1804 + # Return sync affected IDs to skip in update_from_output
1805 + return sync_failed_req_ids
vllm/v1/engine/__init__.py
+6 -3
@@ -19,24 +19,27 @@ from vllm.v1.serial_utils import UtilityResult
19
20 # These are possible values of RequestOutput.finish_reason,
21 # so form part of the external API.
22 -FINISH_REASON_STRINGS = ("stop", "length", "abort")
22 +FINISH_REASON_STRINGS = ("stop", "length", "abort", "error")
23
24
25 class FinishReason(enum.IntEnum):
26 """
27 - Reason a request finished - stop, length, or abort.
27 + Reason a request finished - stop, length, abort, or error.
28
29 Int rather than Str for more compact serialization.
30
31 stop - a stop string was emitted
32 length - max_tokens was consumed, or max_model_len was reached
33 - abort - aborted for another reason
33 + abort - aborted by client
34 + error - retryable request-level internal error (e.g., KV load failure).
35 + Invariant: always converted to 500 Internal Server Error.
36
37 """
38
39 STOP = 0
40 LENGTH = 1
41 ABORT = 2
42 + ERROR = 3
43
44 def __str__(self):
45 return FINISH_REASON_STRINGS[self.value]
vllm/v1/request.py
+2
@@ -255,6 +255,7 @@ class RequestStatus(enum.IntEnum):
255 FINISHED_LENGTH_CAPPED = enum.auto()
256 FINISHED_ABORTED = enum.auto()
257 FINISHED_IGNORED = enum.auto()
258 + FINISHED_ERROR = enum.auto()
259
260 def __str__(self):
261 return self.name
@@ -277,4 +278,5 @@ _FINISHED_REASON_MAP = {
278 RequestStatus.FINISHED_LENGTH_CAPPED: FinishReason.LENGTH,
279 RequestStatus.FINISHED_ABORTED: FinishReason.ABORT,
280 RequestStatus.FINISHED_IGNORED: FinishReason.LENGTH,
281 + RequestStatus.FINISHED_ERROR: FinishReason.ERROR,
282 }