@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dc264bcea1

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, 23:45 UTC

[BugFix] Eagerly abort cancelled final-step requests (#29987)

Currently, when requests are cancelled while executing their final step, "completion" is handled based on normal stop processing (e.g. length or stop token), so the abort has no effect. This is typically not a problem, but when a kv connector is involved it thinks the request completed successfully rather than being aborted. This is problematic for disaggregated prefill which will free kv cache blocks if the request was aborted but not if it completed successfully—since the cancelled request will never be sent to the decode side, kv cache blocks remain pinned until the fall-back timeout expires. The problem is exacerbated when many requests are cancelled and/or there are large prefills whose forward pass takes a long time (since the window is bigger). This PR fixes the problem by processing pending aborts immediately prior to processing model output each step; we process only aborts, not new requests, since it's preferable for latency to process model outputs before new incoming requests. Fixes #26400. Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Dec 5, 2025 at 09:28 UTC dc264bcea17b279173c7cbb99645c35a56bed778

3 files changed +352 -6

tests/v1/engine/test_abort_final_step.py new

+311

@@ -0,0 +1,311 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +"""

5 +Test for the fix in PR #29987: Eagerly abort cancelled final-step requests.

6 +

7 +This test verifies that when a request is aborted during its final execution

8 +step (when it would naturally complete), it is properly marked as aborted

9 +rather than being treated as normally completed.

10 +

11 +The test uses a dummy KV connector to verify that the connector receives

12 +the correct finish status (FINISHED_ABORTED, not FINISHED_LENGTH_CAPPED).

13 +"""

14 +

15 +import asyncio

16 +import tempfile

17 +import time

18 +from pathlib import Path

19 +from typing import Any

20 +from unittest.mock import patch

21 +

22 +import pytest

23 +

24 +from vllm import SamplingParams

25 +from vllm.config import KVTransferConfig, VllmConfig

26 +from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory

27 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (

28 + KVConnectorBase_V1,

29 + KVConnectorMetadata,

30 + KVConnectorRole,

31 +)

32 +from vllm.engine.arg_utils import AsyncEngineArgs

33 +from vllm.platforms import current_platform

34 +from vllm.sampling_params import RequestOutputKind

35 +from vllm.utils.torch_utils import set_default_torch_num_threads

36 +from vllm.v1.core.sched.output import SchedulerOutput

37 +from vllm.v1.engine.async_llm import AsyncLLM

38 +from vllm.v1.kv_cache_interface import KVCacheConfig

39 +from vllm.v1.request import Request

40 +

41 +if not current_platform.is_cuda():

42 + pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True)

43 +

44 +TEXT_PROMPT = "Hello"

45 +

46 +

47 +class DummyKVConnectorMetadata(KVConnectorMetadata):

48 + """Dummy metadata for the test connector."""

49 +

50 + def __init__(self):

51 + self.requests: list = []

52 +

53 +

54 +class DummyKVConnector(KVConnectorBase_V1):

55 + """

56 + Dummy KV connector that captures request finish statuses to a file.

57 + This is used to verify the fix - without the fix, a request aborted

58 + during its final step would be captured as FINISHED_LENGTH_CAPPED

59 + instead of FINISHED_ABORTED.

60 +

61 + The connector runs in a separate process, so we write statuses to a file

62 + that can be read by the test process.

63 + """

64 +

65 + def __init__(

66 + self,

67 + vllm_config: VllmConfig,

68 + role: KVConnectorRole,

69 + kv_cache_config: KVCacheConfig | None = None,

70 + ):

71 + super().__init__(vllm_config, role, kv_cache_config)

72 + # Get the status file path from extra config

73 + extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config or {}

74 + self.status_file = extra_config.get("status_file")

75 + # Log that we were initialized

76 + if self.status_file:

77 + try:

78 + with open(self.status_file, "a") as f:

79 + f.write(f"INIT:{role.name}\n")

80 + except Exception:

81 + pass

82 +

83 + def get_num_new_matched_tokens(

84 + self,

85 + request: Request,

86 + num_computed_tokens: int,

87 + ) -> tuple[int | None, bool]:

88 + return (0, False)

89 +

90 + def update_state_after_alloc(

91 + self,

92 + request: Request,

93 + blocks: Any,

94 + num_external_tokens: int,

95 + ):

