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

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 15, 13:30 UTC

[V1][PP] Optimization: continue scheduling prefill chunks (#17080)

Signed-off-by: Rui Qiao <ruisearch42@gmail.com>

Rui Qiao committed Apr 24, 2025 at 05:27 UTC c0dfd97519c8ff18b6f9c7d86010294ca3510edf

5 files changed +128 -74

tests/v1/core/test_scheduler.py

-4

@@ -437,7 +437,6 @@ def test_stop_via_update_from_output():

437 req.num_computed_tokens = req.num_tokens

438 scheduler.requests[req.request_id] = req

439 scheduler.running.append(req)

440 - scheduler.scheduled_req_ids.add(req.request_id)

440

441 scheduler_output = SchedulerOutput(scheduled_new_reqs=[],

442 scheduled_cached_reqs=[],

@@ -489,7 +488,6 @@ def test_stop_via_update_from_output():

488 req.num_computed_tokens = req.num_tokens

489 scheduler.requests[req.request_id] = req

490 scheduler.running.append(req)

492 - scheduler.scheduled_req_ids.add(req.request_id)

491

492 scheduler_output = SchedulerOutput(scheduled_new_reqs=[],

493 scheduled_cached_reqs=[],

@@ -539,7 +537,6 @@ def test_stop_via_update_from_output():

537 req.num_computed_tokens = req.num_tokens

538 scheduler.requests[req.request_id] = req

539 scheduler.running.append(req)

542 - scheduler.scheduled_req_ids.add(req.request_id)

540

541 scheduler_output = SchedulerOutput(scheduled_new_reqs=[],

542 scheduled_cached_reqs=[],

@@ -589,7 +586,6 @@ def test_stop_via_update_from_output():

586 requests[0].num_computed_tokens = requests[0].num_tokens

587 scheduler.requests[requests[0].request_id] = requests[0]

588 scheduler.running.append(requests[0])

592 - scheduler.scheduled_req_ids.add(requests[0].request_id)

589

590 scheduler_output = SchedulerOutput(

591 scheduled_new_reqs=[],

tests/v1/engine/test_engine_core.py

+84 -22

@@ -1,10 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import copy

4 -import threading

4 import time

5 import uuid

7 -from concurrent.futures import Future

6 +from concurrent.futures import Future, ThreadPoolExecutor

8 import pytest

9 from transformers import AutoTokenizer

@@ -244,33 +243,33 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

243 self, kv_cache_configs: list[KVCacheConfig]) -> None:

244 super().initialize_from_config(kv_cache_configs)

245

247 - # This executor actually can only run 1 batch at a time

248 - self.semaphore = threading.Semaphore(1)

246 + # Create a thread pool with a single worker

247 + self.thread_pool = ThreadPoolExecutor(max_workers=1)

248

249 def execute_model(

250 self,

251 scheduler_output,

252 ) -> Future[ModelRunnerOutput]:

253 """Make execute_model non-blocking."""

255 - future: Future[ModelRunnerOutput] = Future()

254

257 - def _thread_wrapper(scheduler_output, future):

258 - with self.semaphore:

259 - output = self.collective_rpc("execute_model",

260 - args=(scheduler_output, ))

261 - # Make a copy because output[0] may be reused

262 - # by the next batch.

263 - output = copy.deepcopy(output[0])

264 - future.set_result(output)

255 + def _execute():

256 + output = self.collective_rpc("execute_model",

257 + args=(scheduler_output, ))

258 + # Make a copy because output[0] may be reused

259 + # by the next batch.

260 + return copy.deepcopy(output[0])

261

266 - threading.Thread(target=_thread_wrapper,

267 - args=(scheduler_output, future)).start()

268 - return future

262 + # Use the thread pool instead of creating a new thread

263 + return self.thread_pool.submit(_execute)

264

265 @property

266 def max_concurrent_batches(self) -> int:

267 return 2

268

269 + def shutdown(self):

270 + if hasattr(self, 'thread_pool'):

271 + self.thread_pool.shutdown(wait=False)

272 +

273 with monkeypatch.context() as m:

274 m.setenv("VLLM_USE_V1", "1")

275

@@ -299,14 +298,77 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

298 # Schedule Batch 1: (10, req0)

299 assert engine_core.step_with_batch_queue() is None

300 assert engine_core.batch_queue.qsize() == 1

301 + scheduler_output = engine_core.batch_queue.queue[-1][1]

302 + assert scheduler_output.num_scheduled_tokens[0] == 10

303 + # num_computed_tokens should have been updated immediately.

304 + assert engine_core.scheduler.requests[

305 + req0.request_id].num_computed_tokens == 10

306 +

307 + # Schedule Batch 2: (2, req0), (8, req1)

308 assert engine_core.step_with_batch_queue() is None

309 assert engine_core.batch_queue.qsize() == 2

310 + scheduler_output = engine_core.batch_queue.queue[-1][1]

311 + assert scheduler_output.num_scheduled_tokens[0] == 2

312 + assert scheduler_output.num_scheduled_tokens[1] == 8

313 + # num_computed_tokens should have been updated immediately.

314 + assert engine_core.scheduler.requests[0].num_computed_tokens == 12

315 + assert engine_core.scheduler.requests[1].num_computed_tokens == 8

316 +

317 assert engine_core.scheduler.get_num_unfinished_requests() == 2

318

306 - # Loop through both requests.

307 - while engine_core.scheduler.get_num_unfinished_requests() == 2:

308 - engine_core.step_with_batch_queue()

319 + # Batch queue is full. Finish Batch 1.

320 + engine_core.step_with_batch_queue()

321 +

322 + # Schedule Batch 3: (4, req1). Note that req0 cannot be scheduled

323 + # because it is in the decoding stage now.

324 + engine_core.step_with_batch_queue()

325 + assert engine_core.batch_queue.qsize() == 2

326 + scheduler_output = engine_core.batch_queue.queue[-1][1]

327 + assert scheduler_output.num_scheduled_tokens[1] == 4

328

310 - # Reaching here when got the result of the first request.

311 - while engine_core.scheduler.get_num_unfinished_requests() == 1:

312 - engine_core.step_with_batch_queue()

329 + # Batch queue is full. Finish Batch 2. Get first token of req0.

330 + output = engine_core.step_with_batch_queue()

331 + assert output is not None

332 + assert len(output.outputs) == 1

333 + assert engine_core.scheduler.requests[req0.request_id].num_tokens == 13

334 +

335 + # Schedule Batch 4: (1, req0).

336 + engine_core.step_with_batch_queue()

337 + assert engine_core.batch_queue.qsize() == 2

338 + scheduler_output = engine_core.batch_queue.queue[-1][1]

339 + assert scheduler_output.num_scheduled_tokens[0] == 1

340 +

341 + # Batch queue is full. Finish Batch 3. Get first token of req1.

342 + output = engine_core.step_with_batch_queue()

343 + assert output is not None

344 + assert len(output.outputs) == 1

345 + assert engine_core.scheduler.requests[req1.request_id].num_tokens == 13

346 +

347 + # Schedule Batch 5: (1, req1).

348 + engine_core.step_with_batch_queue()

349 + assert engine_core.batch_queue.qsize() == 2

350 + scheduler_output = engine_core.batch_queue.queue[-1][1]

351 + assert scheduler_output.num_scheduled_tokens[1] == 1

352 +

353 + # Loop until req0 is finished.

354 + step = 0

355 + req_id = 0

356 + expected_num_tokens = [

357 + engine_core.scheduler.requests[0].num_tokens + 1,

358 + engine_core.scheduler.requests[1].num_tokens + 1,

359 + ]

360 + while engine_core.scheduler.get_num_unfinished_requests() == 2:

361 + output = engine_core.step_with_batch_queue()

362 + if step % 2 == 0:

363 + # Even steps consumes an output.

364 + assert output is not None

365 + assert len(output.outputs) == 1

366 + if req_id in engine_core.scheduler.requests:

367 + assert engine_core.scheduler.requests[

368 + req_id].num_tokens == expected_num_tokens[req_id]

369 + expected_num_tokens[req_id] += 1

370 + req_id = (req_id + 1) % 2

371 + else:

372 + # Odd steps schedules a new batch.

373 + assert output is None

374 + step += 1

vllm/v1/core/sched/interface.py

-5

@@ -117,11 +117,6 @@ class SchedulerInterface(ABC):

117 not yet returned in SchedulerOutputs."""

118 return self.has_unfinished_requests() or self.has_finished_requests()

119

120 - @abstractmethod

121 - def get_num_unscheduled_requests(self) -> int:

122 - """Number of requests that are not being processed by the executor."""

123 - raise NotImplementedError

124 -

120 @abstractmethod

121 def reset_prefix_cache(self) -> bool:

122 """Reset the prefix cache for KV cache.

vllm/v1/core/sched/scheduler.py

+32 -35

@@ -3,7 +3,7 @@

3 from __future__ import annotations

5 import time

6 -from collections import deque

6 +from collections import defaultdict, deque

7 from collections.abc import Iterable

8 from typing import Optional, Union

@@ -88,9 +88,6 @@ class Scheduler(SchedulerInterface):

88 # Priority queues for requests.

89 self.waiting: deque[Request] = deque()

90 self.running: list[Request] = []

91 - # The requests that have been scheduled and are being executed

92 - # by the executor.

93 - self.scheduled_req_ids: set[str] = set()

92 # The request IDs that are finished in between the previous and the

93 # current steps. This is used to notify the workers about the finished

@@ -100,8 +97,9 @@ class Scheduler(SchedulerInterface):

98 # OPTIMIZATION: Cache the CachedRequestData objects to avoid creating

99 # them at each scheduling step.

103 - # Request id -> CachedRequestData

104 - self._cached_reqs_data: dict[str, CachedRequestData] = {}

100 + # Request id -> deque of CachedRequestData

101 + self._cached_reqs_data: dict[

102 + str, deque[CachedRequestData]] = defaultdict(deque)

103

104 # Encoder-related.

105 # Calculate encoder cache size if applicable

@@ -171,10 +169,6 @@ class Scheduler(SchedulerInterface):

169 req_index = 0

170 while req_index < len(self.running) and token_budget > 0:

171 request = self.running[req_index]

174 - if request.request_id in self.scheduled_req_ids:

175 - # This request has already been scheduled.

176 - req_index += 1

177 - continue

172

173 num_new_tokens = (request.num_tokens_with_spec -

174 request.num_computed_tokens)

@@ -183,33 +177,35 @@ class Scheduler(SchedulerInterface):

177 num_new_tokens = (

178 self.scheduler_config.long_prefill_token_threshold)

179 num_new_tokens = min(num_new_tokens, token_budget)

186 - assert num_new_tokens > 0

180

181 # Make sure the input position does not exceed the max model len.

182 # This is necessary when using spec decoding.

183 num_new_tokens = min(

184 num_new_tokens,

185 self.max_model_len - request.num_computed_tokens)

193 - assert num_new_tokens > 0

186

187 # Schedule encoder inputs.

188 + encoder_inputs_to_schedule = None

189 + new_encoder_budget = encoder_budget

190 if request.has_encoder_inputs:

191 (encoder_inputs_to_schedule, num_new_tokens,

192 new_encoder_budget) = self._try_schedule_encoder_inputs(

193 request, request.num_computed_tokens, num_new_tokens,

194 encoder_budget)

201 - if num_new_tokens == 0:

202 - # The request cannot be scheduled because the encoder budget

203 - # or the encoder cache is exhausted.

204 - # NOTE(woosuk): By using `continue` instead of `break` here,

205 - # we intentionally relax the strict FCFS scheduling policy

206 - # to allow lower-priority requests to be scheduled when a

207 - # higher-priority request is blocked by encoder constraints.

208 - req_index += 1

209 - continue

210 - else:

211 - encoder_inputs_to_schedule = None

212 - new_encoder_budget = encoder_budget

195 +

196 + if num_new_tokens == 0:

197 + # The request cannot be scheduled because one of the following

198 + # reasons:

199 + # 1. No new tokens to schedule. This may happen when PP>1 and

200 + # we have already scheduled all prompt tokens but they are

201 + # not finished yet.

202 + # 2. The encoder budget is exhausted.

203 + # 3. The encoder cache is exhausted.

204 + # NOTE(woosuk): Here, by doing `continue` instead of `break`,

205 + # we do not strictly follow the FCFS scheduling policy and

206 + # allow the lower-priority requests to be scheduled.

207 + req_index += 1

208 + continue

209

210 while True:

211 new_blocks = self.kv_cache_manager.allocate_slots(

@@ -243,7 +239,6 @@ class Scheduler(SchedulerInterface):

239

240 # Schedule the request.

241 scheduled_running_reqs.append(request)

246 - self.scheduled_req_ids.add(request.request_id)

242 if request.use_structured_output:

243 # PERF: in case of chunked prefill,

244 # request might not include any new tokens.

@@ -382,7 +377,6 @@ class Scheduler(SchedulerInterface):

377 request.request_id] = req_index

378 req_index += 1

379 self.running.append(request)

385 - self.scheduled_req_ids.add(request.request_id)

380 if self.log_stats:

381 request.record_event(EngineCoreEventType.SCHEDULED,

382 scheduled_timestamp)

@@ -521,18 +515,21 @@ class Scheduler(SchedulerInterface):

515 num_regular_tokens = num_scheduled_tokens - num_scheduled_spec_tokens

516 new_token_ids = request.all_token_ids[

517 num_computed_tokens:num_computed_tokens + num_regular_tokens]

524 - req_data = self._cached_reqs_data.get(request.request_id)

525 - if req_data is not None:

518 +

519 + req_data_queue = self._cached_reqs_data.get(request.request_id)

520 + if req_data_queue:

521 + req_data = req_data_queue.popleft()

522 req_data.resumed_from_preemption = resumed_from_preemption

523 req_data.new_token_ids = new_token_ids

524 req_data.new_block_ids = new_block_ids

525 req_data.num_computed_tokens = num_computed_tokens

526 else:

527 + # No cached request data, or all cached request data has been

528 + # used by the scheduled requests.

529 req_data = CachedRequestData.from_request(request,

530 resumed_from_preemption,

531 new_token_ids,

532 new_block_ids)

535 - self._cached_reqs_data[request.request_id] = req_data

533 return req_data

534

535 def _try_schedule_encoder_inputs(

@@ -561,6 +558,8 @@ class Scheduler(SchedulerInterface):

558 Note that num_computed_tokens includes both locally cached

559 blocks and externally cached blocks (via KVConnector).

560 """

561 + if num_new_tokens == 0 or not request.has_encoder_inputs:

562 + return [], num_new_tokens, encoder_budget

563 encoder_inputs_to_schedule: list[int] = []

564 mm_positions = request.mm_positions

565 assert mm_positions is not None

@@ -728,10 +727,13 @@ class Scheduler(SchedulerInterface):

727 # Invariant: EngineCore returns no partial prefill outputs.

728 assert not prompt_logprobs_tensors

729

731 - self.scheduled_req_ids.remove(req_id)

730 if not stopped:

731 new_running.append(request)

732

733 + # Return the cached request data to the queue so they can be reused.

734 + for req_data in scheduler_output.scheduled_cached_reqs:

735 + self._cached_reqs_data[req_data.req_id].append(req_data)

736 +

737 self.running = new_running

738 engine_core_outputs = EngineCoreOutputs(

739 outputs=outputs,

@@ -774,7 +776,6 @@ class Scheduler(SchedulerInterface):

776

777 if request.status == RequestStatus.RUNNING:

778 self.running.remove(request)

777 - self.scheduled_req_ids.discard(request.request_id)

779 else:

780 self.waiting.remove(request)

781 request.status = finished_status

@@ -795,10 +796,6 @@ class Scheduler(SchedulerInterface):

796 def has_finished_requests(self) -> bool:

797 return len(self.finished_req_ids) > 0

798

798 - def get_num_unscheduled_requests(self) -> int:

799 - """Number of requests that are not being processed by the executor."""

800 - return self.get_num_unfinished_requests() - len(self.scheduled_req_ids)

801 -

799 def reset_prefix_cache(self) -> bool:

800 return self.kv_cache_manager.reset_prefix_cache()

801

vllm/v1/engine/core.py

+12 -8

@@ -210,10 +210,10 @@ class EngineCore:

210 Note that if nothing to output in this step, None is returned.

211

212 The execution flow is as follows:

213 - 1. Try to schedule a new batch if there are unscheduled requests

214 - and the job queue is not full. If a new batch is scheduled, directly

215 - return an empty engine core output. In other words, we won't check

216 - and return model outputs before the batch queue is full.

213 + 1. Try to schedule a new batch if the batch queue is not full.

214 + If a new batch is scheduled, directly return an empty engine core

215 + output. In other words, fulfilling the batch queue has a higher priority

216 + than getting model outputs.

217 2. If there is no new scheduled batch, meaning that the batch queue

218 is full or no other requests can be scheduled, we block until the first

219 batch in the job queue is finished.

@@ -223,10 +223,10 @@ class EngineCore:

223

224 engine_core_outputs = None

225 scheduler_output = None

226 - # If there are unscheduled requests and the job queue

227 - # is not full, schedule a new batch. Note that this is not blocking.

228 - if (self.scheduler.get_num_unscheduled_requests() > 0

229 - and not self.batch_queue.full()):

226 + # Try to schedule a new batch if the batch queue is not full, but

227 + # the scheduler may return an empty batch if all requests are scheduled.

228 + # Note that this is not blocking.

229 + if not self.batch_queue.full():

230 scheduler_output = self.scheduler.schedule()

231 if scheduler_output.total_num_scheduled_tokens > 0:

232 future = self.model_executor.execute_model(scheduler_output)

@@ -238,6 +238,10 @@ class EngineCore:

238

239 # If no more requests can be scheduled and the job queue is not empty,

240 # block until the first batch in the job queue is finished.

241 + # TODO(comaniac): Ideally we should peek the first batch in the

242 + # job queue to check if it's finished before scheduling a new batch,

243 + # but peeking the first element in a queue is not thread-safe,

244 + # so we need more work.

245 if not scheduled_batch and not self.batch_queue.empty():

246 future, scheduler_output = self.batch_queue.get_nowait()

247 # Blocking until the first result is available.