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

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

[Core] Add sleep level 0 mode with enqueue/wait pattern (#33195)

Signed-off-by: Jaewon Lee <jaewon@meta.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>

Jaewon committed Feb 12, 2026 at 16:16 UTC aa181c923bf83b6f8c4ce5613492a6b410c0c535

5 files changed +167 -26

vllm/entrypoints/llm.py

+112 -17

@@ -458,6 +458,93 @@ class LLM:

458

459 return self.engine_class.validate_outputs(outputs, RequestOutput)

460

461 + def enqueue(

462 + self,

463 + prompts: PromptType | Sequence[PromptType],

464 + sampling_params: SamplingParams | Sequence[SamplingParams] | None = None,

465 + lora_request: list[LoRARequest] | LoRARequest | None = None,

466 + priority: list[int] | None = None,

467 + use_tqdm: bool | Callable[..., tqdm] = True,

468 + tokenization_kwargs: dict[str, Any] | None = None,

469 + ) -> list[str]:

470 + """Enqueue prompts for generation without waiting for completion.

471 +

472 + This method adds requests to the engine queue but does not start

473 + processing them. Use wait_for_completion() to process the queued

474 + requests and get results.

475 +

476 + Args:

477 + prompts: The prompts to the LLM. See generate() for details.

478 + sampling_params: The sampling parameters for text generation.

479 + lora_request: LoRA request to use for generation, if any.

480 + priority: The priority of the requests, if any.

481 + use_tqdm: If True, shows a tqdm progress bar while adding requests.

482 + tokenization_kwargs: Overrides for `tokenizer.encode`.

483 +

484 + Returns:

485 + A list of request IDs for the enqueued requests.

486 + """

487 + model_config = self.model_config

488 + runner_type = model_config.runner_type

489 + if runner_type != "generate":

490 + raise ValueError("LLM.enqueue() is only supported for generative models.")

491 +

492 + if sampling_params is None:

493 + sampling_params = self.get_default_sampling_params()

494 +

495 + # Use the same preprocessing as _run_completion

496 + seq_prompts = prompt_to_seq(prompts)

497 + seq_params = self._params_to_seq(sampling_params, len(seq_prompts))

498 +

499 + if any(param.truncate_prompt_tokens is not None for param in seq_params):

500 + engine_prompts: Sequence[DictPrompt | TokPrompt] = [

501 + engine_prompt

502 + for prompt, param in zip(seq_prompts, seq_params)

503 + for engine_prompt in self._preprocess_completion(

504 + [prompt],

505 + tokenization_kwargs=merge_kwargs(

506 + tokenization_kwargs,

507 + dict(truncate_prompt_tokens=param.truncate_prompt_tokens),

508 + ),

509 + )

510 + ]

511 + else:

512 + engine_prompts = self._preprocess_completion(

513 + seq_prompts,

514 + tokenization_kwargs=tokenization_kwargs,

515 + )

516 +

517 + request_ids = self._validate_and_add_requests(

518 + prompts=engine_prompts,

519 + params=seq_params,

520 + use_tqdm=use_tqdm,

521 + lora_request=self._get_modality_specific_lora_reqs(

522 + engine_prompts, lora_request

523 + ),

524 + tokenization_kwargs=tokenization_kwargs,

525 + priority=priority,

526 + )

527 +

528 + return request_ids

529 +

530 + def wait_for_completion(

531 + self,

532 + use_tqdm: bool | Callable[..., tqdm] = True,

533 + ) -> list[RequestOutput]:

534 + """Wait for all enqueued requests to complete and return results.

535 +

536 + This method processes all requests currently in the engine queue

537 + and returns their outputs. Use after enqueue() to get results.

538 +

539 + Args:

540 + use_tqdm: If True, shows a tqdm progress bar.

541 +

542 + Returns:

543 + A list of RequestOutput objects for all completed requests.

544 + """

545 + outputs = self._run_engine(use_tqdm=use_tqdm)

546 + return self.engine_class.validate_outputs(outputs, RequestOutput)

547 +

548 def _get_modality_specific_lora_reqs(

549 self,

550 prompts: Sequence[DictPrompt | TokPrompt],

@@ -1618,19 +1705,22 @@ class LLM:

1705 during the sleep period, before `wake_up` is called.

1706

1707 Args:

1621 - level: The sleep level. Level 1 sleep will offload the model

1622 - weights and discard the kv cache. The content of kv cache

1623 - is forgotten. Level 1 sleep is good for sleeping and waking

1624 - up the engine to run the same model again. The model weights

1625 - are backed up in CPU memory. Please make sure there's enough

1626 - CPU memory to store the model weights. Level 2 sleep will

1627 - discard both the model weights and the kv cache. The content

1628 - of both the model weights and kv cache is forgotten. Level 2

1629 - sleep is good for sleeping and waking up the engine to run a

1630 - different model or update the model, where previous model

1631 - weights are not needed. It reduces CPU memory pressure.

1708 + level: The sleep level.

1709 + - Level 0: Pause scheduling but continue accepting requests.

1710 + Requests are queued but not processed.

1711 + - Level 1: Offload model weights to CPU, discard KV cache.

1712 + The content of kv cache is forgotten. Good for

1713 + sleeping and waking up the engine to run the same

1714 + model again. Please make sure there's enough CPU

1715 + memory to store the model weights.

1716 + - Level 2: Discard all GPU memory (weights + KV cache).

1717 + Good for sleeping and waking up the engine to run

1718 + a different model or update the model, where

1719 + previous model weights are not needed. It reduces

1720 + CPU memory pressure.

1721 """

1633 - self.reset_prefix_cache()

1722 + if level > 0:

1723 + self.reset_prefix_cache()

1724 self.llm_engine.sleep(level=level)

1725

1726 def wake_up(self, tags: list[str] | None = None):

@@ -1641,9 +1731,10 @@ class LLM:

1731 Args:

1732 tags: An optional list of tags to reallocate the engine memory

1733 for specific memory allocations. Values must be in

1644 - `("weights", "kv_cache")`. If None, all memory is reallocated.

1645 - wake_up should be called with all tags (or None) before the

1646 - engine is used again.

1734 + `("weights", "kv_cache", "scheduling")`. If None, all memory

1735 + is reallocated. wake_up should be called with all tags

1736 + (or None) before the engine is used again.

1737 + Use tags=["scheduling"] to resume from level 0 sleep.

1738 """

1739 self.llm_engine.wake_up(tags)

1740

@@ -1810,7 +1901,7 @@ class LLM:

1901 lora_request: Sequence[LoRARequest | None] | LoRARequest | None,

1902 tokenization_kwargs: dict[str, Any] | None = None,

1903 priority: list[int] | None = None,

1813 - ) -> None:

1904 + ) -> list[str]:

