From c013d32c758699fbe5804af1b9d9408acd6cb8b7 Mon Sep 17 00:00:00 2001 From: Zedong Peng Date: Wed, 10 Apr 2024 12:30:03 +0800 Subject: [PATCH] [Benchmark] Add cpu options to bench scripts (#3915) --- benchmarks/benchmark_latency.py | 4 ++-- benchmarks/benchmark_throughput.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/benchmarks/benchmark_latency.py b/benchmarks/benchmark_latency.py index e2d358ea..91510daf 100644 --- a/benchmarks/benchmark_latency.py +++ b/benchmarks/benchmark_latency.py @@ -169,8 +169,8 @@ if __name__ == '__main__': "--device", type=str, default="cuda", - choices=["cuda"], - help='device type for vLLM execution, supporting CUDA only currently.') + choices=["cuda", "cpu"], + help='device type for vLLM execution, supporting CUDA and CPU.') parser.add_argument('--block-size', type=int, default=16, diff --git a/benchmarks/benchmark_throughput.py b/benchmarks/benchmark_throughput.py index d6bf18c8..e7133827 100644 --- a/benchmarks/benchmark_throughput.py +++ b/benchmarks/benchmark_throughput.py @@ -329,8 +329,8 @@ if __name__ == "__main__": "--device", type=str, default="cuda", - choices=["cuda"], - help='device type for vLLM execution, supporting CUDA only currently.') + choices=["cuda", "cpu"], + help='device type for vLLM execution, supporting CUDA and CPU.') parser.add_argument( "--enable-prefix-caching", action='store_true',