Repository: kaito-project/kaito

Issue:
Support multi-tiered KV caching strategies
**Is your feature request related to a problem? Please describe.**
This is essential for optimizing memory usage and inference speed

**Describe the solution you'd like**

**Describe alternatives you've considered**

**Additional context**

Discussion:
@zhuangqh: This functionality is related to LMCache/mooncake project. 
https://github.com/LMCache/LMCache
https://github.com/kvcache-ai/Mooncake

We have some thought about that. There are two choices to leverage LMCache
1. standalone LMCache instance for each node locally. Offload KV Cache to CPU.
2. centric LMCache for remote access. Shared Remote KV Cache.

Considering centric LMCache requires high-perf RDMA to balance the recalculation and transmission latency and RDMA is not common for most GPU sku, I preferred to choose opt1 at first.@kris-gaudel: Hi I'm interested in contributing, are there any other vLLM related tasks I could work on?@andyzhangx: https://github.com/kaito-project/kaito/pull/1326 would add cpu memory offload feature, and I have added comparison tests in that PR description. Next I will work on offload to nvme disk feature.

Solution patch:
diff --git a/presets/workspace/dependencies/requirements.txt b/presets/workspace/dependencies/requirements.txt
index 0a73c61..69ba063 100644
--- a/presets/workspace/dependencies/requirements.txt
+++ b/presets/workspace/dependencies/requirements.txt
@@ -1,7 +1,7 @@
 # Dependencies for TFS
 
 # Core Dependencies
-vllm==0.9.0
+vllm==0.9.1
 transformers == 4.51.3
 torch==2.7.0
 accelerate==1.0.0
@@ -14,6 +14,7 @@ numpy<3.0,>=1.25.0
 sentencepiece==0.2.0
 jinja2>=3.1.0
 starlette==0.47.2
+lmcache==0.3.2
 
 # For accessing Ray dashboard. ray[default] version should be consistent with the one in vllm/vllm-openai:<vllm-version>.
 # Check with `docker run --entrypoint "" vllm/vllm-openai:<vllm-version> pip freeze | grep ray`
diff --git a/presets/workspace/inference/vllm/inference_api.py b/presets/workspace/inference/vllm/inference_api.py
index 074e7d2..f87e36b 100644
--- a/presets/workspace/inference/vllm/inference_api.py
+++ b/presets/workspace/inference/vllm/inference_api.py
@@ -20,9 +20,11 @@ from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any
 
+import psutil
 import torch
 import uvloop
 import vllm.entrypoints.openai.api_server as api_server
+import vllm.envs as envs
 import yaml
 from vllm.engine.llm_engine import EngineArgs, LLMEngine, VllmConfig
 from vllm.entrypoints.openai.serving_models import LoRAModulePath
@@ -68,6 +70,12 @@ class KAITOArgumentParser(argparse.ArgumentParser):
             type=int,
             help="Maximum number of steps to find the max available seq len fitting in the GPU memory.",
         )
+        self.add_argument(
+            "--kaito-kv-cache-cpu-memory-utilization",
+            type=float,
+            default=0.5,
+            help="KV cache CPU memory utilization.",
+        )
 
     def _reset_vllm_defaults(self):
         local_rank = int(os.environ.get("LOCAL_RANK", 0))  # Default to 0 if not set
@@ -100,6 +108,10 @@ class KAITOArgumentParser(argparse.ArgumentParser):
             file_config = KaitoConfig.from_yaml(kaito_args.kaito_config_file)
             if kaito_args.kaito_max_probe_steps is None:
                 kaito_args.kaito_max_probe_steps = file_config.max_probe_steps
+            if kaito_args.kaito_kv_cache_cpu_memory_utilization is None:
+                kaito_args.kaito_kv_cache_cpu_memory_utilization = (
+                    file_config.kv_cache_cpu_memory_utilization
+                )
 
             for key, value in file_config.vllm.items():
                 runtime_args.append(f"--{key}")
@@ -124,6 +136,9 @@ class KaitoConfig:
     # Maximum number of steps to find the max available seq len fitting in the GPU memory.
     max_probe_steps: int
 
