@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 358 lines 12.1 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Benchmark and regression-test pinned (page-locked) CPU memory for vLLM.
4
5 Verifies that enabling pinned memory does not regress throughput or latency
6 compared to unpinned memory. Each condition runs in an isolated ``spawn``
7 subprocess so both start from a cold CUDA context, giving an unbiased
8 comparison.
9
10 Usage
11 -----
12 Run all tests with the default model::
13
14 python benchmarks/benchmark_pin_memory.py -v
15
16 Override the model and optional max-model-len::
17
18 python benchmarks/benchmark_pin_memory.py --model unsloth/Qwen3-1.7B -v
19 python benchmarks/benchmark_pin_memory.py --model unsloth/Qwen3-1.7B \
20 --max-model-len 8192 -v
21
22 Run only throughput or latency tests::
23
24 python benchmarks/benchmark_pin_memory.py -v -k test_throughput
25 python benchmarks/benchmark_pin_memory.py -v -k test_latency
26
27 Run only the v1 or v2 runner variant::
28
29 python benchmarks/benchmark_pin_memory.py -v -k v1
30 python benchmarks/benchmark_pin_memory.py -v -k v2
31
32 Note: on WSL2, v1 runner tests are skipped because pin memory is not available
33 for the v1 runner without cpu_offload_gb. Run on other platforms to exercise v1.
34 """
35
36 import argparse
37 import json
38 import multiprocessing
39 import sys
40 import tempfile
41
42 import pytest
43
44 # Allow up to 2% degradation. Both benchmark runs start from an identical
45 # cold CUDA context (separate spawn subprocesses), so the measured difference
46 # reflects the genuine pin_memory overhead rather than cold/warm ordering bias.
47 _THROUGHPUT_TOLERANCE = 0.98
48 _THROUGHPUT_NUM_REQUESTS = 200
49 _THROUGHPUT_INPUT_LEN = 128
50 _THROUGHPUT_OUTPUT_LEN = 512
51 _THROUGHPUT_MAX_NUM_SEQS = 128
52
53 # Latency benchmark constants — match latency.py defaults.
54 _LATENCY_TOLERANCE = 1.02 # Allow up to 2% latency regression.
55 _LATENCY_BATCH_SIZE = 64
56 _LATENCY_INPUT_LEN = 32
57 _LATENCY_OUTPUT_LEN = 128
58 _LATENCY_WARMUP_ITERS = 5
59 _LATENCY_BENCH_ITERS = 15
60
61 _DEFAULT_MODEL = "unsloth/Qwen3-1.7B"
62 _DEFAULT_MAX_MODEL_LEN = 16384
63
64
65 def _benchmark_args() -> argparse.Namespace:
66 parser = argparse.ArgumentParser(add_help=False)
67 parser.add_argument("--model", default=_DEFAULT_MODEL)
68 parser.add_argument("--max-model-len", type=int, default=_DEFAULT_MAX_MODEL_LEN)
69 args, _ = parser.parse_known_args()
70 return args
71
72
73 @pytest.fixture
74 def model() -> str:
75 return _benchmark_args().model
76
77
78 @pytest.fixture
79 def max_model_len() -> int:
80 return _benchmark_args().max_model_len
81
82
83 def _skip_if_pin_memory_not_available(engine_args_kwargs: dict) -> None:
84 """Skip the current pytest test if pin_memory is unavailable for this config."""
85 import vllm.utils.platform_utils as pu
86 from vllm.config import set_current_vllm_config
87 from vllm.engine.arg_utils import EngineArgs
88
89 vllm_config = EngineArgs(**engine_args_kwargs).create_engine_config()
90 with set_current_vllm_config(vllm_config):
91 pu.is_pin_memory_available.cache_clear()
92 if not pu.is_pin_memory_available():
93 import os
94
95 runner = "v2" if os.environ.get("VLLM_USE_V2_MODEL_RUNNER") == "1" else "v1"
96 model = engine_args_kwargs.get("model", "unknown")
97 print(
98 f"\033[33mSKIP: pin_memory not available for "
99 f"{runner} runner, model={model}\033[0m"
100 )
101 pytest.skip("pin_memory not available for this configuration")
102
103
104 def _throughput_worker(
105 pin: bool,
106 engine_args_kwargs: dict,
107 q: "multiprocessing.Queue[float]",
108 v2_mode: bool = False,
109 ) -> None:
110 """Run throughput benchmark in a fresh spawn subprocess.
111
112 Delegates to vllm/benchmarks/throughput.py main() using the random dataset,
113 so the methodology matches the official benchmark. Results are written to a
114 temp JSON file and forwarded through the queue as tokens/s.
115
116 v2_mode: when True, monkeypatches is_uva_available() to always return True
117 so the v2 model runner's UVA buffers remain functional even when pin=False.
118 This isolates the non-UVA pin_memory paths in v2.
119 """
120 import vllm.utils.platform_utils as pu
121 from vllm.platforms import current_platform
122
123 pu.is_pin_memory_available.cache_clear()
124 pu.is_uva_available.cache_clear()
125 type(current_platform).is_pin_memory_available = classmethod(lambda cls: pin)
126 if v2_mode:
127 pu.is_uva_available = lambda: True
128
129 from vllm.benchmarks.throughput import add_cli_args
130 from vllm.benchmarks.throughput import main as throughput_main
131
132 parser = argparse.ArgumentParser()
133 add_cli_args(parser)
134 args = parser.parse_args([])
135
136 for key, val in engine_args_kwargs.items():
137 setattr(args, key, val)
138 args.max_num_seqs = _THROUGHPUT_MAX_NUM_SEQS
139 args.dataset_name = "random"
140 args.input_len = _THROUGHPUT_INPUT_LEN
141 args.output_len = _THROUGHPUT_OUTPUT_LEN
142 # Nullify defaults that conflict with explicit input/output_len.
143 args.random_input_len = None
144 args.random_output_len = None
145 args.random_prefix_len = None
146 args.num_prompts = _THROUGHPUT_NUM_REQUESTS
147 args.seed = 0
148 args.disable_detokenize = True
149
150 with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
151 tmp_path = f.name
152 args.output_json = tmp_path
153
154 throughput_main(args)
155
156 with open(tmp_path) as f:
157 results = json.load(f)
158 q.put(results["tokens_per_second"])
159
160
161 def _run_throughput_benchmark(
162 pin: bool,
163 engine_args_kwargs: dict,
164 v2_mode: bool = False,
165 ) -> float:
166 ctx = multiprocessing.get_context("spawn")
167 q = ctx.Queue()
168 p = ctx.Process(
169 target=_throughput_worker,
170 args=(pin, engine_args_kwargs, q, v2_mode),
171 )
172 p.start()
173 p.join()
174 if p.exitcode != 0:
175 raise RuntimeError(
176 f"Throughput benchmark subprocess (pin={pin}) exited with code {p.exitcode}"
177 )
178 return q.get()
179
180
181 def _latency_worker(
182 pin: bool,
183 engine_args_kwargs: dict,
184 q: "multiprocessing.Queue[dict]",
185 v2_mode: bool = False,
186 ) -> None:
187 """Run latency benchmark in a fresh spawn subprocess.
188
189 Follows latency.py methodology: fixed batch of dummy token IDs, warmup
190 iterations to reach steady state, then timed iterations reduced to avg
191 and percentiles. Results are written to a temp JSON file by latency_main
192 and forwarded through the queue.
193 """
194 import vllm.utils.platform_utils as pu
195 from vllm.platforms import current_platform
196
197 pu.is_pin_memory_available.cache_clear()
198 pu.is_uva_available.cache_clear()
199 type(current_platform).is_pin_memory_available = classmethod(lambda cls: pin)
200 if v2_mode:
201 pu.is_uva_available = lambda: True
202
203 from vllm.benchmarks.latency import add_cli_args
204 from vllm.benchmarks.latency import main as latency_main
205
206 parser = argparse.ArgumentParser()
207 add_cli_args(parser)
208 args = parser.parse_args([])
209
210 for key, val in engine_args_kwargs.items():
211 setattr(args, key, val)
212 args.input_len = _LATENCY_INPUT_LEN
213 args.output_len = _LATENCY_OUTPUT_LEN
214 args.batch_size = _LATENCY_BATCH_SIZE
215 args.num_iters_warmup = _LATENCY_WARMUP_ITERS
216 args.num_iters = _LATENCY_BENCH_ITERS
217 args.profile = False
218 args.disable_detokenize = True
219
220 with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
221 tmp_path = f.name
222 args.output_json = tmp_path
223
224 latency_main(args)
225
226 with open(tmp_path) as f:
227 results = json.load(f)
228 q.put(results)
229
230
231 def _run_latency_benchmark(
232 pin: bool,
233 engine_args_kwargs: dict,
234 v2_mode: bool = False,
235 ) -> dict:
236 ctx = multiprocessing.get_context("spawn")
237 q = ctx.Queue()
238 p = ctx.Process(
239 target=_latency_worker,
240 args=(pin, engine_args_kwargs, q, v2_mode),
241 )
242 p.start()
243 p.join()
244 if p.exitcode != 0:
245 raise RuntimeError(
246 f"Latency benchmark subprocess (pin={pin}) exited with code {p.exitcode}"
247 )
248 return q.get()
249
250
251 @pytest.mark.parametrize(
252 "test_v2_runner",
253 [
254 pytest.param(False, id="v1"),
255 pytest.param(True, id="v2"),
256 ],
257 )
258 class TestPinnedMemory:
259 """Verify pinned memory yields >= throughput vs unpinned via real vLLM inference."""
260
261 def test_throughput(self, monkeypatch, test_v2_runner, model, max_model_len):
262 """Benchmark throughput with pin_memory forced on then off.
263
264 Delegates to vllm/benchmarks/throughput.py main() with the random
265 dataset. Each condition runs in an isolated spawn subprocess so both
266 start from a cold CUDA context, giving an unbiased comparison.
267 """
268 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
269 monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if test_v2_runner else "0")
270
271 engine_args_kwargs = dict(
272 model=model,
273 gpu_memory_utilization=0.88,
274 max_model_len=max_model_len,
275 enable_prefix_caching=False,
276 )
277
278 _skip_if_pin_memory_not_available(engine_args_kwargs)
279
280 unpinned_tps = _run_throughput_benchmark(
281 False, engine_args_kwargs, v2_mode=test_v2_runner
282 )
283 pinned_tps = _run_throughput_benchmark(
284 True, engine_args_kwargs, v2_mode=test_v2_runner
285 )
286
287 pct_diff = (pinned_tps - unpinned_tps) / unpinned_tps * 100
288 runner = "v2" if test_v2_runner else "v1"
289 print(
290 f"\n=== Throughput results ({runner} runner, {model}) ==="
291 f"\npin_memory=True: {pinned_tps:.1f} tok/s"
292 f"\npin_memory=False: {unpinned_tps:.1f} tok/s"
293 f"\nDifference: {pct_diff:+.1f}% (pinned vs unpinned)"
294 )
295
296 assert pinned_tps >= unpinned_tps * _THROUGHPUT_TOLERANCE, (
297 f"Pinned throughput ({pinned_tps:.1f} tok/s) fell more than "
298 f"{(1.0 - _THROUGHPUT_TOLERANCE) * 100:.1f}% below "
299 f"unpinned ({unpinned_tps:.1f} tok/s)."
300 )
301
302 def test_latency(self, monkeypatch, test_v2_runner, model, max_model_len):
303 """Benchmark per-batch latency with pin_memory forced on then off.
304
305 Follows vllm/benchmarks/latency.py: fixed dummy-token batch, warmup
306 iterations to reach steady state, then timed iterations reduced to avg
307 and percentiles. Subprocesses run serially so each gets a cold CUDA
308 context without GPU memory pressure from the other run.
309 """
310 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
311 monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if test_v2_runner else "0")
312
313 engine_args_kwargs = dict(
314 model=model,
315 gpu_memory_utilization=0.88,
316 max_model_len=max_model_len,
317 enable_prefix_caching=False,
318 )
319
320 _skip_if_pin_memory_not_available(engine_args_kwargs)
321
322 unpinned = _run_latency_benchmark(
323 False, engine_args_kwargs, v2_mode=test_v2_runner
324 )
325 pinned = _run_latency_benchmark(
326 True, engine_args_kwargs, v2_mode=test_v2_runner
327 )
328
329 pct_diff = (
330 (pinned["avg_latency"] - unpinned["avg_latency"])
331 / unpinned["avg_latency"]
332 * 100
333 )
334 runner = "v2" if test_v2_runner else "v1"
335 print(
336 f"\n=== Latency results ({runner} runner, {model}) ==="
337 f"\npin_memory=True: avg={pinned['avg_latency']:.3f}s"
338 f" p50={pinned['percentiles']['50']:.3f}s"
339 f" p99={pinned['percentiles']['99']:.3f}s"
340 f"\npin_memory=False: avg={unpinned['avg_latency']:.3f}s"
341 f" p50={unpinned['percentiles']['50']:.3f}s"
342 f" p99={unpinned['percentiles']['99']:.3f}s"
343 f"\nDifference: {pct_diff:+.1f}% (pinned vs unpinned)"
344 )
345
346 assert pinned["avg_latency"] <= unpinned["avg_latency"] * _LATENCY_TOLERANCE, (
347 f"Pinned avg latency ({pinned['avg_latency']:.3f}s) exceeded "
348 f"unpinned ({unpinned['avg_latency']:.3f}s) by more than "
349 f"{(_LATENCY_TOLERANCE - 1.0) * 100:.1f}%."
350 )
351
352
353 if __name__ == "__main__":
354 _parser = argparse.ArgumentParser(add_help=False)
355 _parser.add_argument("--model", default=_DEFAULT_MODEL)
356 _parser.add_argument("--max-model-len", type=int, default=_DEFAULT_MAX_MODEL_LEN)
357 _, _remaining = _parser.parse_known_args()
358 sys.exit(pytest.main([__file__] + _remaining))