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

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 14, 05:30 UTC

[Perf][V1] Fully overlap model execution (#23569)

Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai>

Benjamin Chislett committed Sep 5, 2025 at 21:20 UTC cee182b2970b56d3f9d5aa0017ce165cbdb788f6

5 files changed +252 -31

vllm/v1/executor/multiproc_executor.py

+45 -5

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

3 import multiprocessing

4 import os

5 import pickle

6 +import queue

7 import signal

8 import threading

9 import time

@@ -33,7 +34,8 @@ from vllm.utils import (decorate_logs, get_distributed_init_method,

34 get_loopback_ip, get_mp_context, get_open_port,

35 set_process_title)

36 from vllm.v1.executor.abstract import Executor, FailureCallback

36 -from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput

37 +from vllm.v1.outputs import (AsyncModelRunnerOutput, DraftTokenIds,

38 + ModelRunnerOutput)

39 from vllm.worker.worker_base import WorkerWrapperBase

41 logger = init_logger(__name__)

@@ -414,6 +416,16 @@ class WorkerProc:

416 # Initializes a message queue for sending the model output

417 self.worker_response_mq = MessageQueue(1, 1)

418

419 + scheduler_config = vllm_config.scheduler_config

420 + self.use_async_scheduling = scheduler_config.async_scheduling

421 + if self.use_async_scheduling:

422 + self.async_output_queue: queue.Queue = queue.Queue()

423 + self.async_output_copy_thread = Thread(

424 + target=self.async_output_busy_loop,

425 + daemon=True,

426 + name="WorkerAsyncOutputCopy")

427 + self.async_output_copy_thread.start()

428 +

429 # Initialize device and loads weights

430 self.worker.init_device()

431 self.worker.load_model()

@@ -595,6 +607,36 @@ class WorkerProc:

607 SUCCESS = auto()

608 FAILURE = auto()

609

610 + def enqueue_output(self, output: Any):

611 + """Prepares output from the worker and enqueues it to the

612 + worker_response_mq. If the output is an Exception, it is

613 + converted to a FAILURE response.

614 + """

615 + if isinstance(output, AsyncModelRunnerOutput):

616 + output = output.get_output()

617 +

618 + if isinstance(output, Exception):

619 + result = (WorkerProc.ResponseStatus.FAILURE, str(output))

620 + else:

621 + result = (WorkerProc.ResponseStatus.SUCCESS, output)

622 + self.worker_response_mq.enqueue(result)

623 +

624 + def handle_output(self, output: Any):

625 + """Handles output from the worker. If async scheduling is enabled,

626 + it is passed to the async_output_busy_loop thread. Otherwise, it is

627 + enqueued directly to the worker_response_mq.

628 + """

629 + if self.use_async_scheduling:

630 + self.async_output_queue.put(output)

631 + else:

632 + self.enqueue_output(output)

633 +

634 + def async_output_busy_loop(self):

635 + """Entrypoint for the thread which handles outputs asynchronously."""

636 + while True:

637 + output = self.async_output_queue.get()

638 + self.enqueue_output(output)

639 +

640 def worker_busy_loop(self):

641 """Main busy loop for Multiprocessing Workers"""

642 while True:

@@ -614,10 +656,8 @@ class WorkerProc:

656 # exception might not be serializable, so we convert it to

657 # string, only for logging purpose.

658 if output_rank is None or self.rank == output_rank:

617 - self.worker_response_mq.enqueue(

618 - (WorkerProc.ResponseStatus.FAILURE, str(e)))

659 + self.handle_output(e)

660 continue

661

662 if output_rank is None or self.rank == output_rank:

622 - self.worker_response_mq.enqueue(

623 - (WorkerProc.ResponseStatus.SUCCESS, output))

663 + self.handle_output(output)

vllm/v1/outputs.py

+15

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

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

4 +from abc import ABC, abstractmethod

5 from dataclasses import dataclass

6 from typing import NamedTuple, Optional

@@ -114,6 +115,20 @@ class ModelRunnerOutput:

115 num_nans_in_logits: Optional[dict[str, int]] = None

116

117

118 +# ModelRunnerOutput wrapper for async scheduling.

119 +class AsyncModelRunnerOutput(ABC):

120 +

121 + @abstractmethod

122 + def get_output(self) -> ModelRunnerOutput:

123 + """Get the ModelRunnerOutput for this async output.

124 +

125 + This is a blocking call that waits until the results are ready, which

126 + might involve copying device tensors to the host.

127 + This method should only be called once per AsyncModelRunnerOutput.

128 + """

129 + pass

130 +

131 +

132 @dataclass

133 class DraftTokenIds:

134

vllm/v1/worker/gpu_input_batch.py

@@ -250,6 +250,11 @@ class InputBatch:

250

251 self.pooling_params: dict[str, PoolingParams] = {}

252

253 + # Cached reference to the GPU tensor of previously sampled tokens

254 + self.prev_sampled_token_ids: Optional[torch.Tensor] = None

255 + self.prev_sampled_token_ids_invalid_indices: Optional[set[int]] = None

256 + self.prev_req_id_to_index: Optional[dict[str, int]] = None

257 +

258 @property

259 def req_ids(self) -> list[str]:

260 # None elements should only be present transiently

vllm/v1/worker/gpu_model_runner.py

+182 -21

@@ -67,8 +67,8 @@ from vllm.v1.kv_cache_interface import (AttentionSpec,

67 FullAttentionSpec, KVCacheConfig,

68 KVCacheGroupSpec, KVCacheSpec,

69 MambaSpec, SlidingWindowSpec)

70 -from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, DraftTokenIds,

71 - LogprobsTensors, ModelRunnerOutput)

