A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 12:15 UTC
[Misc] refactor example - cpu_offload_lmcache (#17460)
Signed-off-by: reidliu41 <reid201711@gmail.com> Co-authored-by: reidliu41 <reid201711@gmail.com>
Reid committed
May 1, 2025 at 23:05 UTC
7423cf0a9b24ce97f2fee99bcdb9e3abc5b189f7
3 files changed
+64
-70
examples/lmcache/README.md
+2
-2
@@ -44,8 +44,8 @@ The main script generates several log files:
44
45
## 2. CPU Offload Examples
46
47
-- `cpu_offload_lmcache_v0.py` - CPU offloading implementation for vLLM v0
48
-- `cpu_offload_lmcache_v1.py` - CPU offloading implementation for vLLM v1
47
+- `python cpu_offload_lmcache.py -v v0` - CPU offloading implementation for vLLM v0
48
+- `python cpu_offload_lmcache.py -v v1` - CPU offloading implementation for vLLM v1
49
50
## 3. KV Cache Sharing
51
examples/lmcache/cpu_offload_lmcache.py
renamed
+62
-11
@@ -1,22 +1,37 @@
1
# SPDX-License-Identifier: Apache-2.0
2
"""
3
This file demonstrates the example usage of cpu offloading
4
-with LMCache.
4
+with LMCache in vLLM v1 or v0.
5
+
6
+Usage:
7
+
8
+ Specify vLLM version
9
+
10
+ -v v0 : Use LMCacheConnector
11
+ model = mistralai/Mistral-7B-Instruct-v0.2
12
+ (Includes enable_chunked_prefill = True)
13
+
14
+ -v v1 : Use LMCacheConnectorV1 (default)
15
+ model = meta-llama/Meta-Llama-3.1-8B-Instruct
16
+ (Without enable_chunked_prefill)
17
18
Note that `lmcache` is needed to run this example.
19
Requirements: Linux, Python: 3.10 or higher, CUDA: 12.1
20
Learn more about LMCache environment setup, please refer to:
21
https://docs.lmcache.ai/getting_started/installation.html
22
"""
23
+import argparse
24
import contextlib
25
import os
26
import time
27
+from dataclasses import asdict
28
29
from lmcache.experimental.cache_engine import LMCacheEngineBuilder
30
from lmcache.integration.vllm.utils import ENGINE_NAME
31
32
from vllm import LLM, SamplingParams
33
from vllm.config import KVTransferConfig
34
+from vllm.engine.arg_utils import EngineArgs
35
36
37
def setup_environment_variables():
@@ -32,18 +47,32 @@ def setup_environment_variables():
47
48
49
@contextlib.contextmanager
35
-def build_llm_with_lmcache():
36
- ktc = KVTransferConfig.from_cli(
37
- '{"kv_connector":"LMCacheConnector", "kv_role":"kv_both"}')
50
+def build_llm_with_lmcache(lmcache_connector: str, model: str,
51
+ vllm_version: str):
52
+ ktc = KVTransferConfig(
53
+ kv_connector=lmcache_connector,
54
+ kv_role="kv_both",
55
+ )
56
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
57
# memory. Reduce the value if your GPU has less memory.
58
# Note: LMCache supports chunked prefill (see vLLM#14505, LMCache#392).
41
- llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
42
- kv_transfer_config=ktc,
43
- max_model_len=8000,
44
- enable_chunked_prefill=True,
45
- gpu_memory_utilization=0.8)
46
-
59
+ if vllm_version == "v0":
60
+ llm_args = EngineArgs(
61
+ model=model,
62
+ kv_transfer_config=ktc,
63
+ max_model_len=8000,
64
+ gpu_memory_utilization=0.8,
65
+ enable_chunked_prefill=True, # Only in v0
66
+ )
67
+ else:
68
+ llm_args = EngineArgs(
69
+ model=model,
70
+ kv_transfer_config=ktc,
71
+ max_model_len=8000,
72
+ gpu_memory_utilization=0.8,
73
+ )
74
+
75
+ llm = LLM(**asdict(llm_args))
76
try:
77
yield llm
78
finally:
@@ -57,6 +86,9 @@ def print_output(
86
sampling_params: SamplingParams,
87
req_str: str,
88
):
89
+ # Should be able to see logs like the following:
90
+ # `LMCache INFO: Storing KV cache for 6006 out of 6006 tokens for request 0`
91
+ # This indicates that the KV cache has been stored in LMCache.
92
start = time.time()
93
outputs = llm.generate(prompt, sampling_params)
94
print("-" * 50)
@@ -68,10 +100,29 @@ def print_output(
100
print("-" * 50)
101
102
103
+def parse_args():
104
+ parser = argparse.ArgumentParser()
105
+ parser.add_argument("-v",
106
+ "--version",
107
+ choices=["v0", "v1"],
108
+ default="v1",
109
+ help="Specify vLLM version (default: v1)")
110
+ return parser.parse_args()
111
+
112
+
113
def main():
114
+ args = parse_args()
115
+
116
+ if args.version == "v0":
117
+ lmcache_connector = "LMCacheConnector"
118
+ model = "mistralai/Mistral-7B-Instruct-v0.2"
119
+ else:
120
+ lmcache_connector = "LMCacheConnectorV1"
121
+ model = "meta-llama/Meta-Llama-3.1-8B-Instruct"
122
+
123
setup_environment_variables()
124
74
- with build_llm_with_lmcache() as llm:
125
+ with build_llm_with_lmcache(lmcache_connector, model, args.version) as llm:
126
127
# This example script runs two requests with a shared prefix.
128
# Define the shared prompt and specific prompts
examples/lmcache/cpu_offload_lmcache_v1.py
deleted
-57
@@ -1,57 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-"""
3
-This file demonstrates the example usage of cpu offloading
4
-with LMCache in vLLM v1.
5
-
6
-Note that lmcache needs to be installed to run this example.
7
-Learn more about LMCache in https://github.com/LMCache/LMCache.
8
-"""
9
-import os
10
-
11
-from lmcache.experimental.cache_engine import LMCacheEngineBuilder
12
-from lmcache.integration.vllm.utils import ENGINE_NAME
13
-
14
-from vllm import LLM, SamplingParams
15
-from vllm.config import KVTransferConfig
16
-
17
-# LMCache-related environment variables
18
-# Use experimental features in LMCache
19
-os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
20
-# LMCache is set to use 256 tokens per chunk
21
-os.environ["LMCACHE_CHUNK_SIZE"] = "256"
22
-# Enable local CPU backend in LMCache
23
-os.environ["LMCACHE_LOCAL_CPU"] = "True"
24
-# Set local CPU memory limit to 5.0 GB
25
-os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = "5.0"
26
-
27
-# This example script runs two requests with a shared prefix.
28
-shared_prompt = "Hello, how are you?" * 1000
29
-first_prompt = [
30
- shared_prompt + "Hello, my name is",
31
-]
32
-second_prompt = [
33
- shared_prompt + "Tell me a very long story",
34
-]
35
-
36
-sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
37
-
38
-ktc = KVTransferConfig.from_cli(
39
- '{"kv_connector":"LMCacheConnectorV1", "kv_role":"kv_both"}')
40
-# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
41
-# memory. Reduce the value if your GPU has less memory.
42
-# Note that LMCache is not compatible with chunked prefill for now.
43
-llm = LLM(model="meta-llama/Meta-Llama-3.1-8B-Instruct",
44
- kv_transfer_config=ktc,
45
- max_model_len=8000,
46
- gpu_memory_utilization=0.8)
47
-
48
-# Should be able to see logs like the following:
49
-# `LMCache INFO: Storing KV cache for 6006 out of 6006 tokens for request 0`
50
-# This indicates that the KV cache has been stored in LMCache.
51
-outputs = llm.generate(first_prompt, sampling_params)
52
-for output in outputs:
53
- generated_text = output.outputs[0].text
54
- print(f"Generated text: {generated_text!r}")
55
-
56
-# Clean up lmcache backend
57
-LMCacheEngineBuilder.destroy(ENGINE_NAME)