f849ee739c · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -0,0 +1,380 @@

1 +#!/usr/bin/env python3

2 +# SPDX-License-Identifier: Apache-2.0

3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +"""

5 +Benchmark to measure the performance overhead of VLLM_BATCH_INVARIANT mode.

6 +

7 +This benchmark runs the same workload twice:

8 +1. With VLLM_BATCH_INVARIANT=0 (baseline)

9 +2. With VLLM_BATCH_INVARIANT=1 (batch invariant mode)

10 +

11 +And reports the timing and throughput metrics for comparison.

12 +

13 +Environment variables:

14 + VLLM_BENCH_MODEL: Model to benchmark (default: "Qwen/Qwen3-1.7B")

15 + VLLM_BENCH_TP_SIZE: Tensor parallel size (default: 1, use 8 for deepseek)

16 + VLLM_BENCH_BATCH_SIZE: Max batch size (default: 128)

17 + VLLM_BENCH_NUM_TRIALS: Number of trials to run (default: 5)

18 + VLLM_BENCH_MIN_PROMPT: Min prompt length in words (default: 1024)

19 + VLLM_BENCH_MAX_PROMPT: Max prompt length in words (default: 2048)

20 + VLLM_BENCH_MAX_TOKENS: Max tokens to generate (default: 128)

21 + VLLM_BENCH_TEMPERATURE: Temperature for sampling (default: 0.0)

22 + VLLM_BENCH_GPU_MEMORY_UTILIZATION: GPU memory utilization (default: 0.4)

23 + VLLM_BENCH_MAX_MODEL_LEN: Max model length (default: 5120)

24 + VLLM_BENCH_BACKEND: Attention backend (default: FLASH_ATTN)

25 +

26 +Example usage:

27 + # Benchmark qwen3 (default)

28 + python benchmarks/benchmark_batch_invariance.py

29 +

30 + # Benchmark deepseek with 8 GPUs

31 + VLLM_BENCH_MODEL="deepseek-ai/DeepSeek-V3" VLLM_BENCH_TP_SIZE=8 \\

32 + python benchmarks/benchmark_batch_invariance.py

33 +

34 + # Quick test with fewer trials

35 + VLLM_BENCH_NUM_TRIALS=2 VLLM_BENCH_BATCH_SIZE=32 \\

36 + python benchmarks/benchmark_batch_invariance.py

37 +"""

38 +

39 +import contextlib

40 +import os

41 +import random

42 +import time

43 +

44 +from vllm import LLM, SamplingParams

45 +from vllm.platforms import current_platform

46 +

47 +

48 +def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:

49 + """Generate a random prompt for benchmarking."""

50 + prompt_templates = [

51 + "Question: What is the capital of France?\nAnswer: The capital of France is",

52 + "Q: How does photosynthesis work?\nA: Photosynthesis is the process by which",

53 + "User: Can you explain quantum mechanics?\nAssistant: Quantum mechanics is",

54 + "Once upon a time in a distant galaxy, there lived",

55 + "The old man walked slowly down the street, remembering",

56 + "In the year 2157, humanity finally discovered",

57 + "To implement a binary search tree in Python, first we need to",

58 + "The algorithm works by iterating through the array and",

59 + "Here's how to optimize database queries using indexing:",

60 + "The Renaissance was a period in European history that",

61 + "Climate change is caused by several factors including",

62 + "The human brain contains approximately 86 billion neurons which",

63 + "I've been thinking about getting a new laptop because",

64 + "Yesterday I went to the store and bought",

65 + "My favorite thing about summer is definitely",

66 + ]

67 +

68 + base_prompt = random.choice(prompt_templates)

69 +

70 + if max_words < min_words:

71 + max_words = min_words

72 + target_words = random.randint(min_words, max_words)

73 +

74 + if target_words > 50:

