pre-commit

Author: gshtras

benchmarks/profiling/benchmark_latency.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 """Benchmark the latency of processing a single batch of requests."""
+
 import argparse
 import dataclasses
 import json
@@ -26,6 +27,7 @@ def main(args: argparse.Namespace):
     @contextmanager
     def rpd_profiler_context():
         from rpdTracerControl import rpdTracerControl as rpd
+
         llm.start_profile()
         yield
         llm.stop_profile()
@@ -39,15 +41,16 @@ def torch_profiler_context(profile_result_dir: Optional[str] = None):
                 torch.profiler.ProfilerActivity.CUDA,
             ],
             on_trace_ready=torch.profiler.tensorboard_trace_handler(
-                str(profile_result_dir)))
+                str(profile_result_dir)
+            ),
+        )
         p.start()
         try:
             with torch.no_grad():
                 yield p
         finally:
             p.stop()
-            print(p.key_averages().table(sort_by="self_cuda_time_total",
-                                         row_limit=-1))
+            print(p.key_averages().table(sort_by="self_cuda_time_total", row_limit=-1))
 
     def get_profiling_context(profile_result_dir: Optional[str] = None):
         if args.profile_torch:
@@ -58,15 +61,16 @@ def get_profiling_context(profile_result_dir: Optional[str] = None):
             return nullcontext()
 
     if args.profile_torch or args.profile_rpd:
-        profile_result_dir = Path(args.profile_result_dir
-                                  or "./vllm_benchmark_latency_result")
+        profile_result_dir = Path(
+            args.profile_result_dir or "./vllm_benchmark_latency_result"
+        )
         profile_result_dir.mkdir(parents=True, exist_ok=True)
         name = os.path.basename(os.path.normpath(args.model))
         model_trace_name = (
             f"{name}_in_{args.input_len}_out_{args.output_len}_"
-            f"batch_{args.batch_size}_tp_{args.tensor_parallel_size}")
-        print(
-            f"Profiling (results will be saved to '{profile_result_dir}')...")
+            f"batch_{args.batch_size}_tp_{args.tensor_parallel_size}"
+        )
+        print(f"Profiling (results will be saved to '{profile_result_dir}')...")
         if args.profile_rpd:
             profile_result_dir /= f"{model_trace_name}.rpd"
             os.environ["VLLM_RPD_PROFILER_DIR"] = str(profile_result_dir)
@@ -85,26 +89,25 @@ def get_profiling_context(profile_result_dir: Optional[str] = None):
         max_tokens=args.output_len,
     )
     print(sampling_params)
-    dummy_prompt_token_ids = np.random.randint(10000,
-                                               size=(args.batch_size,
-                                                     args.input_len))
-    dummy_prompts: list[PromptType] = [{
-        "prompt_token_ids": batch
-    } for batch in dummy_prompt_token_ids.tolist()]
+    dummy_prompt_token_ids = np.random.randint(
+        10000, size=(args.batch_size, args.input_len)
+    )
+    dummy_prompts: list[PromptType] = [
+        {"prompt_token_ids": batch} for batch in dummy_prompt_token_ids.tolist()
+    ]
 
     def llm_generate():
         if not args.use_beam_search:
-            llm.generate(dummy_prompts,
-                         sampling_params=sampling_params,
-                         use_tqdm=False)
+            llm.generate(dummy_prompts, sampling_params=sampling_params, use_tqdm=False)
         else:
             llm.beam_search(
                 dummy_prompts,
                 BeamSearchParams(
                     beam_width=args.n,
                     max_tokens=args.output_len,
                     ignore_eos=True,
-                ))
+                ),
+            )
 
     def run_to_completion(profile_dir: Optional[str] = None):
         if profile_dir:
@@ -132,9 +135,9 @@ def run_to_completion(profile_dir: Optional[str] = None):
     latencies = np.array(latencies)
     percentages = [10, 25, 50, 75, 90, 99]
     percentiles = np.percentile(latencies, percentages)
-    print(f'Avg latency: {np.mean(latencies)} seconds')
+    print(f"Avg latency: {np.mean(latencies)} seconds")
     for percentage, percentile in zip(percentages, percentiles):
-        print(f'{percentage}% percentile latency: {percentile} seconds')
+        print(f"{percentage}% percentile latency: {percentile} seconds")
 
     # Output JSON results if specified
     if args.output_json:
@@ -147,45 +150,52 @@ def run_to_completion(profile_dir: Optional[str] = None):
             json.dump(results, f, indent=4)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = FlexibleArgumentParser(
-        description='Benchmark the latency of processing a single batch of '
-        'requests till completion.')
-    parser.add_argument('--input-len', type=int, default=32)
-    parser.add_argument('--output-len', type=int, default=128)
-    parser.add_argument('--batch-size', type=int, default=8)
-    parser.add_argument('--n',
-                        type=int,
-                        default=1,
-                        help='Number of generated sequences per prompt.')
-    parser.add_argument('--use-beam-search', action='store_true')
-    parser.add_argument('--num-iters-warmup',
-                        type=int,
-                        default=10,
-                        help='Number of iterations to run for warmup.')
-    parser.add_argument('--num-iters',
-                        type=int,
-                        default=30,
-                        help='Number of iterations to run.')
+        description="Benchmark the latency of processing a single batch of "
+        "requests till completion."
+    )
+    parser.add_argument("--input-len", type=int, default=32)
+    parser.add_argument("--output-len", type=int, default=128)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument(
+        "--n", type=int, default=1, help="Number of generated sequences per prompt."
+    )
+    parser.add_argument("--use-beam-search", action="store_true")
+    parser.add_argument(
+        "--num-iters-warmup",
+        type=int,
+        default=10,
+        help="Number of iterations to run for warmup.",
+    )
     parser.add_argument(
-        '--profile-torch',
-        action='store_true',
-        help='profile the generation process of a single batch')
+        "--num-iters", type=int, default=30, help="Number of iterations to run."
+    )
     parser.add_argument(
-        '--profile-rpd',
-        action='store_true',
-        help='profile the generation process of a single batch')
+        "--profile-torch",
+        action="store_true",
+        help="profile the generation process of a single batch",
+    )
     parser.add_argument(
-        '--profile-result-dir',
+        "--profile-rpd",
+        action="store_true",
+        help="profile the generation process of a single batch",
+    )
+    parser.add_argument(
+        "--profile-result-dir",
         type=str,
-        default=os.getenv('VLLM_RPD_PROFILER_DIR', default=None),
-        help=('path to save the profiler output. Can be visualized '
-              'with ui.perfetto.dev or Tensorboard.'))
+        default=os.getenv("VLLM_RPD_PROFILER_DIR", default=None),
+        help=(
+            "path to save the profiler output. Can be visualized "
+            "with ui.perfetto.dev or Tensorboard."
+        ),
+    )
     parser.add_argument(
-        '--output-json',
+        "--output-json",
         type=str,
         default=None,
-        help='Path to save the latency results in JSON format.')
+        help="Path to save the latency results in JSON format.",
+    )
 
     parser = EngineArgs.add_cli_args(parser)
     args = parser.parse_args()