96 + pass

97 +

98 + def build_connector_meta(

99 + self, scheduler_output: SchedulerOutput

100 + ) -> KVConnectorMetadata:

101 + return DummyKVConnectorMetadata()

102 +

103 + def request_finished(

104 + self,

105 + request: Request,

106 + block_ids: list[int],

107 + ) -> tuple[bool, dict[str, Any] | None]:

108 + """Capture the request status when finished by writing to a file."""

109 + if self.status_file:

110 + try:

111 + with open(self.status_file, "a") as f:

112 + # Write the status name (e.g., "FINISHED_ABORTED")

113 + f.write(f"{request.status.name}\n")

114 + except Exception as e:

115 + # Log but don't fail - this is just test instrumentation

116 + print(f"[DummyKVConnector] Failed to write status: {e}")

117 + return False, None

118 +

119 + def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None:

120 + pass

121 +

122 + def wait_for_layer_load(self, layer_name: str) -> None:

123 + pass

124 +

125 + def save_kv_layer(

126 + self,

127 + layer_name: str,

128 + kv_layer: Any,

129 + attn_metadata: Any,

130 + **kwargs: Any,

131 + ) -> None:

132 + pass

133 +

134 + def wait_for_save(self):

135 + pass

136 +

137 +

138 +# Register the dummy connector

139 +KVConnectorFactory.register_connector(

140 + "DummyKVConnector", __name__, DummyKVConnector.__name__

141 +)

142 +

143 +

144 +@pytest.mark.parametrize("async_scheduling", [False, True])

145 +@pytest.mark.asyncio

146 +async def test_abort_during_final_step(async_scheduling: bool):

147 + """

148 + Test that a request aborted during its final execution step is treated as

149 + aborted rather than completed.

150 +

151 + This test:

152 + 1. Monkeypatches execute_model to wait for a file to be deleted

153 + 2. Configures a dummy KV connector to capture finish statuses

154 + 3. Starts a request with max_tokens=1 (will complete on first decode step)

155 + 4. Aborts the request, then deletes the file to unblock execute_model

156 + 5. Verifies the KV connector received FINISHED_ABORTED not FINISHED_LENGTH_CAPPED

157 +

158 + See https://github.com/vllm-project/vllm/pull/29987.

159 +

160 + Without the fix, the KV connector would see FINISHED_LENGTH_CAPPED because

161 + update_from_output() would mark the request as completed before processing

162 + the abort. This causes KV cache blocks to not be freed properly in

163 + disaggregated prefill scenarios.

164 +

165 + With the fix, _process_aborts_queue() runs before update_from_output(), so the

166 + abort takes precedence and the KV connector sees FINISHED_ABORTED.

167 + """

168 +

169 + # Create three temporary files:

170 + # 1. ready_file: deleted by execute_model to signal it has started

171 + # 2. block_file: execute_model waits for this to be deleted

172 + # 3. status_file: KV connector writes finish statuses here

173 + with tempfile.NamedTemporaryFile(delete=False) as f:

174 + ready_file = Path(f.name)

175 + with tempfile.NamedTemporaryFile(delete=False) as f2:

176 + block_file = Path(f2.name)

177 + with tempfile.NamedTemporaryFile(delete=False, mode="w") as f3:

178 + status_file = Path(f3.name)

179 +

180 + try:

181 + # Get the original execute_model method

182 + from vllm.v1.worker.gpu_worker import Worker

183 +

184 + original_execute_model = Worker.execute_model

185 +

186 + def execute_model_with_wait(self, scheduler_output):

187 + # Signal that execute_model has been called by deleting ready_file

188 + if ready_file.exists():

189 + ready_file.unlink()

190 +

191 + # Wait for the block file to be deleted (triggered from test after abort)

192 + # This runs in the worker process (after fork), so we poll the filesystem

193 + while block_file.exists():

194 + time.sleep(0.01)

195 + return original_execute_model(self, scheduler_output)

196 +

197 + # Patch execute_model to inject the wait

198 + # This happens before the worker process is forked, so the patch applies there

199 + with patch.object(Worker, "execute_model", execute_model_with_wait):

200 + request_id = "test-abort-final-step"

201 +

202 + # Configure engine with dummy KV connector

203 + # Pass the status file path so the connector can write to it