70 +from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput,

71 + DraftTokenIds, LogprobsTensors, ModelRunnerOutput)

72 from vllm.v1.pool.metadata import PoolingMetadata

73 from vllm.v1.sample.logits_processor import LogitsProcessors, build_logitsprocs

74 from vllm.v1.sample.metadata import SamplingMetadata

@@ -100,6 +100,53 @@ else:

100 logger = init_logger(__name__)

101

102

103 +# Wrapper for ModelRunnerOutput to support overlapped execution.

104 +class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

105 +

106 + def __init__(

107 + self,

108 + model_runner_output: ModelRunnerOutput,

109 + sampled_token_ids: torch.Tensor,

110 + invalid_req_indices: list[int],

111 + async_output_copy_stream: torch.cuda.Stream,

112 + ):

113 + self._model_runner_output = model_runner_output

114 + self._invalid_req_indices = invalid_req_indices

115 +

116 + # Event on the copy stream so we can synchronize the non-blocking copy.

117 + self._async_copy_ready_event = torch.cuda.Event()

118 +

119 + # Keep a reference to the device tensor to avoid it being

120 + # deallocated until we finish copying it to the host.

121 + self._sampled_token_ids = sampled_token_ids

122 +

123 + # Initiate the copy on a separate stream, but do not synchronize it.

124 + default_stream = torch.cuda.current_stream()

125 + with torch.cuda.stream(async_output_copy_stream):

126 + async_output_copy_stream.wait_stream(default_stream)

127 + self._sampled_token_ids_cpu = self._sampled_token_ids.to(

128 + 'cpu', non_blocking=True)

129 + self._async_copy_ready_event.record()

130 +

131 + def get_output(self) -> ModelRunnerOutput:

132 + """Copy the device tensors to the host and return a ModelRunnerOutput.

133 +

134 + This function blocks until the copy is finished.

135 + """

136 + self._async_copy_ready_event.synchronize()

137 +

138 + # Release the device tensor once the copy has completed

139 + del self._sampled_token_ids

140 +

141 + valid_sampled_token_ids = self._sampled_token_ids_cpu.tolist()

142 + for i in self._invalid_req_indices:

143 + valid_sampled_token_ids[i].clear()

144 +

145 + output = self._model_runner_output

146 + output.sampled_token_ids = valid_sampled_token_ids

147 + return output

148 +

149 +

150 class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

151

152 def __init__(

@@ -230,6 +277,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

277 is_pooling_model=self.is_pooling_model,

278 )

