[Core] Make parallel sampling (n>1) reqs admission atomic (#53936)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Sep 18, 2026 at 18:08 UTC
a7156060c1fc17c5b31d46d3760829bed8ac1a58
2 files changed
+107
-7
tests/v1/engine/test_admission_control.py
+90
@@ -15,6 +15,7 @@ These tests cover:
15
import argparse
16
import asyncio
17
import multiprocessing
18
+from collections.abc import Awaitable, Callable
19
from http import HTTPStatus
20
from types import SimpleNamespace
21
from unittest.mock import AsyncMock, MagicMock
@@ -35,6 +36,7 @@ from vllm.exceptions import (
36
from vllm.pooling_params import PoolingParams
37
from vllm.sampling_params import SamplingParams
38
from vllm.utils.argparse_utils import human_readable_int
39
+from vllm.v1.engine import EngineCoreRequest
40
from vllm.v1.engine.admission_control import SharedAdmissionStats
41
from vllm.v1.engine.async_llm import AsyncLLM
42
from vllm.v1.engine.output_processor import OutputProcessor
@@ -78,6 +80,45 @@ def _make_output_processor(**request_states) -> OutputProcessor:
80
return op
81
82
83
+def _make_engine_request(request_id: str, n: int) -> EngineCoreRequest:
84
+ return EngineCoreRequest(
85
+ request_id=request_id,
86
+ prompt_token_ids=[1],
87
+ mm_features=None,
88
+ sampling_params=SamplingParams(n=n),
89
+ pooling_params=None,
90
+ arrival_time=0,
91
+ lora_request=None,
92
+ cache_salt=None,
93
+ data_parallel_rank=None,
94
+ )
95
+
96
+
97
+def _make_request_test_llm(
98
+ max_num_queued_reqs: int,
99
+ add_request_async: Callable[[EngineCoreRequest], Awaitable[None]],
100
+) -> AsyncLLM:
101
+ llm = _make_async_llm(max_num_queued_reqs=max_num_queued_reqs)
102
+ llm.output_processor = OutputProcessor(None, log_stats=False)
103
+ llm.engine_core = SimpleNamespace(
104
+ resources=SimpleNamespace(engine_dead=False),
105
+ add_request_async=AsyncMock(side_effect=add_request_async),
106
+ abort_requests_async=AsyncMock(),
107
+ shutdown=MagicMock(),
108
+ )
109
+ llm.vllm_config = SimpleNamespace(
110
+ cache_config=SimpleNamespace(kv_sharing_fast_prefill=False)
111
+ )
112
+ llm.output_handler = None
113
+ llm.log_requests = False
114
+ llm._run_output_handler = MagicMock()
115
+ llm.input_processor = MagicMock()
116
+ llm.input_processor.assign_request_id.side_effect = lambda request: setattr(
117
+ request, "external_req_id", request.request_id
118
+ )
119
+ return llm
120
+
121
+
122
def _make_scheduler_config(**kwargs) -> SchedulerConfig:
123
return SchedulerConfig(
124
runner_type="generate",
@@ -340,6 +381,55 @@ async def test_concurrent_single_request_admission_respects_limit():
381
assert len(request_states) == 1
382
383
384
+@pytest.mark.asyncio
385
+@pytest.mark.parametrize("first_n", [1, 3])
386
+async def test_parallel_admission_is_all_or_nothing(first_n: int):
387
+ """An n>1 request reserves either all its capacity or none of it."""
388
+
389
+ async def add_request_async(_request):
390
+ await asyncio.sleep(0)
391
+
392
+ llm = _make_request_test_llm(3, add_request_async)
393
+ second_n = 3 if first_n == 1 else 1
394
+ first = _make_engine_request("first", first_n)
395
+ second = _make_engine_request("second", second_n)
396
+ results = await asyncio.gather(
397
+ llm.add_request("first", first, first.params),
398
+ llm.add_request("second", second, second.params),
399
+ return_exceptions=True,
400
+ )
401
+
402
+ assert not isinstance(results[0], Exception)
403
+ assert isinstance(results[1], QueueOverflowError)
404
+ assert llm.output_processor.get_num_unfinished_requests() == first_n
405
+ assert llm.engine_core.add_request_async.await_count == first_n
406
+
407
+
408
+@pytest.mark.asyncio
409
+async def test_parallel_admission_cancellation_cleans_up_all_children():
410
+ """Cancellation during core submission must release every reserved slot."""
411
+ first_submission_started = asyncio.Event()
412
+
413
+ async def add_request_async(_request):
414
+ first_submission_started.set()
415
+ await asyncio.Event().wait()
416
+
417
+ llm = _make_request_test_llm(3, add_request_async)
418
+ request = _make_engine_request("parallel", 3)
419
+ output = llm.generate(request, request.params, request.request_id)
420
+ generate_task = asyncio.create_task(anext(output))
421
+
422
+ await first_submission_started.wait()
423
+ generate_task.cancel()
424
+ with pytest.raises(asyncio.CancelledError):
425
+ await generate_task
426
+
427
+ assert llm.engine_core.add_request_async.await_count == 1
428
+ assert llm.output_processor.get_num_unfinished_requests() == 0
429
+ aborted_ids = llm.engine_core.abort_requests_async.await_args.args[0]
430
+ assert set(aborted_ids) == {"0_parallel", "1_parallel", "2_parallel"}
431
+
432
+
433
# ---------------------------------------------------------------------------
434
# create_error_response integration
435
# ---------------------------------------------------------------------------
vllm/v1/engine/async_llm.py
+17
-7
@@ -401,10 +401,6 @@ class AsyncLLM(EngineClient):
401
"prompt logprobs"
402
)
403
404
- if isinstance(params, SamplingParams) and params.n > 1:
405
- # TODO (NickLucche) Batch check admission check for all n requests
406
- self.check_admission(params.n, request_id)
407
-
404
if isinstance(prompt, AsyncGenerator):
405
if reasoning_ended is not None or reasoning_parser_kwargs is not None:
406
raise NotImplementedError
@@ -500,14 +496,28 @@ class AsyncLLM(EngineClient):
496
497
# Fan out child requests (for n>1).
498
parent_request = ParentRequest(request)
499
+ child_requests: list[EngineCoreRequest] = []
500
for idx in range(parent_params.n):
501
request_id, child_params = parent_request.get_child_info(idx)
502
child_request = request if idx == parent_params.n - 1 else copy(request)
503
child_request.request_id = request_id
504
child_request.sampling_params = child_params
508
- await self._add_request(
509
- child_request, prompt_text, parent_request, idx, queue
510
- )
505
+ child_requests.append(child_request)
506
+
507
+ self.check_admission(parent_params.n, parent_request.request_id)
508
+ try:
509
+ for idx, child_request in enumerate(child_requests):
510
+ self.output_processor.add_request(
511
+ child_request, prompt_text, parent_request, idx, queue
512
+ )
513
+
514
+ for child_request in child_requests:
515
+ await self.engine_core.add_request_async(child_request)
516
+ if self.log_requests:
517
+ logger.info("Added request %s.", child_request.request_id)
518
+ except BaseException:
519
+ await self.abort(parent_request.request_id, internal=True)
520
+ raise
521
return queue
522
523
async def _add_request(