204 + kv_transfer_config = KVTransferConfig(

205 + kv_connector="DummyKVConnector",

206 + kv_role="kv_both",

207 + kv_connector_extra_config={"status_file": str(status_file)},

208 + )

209 + engine_args = AsyncEngineArgs(

210 + model="meta-llama/Llama-3.2-1B-Instruct",

211 + enforce_eager=True,

212 + async_scheduling=async_scheduling,

213 + kv_transfer_config=kv_transfer_config,

214 + )

215 +

216 + with set_default_torch_num_threads(1):

217 + engine = AsyncLLM.from_engine_args(engine_args)

218 +

219 + try:

220 + # Create a request that will complete after just 1 token

221 + sampling_params = SamplingParams(

222 + max_tokens=1,

223 + ignore_eos=True,

224 + output_kind=RequestOutputKind.DELTA,

225 + )

226 +

227 + # Start generation in a task

228 + outputs = []

229 +

230 + async def generate():

231 + async for output in engine.generate(

232 + request_id=request_id,

233 + prompt=TEXT_PROMPT,

234 + sampling_params=sampling_params,

235 + ):

236 + outputs.append(output)

237 +

238 + gen_task = asyncio.create_task(generate())

239 +

240 + # Wait for execute_model to signal it has started (with timeout)

241 + timeout = 5.0 # 5 second timeout

242 + start_time = time.time()

243 + while ready_file.exists():

244 + if time.time() - start_time > timeout:

245 + raise TimeoutError(

246 + "Timeout waiting for execute_model to start. "

247 + "The monkeypatch may not be working correctly, "

248 + "for example if spawn was used instead of fork."

249 + )

250 + await asyncio.sleep(0.01)

251 +

252 + # Abort the request while execute_model is blocked

253 + await engine.abort(request_id)

254 +

255 + # Now unblock execute_model by deleting the file

256 + # The abort should be processed before the model output

257 + block_file.unlink()

258 +

259 + # Wait for generation to complete

260 + await gen_task

261 +

262 + # Give the scheduler a moment to finish cleanup

263 + await asyncio.sleep(0.1)

264 +

265 + # Verify we got output

266 + assert len(outputs) > 0, "Should have received at least one output"

267 +

268 + # The final output should have finish_reason="abort"

269 + final_output = outputs[-1]

270 + assert final_output.finished, (

271 + "Final output should be marked as finished"

272 + )

273 + assert final_output.outputs[0].finish_reason == "abort", (

274 + f"Expected finish_reason='abort' but got "

275 + f"'{final_output.outputs[0].finish_reason}'. "

276 + )

277 +

278 + with open(status_file) as f4:

279 + status_lines = f4.read().strip().split("\n")

280 + # Filter for actual finish statuses (not INIT or empty lines)

281 + captured_statuses = [

282 + line

283 + for line in status_lines

284 + if line and line.startswith("FINISHED_")

285 + ]

286 +

287 + assert len(captured_statuses) >= 1, (

288 + f"Expected at least 1 captured finish status, got "

289 + f"{len(captured_statuses)}. File content: {status_lines}"

290 + )

291 +

292 + assert "FINISHED_ABORTED" in captured_statuses, (

293 + f"KV connector should see FINISHED_ABORTED but got "

294 + f"{captured_statuses}. "

295 + )

296 +

297 + # Verify cleanup

298 + assert not engine.output_processor.has_unfinished_requests()

299 +

300 + finally:

301 + # Shutdown the engine

302 + engine.shutdown()

303 +

304 + finally:

305 + # Clean up temporary files if they still exist

306 + if ready_file.exists():

307 + ready_file.unlink()

308 + if block_file.exists():

309 + block_file.unlink()

310 + if status_file.exists():

311 + status_file.unlink()

vllm/v1/engine/core.py

+34 -3

@@ -204,6 +204,8 @@ class EngineCore:

204 )

205 self.async_scheduling = vllm_config.scheduler_config.async_scheduling

206

207 + self.aborts_queue = queue.Queue[list[str]]()

208 +

209 # Mark the startup heap as static so that it's ignored by GC.

210 # Reduces pause times of oldest generation collections.

211 freeze_gc_heap()

@@ -347,6 +349,9 @@ class EngineCore:

349 if model_output is None:

350 model_output = self.model_executor.sample_tokens(grammar_output)