279

280 + self.use_async_scheduling = self.scheduler_config.async_scheduling

281 + self.async_output_copy_stream = torch.cuda.Stream() if \

282 + self.use_async_scheduling else None

283 +

284 # TODO(woosuk): Provide an option to tune the max cudagraph batch size.

285 # The convention is different.

286 # self.cudagraph_batch_sizes sorts in ascending order.

@@ -654,6 +705,73 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

705

706 return cu_num_tokens, arange

707

708 + def _prepare_input_ids(self, total_num_scheduled_tokens: int,

709 + cu_num_tokens: np.ndarray) -> None:

710 + """Prepare the input IDs for the current batch.

711 +

712 + Carefully handles the `prev_sampled_token_ids` which can be cached

713 + from the previous engine iteration, in which case those tokens on the

714 + GPU need to be copied into the corresponding slots into input_ids."""

715 +

716 + if self.input_batch.prev_sampled_token_ids is None:

717 + # Normal scheduling case

718 + self.input_ids.copy_to_gpu(total_num_scheduled_tokens)

719 + return

720 +

721 + # Async scheduling case, where some decode requests from the previous

722 + # iteration won't have entries in input_ids_cpu and need to be copied

723 + # on the GPU from prev_sampled_token_ids.

724 + prev_req_id_to_index = self.input_batch.prev_req_id_to_index

725 + assert prev_req_id_to_index is not None

726 + flattened_indices = []

727 + prev_common_req_indices = []

728 + indices_match = True

729 + max_flattened_index = -1

730 + for req_id, cur_index in self.input_batch.req_id_to_index.items():

731 + if (prev_index := prev_req_id_to_index.get(req_id)) is not None:

732 + prev_common_req_indices.append(prev_index)

733 + # We need to compute the flattened input_ids index of the

734 + # last token in each common request.

735 + flattened_index = cu_num_tokens[cur_index].item() - 1

736 + flattened_indices.append(flattened_index)

737 + indices_match &= (prev_index == flattened_index)

738 + max_flattened_index = max(max_flattened_index, flattened_index)

739 + num_commmon_tokens = len(flattened_indices)

740 + if num_commmon_tokens < total_num_scheduled_tokens:

741 + # If not all requests are decodes from the last iteration,

742 + # We need to copy the input_ids_cpu to the GPU first.

743 + self.input_ids.copy_to_gpu(total_num_scheduled_tokens)

744 + if num_commmon_tokens == 0:

745 + # No requests in common with the previous iteration

746 + # So input_ids_cpu will have all the input ids.

747 + return

748 + if indices_match and max_flattened_index == (num_commmon_tokens - 1):

749 + # Common-case optimization: the batch is unchanged

750 + # and no reordering happened.

751 + # The indices are both the same permutation of 0..N-1 so

752 + # we can copy directly using a single slice.

753 + self.input_ids.gpu[:num_commmon_tokens].copy_(

754 + self.input_batch.prev_sampled_token_ids[:num_commmon_tokens,

755 + 0],

756 + non_blocking=True)

757 + return

758 + # Upload the index tensors asynchronously

759 + # so the scatter can be non-blocking.

760 + input_ids_index_tensor = torch.tensor(flattened_indices,

761 + dtype=torch.int64,

762 + pin_memory=self.pin_memory).to(

763 + self.device,

764 + non_blocking=True)

765 + prev_common_req_indices_tensor = torch.tensor(

766 + prev_common_req_indices,

767 + dtype=torch.int64,

768 + pin_memory=self.pin_memory).to(self.device, non_blocking=True)

769 + self.input_ids.gpu.scatter_(

770 + dim=0,

771 + index=input_ids_index_tensor,

772 + src=self.input_batch.prev_sampled_token_ids[

773 + prev_common_req_indices_tensor, 0])

774 +

