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

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 5, 12:15 UTC

[Feat][Core] safely abort requests when FSM fails to advance (#38663)

Signed-off-by: walterbm <walter.beller.morales@gmail.com>

Walter Beller-Morales committed Apr 6, 2026 at 11:00 UTC e69a265135ef48312d78130f64b7bfce4cd81a37

3 files changed +164 -13

tests/v1/core/test_async_scheduler.py

+66 -1

@@ -1,10 +1,12 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 from collections import deque

4 +from unittest.mock import Mock

6 import pytest

7 -from vllm.v1.core.sched.output import SchedulerOutput

8 +from vllm.v1.core.sched.async_scheduler import AsyncScheduler

9 +from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput

10 from vllm.v1.outputs import ModelRunnerOutput

11 from vllm.v1.request import RequestStatus

12 from vllm.v1.utils import ConstantList

@@ -247,3 +249,66 @@ def test_prefix_caching_for_multi_turn():

249 # requests.

250 for req in next_turn_requests:

251 assert req.num_cached_tokens == req.num_prompt_tokens // BLOCK_SIZE * BLOCK_SIZE

252 +

253 +

254 +def test_abort_request_when_structured_output_fsm_cannot_advance():

255 + scheduler = object.__new__(AsyncScheduler)

256 + request = create_requests(num_requests=1, num_tokens=1)[0]

257 + request.structured_output_request = Mock()

258 + request.structured_output_request.grammar = Mock()

259 + request.structured_output_request.grammar.accept_tokens.return_value = False

260 + request.status = RequestStatus.RUNNING

261 + request.num_computed_tokens = request.num_tokens

262 + request.num_output_placeholders = 1

263 +

264 + scheduler.perf_metrics = None

265 + scheduler.connector = None

266 + scheduler.structured_output_manager = Mock()

267 + scheduler.structured_output_manager.should_advance.return_value = True

268 + scheduler.requests = {request.request_id: request}

269 + scheduler.running = [request]

270 + scheduler.waiting = Mock()

271 + scheduler.kv_cache_manager = Mock()

272 + scheduler.kv_cache_manager.take_events.return_value = None

273 + scheduler.kv_event_publisher = Mock()

274 + scheduler.finished_req_ids = set()

275 + scheduler.finished_req_ids_dict = None

276 + scheduler.vllm_config = Mock()

277 + scheduler.vllm_config.model_config.enable_return_routed_experts = False

278 + scheduler.recompute_kv_load_failures = False

279 + scheduler.make_stats = Mock(return_value=None)

280 + scheduler.max_model_len = 128

281 +

282 + def free_request(req, delay_free_blocks=False):

283 + scheduler.finished_req_ids.add(req.request_id)

284 + scheduler.requests.pop(req.request_id, None)

285 + return None

286 +

287 + scheduler._free_request = Mock(side_effect=free_request)

288 +

289 + output = SchedulerOutput(

290 + scheduled_new_reqs=[],

291 + scheduled_cached_reqs=CachedRequestData.make_empty(),

292 + num_scheduled_tokens={request.request_id: 1},

293 + total_num_scheduled_tokens=1,

294 + scheduled_encoder_inputs={},

295 + scheduled_spec_decode_tokens={},

296 + num_common_prefix_blocks=[],

297 + finished_req_ids=set(),

298 + free_encoder_mm_hashes=[],

299 + )

300 + model_runner_output = ModelRunnerOutput(

301 + req_ids=[request.request_id],

302 + req_id_to_index={request.request_id: 0},

303 + sampled_token_ids=[[123]],

304 + logprobs=None,

305 + prompt_logprobs_dict={},

306 + pooler_output=[],

307 + )

308 +

309 + scheduler.update_from_output(output, model_runner_output)

310 +

311 + assert request.resumable is False

312 + assert request.status == RequestStatus.FINISHED_ERROR

313 + assert request.request_id not in scheduler.requests

314 + assert not scheduler.running

tests/v1/core/test_scheduler.py

+81

@@ -26,6 +26,7 @@ from vllm.v1.core.encoder_cache_manager import EncoderCacheManager

26 from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash

27 from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput

28 from vllm.v1.core.sched.scheduler import Scheduler

29 +from vllm.v1.engine import FinishReason

30 from vllm.v1.kv_cache_interface import (

31 FullAttentionSpec,

32 KVCacheConfig,

@@ -2463,6 +2464,86 @@ def test_schedule_skip_tokenizer_init_structured_output_request():

2464 assert len(scheduler.skipped_waiting) == 1

2465

2466

2467 +def test_abort_request_when_structured_output_fsm_cannot_advance():

2468 + scheduler = object.__new__(Scheduler)

2469 + sampling_params = SamplingParams(ignore_eos=True, max_tokens=4)

2470 + sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)

2471 +

2472 + request = Request(

2473 + request_id="0",

2474 + prompt_token_ids=[0, 1],

2475 + mm_features=None,

2476 + sampling_params=sampling_params,

2477 + pooling_params=None,

2478 + )

2479 + request.structured_output_request = Mock()

2480 + request.structured_output_request.grammar = Mock()

2481 + request.structured_output_request.grammar.accept_tokens.return_value = False

2482 + request.status = RequestStatus.RUNNING

2483 + request.num_computed_tokens = request.num_tokens

2484 +

2485 + scheduler.perf_metrics = None

2486 + scheduler.connector = None

2487 + scheduler.structured_output_manager = Mock()

2488 + scheduler.structured_output_manager.should_advance.return_value = True

2489 + scheduler.requests = {request.request_id: request}

2490 + scheduler.running = [request]

2491 + scheduler.waiting = Mock()

2492 + scheduler.kv_cache_manager = Mock()

2493 + scheduler.kv_cache_manager.take_events.return_value = None

2494 + scheduler.kv_event_publisher = Mock()

2495 + scheduler.finished_req_ids = set()

2496 + scheduler.finished_req_ids_dict = None

2497 + scheduler.vllm_config = Mock()

2498 + scheduler.vllm_config.model_config.enable_return_routed_experts = False

2499 + scheduler.recompute_kv_load_failures = False

2500 + scheduler.make_stats = Mock(return_value=None)

2501 + scheduler.max_model_len = 128

2502 +

2503 + def free_request(req: Request, delay_free_blocks: bool = False):

2504 + scheduler.finished_req_ids.add(req.request_id)

2505 + scheduler.requests.pop(req.request_id, None)

2506 + return None

2507 +

2508 + scheduler._free_request = Mock(side_effect=free_request)

2509 +

2510 + output = SchedulerOutput(

2511 + scheduled_new_reqs=[],

2512 + scheduled_cached_reqs=CachedRequestData.make_empty(),

2513 + num_scheduled_tokens={request.request_id: 1},

2514 + total_num_scheduled_tokens=1,

2515 + scheduled_encoder_inputs={},

2516 + scheduled_spec_decode_tokens={},

2517 + num_common_prefix_blocks=[],

2518 + finished_req_ids=set(),

2519 + free_encoder_mm_hashes=[],

2520 + )

2521 +

2522 + model_runner_output = ModelRunnerOutput(

2523 + req_ids=[request.request_id],

2524 + req_id_to_index={request.request_id: 0},

2525 + sampled_token_ids=[[123]],

2526 + logprobs=None,

2527 + prompt_logprobs_dict={},

2528 + pooler_output=[],

2529 + )

2530 + engine_core_outputs = scheduler.update_from_output(output, model_runner_output)

2531 +

2532 + request.structured_output_request.grammar.accept_tokens.assert_called_once_with(

2533 + request.request_id, [123]

2534 + )

2535 + assert request.resumable is False

2536 + assert request.status == RequestStatus.FINISHED_ERROR

2537 + assert request.request_id not in scheduler.requests

2538 + assert not scheduler.running

2539 + scheduler._free_request.assert_called_once_with(request)

2540 + assert len(engine_core_outputs[0].outputs) == 1

2541 + engine_core_output = engine_core_outputs[0].outputs[0]

2542 + assert engine_core_output.request_id == request.request_id

2543 + assert engine_core_output.new_token_ids == [123]

2544 + assert engine_core_output.finish_reason == FinishReason.ERROR

2545 +

2546 +

2547 @pytest.mark.parametrize(

2548 "use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]

2549 )

vllm/v1/core/sched/scheduler.py

+17 -12

@@ -1406,6 +1406,23 @@ class Scheduler(SchedulerInterface):

1406 request.status = RequestStatus.FINISHED_STOPPED

1407 stopped = True

1408

1409 + if new_token_ids and self.structured_output_manager.should_advance(request):

1410 + struct_output_request = request.structured_output_request

1411 + assert struct_output_request is not None

1412 + assert struct_output_request.grammar is not None

1413 + if not struct_output_request.grammar.accept_tokens( # type: ignore[union-attr]

1414 + req_id, new_token_ids

1415 + ):

1416 + logger.error(

1417 + "Unexpected: grammar rejected tokens %s for request %s. "

1418 + "Terminating request.",

1419 + new_token_ids,

1420 + req_id,

1421 + )

1422 + request.status = RequestStatus.FINISHED_ERROR

1423 + request.resumable = False

1424 + stopped = True

1425 +

1426 routed_experts = None

1427 finish_reason = None

1428 if stopped:

@@ -1431,18 +1448,6 @@ class Scheduler(SchedulerInterface):

1448 ):

1449 new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))

1450

1434 - if new_token_ids and self.structured_output_manager.should_advance(request):

1435 - struct_output_request = request.structured_output_request

1436 - assert struct_output_request is not None

1437 - assert struct_output_request.grammar is not None

1438 - ok = struct_output_request.grammar.accept_tokens(req_id, new_token_ids)

1439 - if not ok:

1440 - logger.warning(

1441 - "Unexpected: grammar rejected tokens %s for request %s.",

1442 - new_token_ids,

1443 - req_id,

1444 - )

1445 -

1451 if num_nans_in_logits is not None and req_id in num_nans_in_logits:

1452 request.num_nans_in_logits = num_nans_in_logits[req_id]

1453