From bd0e7802e09e40060b857b85227fd93a6a739467 Mon Sep 17 00:00:00 2001 From: Zhuohan Li Date: Mon, 3 Jun 2024 19:36:41 -0700 Subject: [PATCH] [Bugfix] Add warmup for prefix caching example (#5235) --- examples/offline_inference_with_prefix.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/offline_inference_with_prefix.py b/examples/offline_inference_with_prefix.py index 166e9854..04c28437 100644 --- a/examples/offline_inference_with_prefix.py +++ b/examples/offline_inference_with_prefix.py @@ -51,8 +51,10 @@ for output in outputs: print("-" * 80) -# The llm.generate call will batch all prompts and send the batch at once -# if resources allow. +# Warmup so that the shared prompt's KV cache is computed. +prefix_cached_llm.generate(generating_prompts[0], sampling_params) + +# Generate with prefix caching. start_time_cached = time() outputs = prefix_cached_llm.generate(generating_prompts, sampling_params) duration_cached = time() - start_time_cached