775 def _prepare_inputs(

776 self,

777 scheduler_output: "SchedulerOutput",

@@ -740,7 +858,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

858 max_seq_len = self.seq_lens.np[:num_reqs].max().item()

859

860 # Copy the tensors to the GPU.

743 - self.input_ids.copy_to_gpu(total_num_scheduled_tokens)

861 + self._prepare_input_ids(total_num_scheduled_tokens, cu_num_tokens)

862 +

863 if self.uses_mrope:

864 # Only relevant for models using M-RoPE (e.g, Qwen2-VL)

865 self.mrope_positions.gpu[:, :total_num_scheduled_tokens].copy_(

@@ -1458,7 +1577,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1577 self,

1578 scheduler_output: "SchedulerOutput",

1579 intermediate_tensors: Optional[IntermediateTensors] = None,

1461 - ) -> Union[ModelRunnerOutput, IntermediateTensors]:

1580 + ) -> Union[ModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]:

1581 self._update_states(scheduler_output)

1582 if not scheduler_output.total_num_scheduled_tokens:

1583 if not has_kv_transfer_group():

@@ -1673,6 +1792,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1792 # so that we could clear the sampled tokens before returning.

1793 discard_sampled_tokens_req_indices.append(i)

1794

1795 + # Copy some objects so they don't get modified after returning.

1796 + # This is important when using async scheduling.

1797 + req_ids_output_copy = self.input_batch.req_ids.copy()

1798 + req_id_to_index_output_copy = \

1799 + self.input_batch.req_id_to_index.copy()

1800 +

1801 # NOTE: GPU -> CPU Sync happens here.

1802 # Move as many CPU operations as possible before this sync point.

1803 logprobs_tensors = sampler_output.logprobs_tensors

@@ -1685,21 +1810,41 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1810 scheduler_output.num_scheduled_tokens,

1811 )

1812

1688 - # Get the valid generated tokens.

1813 + num_sampled_tokens = sampler_output.sampled_token_ids.shape[0]

1814 sampled_token_ids = sampler_output.sampled_token_ids

1690 - max_gen_len = sampled_token_ids.shape[-1]

1691 - if max_gen_len == 1:

1692 - # No spec decode tokens.

1693 - valid_sampled_token_ids = self._to_list(sampled_token_ids)

1815 + if not self.use_async_scheduling:

1816 + # Get the valid generated tokens.

1817 + max_gen_len = sampled_token_ids.shape[-1]

1818 + if max_gen_len == 1:

1819 + # No spec decode tokens.

1820 + valid_sampled_token_ids = self._to_list(sampled_token_ids)

1821 + else:

1822 + # Includes spec decode tokens.

1823 + valid_sampled_token_ids = self.rejection_sampler.parse_output(

1824 + sampled_token_ids,

1825 + self.input_batch.vocab_size,

1826 + )

1827 + # Mask out the sampled tokens that should not be sampled.

1828 + for i in discard_sampled_tokens_req_indices:

1829 + valid_sampled_token_ids[i].clear()

1830 else:

1695 - # Includes spec decode tokens.

1696 - valid_sampled_token_ids = self.rejection_sampler.parse_output(

1697 - sampled_token_ids,

1698 - self.input_batch.vocab_size,

1699 - )

1700 - # Mask out the sampled tokens that should not be sampled.

1701 - for i in discard_sampled_tokens_req_indices:

1702 - valid_sampled_token_ids[i].clear()

1831 + valid_sampled_token_ids = []

1832 + invalid_req_indices = list(discard_sampled_tokens_req_indices)

1833 + invalid_req_indices_set = set(invalid_req_indices)

1834 + assert sampled_token_ids.shape[-1] == 1

1835 +

1836 + # Cache the sampled tokens on the GPU and avoid CPU sync.

1837 + # These will be copied into input_ids in the next step

1838 + # when preparing inputs.

1839 + self.input_batch.prev_sampled_token_ids = \

1840 + sampled_token_ids

1841 + self.input_batch.prev_sampled_token_ids_invalid_indices = \

1842 + invalid_req_indices_set

1843 + self.input_batch.prev_req_id_to_index = {

1844 + req_id: i

1845 + for i, req_id in enumerate(self.input_batch.req_ids)

1846 + if i not in invalid_req_indices_set

1847 + }