1905 num_requests = len(prompts)

1906 seq_params = self._params_to_seq(params, num_requests)

1907 seq_lora_requests = self._lora_request_to_seq(lora_request, num_requests)

@@ -1844,6 +1935,8 @@ class LLM:

1935 self.llm_engine.abort_request(added_request_ids, internal=True)

1936 raise e

1937

1938 + return added_request_ids

1939 +

1940 def _add_request(

1941 self,

1942 prompt: PromptType | DictPrompt | TokPrompt,

@@ -1895,7 +1988,9 @@ class LLM:

1988 return engine_request.request_id

1989

1990 def _run_engine(

1898 - self, *, use_tqdm: bool | Callable[..., tqdm] = True

1991 + self,

1992 + *,

1993 + use_tqdm: bool | Callable[..., tqdm] = True,

1994 ) -> list[RequestOutput | PoolingRequestOutput]:

1995 # Initialize tqdm.

1996 if use_tqdm:

vllm/v1/engine/async_llm.py

+2 -1

@@ -938,7 +938,8 @@ class AsyncLLM(EngineClient):

938 await self.engine_core.reset_encoder_cache_async()

939

940 async def sleep(self, level: int = 1) -> None:

941 - await self.reset_prefix_cache()

941 + if level > 0:

942 + await self.reset_prefix_cache()

943 await self.engine_core.sleep_async(level)

944

945 if self.logger_manager is not None:

vllm/v1/engine/core.py

+46 -6

@@ -614,13 +614,43 @@ class EngineCore:

614 self.model_executor.reset_encoder_cache()

615

616 def sleep(self, level: int = 1):

617 - self.model_executor.sleep(level)

617 + """Put the engine to sleep at the specified level.

618 +

619 + Args:

620 + level: Sleep level.

621 + - Level 0: Pause scheduling only. Requests are still accepted

622 + but not processed. No GPU memory changes.

623 + - Level 1: Offload model weights to CPU, discard KV cache.

624 + - Level 2: Discard all GPU memory.

625 + """

626 + if level == 0:

627 + # Level 0: Just pause scheduling, don't touch GPU

628 + self.pause_scheduler()

629 + else:

630 + # Level 1+: Delegate to executor for GPU memory management

631 + self.model_executor.sleep(level)

632

633 def wake_up(self, tags: list[str] | None = None):

620 - self.model_executor.wake_up(tags)

634 + """Wake up the engine from sleep.

635 +

636 + Args:

637 + tags: Tags to wake up. Use ["scheduling"] for level 0 wake up.

638 + """

639 + if tags is not None and "scheduling" in tags:

640 + # Level 0 wake up: Resume scheduling

641 + self.resume_scheduler()

642 + # Remove "scheduling" from tags if there are other tags to process

643 + remaining_tags = [t for t in tags if t != "scheduling"]

644 + if remaining_tags:

645 + self.model_executor.wake_up(remaining_tags)

646 + else:

647 + # Full wake up

648 + self.resume_scheduler()

649 + self.model_executor.wake_up(tags)

650

651 def is_sleeping(self) -> bool:

623 - return self.model_executor.is_sleeping

652 + """Check if engine is sleeping at any level."""

653 + return self._scheduler_paused or self.model_executor.is_sleeping

654

655 def execute_dummy_batch(self):

656 self.model_executor.execute_dummy_batch()

@@ -1023,7 +1053,13 @@ class EngineCoreProc(EngineCore):

1053 # 1) Poll the input queue until there is work to do.

1054 self._process_input_queue()

1055 # 2) Step the engine core and return the outputs.

1026 - self._process_engine_step()

1056 + # Skip if scheduling is paused (level 0 sleep)

1057 + if not self._scheduler_paused:

1058 + self._process_engine_step()

1059 + else:

1060 + # When scheduling is paused, still need to check for wake up

1061 + # by processing any utility requests that might resume scheduling

1062 + pass

1063

1064 def _process_input_queue(self):

1065 """Exits when an engine step needs to be performed."""

@@ -1031,7 +1067,7 @@ class EngineCoreProc(EngineCore):

1067 waited = False

1068 while (

1069 not self.engines_running

1034 - and not self.scheduler.has_requests()

1070 + and (not self.scheduler.has_requests() or self._scheduler_paused)

1071 and not self.batch_queue

1072 and not self._scheduler_paused

1073 ):

@@ -1414,11 +1450,15 @@ class DPEngineCoreProc(EngineCoreProc):

1450 # 1) Poll the input queue until there is work to do.

1451 self._process_input_queue()

1452

1453 + # Skip processing if scheduling is paused (level 0 sleep)

1454 + if self._scheduler_paused:

1455 + continue

1456 +

1457 # 2) Step the engine core.

1458 executed = self._process_engine_step()

1459 self._maybe_publish_request_counts()

1420 -

1460 local_unfinished_reqs = self.scheduler.has_unfinished_requests()

1461 +

1462 if not executed:

1463 if not local_unfinished_reqs and not self.engines_running:

1464 # All engines are idle.

vllm/v1/engine/core_client.py

+2 -1

@@ -194,7 +194,7 @@ class EngineCoreClient(ABC):

194 raise NotImplementedError

195

196 def dp_engines_running(self) -> bool:

197 - """Returns True id data parallel engines are collectively in a

197 + """Returns True if data parallel engines are collectively in a

198 running state."""

199 raise NotImplementedError

200

@@ -724,6 +724,7 @@ class SyncMPClient(MPClient):

724 # it is forwarded to the outputs_queue so we can raise it

725 # from this (run_output_handler) task to shut down the server.

726 outputs = self.outputs_queue.get()

727 +

728 if isinstance(outputs, Exception):

729 raise self._format_exception(outputs) from None

730 if outputs.wave_complete is not None:

vllm/v1/engine/llm_engine.py

+5 -1

@@ -312,7 +312,11 @@ class LLMEngine:

312

313 # 4) Record stats

314 with record_function_or_nullcontext("llm_engine step: record_stats"):

315 - if self.logger_manager is not None and outputs.scheduler_stats is not None:

315 + if (

316 + self.logger_manager is not None

317 + and outputs.scheduler_stats is not None

318 + and len(outputs.outputs) > 0

319 + ):

320 self.logger_manager.record(

321 scheduler_stats=outputs.scheduler_stats,

322 iteration_stats=iteration_stats,