| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from vllm import LLM, SamplingParams |
| 5 | from vllm.distributed import cleanup_dist_env_and_memory |
| 6 | |
| 7 | # NOTE: This is just a running example. For benchmarking purpose, |
| 8 | # please see benchmarks/benchmark_prefix_caching.py |
| 9 | |
| 10 | # Common prefix. |
| 11 | prefix = ( |
| 12 | "You are an expert school principal, skilled in effectively managing " |
| 13 | "faculty and staff. Draft 10-15 questions for a potential first grade " |
| 14 | "Head Teacher for my K-12, all-girls', independent school that emphasizes " |
| 15 | "community, joyful discovery, and life-long learning. The candidate is " |
| 16 | "coming in for a first-round panel interview for a 8th grade Math " |
| 17 | "teaching role. They have 5 years of previous teaching experience " |
| 18 | "as an assistant teacher at a co-ed, public school with experience " |
| 19 | "in middle school math teaching. Based on these information, fulfill " |
| 20 | "the following paragraph: " |
| 21 | ) |
| 22 | |
| 23 | # Sample prompts. |
| 24 | prompts = [ |
| 25 | "Hello, my name is", |
| 26 | "The president of the United States is", |
| 27 | "The capital of France is", |
| 28 | "The future of AI is", |
| 29 | ] |
| 30 | |
| 31 | generating_prompts = [prefix + prompt for prompt in prompts] |
| 32 | |
| 33 | # Create a sampling params object. |
| 34 | sampling_params = SamplingParams(temperature=0.0) |
| 35 | |
| 36 | |
| 37 | def main(): |
| 38 | # Create an LLM without prefix caching as a baseline. |
| 39 | regular_llm = LLM(model="facebook/opt-125m", gpu_memory_utilization=0.4) |
| 40 | |
| 41 | print("Results without `enable_prefix_caching`") |
| 42 | |
| 43 | # ruff: noqa: E501 |
| 44 | # Generate texts from the prompts. The output is a list of RequestOutput objects |
| 45 | # that contain the prompt, generated text, and other information. |
| 46 | outputs = regular_llm.generate(generating_prompts, sampling_params) |
| 47 | |
| 48 | regular_generated_texts = [] |
| 49 | # Print the outputs. |
| 50 | print("-" * 50) |
| 51 | for output in outputs: |
| 52 | prompt = output.prompt |
| 53 | generated_text = output.outputs[0].text |
| 54 | regular_generated_texts.append(generated_text) |
| 55 | print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}") |
| 56 | print("-" * 50) |
| 57 | |
| 58 | # Destroy the LLM object and free up the GPU memory. |
| 59 | del regular_llm |
| 60 | cleanup_dist_env_and_memory() |
| 61 | |
| 62 | # Create an LLM with prefix caching enabled. |
| 63 | prefix_cached_llm = LLM( |
| 64 | model="facebook/opt-125m", |
| 65 | enable_prefix_caching=True, |
| 66 | gpu_memory_utilization=0.4, |
| 67 | ) |
| 68 | |
| 69 | # Warmup so that the shared prompt's KV cache is computed. |
| 70 | prefix_cached_llm.generate(generating_prompts[0], sampling_params) |
| 71 | |
| 72 | # Generate with prefix caching. |
| 73 | outputs = prefix_cached_llm.generate(generating_prompts, sampling_params) |
| 74 | |
| 75 | print("Results with `enable_prefix_caching`") |
| 76 | |
| 77 | cached_generated_texts = [] |
| 78 | # Print the outputs. You should see the same outputs as before. |
| 79 | print("-" * 50) |
| 80 | for output in outputs: |
| 81 | prompt = output.prompt |
| 82 | generated_text = output.outputs[0].text |
| 83 | cached_generated_texts.append(generated_text) |
| 84 | print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}") |
| 85 | print("-" * 50) |
| 86 | |
| 87 | # Compare the results and display the speedup |
| 88 | generated_same = all( |
| 89 | [ |
| 90 | regular_generated_texts[i] == cached_generated_texts[i] |
| 91 | for i in range(len(prompts)) |
| 92 | ] |
| 93 | ) |
| 94 | print(f"Generated answers are the same: {generated_same}") |
| 95 | |
| 96 | |
| 97 | if __name__ == "__main__": |
| 98 | main() |