Add HPU support to benchmark_latency and benchmark_throughput (#49)

HabanaAI · Jun 5, 2024 · 9bb5d20 · 9bb5d20
1 parent 1c5d12e
commit 9bb5d20
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/benchmarks/benchmark_latency.py b/benchmarks/benchmark_latency.py
@@ -170,8 +170,8 @@ def run_to_completion(profile_dir: Optional[str] = None):
         "--device",
         type=str,
         default="cuda",
-        choices=["cuda", "cpu"],
-        help='device type for vLLM execution, supporting CUDA and CPU.')
+        choices=["cuda", "cpu", "hpu"],
+        help='device type for vLLM execution, supporting CUDA, CPU and HPU.')
     parser.add_argument('--block-size',
                         type=int,
                         default=16,

diff --git a/benchmarks/benchmark_throughput.py b/benchmarks/benchmark_throughput.py
@@ -334,8 +334,8 @@ def main(args: argparse.Namespace):
         "--device",
         type=str,
         default="cuda",
-        choices=["cuda", "cpu"],
-        help='device type for vLLM execution, supporting CUDA and CPU.')
+        choices=["cuda", "cpu", "hpu"],
+        help='device type for vLLM execution, supporting CUDA, CPU and HPU.')
     parser.add_argument(
         "--enable-prefix-caching",
         action='store_true',