From 55fe1536f7a8858ec04f763a24021bc0a380c109 Mon Sep 17 00:00:00 2001 From: Andy Dai <76841985+Imss27@users.noreply.github.com> Date: Fri, 4 Oct 2024 14:58:57 -0700 Subject: [PATCH] [Misc] Add random seed for prefix cache benchmark (#9081) Signed-off-by: Alvant --- benchmarks/benchmark_prefix_caching.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/benchmarks/benchmark_prefix_caching.py b/benchmarks/benchmark_prefix_caching.py index 3e90fdfb78e10..d0e5b73bb4b90 100644 --- a/benchmarks/benchmark_prefix_caching.py +++ b/benchmarks/benchmark_prefix_caching.py @@ -113,7 +113,7 @@ def repeat_and_sort_requests(requests: List[Tuple[str, int, int]], def main(args): tokenizer = get_tokenizer(args.model, trust_remote_code=True) input_length_range = tuple(map(int, args.input_length_range.split(':'))) - + random.seed(args.seed) if args.dataset_path is not None: print(f"Start to sample {args.num_prompts} prompts" "from {args.dataset_path}") @@ -194,5 +194,7 @@ def main(args): default='128:256', help='Range of input lengths for sampling prompts,' 'specified as "min:max" (e.g., "128:256").') + parser.add_argument("--seed", type=int, default=0, + help='Random seed for reproducibility') args = parser.parse_args() main(args)