+    # Optional: CPU memory utilization for the vllm engine in kv cache offload mode. (default: 0.5, set to 0 to disable)
+    kv_cache_cpu_memory_utilization: float
+
     @staticmethod
     def from_yaml(yaml_file: str) -> "KaitoConfig":
         with open(yaml_file) as file:
@@ -131,6 +146,9 @@ class KaitoConfig:
         return KaitoConfig(
             vllm=config_data.get("vllm", {}),
             max_probe_steps=config_data.get("max_probe_steps", 6),
+            kv_cache_cpu_memory_utilization=config_data.get(
+                "kv_cache_cpu_memory_utilization", 0.5
+            ),
         )
 
     def to_yaml(self) -> str:
@@ -274,6 +292,50 @@ def try_get_max_available_seq_len(args: argparse.Namespace) -> int | None:
         return None
 
 
+def set_kv_cache_offloading_if_appliable(args: argparse.Namespace) -> None:
+    """
+    Set KV cache offloading to CPU RAM if applicable.
+    This is only applicable when VLLM_USE_V1 is enabled and
+    kaito_kv_cache_cpu_memory_utilization is set.
+    """
+    if (
+        not envs.is_set("VLLM_USE_V1")
+        and args.kaito_kv_cache_cpu_memory_utilization > 0
+    ):
+        logger.info(
+            f"VLLM_USE_V1 is not set, but kaito_kv_cache_cpu_memory_utilization is set as {args.kaito_kv_cache_cpu_memory_utilization}, "
+            "run create_engine_config to check whether VLLM_USE_V1 should be set."
+        )
+        EngineArgs.from_cli_args(copy.deepcopy(args)).create_engine_config()
+
+    if (
+        envs.is_set("VLLM_USE_V1")
+        and envs.VLLM_USE_V1
+        and args.kaito_kv_cache_cpu_memory_utilization > 0
+    ):
+        os.environ["LMCACHE_CHUNK_SIZE"] = "256"
+        os.environ["LMCACHE_LOCAL_CPU"] = "True"
+        available_memory_gb = (
+            psutil.virtual_memory().total - psutil.virtual_memory().used
+        ) / (1024**3)
+        logger.info(
+            f"VLLM_USE_V1 is set, Offload KV cache to CPU RAM, size limit: {available_memory_gb} * {args.kaito_kv_cache_cpu_memory_utilization} GB"
+        )
+        os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = (
+            f"{available_memory_gb * args.kaito_kv_cache_cpu_memory_utilization}"
+        )
+
+        if args.kv_transfer_config is None:
+            args.kv_transfer_config = {
+                "kv_connector": "LMCacheConnectorV1",
+                "kv_role": "kv_both",
+            }
+    else:
+        logger.info(
+            "VLLM_USE_V1 or kv_cache_cpu_memory_utilization is not set, do not use KV cache offload to CPU RAM."
+        )
+
+
 if __name__ == "__main__":
     parser = KAITOArgumentParser(description="KAITO wrapper of vLLM serving server")
     args = parser.parse_args()
@@ -287,6 +349,8 @@ if __name__ == "__main__":
     if max_available_seq_len is not None:
         args.max_model_len = max_available_seq_len
 
+    set_kv_cache_offloading_if_appliable(args)
+
     # Run the serving server
     logger.info(f"Starting server on port {args.port}")
     # See https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html for more



Test patch:
diff --git a/presets/workspace/inference/vllm/tests/test_vllm_inference_api.py b/presets/workspace/inference/vllm/tests/test_vllm_inference_api.py
index f0e00f3..90bc774 100644
--- a/presets/workspace/inference/vllm/tests/test_vllm_inference_api.py
+++ b/presets/workspace/inference/vllm/tests/test_vllm_inference_api.py
@@ -71,6 +71,7 @@ def setup_server(request, tmp_path_factory, autouse=True):
     kaito_config = KaitoConfig(
         vllm={"max-model-len": TEST_MODEL_LEN, "served-model-name": TEST_MODEL_NAME},
         max_probe_steps=0,
+        kv_cache_cpu_memory_utilization=0.6,
     )
     with open(config_file, "w") as f:
         f.write(kaito_config.to_yaml())