75 + padding_text = (

76 + " This is an interesting topic that deserves more explanation. "

77 + * (target_words // 50)

78 + )

79 + base_prompt = base_prompt + padding_text

80 +

81 + return base_prompt

82 +

83 +

84 +def run_benchmark_with_batch_invariant(

85 + model: str,

86 + tp_size: int,

87 + max_batch_size: int,

88 + num_trials: int,

89 + min_prompt: int,

90 + max_prompt: int,

91 + max_tokens: int,

92 + temperature: float,

93 + gpu_mem_util: float,

94 + max_model_len: int,

95 + backend: str,

96 + batch_invariant: bool,

97 + seed: int = 12345,

98 +) -> dict:

99 + """

100 + Run the benchmark with the specified configuration.

101 +

102 + Returns a dict with timing and throughput metrics.

103 + """

104 + random.seed(seed)

105 +

106 + # Set environment variables

107 + os.environ["VLLM_ATTENTION_BACKEND"] = backend

108 + if batch_invariant:

109 + os.environ["VLLM_BATCH_INVARIANT"] = "1"

110 + else:

111 + os.environ["VLLM_BATCH_INVARIANT"] = "0"

112 +

113 + print(f"\n{'=' * 80}")

114 + print(f"BENCHMARK: VLLM_BATCH_INVARIANT={int(batch_invariant)}")

115 + print(f" Model: {model}")

116 + print(f" TP Size: {tp_size}")

117 + print(f" Backend: {backend}")

118 + print(f" Max Batch Size: {max_batch_size}")

119 + print(f" Trials: {num_trials}")

120 + print(f" Max Tokens: {max_tokens}")

121 + print(f"{'=' * 80}\n")

122 +

123 + sampling = SamplingParams(

124 + temperature=temperature,

125 + top_p=0.95,

126 + max_tokens=max_tokens,

127 + seed=20240919,

128 + )

129 +

130 + needle_prompt = "There once was a "

131 +

132 + llm = None

133 + try:

134 + # Create LLM engine

135 + start_init = time.perf_counter()

136 + llm = LLM(

137 + model=model,

138 + max_num_seqs=max_batch_size,

139 + gpu_memory_utilization=gpu_mem_util,

140 + max_model_len=max_model_len,

141 + dtype="bfloat16",

142 + tensor_parallel_size=tp_size,

143 + enable_prefix_caching=False,

144 + )

145 + init_time = time.perf_counter() - start_init

146 + print(f"Engine initialization time: {init_time:.2f}s\n")

147 +

148 + # Generate baseline

149 + print("Generating baseline (warmup)...")

150 + baseline_out = llm.generate([needle_prompt], sampling)

151 + assert len(baseline_out) == 1

152 + baseline_text = baseline_out[0].outputs[0].text

153 + print(f"Baseline output: '{baseline_text[:50]}...'\n")

154 +

155 + # Run trials and measure timing

156 + trial_times: list[float] = []

157 + total_tokens = 0

158 + total_prompts = 0

159 +

160 + for trial in range(num_trials):

161 + # Create a batch

162 + prompts: list[str] = []

163 + batch_size = random.randint(max_batch_size // 2, max_batch_size)

164 + needle_pos = random.randint(0, batch_size - 1)

165 + for i in range(batch_size):

166 + if i == needle_pos:

167 + prompts.append(needle_prompt)

168 + else:

169 + prompts.append(_random_prompt(min_prompt, max_prompt))

170 +

171 + # Measure time for this trial

172 + start_time = time.perf_counter()

173 + outputs = llm.generate(prompts, sampling)

174 + trial_time = time.perf_counter() - start_time

175 +

176 + trial_times.append(trial_time)

177 + total_prompts += len(prompts)

178 +

179 + # Count tokens

180 + for output in outputs:

181 + if output.outputs:

182 + total_tokens += len(output.outputs[0].token_ids)

183 +

184 + print(

185 + f"Trial {trial + 1}/{num_trials}: "

186 + f"batch_size={batch_size}, "

187 + f"time={trial_time:.2f}s"

188 + )

189 +

190 + # Verify needle output still matches

191 + needle_output = outputs[needle_pos]

192 + assert needle_output.prompt == needle_prompt

193 +

194 + # Compute statistics

195 + avg_time = sum(trial_times) / len(trial_times)

196 + min_time = min(trial_times)

197 + max_time = max(trial_times)

198 + throughput = total_tokens / sum(trial_times)

199 + prompts_per_sec = total_prompts / sum(trial_times)

200 +

201 + print(f"\n{'=' * 80}")

202 + print("RESULTS:")

203 + print(f" Average time per trial: {avg_time:.2f}s")

204 + print(f" Min time: {min_time:.2f}s")

205 + print(f" Max time: {max_time:.2f}s")

206 + print(f" Total tokens generated: {total_tokens}")

207 + print(f" Total prompts processed: {total_prompts}")

208 + print(f" Throughput: {throughput:.2f} tokens/s")

209 + print(f" Prompts/s: {prompts_per_sec:.2f}")

210 + print(f"{'=' * 80}\n")

211 +

212 + return {

213 + "init_time": init_time,

214 + "avg_time": avg_time,

215 + "min_time": min_time,

216 + "max_time": max_time,

217 + "total_tokens": total_tokens,

218 + "total_prompts": total_prompts,

219 + "throughput": throughput,

220 + "prompts_per_sec": prompts_per_sec,

221 + "trial_times": trial_times,

222 + }

223 +

224 + finally:

225 + # Cleanup

226 + if llm is not None:

227 + with contextlib.suppress(Exception):

228 + llm.shutdown()

229 +

230 +

231 +def main():

232 + # Check platform support

233 + if not (current_platform.is_cuda() and current_platform.has_device_capability(90)):

234 + print("ERROR: Requires CUDA and >= Hopper (SM90)")

235 + print(f"Current platform: {current_platform.device_type}")

236 + if current_platform.is_cuda():

237 + print(f"Device capability: {current_platform.get_device_capability()}")

238 + return 1

239 +

240 + # Read configuration from environment

241 + model = os.getenv("VLLM_BENCH_MODEL", "Qwen/Qwen3-1.7B")

242 + tp_size = int(os.getenv("VLLM_BENCH_TP_SIZE", "1"))

243 + max_batch_size = int(os.getenv("VLLM_BENCH_BATCH_SIZE", "128"))

244 + num_trials = int(os.getenv("VLLM_BENCH_NUM_TRIALS", "5"))

245 + min_prompt = int(os.getenv("VLLM_BENCH_MIN_PROMPT", "1024"))

246 + max_prompt = int(os.getenv("VLLM_BENCH_MAX_PROMPT", "2048"))

247 + max_tokens = int(os.getenv("VLLM_BENCH_MAX_TOKENS", "128"))

248 + temperature = float(os.getenv("VLLM_BENCH_TEMPERATURE", "0.0"))

249 + gpu_mem_util = float(os.getenv("VLLM_BENCH_GPU_MEMORY_UTILIZATION", "0.4"))

250 + max_model_len = int(os.getenv("VLLM_BENCH_MAX_MODEL_LEN", "5120"))

251 + backend = os.getenv("VLLM_BENCH_BACKEND", "FLASH_ATTN")

252 +

253 + print("\n" + "=" * 80)

254 + print("VLLM BATCH INVARIANCE BENCHMARK")

255 + print("=" * 80)

256 + print("\nConfiguration:")

257 + print(f" Model: {model}")

258 + print(f" Tensor Parallel Size: {tp_size}")

259 + print(f" Attention Backend: {backend}")

260 + print(f" Max Batch Size: {max_batch_size}")

261 + print(f" Number of Trials: {num_trials}")

262 + print(f" Prompt Length Range: {min_prompt}-{max_prompt} words")

263 + print(f" Max Tokens to Generate: {max_tokens}")

264 + print(f" Temperature: {temperature}")

265 + print(f" GPU Memory Utilization: {gpu_mem_util}")

266 + print(f" Max Model Length: {max_model_len}")

267 + print("=" * 80)

268 +

269 + # Run benchmark WITHOUT batch invariance (baseline)

270 + print("\n" + "=" * 80)

271 + print("PHASE 1: Running WITHOUT batch invariance (baseline)")

272 + print("=" * 80)

273 + baseline_results = run_benchmark_with_batch_invariant(

274 + model=model,

275 + tp_size=tp_size,

276 + max_batch_size=max_batch_size,

277 + num_trials=num_trials,

278 + min_prompt=min_prompt,

279 + max_prompt=max_prompt,

280 + max_tokens=max_tokens,

281 + temperature=temperature,

282 + gpu_mem_util=gpu_mem_util,

283 + max_model_len=max_model_len,

284 + backend=backend,

285 + batch_invariant=False,

286 + )

287 +

288 + # Run benchmark WITH batch invariance

289 + print("\n" + "=" * 80)

290 + print("PHASE 2: Running WITH batch invariance")

291 + print("=" * 80)

292 + batch_inv_results = run_benchmark_with_batch_invariant(

293 + model=model,

294 + tp_size=tp_size,

295 + max_batch_size=max_batch_size,

296 + num_trials=num_trials,

297 + min_prompt=min_prompt,

298 + max_prompt=max_prompt,

299 + max_tokens=max_tokens,

300 + temperature=temperature,

301 + gpu_mem_util=gpu_mem_util,

302 + max_model_len=max_model_len,

303 + backend=backend,

304 + batch_invariant=True,

305 + )

306 +

307 + # Compare results

308 + print("\n" + "=" * 80)

309 + print("COMPARISON: Batch Invariance vs Baseline")

310 + print("=" * 80)

311 +

312 + init_overhead_pct = (

313 + (batch_inv_results["init_time"] - baseline_results["init_time"])

314 + / baseline_results["init_time"]

315 + * 100

316 + )

317 + time_overhead_pct = (

318 + (batch_inv_results["avg_time"] - baseline_results["avg_time"])

319 + / baseline_results["avg_time"]

320 + * 100

321 + )

322 + throughput_change_pct = (

323 + (batch_inv_results["throughput"] - baseline_results["throughput"])

324 + / baseline_results["throughput"]

325 + * 100

326 + )

327 +

328 + print("\nInitialization Time:")

329 + print(f" Baseline: {baseline_results['init_time']:.2f}s")

330 + print(f" Batch Invariant: {batch_inv_results['init_time']:.2f}s")

331 + print(f" Overhead: {init_overhead_pct:+.2f}%")

332 +

333 + print("\nAverage Trial Time:")

334 + print(f" Baseline: {baseline_results['avg_time']:.2f}s")

335 + print(f" Batch Invariant: {batch_inv_results['avg_time']:.2f}s")

336 + print(f" Overhead: {time_overhead_pct:+.2f}%")

337 +

338 + print("\nThroughput (tokens/s):")

339 + print(f" Baseline: {baseline_results['throughput']:.2f}")

340 + print(f" Batch Invariant: {batch_inv_results['throughput']:.2f}")

341 + print(f" Change: {throughput_change_pct:+.2f}%")

342 +

343 + print("\nPrompts/s:")

344 + print(f" Baseline: {baseline_results['prompts_per_sec']:.2f}")

345 + print(f" Batch Invariant: {batch_inv_results['prompts_per_sec']:.2f}")

346 +

347 + print("\n" + "=" * 80)

348 + print("SUMMARY")

349 + print("=" * 80)

350 + if time_overhead_pct > 0:

351 + print(

352 + f"Batch invariance mode adds approximately {time_overhead_pct:.1f}% "

353 + "overhead"

354 + )

355 + else:

356 + print(

357 + f"Batch invariance mode is approximately {-time_overhead_pct:.1f}% "

358 + "faster (unexpected!)"

359 + )

360 +

361 + if abs(throughput_change_pct) < 1.0:

362 + print("Throughput difference is negligible (< 1%)")

363 + elif throughput_change_pct < 0:

364 + print(

365 + f"Throughput decreased by {-throughput_change_pct:.1f}% "

366 + "with batch invariance"

367 + )

368 + else:

369 + print(

370 + f"Throughput increased by {throughput_change_pct:.1f}% "

371 + "with batch invariance (unexpected!)"

372 + )

373 +

374 + print("=" * 80 + "\n")

375 +

376 + return 0

377 +

378 +

379 +if __name__ == "__main__":

380 + exit(main())

Adding a benchmark for batch invariance (#28161)