1848

1849 # Cache the sampled tokens in the model runner, so that the scheduler

1850 # doesn't need to send them back.

@@ -1707,7 +1852,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1852 # the sampled tokens back, because there's no direct communication

1853 # between the first-stage worker and the last-stage worker.

1854 req_ids = self.input_batch.req_ids

1710 - for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):

1855 + for req_idx in range(num_sampled_tokens):

1856 + if self.use_async_scheduling:

1857 + sampled_ids = [-1] if \

1858 + req_idx not in invalid_req_indices_set else None

1859 + else:

1860 + sampled_ids = valid_sampled_token_ids[req_idx]

1861 if not sampled_ids:

1862 continue

1863

@@ -1722,6 +1872,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1872 start_idx:end_idx] = sampled_ids

1873 self.input_batch.num_tokens_no_spec[req_idx] = end_idx

1874 self.input_batch.num_tokens[req_idx] = end_idx

1875 +

1876 req_id = req_ids[req_idx]

1877 req_state = self.requests[req_id]

1878 req_state.output_token_ids.extend(sampled_ids)

@@ -1741,9 +1892,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1892

1893 self.eplb_step()

1894

1744 - return ModelRunnerOutput(

1745 - req_ids=self.input_batch.req_ids,

1746 - req_id_to_index=self.input_batch.req_id_to_index,

1895 + output = ModelRunnerOutput(

1896 + req_ids=req_ids_output_copy,

1897 + req_id_to_index=req_id_to_index_output_copy,

1898 sampled_token_ids=valid_sampled_token_ids,

1899 logprobs=logprobs_lists,

1900 prompt_logprobs_dict=prompt_logprobs_dict,

@@ -1752,6 +1903,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1903 num_nans_in_logits=num_nans_in_logits,

1904 )

1905

1906 + if not self.use_async_scheduling:

1907 + return output

1908 +

1909 + return AsyncGPUModelRunnerOutput(

1910 + model_runner_output=output,

1911 + sampled_token_ids=sampled_token_ids,

1912 + invalid_req_indices=invalid_req_indices,

1913 + async_output_copy_stream=self.async_output_copy_stream,

1914 + )

1915 +

1916 def take_draft_token_ids(self) -> Optional[DraftTokenIds]:

1917 if self._draft_token_ids is None:

1918 return None

vllm/v1/worker/gpu_worker.py

+5 -5

@@ -5,7 +5,7 @@ import copy

5 import gc

6 import os

7 from contextlib import AbstractContextManager, nullcontext

8 -from typing import TYPE_CHECKING, Any, Optional

8 +from typing import TYPE_CHECKING, Any, Optional, Union

10 import torch

11 import torch.distributed

@@ -28,8 +28,8 @@ from vllm.tasks import SupportedTask

28 from vllm.utils import GiB_bytes, MemorySnapshot, memory_profiling

29 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType

30 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec

31 -from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, DraftTokenIds,

32 - ModelRunnerOutput)

31 +from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput,

32 + DraftTokenIds, ModelRunnerOutput)

33 from vllm.v1.utils import report_usage_stats

34 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

35 from vllm.v1.worker.worker_base import WorkerBase

@@ -355,7 +355,7 @@ class Worker(WorkerBase):

355 def execute_model(

356 self,

357 scheduler_output: "SchedulerOutput",

358 - ) -> Optional[ModelRunnerOutput]:

358 + ) -> Optional[Union[ModelRunnerOutput, AsyncModelRunnerOutput]]:

359 intermediate_tensors = None

360 forward_pass = scheduler_output.total_num_scheduled_tokens > 0

361 if forward_pass and not get_pp_group().is_first_rank:

@@ -365,7 +365,7 @@ class Worker(WorkerBase):

365

366 output = self.model_runner.execute_model(scheduler_output,

367 intermediate_tensors)

368 - if isinstance(output, ModelRunnerOutput):

368 + if isinstance(output, (ModelRunnerOutput, AsyncModelRunnerOutput)):

369 return output

370

371 assert isinstance(output, IntermediateTensors)