351

352 + # Before processing the model output, process any aborts that happened

353 + # during the model execution.

354 + self._process_aborts_queue()

355 engine_core_outputs = self.scheduler.update_from_output(

356 scheduler_output, model_output

357 )

@@ -440,6 +445,9 @@ class EngineCore:

445 with self.log_error_detail(scheduler_output):

446 model_output = future.result()

447

448 + # Before processing the model output, process any aborts that happened

449 + # during the model execution.

450 + self._process_aborts_queue()

451 engine_core_outputs = self.scheduler.update_from_output(

452 scheduler_output, model_output

453 )

@@ -458,6 +466,18 @@ class EngineCore:

466

467 return engine_core_outputs, model_executed

468

469 + def _process_aborts_queue(self):

470 + if not self.aborts_queue.empty():

471 + request_ids = []

472 + while not self.aborts_queue.empty():

473 + ids = self.aborts_queue.get_nowait()

474 + if isinstance(ids, str):

475 + # Should be a list here, but also handle string just in case.

476 + ids = (ids,)

477 + request_ids.extend(ids)

478 + # More efficient to abort all as a single batch.

479 + self.abort_requests(request_ids)

480 +

481 def shutdown(self):

482 self.structured_output_manager.clear_backend()

483 if self.model_executor:

@@ -871,9 +891,13 @@ class EngineCoreProc(EngineCore):

891 and not self.scheduler.has_requests()

892 and not self.batch_queue

893 ):

874 - if logger.isEnabledFor(DEBUG) and self.input_queue.empty():

875 - logger.debug("EngineCore waiting for work.")

876 - waited = True

894 + if self.input_queue.empty():

895 + # Drain aborts queue; all aborts are also processed via input_queue.

896 + with self.aborts_queue.mutex:

897 + self.aborts_queue.queue.clear()

898 + if logger.isEnabledFor(DEBUG):

899 + logger.debug("EngineCore waiting for work.")

900 + waited = True

901 req = self.input_queue.get()

902 self._handle_client_request(*req)

903

@@ -1027,6 +1051,13 @@ class EngineCoreProc(EngineCore):

1051 else:

1052 request = generic_decoder.decode(data_frames)

1053

1054 + if request_type == EngineCoreRequestType.ABORT:

1055 + # Aborts are added to *both* queues, allows us to eagerly

1056 + # process aborts while also ensuring ordering in the input

1057 + # queue to avoid leaking requests. This is ok because

1058 + # aborting in the scheduler is idempotent.

1059 + self.aborts_queue.put_nowait(request)

1060 +

1061 # Push to input queue for core busy loop.

1062 self.input_queue.put_nowait((request_type, request))

1063

vllm/v1/worker/gpu_worker.py

+7 -3

@@ -51,7 +51,6 @@ from vllm.v1.outputs import (

51 ModelRunnerOutput,

52 )

53 from vllm.v1.utils import report_usage_stats

54 -from vllm.v1.worker.gpu_model_runner import GPUModelRunner

54 from vllm.v1.worker.utils import is_residual_scattered_for_sp

55 from vllm.v1.worker.worker_base import WorkerBase

@@ -59,6 +58,7 @@ logger = init_logger(__name__)

59 if TYPE_CHECKING:

60 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig

61 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner

64 class Worker(WorkerBase):

@@ -259,7 +259,11 @@ class Worker(WorkerBase):

259 self.vllm_config, self.device

260 )

261 else:

262 - self.model_runner = GPUModelRunner(self.vllm_config, self.device)

262 + from vllm.v1.worker.gpu_model_runner import (

263 + GPUModelRunner as GPUModelRunnerV1,

264 + )

265 +

266 + self.model_runner = GPUModelRunnerV1(self.vllm_config, self.device)

267

268 if self.rank == 0:

269 # If usage stat is enabled, collect relevant info.

@@ -556,7 +560,7 @@ class Worker(WorkerBase):

560 and forward_pass

561 ):

562 # currently only supported by V1 GPUModelRunner

559 - assert isinstance(self.model_runner, GPUModelRunner)

563 + assert not self.use_v2_model_runner

564 num_scheduled_tokens_np = np.array(

565 list(scheduler_output.num_scheduled_tokens.values()),

566 dtype=np.int32,