| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """This file demonstrates the example usage of CPU offloading |
| 4 | with LMCache in vLLM v1. |
| 5 | |
| 6 | Note that `lmcache` is needed to run this example. |
| 7 | Requirements: |
| 8 | https://docs.lmcache.ai/getting_started/installation.html#prerequisites |
| 9 | Learn more about LMCache environment setup, please refer to: |
| 10 | https://docs.lmcache.ai/getting_started/installation.html |
| 11 | """ |
| 12 | |
| 13 | import contextlib |
| 14 | import os |
| 15 | import time |
| 16 | from dataclasses import asdict |
| 17 | |
| 18 | from lmcache.integration.vllm.utils import ENGINE_NAME |
| 19 | from lmcache.v1.cache_engine import LMCacheEngineBuilder |
| 20 | |
| 21 | from vllm import LLM, SamplingParams |
| 22 | from vllm.config import KVTransferConfig |
| 23 | from vllm.engine.arg_utils import EngineArgs |
| 24 | |
| 25 | |
| 26 | def setup_environment_variables(): |
| 27 | # LMCache-related environment variables |
| 28 | # LMCache is set to use 256 tokens per chunk |
| 29 | os.environ["LMCACHE_CHUNK_SIZE"] = "256" |
| 30 | # Enable local CPU backend in LMCache |
| 31 | os.environ["LMCACHE_LOCAL_CPU"] = "True" |
| 32 | # Set local CPU memory limit to 5.0 GB |
| 33 | os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = "5.0" |
| 34 | |
| 35 | |
| 36 | @contextlib.contextmanager |
| 37 | def build_llm_with_lmcache(model: str): |
| 38 | ktc = KVTransferConfig( |
| 39 | kv_connector="LMCacheConnectorV1", |
| 40 | kv_role="kv_both", |
| 41 | ) |
| 42 | # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB |
| 43 | # memory. Reduce the value if your GPU has less memory. |
| 44 | # Note: LMCache supports chunked prefill (see vLLM#14505, LMCache#392). |
| 45 | llm_args = EngineArgs( |
| 46 | model=model, |
| 47 | kv_transfer_config=ktc, |
| 48 | max_model_len=8000, |
| 49 | gpu_memory_utilization=0.8, |
| 50 | ) |
| 51 | |
| 52 | llm = LLM(**asdict(llm_args)) |
| 53 | try: |
| 54 | yield llm |
| 55 | finally: |
| 56 | # Clean up lmcache backend |
| 57 | LMCacheEngineBuilder.destroy(ENGINE_NAME) |
| 58 | |
| 59 | |
| 60 | def print_output( |
| 61 | llm: LLM, |
| 62 | prompt: list[str], |
| 63 | sampling_params: SamplingParams, |
| 64 | req_str: str, |
| 65 | ): |
| 66 | # Should be able to see logs like the following: |
| 67 | # `LMCache INFO: Storing KV cache for 6006 out of 6006 tokens for request 0` |
| 68 | # This indicates that the KV cache has been stored in LMCache. |
| 69 | start = time.time() |
| 70 | outputs = llm.generate(prompt, sampling_params) |
| 71 | print("-" * 50) |
| 72 | for output in outputs: |
| 73 | generated_text = output.outputs[0].text |
| 74 | print(f"Generated text: {generated_text!r}") |
| 75 | print(f"Generation took {time.time() - start:.2f} seconds, {req_str} request done.") |
| 76 | print("-" * 50) |
| 77 | |
| 78 | |
| 79 | def main(): |
| 80 | model = "meta-llama/Meta-Llama-3.1-8B-Instruct" |
| 81 | setup_environment_variables() |
| 82 | with build_llm_with_lmcache(model) as llm: |
| 83 | # This example script runs two requests with a shared prefix. |
| 84 | # Define the shared prompt and specific prompts |
| 85 | shared_prompt = "Hello, how are you?" * 1000 |
| 86 | first_prompt = [ |
| 87 | shared_prompt + "Hello, my name is", |
| 88 | ] |
| 89 | second_prompt = [ |
| 90 | shared_prompt + "Tell me a very long story", |
| 91 | ] |
| 92 | |
| 93 | sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) |
| 94 | |
| 95 | # Print the first output |
| 96 | print_output(llm, first_prompt, sampling_params, "first") |
| 97 | |
| 98 | time.sleep(1) |
| 99 | |
| 100 | # print the second output |
| 101 | print_output(llm, second_prompt, sampling_params, "second") |
| 102 | |
| 103 | |
| 104 | if __name__ == "__main__": |
| 105 | main() |