From fadf5d332c6e9bb6e552c1ca529511bce0f79802 Mon Sep 17 00:00:00 2001 From: pmaybank <113125070+pmaybank@users.noreply.github.com> Date: Tue, 19 May 2026 23:16:02 -0400 Subject: [PATCH] add enqueue all option to throughput benchmark (#42975) Signed-off-by: Philip Maybank Signed-off-by: pmaybank <113125070+pmaybank@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../llm/test_mm_processor_kwargs.py | 45 +++++++ vllm/benchmarks/throughput.py | 60 ++++++++- vllm/entrypoints/llm.py | 119 +++++++++++++++++- 3 files changed, 217 insertions(+), 7 deletions(-) diff --git a/tests/entrypoints/llm/test_mm_processor_kwargs.py b/tests/entrypoints/llm/test_mm_processor_kwargs.py index 1b0092df01..f970d68aca 100644 --- a/tests/entrypoints/llm/test_mm_processor_kwargs.py +++ b/tests/entrypoints/llm/test_mm_processor_kwargs.py @@ -76,6 +76,51 @@ def test_chat_forwards_mm_processor_kwargs() -> None: ) +def test_enqueue_chat_forwards_mm_processor_kwargs() -> None: + llm = _make_mock_llm() + mm_processor_kwargs = {"do_pan_and_scan": True} + sampling_params = SamplingParams(max_tokens=1) + messages = [{"role": "user", "content": "hello"}] + + llm._add_chat_requests = Mock(return_value=["req-0"]) + + request_ids = llm.enqueue_chat( + messages, + sampling_params=sampling_params, + use_tqdm=False, + mm_processor_kwargs=mm_processor_kwargs, + ) + + assert request_ids == ["req-0"] + assert llm._add_chat_requests.call_args.kwargs["mm_processor_kwargs"] == ( + mm_processor_kwargs + ) + + +def test_run_chat_forwards_mm_processor_kwargs() -> None: + llm = _make_mock_llm() + mm_processor_kwargs = {"num_crops": 8} + sampling_params = SamplingParams(max_tokens=1) + messages = [{"role": "user", "content": "hello"}] + sentinel_output = ["done"] + + llm._add_chat_requests = Mock() + llm._run_engine = Mock(return_value=sentinel_output) + + outputs = llm._run_chat( + messages=messages, + params=sampling_params, + output_type=object, + use_tqdm=False, + mm_processor_kwargs=mm_processor_kwargs, + ) + + assert outputs == sentinel_output + assert llm._add_chat_requests.call_args.kwargs["mm_processor_kwargs"] == ( + mm_processor_kwargs + ) + + def test_run_completion_forwards_mm_processor_kwargs() -> None: llm = _make_mock_llm() mm_processor_kwargs = {"min_pixels": 4 * 28 * 28} diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py index 1c3b728fdb..95a87e6be1 100644 --- a/vllm/benchmarks/throughput.py +++ b/vllm/benchmarks/throughput.py @@ -50,6 +50,7 @@ def run_vllm( engine_args: EngineArgs, do_profile: bool, disable_detokenize: bool = False, + prequeue_requests: bool = False, ) -> tuple[float, list[RequestOutput] | None]: from vllm import LLM, SamplingParams @@ -94,12 +95,29 @@ def run_vllm( outputs = None if not use_beam_search: + if prequeue_requests: + llm.sleep(level=0, mode="abort") + start = time.perf_counter() if do_profile: llm.start_profile() - outputs = llm.generate( - prompts, sampling_params, lora_request=lora_requests, use_tqdm=True - ) + + if prequeue_requests: + try: + llm.enqueue( + prompts, + sampling_params, + lora_request=lora_requests, + use_tqdm=True, + ) + finally: + llm.wake_up(tags=["scheduling"]) + outputs = llm.wait_for_completion(output_type=RequestOutput, use_tqdm=True) + else: + outputs = llm.generate( + prompts, sampling_params, lora_request=lora_requests, use_tqdm=True + ) + if do_profile: llm.stop_profile() end = time.perf_counter() @@ -133,6 +151,7 @@ def run_vllm_chat( engine_args: EngineArgs, do_profile: bool, disable_detokenize: bool = False, + prequeue_requests: bool = False, ) -> tuple[float, list[RequestOutput]]: """ Run vLLM chat benchmark. This function is recommended ONLY for benchmarking @@ -166,12 +185,25 @@ def run_vllm_chat( detokenize=not disable_detokenize, ) ) + if prequeue_requests: + llm.sleep(level=0, mode="abort") + start = time.perf_counter() if do_profile: llm.start_profile() - outputs = llm.chat(prompts, sampling_params, use_tqdm=True) + + if prequeue_requests: + try: + llm.enqueue_chat(prompts, sampling_params, use_tqdm=True) + finally: + llm.wake_up(tags=["scheduling"]) + outputs = llm.wait_for_completion(output_type=RequestOutput, use_tqdm=True) + else: + outputs = llm.chat(prompts, sampling_params, use_tqdm=True) + if do_profile: llm.stop_profile() + end = time.perf_counter() return end - start, outputs @@ -543,6 +575,10 @@ def validate_args(args): valid_backends = {"vllm", "hf", "mii", "vllm-chat"} if args.backend not in valid_backends: raise ValueError(f"Unsupported backend: {args.backend}") + if args.prequeue_requests and args.backend not in {"vllm", "vllm-chat"}: + raise ValueError("--prequeue-requests requires --backend vllm or vllm-chat") + if args.prequeue_requests and args.async_engine: + raise ValueError("--prequeue-requests is not supported with --async-engine") # === Dataset Configuration === if ( @@ -789,6 +825,20 @@ def add_cli_args(parser: argparse.ArgumentParser): default=False, help="Use vLLM async engine rather than LLM class.", ) + parser.add_argument( + "--prequeue-requests", + action="store_true", + default=False, + help=( + "For the vLLM backends, enqueue all requests before allowing the " + "scheduler to process them. This can improve benchmark " + "reproducibility by removing overlap between request rendering " + "and engine scheduling, but may reduce measured throughput. " + "Request rendering is typically fast relative to scheduling and " + "processing; the intended use case of this flag is multimodal " + "benchmarks with time-consuming image rendering." + ), + ) parser.add_argument( "--disable-detokenize", action="store_true", @@ -942,6 +992,7 @@ def main(args: argparse.Namespace): EngineArgs.from_cli_args(args), disable_detokenize=args.disable_detokenize, do_profile=args.profile, + prequeue_requests=args.prequeue_requests, ) elif args.backend == "hf": assert args.tensor_parallel_size == 1 @@ -965,6 +1016,7 @@ def main(args: argparse.Namespace): EngineArgs.from_cli_args(args), disable_detokenize=args.disable_detokenize, do_profile=args.profile, + prequeue_requests=args.prequeue_requests, ) else: raise ValueError(f"Unknown backend: {args.backend}") diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 0e29f38c0b..13e12683da 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -894,6 +894,83 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin): mm_processor_kwargs=mm_processor_kwargs, ) + def enqueue_chat( + self, + messages: list[ChatCompletionMessageParam] + | Sequence[list[ChatCompletionMessageParam]], + sampling_params: SamplingParams | Sequence[SamplingParams] | None = None, + use_tqdm: bool | Callable[..., tqdm] = True, + lora_request: Sequence[LoRARequest] | LoRARequest | None = None, + priority: list[int] | None = None, + chat_template: str | None = None, + chat_template_content_format: ChatTemplateContentFormatOption = "auto", + add_generation_prompt: bool = True, + continue_final_message: bool = False, + tools: list[dict[str, Any]] | None = None, + chat_template_kwargs: dict[str, Any] | None = None, + tokenization_kwargs: dict[str, Any] | None = None, + mm_processor_kwargs: dict[str, Any] | None = None, + ) -> list[str]: + """Enqueue chat conversations for generation without waiting. + + This method renders chat conversations and adds the resulting requests + to the engine queue. Use wait_for_completion() to get results. To + guarantee that all requests are queued before scheduling starts, pause + scheduling with sleep(level=0) before calling this method and resume it + with wake_up(tags=["scheduling"]) afterward. + + Args: + messages: A sequence of conversations or a single conversation. + Each conversation is represented as a list of messages. + sampling_params: The sampling parameters for text generation. + If None, we use the default sampling parameters. + use_tqdm: If `True`, shows a tqdm progress bar while rendering + conversations. + lora_request: LoRA request to use for generation, if any. + priority: The priority of the requests, if any. + chat_template: The template to use for structuring the chat. + chat_template_content_format: The format to render message content. + add_generation_prompt: If True, adds a generation template + to each message. + continue_final_message: If True, continues the final message in + the conversation instead of starting a new one. + tools: Tools to make available to the model, if any. + chat_template_kwargs: Additional kwargs to pass to the chat + template. + tokenization_kwargs: Overrides for `tokenizer.encode`. + mm_processor_kwargs: Overrides for `processor.__call__`. + + Returns: + A list of request IDs for the enqueued requests. + """ + model_config = self.model_config + runner_type = model_config.runner_type + if runner_type != "generate": + raise ValueError( + "LLM.enqueue_chat() is only supported for generative models. " + "Try passing `--runner generate` to use the model as a " + "generative model." + ) + + if sampling_params is None: + sampling_params = self.get_default_sampling_params() + + return self._add_chat_requests( + messages=messages, + params=sampling_params, + use_tqdm=use_tqdm, + lora_request=lora_request, + priority=priority, + chat_template=chat_template, + chat_template_content_format=chat_template_content_format, + chat_template_kwargs=chat_template_kwargs, + add_generation_prompt=add_generation_prompt, + continue_final_message=continue_final_message, + tools=tools, + tokenization_kwargs=tokenization_kwargs, + mm_processor_kwargs=mm_processor_kwargs, + ) + def start_profile(self, profile_prefix: str | None = None) -> None: """Start profiling with optional custom trace prefix. @@ -1095,9 +1172,46 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin): tokenization_kwargs: dict[str, Any] | None = None, mm_processor_kwargs: dict[str, Any] | None = None, ): + self._add_chat_requests( + messages=messages, + params=params, + use_tqdm=use_tqdm, + lora_request=lora_request, + chat_template=chat_template, + chat_template_content_format=chat_template_content_format, + chat_template_kwargs=chat_template_kwargs, + add_generation_prompt=add_generation_prompt, + continue_final_message=continue_final_message, + tools=tools, + tokenization_kwargs=tokenization_kwargs, + mm_processor_kwargs=mm_processor_kwargs, + ) + return self._run_engine(output_type=output_type, use_tqdm=use_tqdm) + + def _add_chat_requests( + self, + messages: list[ChatCompletionMessageParam] + | Sequence[list[ChatCompletionMessageParam]], + params: SamplingParams + | PoolingParams + | Sequence[SamplingParams | PoolingParams], + *, + use_tqdm: bool | Callable[..., tqdm] = True, + lora_request: Sequence[LoRARequest] | LoRARequest | None = None, + priority: list[int] | None = None, + chat_template: str | None = None, + chat_template_content_format: ChatTemplateContentFormatOption = "auto", + add_generation_prompt: bool = True, + continue_final_message: bool = False, + tools: list[dict[str, Any]] | None = None, + chat_template_kwargs: dict[str, Any] | None = None, + tokenization_kwargs: dict[str, Any] | None = None, + mm_processor_kwargs: dict[str, Any] | None = None, + ) -> list[str]: seq_convs = conversation_to_seq(messages) seq_params = self._params_to_seq(params, len(seq_convs)) seq_lora_requests = self._lora_request_to_seq(lora_request, len(seq_convs)) + seq_priority = self._priority_to_seq(priority, len(seq_convs)) # When thinking is enabled or tools are provided, and the model # uses special tokens for structured output (e.g. Gemma4's @@ -1110,7 +1224,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin): if needs_parsing: self._adjust_params_for_parsing(seq_params) - return self._render_and_run_requests( + return self._render_and_add_requests( prompts=( self._preprocess_chat_one( conversation, @@ -1130,9 +1244,8 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin): ) ), params=seq_params, - output_type=output_type, lora_requests=seq_lora_requests, - use_tqdm=use_tqdm, + priorities=seq_priority, ) def _adjust_params_for_parsing(