Offline Inference With Prefix

Offline Inference With Prefix#

Source: examples/offline_inference_with_prefix.py.

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