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

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

break execute_model in gpu_model_runner into sub-functions for custom scopes (#24265)

Co-authored-by: Bangsheng Tang <bangsheng@meta.com>

Bangsheng Tang committed Sep 6, 2025 at 14:02 UTC 848562bd49a91d21f2dc76912d5fcb25ae80c8a8

3 files changed +207 -108

vllm/envs.py

@@ -168,6 +168,7 @@ if TYPE_CHECKING:

168 VLLM_ALLREDUCE_USE_SYMM_MEM: bool = False

169 VLLM_TUNED_CONFIG_FOLDER: Optional[str] = None

170 VLLM_DISABLE_PAD_FOR_CUDAGRAPH: bool = False

171 + VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False

172

173

174 def get_default_cache_root():

@@ -1200,6 +1201,9 @@ environment_variables: dict[str, Callable[[], Any]] = {

1201 "VLLM_TUNED_CONFIG_FOLDER":

1202 lambda: os.getenv("VLLM_TUNED_CONFIG_FOLDER", None),

1203

1204 + # Add optional custom scopes for profiling, disable to avoid overheads

1205 + "VLLM_CUSTOM_SCOPES_FOR_PROFILING":

1206 + lambda: bool(int(os.getenv("VLLM_CUSTOM_SCOPES_FOR_PROFILING", "0"))),

1207 }

1208

1209 # --8<-- [end:env-vars-definition]

vllm/v1/utils.py

+15 -4

@@ -1,17 +1,21 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 import argparse

4 +import contextlib

5 import multiprocessing

6 import time

7 import weakref

8 from collections.abc import Sequence

9 +from contextlib import AbstractContextManager

10 from multiprocessing import connection

11 from multiprocessing.process import BaseProcess

12 from typing import (TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar,

13 Union, overload)

15 import torch

16 +from torch.autograd.profiler import record_function

18 +import vllm.envs as envs

19 from vllm.logger import init_logger

20 from vllm.usage.usage_lib import (UsageContext, is_usage_stats_enabled,

21 usage_message)

@@ -155,7 +159,7 @@ def get_engine_client_zmq_addr(local_only: bool,

159

160 class APIServerProcessManager:

161 """Manages a group of API server processes.

158 -

162 +

163 Handles creation, monitoring, and termination of API server worker

164 processes. Also monitors extra processes to check if they are healthy.

165 """

@@ -172,7 +176,7 @@ class APIServerProcessManager:

176 stats_update_address: Optional[str] = None,

177 ):

178 """Initialize and start API server worker processes.

175 -

179 +

180 Args:

181 target_server_fn: Function to call for each API server process

182 listen_address: Address to listen for client connections

@@ -181,7 +185,7 @@ class APIServerProcessManager:

185 num_servers: Number of API server processes to start

186 input_addresses: Input addresses for each API server

187 output_addresses: Output addresses for each API server

184 - stats_update_address: Optional stats update address

188 + stats_update_address: Optional stats update address

189 """

190 self.listen_address = listen_address

191 self.sock = sock

@@ -225,7 +229,7 @@ def wait_for_completion_or_failure(

229 "CoreEngineActorManager"]] = None,

230 coordinator: Optional["DPCoordinator"] = None) -> None:

231 """Wait for all processes to complete or detect if any fail.

228 -

232 +

233 Raises an exception if any process exits with a non-zero status.

234

235 Args:

@@ -368,3 +372,10 @@ def report_usage_stats(

372 "disable_custom_all_reduce":

373 vllm_config.parallel_config.disable_custom_all_reduce,

374 })

375 +

376 +

377 +def record_function_or_nullcontext(name: str) -> AbstractContextManager:

378 + if envs.VLLM_CUSTOM_SCOPES_FOR_PROFILING:

379 + return record_function(name)

380 + else:

381 + return contextlib.nullcontext()

vllm/v1/worker/gpu_model_runner.py

+188 -104

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

69 KVCacheGroupSpec, KVCacheSpec,

70 MambaSpec, SlidingWindowSpec)

71 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput,

72 - DraftTokenIds, LogprobsTensors, ModelRunnerOutput)

72 + DraftTokenIds, LogprobsLists, LogprobsTensors,

73 + ModelRunnerOutput, SamplerOutput)

74 from vllm.v1.pool.metadata import PoolingMetadata

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

76 from vllm.v1.sample.metadata import SamplingMetadata

@@ -79,7 +80,7 @@ from vllm.v1.spec_decode.eagle import EagleProposer

80 from vllm.v1.spec_decode.medusa import MedusaProposer

81 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

82 from vllm.v1.spec_decode.ngram_proposer import NgramProposer

82 -from vllm.v1.utils import CpuGpuBuffer

83 +from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext

84 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

85 from vllm.v1.worker.kv_connector_model_runner_mixin import (

86 KVConnectorModelRunnerMixin, KVConnectorOutput)

@@ -1587,31 +1588,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1588 kv_connector_output=kv_connector_output,

1589 )

1590

1590 - @torch.inference_mode()

1591 - def execute_model(

1591 + def _preprocess(

1592 self,

1593 scheduler_output: "SchedulerOutput",

1594 intermediate_tensors: Optional[IntermediateTensors] = None,

1595 - ) -> Union[ModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]:

1596 - self._update_states(scheduler_output)

1597 - if not scheduler_output.total_num_scheduled_tokens:

1598 - if not has_kv_transfer_group():

1599 - # Return empty ModelRunnerOutput if there's no work to do.

1600 - return EMPTY_MODEL_RUNNER_OUTPUT

1601 -

1602 - return self.kv_connector_no_forward(scheduler_output,

1603 - self.vllm_config)

1604 -

1605 - if self.cache_config.kv_sharing_fast_prefill:

1606 - assert not self.input_batch.num_prompt_logprobs, (

1607 - "--kv-sharing-fast-prefill produces incorrect logprobs for "

1608 - "prompt tokens, tokens, please disable it when the requests "

1609 - "need prompt logprobs")

1610 -

1611 - # Prepare the decoder inputs.

1612 - (attn_metadata, logits_indices, spec_decode_metadata,

1613 - num_scheduled_tokens_np, spec_decode_common_attn_metadata,

1614 - max_query_len) = self._prepare_inputs(scheduler_output)

1595 + ) -> tuple[int, int, Optional[torch.Tensor], Optional[torch.Tensor],

1596 + Optional[torch.Tensor], torch.Tensor,

1597 + Optional[IntermediateTensors], dict[str, Any]]:

1598

1599 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

1600 if (self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

@@ -1683,75 +1666,21 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1666 intermediate_tensors = self.sync_and_slice_intermediate_tensors(

1667 num_input_tokens, intermediate_tensors, True)

1668

1686 - uniform_decode = (max_query_len == self.uniform_decode_query_len) and (

1687 - num_scheduled_tokens == self.input_batch.num_reqs * max_query_len)

1688 - batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens,

1689 - uniform_decode=uniform_decode)

1690 - cudagraph_runtime_mode, batch_descriptor = \

1691 - self.cudagraph_dispatcher.dispatch(batch_descriptor)

1692 -

1693 - # Run the model.

1694 - # Use persistent buffers for CUDA graphs.

1695 - with set_forward_context(

1696 - attn_metadata,

1697 - self.vllm_config,

1698 - num_tokens=num_input_tokens,

1699 - num_tokens_across_dp=num_tokens_across_dp,

1700 - cudagraph_runtime_mode=cudagraph_runtime_mode,

1701 - batch_descriptor=batch_descriptor,

1702 - ), self.maybe_get_kv_connector_output(

1703 - scheduler_output) as kv_connector_output:

1704 -

1705 - model_output = self.model(

1706 - input_ids=input_ids,

1707 - positions=positions,

1708 - intermediate_tensors=intermediate_tensors,

1709 - inputs_embeds=inputs_embeds,

1710 - **model_kwargs,

1711 - )

1712 -

1713 - if self.use_aux_hidden_state_outputs:

1714 - hidden_states, aux_hidden_states = model_output

1715 - else:

1716 - hidden_states = model_output

1717 - aux_hidden_states = None

1718 -

1719 - # Broadcast PP output for external_launcher (torchrun)

1720 - # to make sure we are synced across pp ranks

1721 - # TODO: Support overlapping mirco-batches

1722 - # https://github.com/vllm-project/vllm/issues/18019

1723 - broadcast_pp_output = \

1724 - self.parallel_config.distributed_executor_backend \

1725 - == "external_launcher" and len(get_pp_group().ranks) > 0

1726 - if not get_pp_group().is_last_rank:

1727 - # For mid-pipeline stages, return the hidden states.

1728 - assert isinstance(hidden_states, IntermediateTensors)

1729 - if not broadcast_pp_output:

1730 - hidden_states.kv_connector_output = kv_connector_output

1731 - return hidden_states

1732 - get_pp_group().send_tensor_dict(hidden_states.tensors,

1733 - all_gather_group=get_tp_group())

1734 - logits = None

1735 - else:

1736 - if self.is_pooling_model:

1737 - return self._pool(hidden_states, num_scheduled_tokens,

1738 - num_scheduled_tokens_np, kv_connector_output)

1739 -

1740 - sample_hidden_states = hidden_states[logits_indices]

1741 - logits = self.model.compute_logits(sample_hidden_states, None)

1742 - if broadcast_pp_output:

1743 - model_output_broadcast_data = {

1744 - "logits": logits.contiguous(),

1745 - } if logits is not None else {}

1746 - model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(

1747 - model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)

1748 - assert model_output_broadcast_data is not None

1749 - logits = model_output_broadcast_data["logits"]

1750 -

1751 - # Apply structured output bitmasks if present

1752 - if scheduler_output.grammar_bitmask is not None:

1753 - self.apply_grammar_bitmask(scheduler_output, logits)

1669 + return (

1670 + num_scheduled_tokens,

1671 + num_input_tokens,

1672 + num_tokens_across_dp,

1673 + input_ids,

1674 + inputs_embeds,

1675 + positions,

1676 + intermediate_tensors,

1677 + model_kwargs,

1678 + )

1679

1680 + def _sample(

1681 + self, logits: Optional[torch.Tensor],

1682 + spec_decode_metadata: Optional[SpecDecodeMetadata]

1683 + ) -> SamplerOutput:

1684 # Sample the next token and get logprobs if needed.

1685 sampling_metadata = self.input_batch.sampling_metadata

1686 if spec_decode_metadata is None:

@@ -1785,6 +1714,21 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1714 )

1715 sampler_output.sampled_token_ids = output_token_ids

1716

1717 + return sampler_output

1718 +

1719 + def _bookkeeping_sync(

1720 + self, scheduler_output: "SchedulerOutput",

1721 + sampler_output: SamplerOutput, logits: Optional[torch.Tensor],

1722 + hidden_states: torch.Tensor, num_scheduled_tokens: int

1723 + ) -> tuple[

1724 + dict[str, int],

1725 + Optional[LogprobsLists],

1726 + list[list[int]],

1727 + dict[str, Optional[LogprobsTensors]],

1728 + list[str],

1729 + dict[str, int],

1730 + list[int],

1731 + ]:

1732 num_nans_in_logits = {}

1733 if envs.VLLM_COMPUTE_NANS_IN_LOGITS:

1734 num_nans_in_logits = self._get_nans_in_logits(logits)

@@ -1827,6 +1771,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1771

1772 num_sampled_tokens = sampler_output.sampled_token_ids.shape[0]

1773 sampled_token_ids = sampler_output.sampled_token_ids

1774 + invalid_req_indices = []

1775 if not self.use_async_scheduling:

1776 # Get the valid generated tokens.

1777 max_gen_len = sampled_token_ids.shape[-1]

@@ -1892,20 +1837,159 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1837 req_state = self.requests[req_id]

1838 req_state.output_token_ids.extend(sampled_ids)

1839

1840 + return (

1841 + num_nans_in_logits,

1842 + logprobs_lists,

1843 + valid_sampled_token_ids,

1844 + prompt_logprobs_dict,

1845 + req_ids_output_copy,

1846 + req_id_to_index_output_copy,

1847 + invalid_req_indices,

1848 + )

1849 +

1850 + @torch.inference_mode()

1851 + def execute_model(

1852 + self,

1853 + scheduler_output: "SchedulerOutput",

1854 + intermediate_tensors: Optional[IntermediateTensors] = None,

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

1856 + with record_function_or_nullcontext("Preprocess"):

1857 + self._update_states(scheduler_output)

1858 + if not scheduler_output.total_num_scheduled_tokens:

1859 + if not has_kv_transfer_group():

1860 + # Return empty ModelRunnerOutput if there's no work to do.

1861 + return EMPTY_MODEL_RUNNER_OUTPUT

1862 + return self.kv_connector_no_forward(scheduler_output,

1863 + self.vllm_config)

1864 + if self.cache_config.kv_sharing_fast_prefill:

1865 + assert not self.input_batch.num_prompt_logprobs, (

1866 + "--kv-sharing-fast-prefill produces incorrect logprobs for "

1867 + "prompt tokens, tokens, please disable it when the requests"

1868 + " need prompt logprobs")

1869 +

1870 + # Prepare the decoder inputs.

1871 + (attn_metadata, logits_indices, spec_decode_metadata,

1872 + num_scheduled_tokens_np, spec_decode_common_attn_metadata,

1873 + max_query_len) = self._prepare_inputs(scheduler_output)

1874 +

1875 + (

1876 + num_scheduled_tokens,

1877 + num_input_tokens,

1878 + num_tokens_across_dp,

1879 + input_ids,

1880 + inputs_embeds,

1881 + positions,

1882 + intermediate_tensors,

1883 + model_kwargs,

1884 + ) = self._preprocess(scheduler_output, intermediate_tensors)

1885 +

1886 + uniform_decode = (max_query_len

1887 + == self.uniform_decode_query_len) and (

1888 + num_scheduled_tokens

1889 + == self.input_batch.num_reqs * max_query_len)

1890 + batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens,

1891 + uniform_decode=uniform_decode)

1892 + cudagraph_runtime_mode, batch_descriptor = \

1893 + self.cudagraph_dispatcher.dispatch(batch_descriptor)

1894 +

1895 + # Run the model.

1896 + # Use persistent buffers for CUDA graphs.

1897 + with (set_forward_context(

1898 + attn_metadata,

1899 + self.vllm_config,

1900 + num_tokens=num_input_tokens,

1901 + num_tokens_across_dp=num_tokens_across_dp,

1902 + cudagraph_runtime_mode=cudagraph_runtime_mode,

1903 + batch_descriptor=batch_descriptor,

1904 + ), record_function_or_nullcontext("Forward"),

1905 + self.maybe_get_kv_connector_output(scheduler_output) as

1906 + kv_connector_output):

1907 + model_output = self.model(

1908 + input_ids=input_ids,

1909 + positions=positions,

1910 + intermediate_tensors=intermediate_tensors,

1911 + inputs_embeds=inputs_embeds,

1912 + **model_kwargs,

1913 + )

1914 +

1915 + with record_function_or_nullcontext("Postprocess"):

1916 + if self.use_aux_hidden_state_outputs:

1917 + hidden_states, aux_hidden_states = model_output

1918 + else:

1919 + hidden_states = model_output

1920 + aux_hidden_states = None

1921 +

1922 + # Broadcast PP output for external_launcher (torchrun)

1923 + # to make sure we are synced across pp ranks

1924 + # TODO: Support overlapping mirco-batches

1925 + # https://github.com/vllm-project/vllm/issues/18019

1926 + broadcast_pp_output = \

1927 + self.parallel_config.distributed_executor_backend \

1928 + == "external_launcher" and len(get_pp_group().ranks) > 0

1929 + if not get_pp_group().is_last_rank:

1930 + # For mid-pipeline stages, return the hidden states.

1931 + assert isinstance(hidden_states, IntermediateTensors)

1932 + if not broadcast_pp_output:

1933 + hidden_states.kv_connector_output = kv_connector_output

1934 + return hidden_states

1935 + get_pp_group().send_tensor_dict(

1936 + hidden_states.tensors, all_gather_group=get_tp_group())

1937 + logits = None

1938 + else:

1939 + if self.is_pooling_model:

1940 + return self._pool(hidden_states, num_scheduled_tokens,

1941 + num_scheduled_tokens_np,

1942 + kv_connector_output)

1943 +

1944 + sample_hidden_states = hidden_states[logits_indices]

1945 + logits = self.model.compute_logits(sample_hidden_states, None)

1946 + if broadcast_pp_output:

1947 + model_output_broadcast_data = {

1948 + "logits": logits.contiguous(),

1949 + } if logits is not None else {}

1950 + model_output_broadcast_data = get_pp_group(

1951 + ).broadcast_tensor_dict(model_output_broadcast_data,

1952 + src=len(get_pp_group().ranks) - 1)

1953 + assert model_output_broadcast_data is not None

1954 + logits = model_output_broadcast_data["logits"]

1955 +

1956 + # Apply structured output bitmasks if present

1957 + if scheduler_output.grammar_bitmask is not None:

1958 + self.apply_grammar_bitmask(scheduler_output, logits)

1959 +

1960 + with record_function_or_nullcontext("Sample"):

1961 + sampler_output = self._sample(logits, spec_decode_metadata)

1962 +

1963 + with record_function_or_nullcontext("Bookkeep"):

1964 + assert isinstance(hidden_states, torch.Tensor)

1965 + (

1966 + num_nans_in_logits,

1967 + logprobs_lists,

1968 + valid_sampled_token_ids,

1969 + prompt_logprobs_dict,

1970 + req_ids_output_copy,

1971 + req_id_to_index_output_copy,

1972 + invalid_req_indices,

1973 + ) = self._bookkeeping_sync(scheduler_output, sampler_output,

1974 + logits, hidden_states,

1975 + num_scheduled_tokens)

1976 +

1977 if self.speculative_config:

1978 assert spec_decode_common_attn_metadata is not None

1897 - self._draft_token_ids = self.propose_draft_token_ids(

1898 - scheduler_output,

1899 - valid_sampled_token_ids,

1900 - sampling_metadata,

1901 - hidden_states,

1902 - sample_hidden_states,

1903 - aux_hidden_states,

1904 - spec_decode_metadata,

1905 - spec_decode_common_attn_metadata,

1906 - )

1979 + with record_function_or_nullcontext("Draft"):

1980 + self._draft_token_ids = self.propose_draft_token_ids(

1981 + scheduler_output,

1982 + valid_sampled_token_ids,

1983 + self.input_batch.sampling_metadata,

1984 + hidden_states,

1985 + sample_hidden_states,

1986 + aux_hidden_states,

1987 + spec_decode_metadata,

1988 + spec_decode_common_attn_metadata,

1989 + )

1990

1908 - self.eplb_step()

1991 + with record_function_or_nullcontext("EPLB"):

1992 + self.eplb_step()

1993

1994 output = ModelRunnerOutput(

1995 req_ids=req_ids_output_copy,

@@ -1923,7 +2007,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2007

2008 return AsyncGPUModelRunnerOutput(

2009 model_runner_output=output,

1926 - sampled_token_ids=sampled_token_ids,

2010 + sampled_token_ids=sampler_output.sampled_token_ids,

2011 invalid_req_indices=invalid_req_indices,

2012 async_output_copy_stream=self.async_output_copy_stream,

2013 )