[Misc] refactor example - cpu_offload_lmcache (#17460)

Signed-off-by: reidliu41 <[email protected]>
Co-authored-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.py -v v0` - CPU offloading implementation for vLLM v0
+- `python cpu_offload_lmcache.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,22 +1,37 @@
 # 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
 Learn more about LMCache environment setup, please refer to:
 https://docs.lmcache.ai/getting_started/installation.html
 """
+import argparse
 import contextlib
 import os
 import time
+from dataclasses import asdict
 
 from lmcache.experimental.cache_engine import LMCacheEngineBuilder
 from lmcache.integration.vllm.utils import ENGINE_NAME
 
 from vllm import LLM, SamplingParams
 from vllm.config import KVTransferConfig
+from vllm.engine.arg_utils import EngineArgs
 
 
 def setup_environment_variables():
@@ -32,18 +47,32 @@ def setup_environment_variables():
 
 
 @contextlib.contextmanager
-def build_llm_with_lmcache():
-    ktc = KVTransferConfig.from_cli(
-        '{"kv_connector":"LMCacheConnector", "kv_role":"kv_both"}')
+def build_llm_with_lmcache(lmcache_connector: str, model: str,
+                           vllm_version: str):
+    ktc = KVTransferConfig(
+        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)
-
+    if vllm_version == "v0":
+        llm_args = EngineArgs(
+            model=model,
+            kv_transfer_config=ktc,
+            max_model_len=8000,
+            gpu_memory_utilization=0.8,
+            enable_chunked_prefill=True,  # Only in v0
+        )
+    else:
+        llm_args = EngineArgs(
+            model=model,
+            kv_transfer_config=ktc,
+            max_model_len=8000,
+            gpu_memory_utilization=0.8,
+        )
+
+    llm = LLM(**asdict(llm_args))
     try:
         yield llm
     finally:
@@ -57,6 +86,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 +100,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