[Misc] refactor example - cpu_offload_lmcache

Signed-off-by: reidliu41 <[email protected]>

Author: reidliu41

examples/lmcache/README.md

@@ -44,8 +44,8 @@ The main script generates several log files:
 
 ## 2. CPU Offload Examples
 
-- `cpu_offload_lmcache_v0.py` - CPU offloading implementation for vLLM v0
-- `cpu_offload_lmcache_v1.py` - CPU offloading implementation for vLLM v1
+- `python cpu_offload_lmcache_v0.py -v v0` - CPU offloading implementation for vLLM v0
+- `python cpu_offload_lmcache_v1.py -v v1` - CPU offloading implementation for vLLM v1
 
 ## 3. KV Cache Sharing
 

examples/lmcache/cpu_offload_lmcache_v0.py renamed to examples/lmcache/cpu_offload_lmcache.py

@@ -1,7 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 """
 This file demonstrates the example usage of cpu offloading
-with LMCache.
+with LMCache in vLLM v1 or v0.
+
+Usage:
+
+    Specify vLLM version
+
+    -v v0 : Use LMCacheConnector
+            model = mistralai/Mistral-7B-Instruct-v0.2
+            (Includes enable_chunked_prefill = True)
+
+    -v v1 : Use LMCacheConnectorV1 (default)
+            model = meta-llama/Meta-Llama-3.1-8B-Instruct
+            (Without enable_chunked_prefill)
 
 Note that `lmcache` is needed to run this example.
 Requirements: Linux, Python: 3.10 or higher, CUDA: 12.1
@@ -11,6 +23,7 @@
 import contextlib
 import os
 import time
+import argparse
 
 from lmcache.experimental.cache_engine import LMCacheEngineBuilder
 from lmcache.integration.vllm.utils import ENGINE_NAME
@@ -32,17 +45,28 @@ def setup_environment_variables():
 
 
 @contextlib.contextmanager
-def build_llm_with_lmcache():
+def build_llm_with_lmcache(
+    lmcache_connector: str,
+    model: str,
+    vllm_version: str
+):
     ktc = KVTransferConfig.from_cli(
-        '{"kv_connector":"LMCacheConnector", "kv_role":"kv_both"}')
+        f'{{"kv_connector":"{lmcache_connector}", "kv_role":"kv_both"}}')
     # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
     # memory. Reduce the value if your GPU has less memory.
     # Note: LMCache supports chunked prefill (see vLLM#14505, LMCache#392).
-    llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
-              kv_transfer_config=ktc,
-              max_model_len=8000,
-              enable_chunked_prefill=True,
-              gpu_memory_utilization=0.8)
+    llm_args = {
+        "model": model,
+        "kv_transfer_config": ktc,
+        "max_model_len": 8000,
+        "gpu_memory_utilization": 0.8,
+    }
+
+    # Only for v0
+    if vllm_version == "v0":
+        llm_args["enable_chunked_prefill"] = True
+
+    llm = LLM(**llm_args)
 
     try:
         yield llm
@@ -57,6 +81,9 @@ def print_output(
     sampling_params: SamplingParams,
     req_str: str,
 ):
+    # Should be able to see logs like the following:
+    # `LMCache INFO: Storing KV cache for 6006 out of 6006 tokens for request 0`
+    # This indicates that the KV cache has been stored in LMCache.
     start = time.time()
     outputs = llm.generate(prompt, sampling_params)
     print("-" * 50)
@@ -68,10 +95,29 @@ def print_output(
     print("-" * 50)
 
 
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v",
+                        "--version",
+                        choices=["v0", "v1"],
+                        default="v1",
+                        help="Specify vLLM version (default: v1)")
+    return parser.parse_args()
+
+
 def main():
+    args = parse_args()
+
+    if args.version == "v0":
+        lmcache_connector = "LMCacheConnector"
+        model = "mistralai/Mistral-7B-Instruct-v0.2"
+    else:
+        lmcache_connector = "LMCacheConnectorV1"
+        model = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+
     setup_environment_variables()
 
-    with build_llm_with_lmcache() as llm:
+    with build_llm_with_lmcache(lmcache_connector, model, args.version) as llm:
 
         # This example script runs two requests with a shared prefix.
         # Define the shared prompt and specific prompts