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, 12:30 UTC
[Misc] Add run one batch script that supports profiling (#32968)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Feb 10, 2026 at 19:29 UTC
ba0511fd80b95d05ffab867cce54f3590e57a7fc
1 file changed
+112
examples/offline_inference/run_one_batch.py
new
+112
@@ -0,0 +1,112 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+from vllm import LLM, EngineArgs
7
+from vllm.config import ProfilerConfig
8
+from vllm.utils.argparse_utils import FlexibleArgumentParser
9
+
10
+DEFAULT_MAX_TOKENS = 16
11
+
12
+
13
+def create_parser() -> FlexibleArgumentParser:
14
+ parser = FlexibleArgumentParser()
15
+ EngineArgs.add_cli_args(parser)
16
+ parser.set_defaults(model="meta-llama/Llama-3.2-1B-Instruct")
17
+
18
+ batch_group = parser.add_argument_group("Batch parameters")
19
+ batch_group.add_argument("--batch-size", type=int, default=1)
20
+ batch_group.add_argument("--prompt-size", type=int, default=128)
21
+ batch_group.add_argument("--prompt-prefix", type=str, default="Hello, my name is")
22
+
23
+ profile_group = parser.add_argument_group("Profiling parameters")
24
+ profile_group.add_argument(
25
+ "--profile",
26
+ choices=["none", "prefill", "decode", "both"],
27
+ default="none",
28
+ )
29
+ profile_group.add_argument(
30
+ "--profile-dir",
31
+ type=str,
32
+ default="",
33
+ help="Required when --profile is not 'none'.",
34
+ )
35
+
36
+ return parser
37
+
38
+
39
+def _build_prompt(prefix: str, prompt_size: int) -> str:
40
+ if prompt_size <= 0:
41
+ return ""
42
+ if not prefix:
43
+ prefix = " "
44
+ if len(prefix) >= prompt_size:
45
+ return prefix[:prompt_size]
46
+ repeat_count = (prompt_size + len(prefix) - 1) // len(prefix)
47
+ return (prefix * repeat_count)[:prompt_size]
48
+
49
+
50
+def _build_profiler_config(
51
+ profile: str, profile_dir: str, max_tokens: int
52
+) -> ProfilerConfig | None:
53
+ if profile == "none":
54
+ return None
55
+ if not profile_dir:
56
+ raise ValueError("--profile-dir must be set when profiling is enabled.")
57
+ if profile == "prefill":
58
+ delay_iterations = 0
59
+ max_iterations = 1
60
+ elif profile == "decode":
61
+ delay_iterations = 1
62
+ max_iterations = max(1, max_tokens)
63
+ else:
64
+ delay_iterations = 0
65
+ max_iterations = 0
66
+
67
+ return ProfilerConfig(
68
+ profiler="torch",
69
+ torch_profiler_dir=profile_dir,
70
+ delay_iterations=delay_iterations,
71
+ max_iterations=max_iterations,
72
+ )
73
+
74
+
75
+def main(args: dict) -> None:
76
+ max_tokens = DEFAULT_MAX_TOKENS
77
+ batch_size = args.pop("batch_size")
78
+ prompt_size = args.pop("prompt_size")
79
+ prompt_prefix = args.pop("prompt_prefix")
80
+ profile = args.pop("profile")
81
+ profile_dir = args.pop("profile_dir")
82
+
83
+ profiler_config = _build_profiler_config(profile, profile_dir, max_tokens)
84
+ if profiler_config is not None:
85
+ args["profiler_config"] = profiler_config
86
+
87
+ llm = LLM(**args)
88
+
89
+ sampling_params = llm.get_default_sampling_params()
90
+ sampling_params.max_tokens = max_tokens
91
+ sampling_params.min_tokens = max_tokens
92
+ sampling_params.ignore_eos = True
93
+
94
+ prompt = _build_prompt(prompt_prefix, prompt_size)
95
+ prompts = [prompt] * batch_size
96
+
97
+ if profile != "none":
98
+ llm.start_profile()
99
+ outputs = llm.generate(prompts, sampling_params)
100
+ if profile != "none":
101
+ llm.stop_profile()
102
+
103
+ print("-" * 50)
104
+ for output in outputs:
105
+ generated_text = output.outputs[0].text
106
+ print(f"Prompt: {output.prompt!r}\nGenerated text: {generated_text!r}")
107
+ print("-" * 50)
108
+
109
+
110
+if __name__ == "__main__":
111
+ parser = create_parser()
112
+ main(vars(parser.parse_args()))