@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 378 lines 13.1 KB
Raw
1 #!/usr/bin/env python3
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 """Benchmark to measure the performance overhead of VLLM_BATCH_INVARIANT mode.
5
6 This benchmark runs the same workload twice:
7 1. With VLLM_BATCH_INVARIANT=0 (baseline)
8 2. With VLLM_BATCH_INVARIANT=1 (batch invariant mode)
9
10 And reports the timing and throughput metrics for comparison.
11
12 Environment variables:
13 VLLM_BENCH_MODEL: Model to benchmark (default: "Qwen/Qwen3-1.7B")
14 VLLM_BENCH_TP_SIZE: Tensor parallel size (default: 1, use 8 for deepseek)
15 VLLM_BENCH_BATCH_SIZE: Max batch size (default: 128)
16 VLLM_BENCH_NUM_TRIALS: Number of trials to run (default: 5)
17 VLLM_BENCH_MIN_PROMPT: Min prompt length in words (default: 1024)
18 VLLM_BENCH_MAX_PROMPT: Max prompt length in words (default: 2048)
19 VLLM_BENCH_MAX_TOKENS: Max tokens to generate (default: 128)
20 VLLM_BENCH_TEMPERATURE: Temperature for sampling (default: 0.0)
21 VLLM_BENCH_GPU_MEMORY_UTILIZATION: GPU memory utilization (default: 0.4)
22 VLLM_BENCH_MAX_MODEL_LEN: Max model length (default: 5120)
23 VLLM_BENCH_BACKEND: Attention backend (default: FLASH_ATTN)
24
25 Example usage:
26 # Benchmark qwen3 (default)
27 python benchmarks/benchmark_batch_invariance.py
28
29 # Benchmark deepseek with 8 GPUs
30 VLLM_BENCH_MODEL="deepseek-ai/DeepSeek-V3" VLLM_BENCH_TP_SIZE=8 \\
31 python benchmarks/benchmark_batch_invariance.py
32
33 # Quick test with fewer trials
34 VLLM_BENCH_NUM_TRIALS=2 VLLM_BENCH_BATCH_SIZE=32 \\
35 python benchmarks/benchmark_batch_invariance.py
36 """
37
38 import contextlib
39 import os
40 import random
41 import time
42
43 from vllm import LLM, SamplingParams
44 from vllm.platforms import current_platform
45
46
47 def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:
48 """Generate a random prompt for benchmarking."""
49 prompt_templates = [
50 "Question: What is the capital of France?\nAnswer: The capital of France is",
51 "Q: How does photosynthesis work?\nA: Photosynthesis is the process by which",
52 "User: Can you explain quantum mechanics?\nAssistant: Quantum mechanics is",
53 "Once upon a time in a distant galaxy, there lived",
54 "The old man walked slowly down the street, remembering",
55 "In the year 2157, humanity finally discovered",
56 "To implement a binary search tree in Python, first we need to",
57 "The algorithm works by iterating through the array and",
58 "Here's how to optimize database queries using indexing:",
59 "The Renaissance was a period in European history that",
60 "Climate change is caused by several factors including",
61 "The human brain contains approximately 86 billion neurons which",
62 "I've been thinking about getting a new laptop because",
63 "Yesterday I went to the store and bought",
64 "My favorite thing about summer is definitely",
65 ]
66
67 base_prompt = random.choice(prompt_templates)
68
69 if max_words < min_words:
70 max_words = min_words
71 target_words = random.randint(min_words, max_words)
72
73 if target_words > 50:
74 padding_text = (
75 " This is an interesting topic that deserves more explanation. "
76 * (target_words // 50)
77 )
78 base_prompt = base_prompt + padding_text
79
80 return base_prompt
81
82
83 def run_benchmark_with_batch_invariant(
84 model: str,
85 tp_size: int,
86 max_batch_size: int,
87 num_trials: int,
88 min_prompt: int,
89 max_prompt: int,
90 max_tokens: int,
91 temperature: float,
92 gpu_mem_util: float,
93 max_model_len: int,
94 backend: str,
95 batch_invariant: bool,
96 seed: int = 12345,
97 ) -> dict:
98 """Run the benchmark with the specified configuration.
99
100 Returns a dict with timing and throughput metrics.
101 """
102 random.seed(seed)
103
104 # Set environment variables
105 if batch_invariant:
106 os.environ["VLLM_BATCH_INVARIANT"] = "1"
107 else:
108 os.environ["VLLM_BATCH_INVARIANT"] = "0"
109
110 print(f"\n{'=' * 80}")
111 print(f"BENCHMARK: VLLM_BATCH_INVARIANT={int(batch_invariant)}")
112 print(f" Model: {model}")
113 print(f" TP Size: {tp_size}")
114 print(f" Backend: {backend}")
115 print(f" Max Batch Size: {max_batch_size}")
116 print(f" Trials: {num_trials}")
117 print(f" Max Tokens: {max_tokens}")
118 print(f"{'=' * 80}\n")
119
120 sampling = SamplingParams(
121 temperature=temperature,
122 top_p=0.95,
123 max_tokens=max_tokens,
124 seed=20240919,
125 )
126
127 needle_prompt = "There once was a "
128
129 llm = None
130 try:
131 # Create LLM engine
132 start_init = time.perf_counter()
133 llm = LLM(
134 model=model,
135 max_num_seqs=max_batch_size,
136 gpu_memory_utilization=gpu_mem_util,
137 max_model_len=max_model_len,
138 dtype="bfloat16",
139 tensor_parallel_size=tp_size,
140 attention_config={"backend": backend},
141 enable_prefix_caching=False,
142 )
143 init_time = time.perf_counter() - start_init
144 print(f"Engine initialization time: {init_time:.2f}s\n")
145
146 # Generate baseline
147 print("Generating baseline (warmup)...")
148 baseline_out = llm.generate([needle_prompt], sampling)
149 assert len(baseline_out) == 1
150 baseline_text = baseline_out[0].outputs[0].text
151 print(f"Baseline output: '{baseline_text[:50]}...'\n")
152
153 # Run trials and measure timing
154 trial_times: list[float] = []
155 total_tokens = 0
156 total_prompts = 0
157
158 for trial in range(num_trials):
159 # Create a batch
160 prompts: list[str] = []
161 batch_size = random.randint(max_batch_size // 2, max_batch_size)
162 needle_pos = random.randint(0, batch_size - 1)
163 for i in range(batch_size):
164 if i == needle_pos:
165 prompts.append(needle_prompt)
166 else:
167 prompts.append(_random_prompt(min_prompt, max_prompt))
168
169 # Measure time for this trial
170 start_time = time.perf_counter()
171 outputs = llm.generate(prompts, sampling)
172 trial_time = time.perf_counter() - start_time
173
174 trial_times.append(trial_time)
175 total_prompts += len(prompts)
176
177 # Count tokens
178 for output in outputs:
179 if output.outputs:
180 total_tokens += len(output.outputs[0].token_ids)
181
182 print(
183 f"Trial {trial + 1}/{num_trials}: "
184 f"batch_size={batch_size}, "
185 f"time={trial_time:.2f}s"
186 )
187
188 # Verify needle output still matches
189 needle_output = outputs[needle_pos]
190 assert needle_output.prompt == needle_prompt
191
192 # Compute statistics
193 avg_time = sum(trial_times) / len(trial_times)
194 min_time = min(trial_times)
195 max_time = max(trial_times)
196 throughput = total_tokens / sum(trial_times)
197 prompts_per_sec = total_prompts / sum(trial_times)
198
199 print(f"\n{'=' * 80}")
200 print("RESULTS:")
201 print(f" Average time per trial: {avg_time:.2f}s")
202 print(f" Min time: {min_time:.2f}s")
203 print(f" Max time: {max_time:.2f}s")
204 print(f" Total tokens generated: {total_tokens}")
205 print(f" Total prompts processed: {total_prompts}")
206 print(f" Throughput: {throughput:.2f} tokens/s")
207 print(f" Prompts/s: {prompts_per_sec:.2f}")
208 print(f"{'=' * 80}\n")
209
210 return {
211 "init_time": init_time,
212 "avg_time": avg_time,
213 "min_time": min_time,
214 "max_time": max_time,
215 "total_tokens": total_tokens,
216 "total_prompts": total_prompts,
217 "throughput": throughput,
218 "prompts_per_sec": prompts_per_sec,
219 "trial_times": trial_times,
220 }
221
222 finally:
223 # Cleanup
224 if llm is not None:
225 with contextlib.suppress(Exception):
226 llm.shutdown()
227
228
229 def main():
230 # Check platform support
231 if not (current_platform.is_cuda() and current_platform.has_device_capability(90)):
232 print("ERROR: Requires CUDA and >= Hopper (SM90)")
233 print(f"Current platform: {current_platform.device_type}")
234 if current_platform.is_cuda():
235 print(f"Device capability: {current_platform.get_device_capability()}")
236 return 1
237
238 # Read configuration from environment
239 model = os.getenv("VLLM_BENCH_MODEL", "Qwen/Qwen3-1.7B")
240 tp_size = int(os.getenv("VLLM_BENCH_TP_SIZE", "1"))
241 max_batch_size = int(os.getenv("VLLM_BENCH_BATCH_SIZE", "128"))
242 num_trials = int(os.getenv("VLLM_BENCH_NUM_TRIALS", "5"))
243 min_prompt = int(os.getenv("VLLM_BENCH_MIN_PROMPT", "1024"))
244 max_prompt = int(os.getenv("VLLM_BENCH_MAX_PROMPT", "2048"))
245 max_tokens = int(os.getenv("VLLM_BENCH_MAX_TOKENS", "128"))
246 temperature = float(os.getenv("VLLM_BENCH_TEMPERATURE", "0.0"))
247 gpu_mem_util = float(os.getenv("VLLM_BENCH_GPU_MEMORY_UTILIZATION", "0.4"))
248 max_model_len = int(os.getenv("VLLM_BENCH_MAX_MODEL_LEN", "5120"))
249 backend = os.getenv("VLLM_BENCH_BACKEND", "FLASH_ATTN")
250
251 print("\n" + "=" * 80)
252 print("VLLM BATCH INVARIANCE BENCHMARK")
253 print("=" * 80)
254 print("\nConfiguration:")
255 print(f" Model: {model}")
256 print(f" Tensor Parallel Size: {tp_size}")
257 print(f" Attention Backend: {backend}")
258 print(f" Max Batch Size: {max_batch_size}")
259 print(f" Number of Trials: {num_trials}")
260 print(f" Prompt Length Range: {min_prompt}-{max_prompt} words")
261 print(f" Max Tokens to Generate: {max_tokens}")
262 print(f" Temperature: {temperature}")
263 print(f" GPU Memory Utilization: {gpu_mem_util}")
264 print(f" Max Model Length: {max_model_len}")
265 print("=" * 80)
266
267 # Run benchmark WITHOUT batch invariance (baseline)
268 print("\n" + "=" * 80)
269 print("PHASE 1: Running WITHOUT batch invariance (baseline)")
270 print("=" * 80)
271 baseline_results = run_benchmark_with_batch_invariant(
272 model=model,
273 tp_size=tp_size,
274 max_batch_size=max_batch_size,
275 num_trials=num_trials,
276 min_prompt=min_prompt,
277 max_prompt=max_prompt,
278 max_tokens=max_tokens,
279 temperature=temperature,
280 gpu_mem_util=gpu_mem_util,
281 max_model_len=max_model_len,
282 backend=backend,
283 batch_invariant=False,
284 )
285
286 # Run benchmark WITH batch invariance
287 print("\n" + "=" * 80)
288 print("PHASE 2: Running WITH batch invariance")
289 print("=" * 80)
290 batch_inv_results = run_benchmark_with_batch_invariant(
291 model=model,
292 tp_size=tp_size,
293 max_batch_size=max_batch_size,
294 num_trials=num_trials,
295 min_prompt=min_prompt,
296 max_prompt=max_prompt,
297 max_tokens=max_tokens,
298 temperature=temperature,
299 gpu_mem_util=gpu_mem_util,
300 max_model_len=max_model_len,
301 backend=backend,
302 batch_invariant=True,
303 )
304
305 # Compare results
306 print("\n" + "=" * 80)
307 print("COMPARISON: Batch Invariance vs Baseline")
308 print("=" * 80)
309
310 init_overhead_pct = (
311 (batch_inv_results["init_time"] - baseline_results["init_time"])
312 / baseline_results["init_time"]
313 * 100
314 )
315 time_overhead_pct = (
316 (batch_inv_results["avg_time"] - baseline_results["avg_time"])
317 / baseline_results["avg_time"]
318 * 100
319 )
320 throughput_change_pct = (
321 (batch_inv_results["throughput"] - baseline_results["throughput"])
322 / baseline_results["throughput"]
323 * 100
324 )
325
326 print("\nInitialization Time:")
327 print(f" Baseline: {baseline_results['init_time']:.2f}s")
328 print(f" Batch Invariant: {batch_inv_results['init_time']:.2f}s")
329 print(f" Overhead: {init_overhead_pct:+.2f}%")
330
331 print("\nAverage Trial Time:")
332 print(f" Baseline: {baseline_results['avg_time']:.2f}s")
333 print(f" Batch Invariant: {batch_inv_results['avg_time']:.2f}s")
334 print(f" Overhead: {time_overhead_pct:+.2f}%")
335
336 print("\nThroughput (tokens/s):")
337 print(f" Baseline: {baseline_results['throughput']:.2f}")
338 print(f" Batch Invariant: {batch_inv_results['throughput']:.2f}")
339 print(f" Change: {throughput_change_pct:+.2f}%")
340
341 print("\nPrompts/s:")
342 print(f" Baseline: {baseline_results['prompts_per_sec']:.2f}")
343 print(f" Batch Invariant: {batch_inv_results['prompts_per_sec']:.2f}")
344
345 print("\n" + "=" * 80)
346 print("SUMMARY")
347 print("=" * 80)
348 if time_overhead_pct > 0:
349 print(
350 f"Batch invariance mode adds approximately {time_overhead_pct:.1f}% "
351 "overhead"
352 )
353 else:
354 print(
355 f"Batch invariance mode is approximately {-time_overhead_pct:.1f}% "
356 "faster (unexpected!)"
357 )
358
359 if abs(throughput_change_pct) < 1.0:
360 print("Throughput difference is negligible (< 1%)")
361 elif throughput_change_pct < 0:
362 print(
363 f"Throughput decreased by {-throughput_change_pct:.1f}% "
364 "with batch invariance"
365 )
366 else:
367 print(
368 f"Throughput increased by {throughput_change_pct:.1f}% "
369 "with batch invariance (unexpected!)"
370 )
371
372 print("=" * 80 + "\n")
373
374 return 0
375
376
377 if __name__ == "__main__":
378 exit(main())