@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5e83a7277f
Fork 0

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, 13:30 UTC

[v1] [P/D] Adding LMCache KV connector for v1 (#16625)

Yihua Cheng committed Apr 25, 2025 at 22:03 UTC 5e83a7277f7892432375d3d41594ebfde086ca4e
12 files changed +793
examples/lmcache/README.md new
+56
@@ -0,0 +1,56 @@
1 +# LMCache Examples
2 +
3 +This folder demonstrates how to use LMCache for disaggregated prefilling, CPU offloading and KV cache sharing.
4 +
5 +## 1. Disaggregated Prefill in vLLM v1
6 +
7 +This example demonstrates how to run LMCache with disaggregated prefill using NIXL on a single node.
8 +
9 +### Prerequisites
10 +
11 +- Install [LMCache](https://github.com/ai-dynamo/lmcache)
12 +- Install [NIXL](https://github.com/ai-dynamo/nixl)
13 +- At least 2 GPUs
14 +- Valid Hugging Face token (HF_TOKEN) for Llama 3.1 8B Instruct.
15 +
16 +### Usage
17 +
18 +Run
19 +`cd disagg_prefill_lmcache_v1`
20 +to get into `disagg_prefill_lmcache_v1` folder, and then run
21 +
22 +```bash
23 +bash disagg_example_nixl.sh
24 +```
25 +
26 +to run disaggregated prefill and benchmark the performance.
27 +
28 +### Components
29 +
30 +#### Server Scripts
31 +- `disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh` - Launches individual vLLM servers for prefill/decode, and also launches the proxy server.
32 +- `disagg_prefill_lmcache_v1/disagg_proxy_server.py` - FastAPI proxy server that coordinates between prefiller and decoder
33 +- `disagg_prefill_lmcache_v1/disagg_example_nixl.sh` - Main script to run the example
34 +
35 +#### Configuration
36 +- `disagg_prefill_lmcache_v1/configs/lmcache-prefiller-config.yaml` - Configuration for prefiller server
37 +- `disagg_prefill_lmcache_v1/configs/lmcache-decoder-config.yaml` - Configuration for decoder server
38 +
39 +#### Log Files
40 +The main script generates several log files:
41 +- `prefiller.log` - Logs from the prefill server
42 +- `decoder.log` - Logs from the decode server
43 +- `proxy.log` - Logs from the proxy server
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
49 +
50 +## 3. KV Cache Sharing
51 +
52 +The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV caches between vLLM v1 instances.
53 +
54 +## 4. Disaggregated Prefill in vLLM v0
55 +
56 +The `disaggregated_prefill_lmcache_v0.py` provides an example of how to run disaggregated prefill in vLLM v0.
examples/lmcache/cpu_offload_lmcache_v0.py renamed
examples/lmcache/cpu_offload_lmcache_v1.py new
+57
@@ -0,0 +1,57 @@
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)
examples/lmcache/disagg_prefill_lmcache_v0.py renamed
examples/lmcache/disagg_prefill_lmcache_v1/configs/lmcache-decoder-config.yaml new
+13
@@ -0,0 +1,13 @@
1 +local_cpu: False
2 +max_local_cpu_size: 0
3 +#local_disk:
4 +max_local_disk_size: 0
5 +remote_serde: NULL
6 +
7 +enable_nixl: True
8 +nixl_role: "receiver"
9 +nixl_peer_host: "localhost"
10 +nixl_peer_port: 55555
11 +nixl_buffer_size: 1073741824 # 1GB
12 +nixl_buffer_device: "cuda"
13 +nixl_enable_gc: True
examples/lmcache/disagg_prefill_lmcache_v1/configs/lmcache-prefiller-config.yaml new
+13
@@ -0,0 +1,13 @@
1 +local_cpu: False
2 +max_local_cpu_size: 0
3 +#local_disk:
4 +max_local_disk_size: 0
5 +remote_serde: NULL
6 +
7 +enable_nixl: True
8 +nixl_role: "sender"
9 +nixl_peer_host: "localhost"
10 +nixl_peer_port: 55555
11 +nixl_buffer_size: 1073741824 # 1GB
12 +nixl_buffer_device: "cuda"
13 +nixl_enable_gc: True
examples/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh new
+136
@@ -0,0 +1,136 @@
1 +#!/bin/bash
2 +
3 +echo "Warning: LMCache disaggregated prefill support for vLLM v1 is experimental and subject to change."
4 +
5 +
6 +PIDS=()
7 +
8 +# Switch to the directory of the current script
9 +cd "$(dirname "${BASH_SOURCE[0]}")"
10 +
11 +check_hf_token() {
12 + if [ -z "$HF_TOKEN" ]; then
13 + echo "HF_TOKEN is not set. Please set it to your Hugging Face token."
14 + exit 1
15 + fi
16 + if [[ "$HF_TOKEN" != hf_* ]]; then
17 + echo "HF_TOKEN is not a valid Hugging Face token. Please set it to your Hugging Face token."
18 + exit 1
19 + fi
20 + echo "HF_TOKEN is set and valid."
21 +}
22 +
23 +check_num_gpus() {
24 + # can you check if the number of GPUs are >=2 via nvidia-smi?
25 + num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
26 + if [ "$num_gpus" -lt 2 ]; then
27 + echo "You need at least 2 GPUs to run disaggregated prefill."
28 + exit 1
29 + else
30 + echo "Found $num_gpus GPUs."
31 + fi
32 +}
33 +
34 +ensure_python_library_installed() {
35 + echo "Checking if $1 is installed..."
36 + python -c "import $1" > /dev/null 2>&1
37 + if [ $? -ne 0 ]; then
38 + if [ "$1" == "nixl" ]; then
39 + echo "$1 is not installed. Please refer to https://github.com/ai-dynamo/nixl for installation."
40 + else
41 + echo "$1 is not installed. Please install it via pip install $1."
42 + fi
43 + exit 1
44 + else
45 + echo "$1 is installed."
46 + fi
47 +}
48 +
49 +cleanup() {
50 + echo "Stopping everything…"
51 + trap - INT TERM # prevent re-entrancy
52 + kill -- -$$ # negative PID == “this whole process-group”
53 + wait # reap children so we don't leave zombies
54 + exit 0
55 +}
56 +
57 +wait_for_server() {
58 + local port=$1
59 + local timeout_seconds=1200
60 + local start_time=$(date +%s)
61 +
62 + echo "Waiting for server on port $port..."
63 +
64 + while true; do
65 + if curl -s "localhost:${port}/v1/completions" > /dev/null; then
66 + return 0
67 + fi
68 +
69 + local now=$(date +%s)
70 + if (( now - start_time >= timeout_seconds )); then
71 + echo "Timeout waiting for server"
72 + return 1
73 + fi
74 +
75 + sleep 1
76 + done
77 +}
78 +
79 +
80 +main() {
81 + check_hf_token
82 + check_num_gpus
83 + ensure_python_library_installed lmcache
84 + ensure_python_library_installed nixl
85 + ensure_python_library_installed pandas
86 + ensure_python_library_installed datasets
87 + ensure_python_library_installed vllm
88 +
89 + trap cleanup INT
90 + trap cleanup USR1
91 + trap cleanup TERM
92 +
93 + echo "Launching prefiller, decoder and proxy..."
94 + echo "Please check prefiller.log, decoder.log and proxy.log for logs."
95 +
96 + bash disagg_vllm_launcher.sh prefiller \
97 + > >(tee prefiller.log) 2>&1 &
98 + prefiller_pid=$!
99 + PIDS+=($prefiller_pid)
100 +
101 + bash disagg_vllm_launcher.sh decoder \
102 + > >(tee decoder.log) 2>&1 &
103 + decoder_pid=$!
104 + PIDS+=($decoder_pid)
105 +
106 + python3 disagg_proxy_server.py \
107 + --host localhost \
108 + --port 9000 \
109 + --prefiller-host localhost \
110 + --prefiller-port 8100 \
111 + --decoder-host localhost \
112 + --decoder-port 8200 \
113 + > >(tee proxy.log) 2>&1 &
114 + proxy_pid=$!
115 + PIDS+=($proxy_pid)
116 +
117 + wait_for_server 8100
118 + wait_for_server 8200
119 + wait_for_server 9000
120 +
121 + echo "All servers are up. Starting benchmark..."
122 +
123 + # begin benchmark
124 + cd ../../../benchmarks/
125 + python benchmark_serving.py --port 9000 --seed $(date +%s) \
126 + --model meta-llama/Llama-3.1-8B-Instruct \
127 + --dataset-name random --random-input-len 7500 --random-output-len 200 \
128 + --num-prompts 200 --burstiness 100 --request-rate 3.6 | tee benchmark.log
129 +
130 + echo "Benchmarking done. Cleaning up..."
131 +
132 + cleanup
133 +
134 +}
135 +
136 +main
\ No newline at end of file
examples/lmcache/disagg_prefill_lmcache_v1/disagg_proxy_server.py new
+193
@@ -0,0 +1,193 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import argparse
4 +import os
5 +import time
6 +from contextlib import asynccontextmanager
7 +
8 +import httpx
9 +import numpy as np
10 +from fastapi import FastAPI, Request
11 +from fastapi.responses import StreamingResponse
12 +
13 +
14 +@asynccontextmanager
15 +async def lifespan(app: FastAPI):
16 + """
17 + Lifespan context manager to handle startup and shutdown events.
18 + """
19 + # Startup: Initialize clients
20 + prefiller_base_url = f'http://{global_args.prefiller_host}:{global_args.prefiller_port}/v1'
21 + decoder_base_url = f'http://{global_args.decoder_host}:{global_args.decoder_port}/v1'
22 +
23 + app.state.prefill_client = httpx.AsyncClient(timeout=None,
24 + base_url=prefiller_base_url)
25 + app.state.decode_client = httpx.AsyncClient(timeout=None,
26 + base_url=decoder_base_url)
27 +
28 + yield
29 +
30 + # Shutdown: Close clients
31 + await app.state.prefill_client.aclose()
32 + await app.state.decode_client.aclose()
33 +
34 +
35 +# Update FastAPI app initialization to use lifespan
36 +app = FastAPI(lifespan=lifespan)
37 +
38 +
39 +class StatsCalculator:
40 +
41 + def __init__(self):
42 + self._stats = []
43 + self._last_log_time = time.time()
44 +
45 + def add(self, value):
46 + self._stats.append(value)
47 + if time.time() - self._last_log_time > 5:
48 + self._log_stats()
49 + self._last_log_time = time.time()
50 +
51 + def _log_stats(self):
52 + # Print average, median, and 99th percentile
53 + np_arr = np.array(self._stats)
54 + output_str = f"\nNum requests: {len(self._stats)}" + \
55 + "\nPrefill node TTFT stats:" + \
56 + f"\n - Average (ms): {np.mean(np_arr)}" + \
57 + f"\n - Median (ms): {np.median(np_arr)}" + \
58 + f"\n - 99th Percentile (ms): {np.percentile(np_arr, 99)}\n"
59 + print("===============================", output_str,
60 + "===============================")
61 +
62 +
63 +stats_calculator = StatsCalculator()
64 +counter = 0
65 +
66 +
67 +def parse_args():
68 + parser = argparse.ArgumentParser()
69 +
70 + parser.add_argument("--port", type=int, default=8000)
71 + parser.add_argument("--host", type=str, default="localhost")
72 + parser.add_argument("--prefiller-host", type=str, default="localhost")
73 + parser.add_argument("--prefiller-port", type=int, default=8100)
74 + parser.add_argument("--decoder-host", type=str, default="localhost")
75 + parser.add_argument("--decoder-port", type=int, default=8200)
76 + args = parser.parse_args()
77 + return args
78 +
79 +
80 +# Initialize variables to hold the persistent clients
81 +app.state.prefill_client = None
82 +app.state.decode_client = None
83 +
84 +
85 +async def send_request_to_service(client: httpx.AsyncClient, endpoint: str,
86 + req_data: dict):
87 + """
88 + Send a request to a service using a persistent client.
89 + """
90 + req_data = req_data.copy()
91 + req_data['max_tokens'] = 1
92 + if 'max_completion_tokens' in req_data:
93 + req_data['max_completion_tokens'] = 1
94 +
95 + headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
96 + response = await client.post(endpoint, json=req_data, headers=headers)
97 + response.raise_for_status()
98 + return response
99 +
100 +
101 +async def stream_service_response(client: httpx.AsyncClient, endpoint: str,
102 + req_data: dict):
103 + """
104 + Asynchronously stream the response from a service using a persistent client.
105 + """
106 + headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
107 + async with client.stream("POST", endpoint, json=req_data,
108 + headers=headers) as response:
109 + response.raise_for_status()
110 + async for chunk in response.aiter_bytes():
111 + yield chunk
112 +
113 +
114 +@app.post("/v1/completions")
115 +async def handle_completions(request: Request):
116 + global counter, stats_calculator
117 + counter += 1
118 +
119 + st = time.time()
120 + try:
121 + req_data = await request.json()
122 +
123 + # Send request to prefill service, ignore the response
124 + await send_request_to_service(app.state.prefill_client, "/completions",
125 + req_data)
126 +
127 + et = time.time()
128 + stats_calculator.add(et - st)
129 +
130 + # Stream response from decode service
131 + async def generate_stream():
132 + async for chunk in stream_service_response(app.state.decode_client,
133 + "/completions",
134 + req_data):
135 + yield chunk
136 +
137 + return StreamingResponse(generate_stream(),
138 + media_type="application/json")
139 +
140 + except Exception as e:
141 + import sys
142 + import traceback
143 + exc_info = sys.exc_info()
144 + print("Error occurred in disagg prefill proxy server"
145 + " - completions endpoint")
146 + print(e)
147 + print("".join(traceback.format_exception(*exc_info)))
148 + raise
149 +
150 +
151 +@app.post("/v1/chat/completions")
152 +async def handle_chat_completions(request: Request):
153 + global counter, stats_calculator
154 + counter += 1
155 +
156 + st = time.time()
157 + try:
158 + req_data = await request.json()
159 +
160 + # Send request to prefill service, ignore the response
161 + await send_request_to_service(app.state.prefill_client,
162 + "/chat/completions", req_data)
163 +
164 + et = time.time()
165 + stats_calculator.add(et - st)
166 +
167 + # Stream response from decode service
168 + async def generate_stream():
169 + async for chunk in stream_service_response(app.state.decode_client,
170 + "/chat/completions",
171 + req_data):
172 + yield chunk
173 +
174 + return StreamingResponse(generate_stream(),
175 + media_type="application/json")
176 +
177 + except Exception as e:
178 + import sys
179 + import traceback
180 + exc_info = sys.exc_info()
181 + print("Error occurred in disagg prefill proxy server "
182 + " - chat completions endpoint")
183 + print(e)
184 + print("".join(traceback.format_exception(*exc_info)))
185 + raise
186 +
187 +
188 +if __name__ == '__main__':
189 + global global_args
190 + global_args = parse_args()
191 +
192 + import uvicorn
193 + uvicorn.run(app, host=global_args.host, port=global_args.port)
examples/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh new
+59
@@ -0,0 +1,59 @@
1 +#!/bin/bash
2 +
3 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
4 +
5 +if [[ $# -lt 1 ]]; then
6 + echo "Usage: $0 <prefiller | decoder> [model]"
7 + exit 1
8 +fi
9 +
10 +if [[ $# -eq 1 ]]; then
11 + echo "Using default model: meta-llama/Llama-3.1-8B-Instruct"
12 + MODEL="meta-llama/Llama-3.1-8B-Instruct"
13 +else
14 + echo "Using model: $2"
15 + MODEL=$2
16 +fi
17 +
18 +
19 +if [[ $1 == "prefiller" ]]; then
20 + # Prefiller listens on port 8100
21 + prefill_config_file=$SCRIPT_DIR/configs/lmcache-prefiller-config.yaml
22 +
23 + UCX_TLS=cuda_ipc,cuda_copy,tcp \
24 + LMCACHE_CONFIG_FILE=$prefill_config_file \
25 + LMCACHE_USE_EXPERIMENTAL=True \
26 + VLLM_ENABLE_V1_MULTIPROCESSING=1 \
27 + VLLM_WORKER_MULTIPROC_METHOD=spawn \
28 + CUDA_VISIBLE_DEVICES=0 \
29 + vllm serve $MODEL \
30 + --port 8100 \
31 + --disable-log-requests \
32 + --enforce-eager \
33 + --kv-transfer-config \
34 + '{"kv_connector":"LMCacheConnectorV1","kv_role":"kv_producer","kv_connector_extra_config": {"discard_partial_chunks": false, "lmcache_rpc_port": "producer1"}}'
35 +
36 +
37 +elif [[ $1 == "decoder" ]]; then
38 + # Decoder listens on port 8200
39 + decode_config_file=$SCRIPT_DIR/configs/lmcache-decoder-config.yaml
40 +
41 + UCX_TLS=cuda_ipc,cuda_copy,tcp \
42 + LMCACHE_CONFIG_FILE=$decode_config_file \
43 + LMCACHE_USE_EXPERIMENTAL=True \
44 + VLLM_ENABLE_V1_MULTIPROCESSING=1 \
45 + VLLM_WORKER_MULTIPROC_METHOD=spawn \
46 + CUDA_VISIBLE_DEVICES=1 \
47 + vllm serve $MODEL \
48 + --port 8200 \
49 + --disable-log-requests \
50 + --enforce-eager \
51 + --kv-transfer-config \
52 + '{"kv_connector":"LMCacheConnectorV1","kv_role":"kv_consumer","kv_connector_extra_config": {"discard_partial_chunks": false, "lmcache_rpc_port": "consumer1"}}'
53 +
54 +
55 +else
56 + echo "Invalid role: $1"
57 + echo "Should be either prefill, decode"
58 + exit 1
59 +fi
examples/lmcache/kv_cache_sharing_lmcache_v1.py new
+130
@@ -0,0 +1,130 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""
3 +This file demonstrates the example usage of remote KV cache sharing
4 +with LMCache.
5 +We will launch 2 vllm instances, and launch an additional LMCache server.
6 +KV cache is transferred in the following manner:
7 +(1) vLLM instance 1 -> LMCache server (KV cache store).
8 +(2) LMCache server -> vLLM instance 2 (KV cache reuse/retrieve).
9 +
10 +Note that lmcache needs to be installed to run this example.
11 +Learn more about LMCache in https://github.com/LMCache/LMCache.
12 +"""
13 +import os
14 +import subprocess
15 +import time
16 +from multiprocessing import Event, Process
17 +
18 +from lmcache.experimental.cache_engine import LMCacheEngineBuilder
19 +from lmcache.integration.vllm.utils import ENGINE_NAME
20 +
21 +from vllm import LLM, SamplingParams
22 +from vllm.config import KVTransferConfig
23 +
24 +# LMCache-related environment variables
25 +# The port to start LMCache server
26 +port = 8100
27 +# Use experimental features in LMCache
28 +os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
29 +# LMCache is set to use 256 tokens per chunk
30 +os.environ["LMCACHE_CHUNK_SIZE"] = "256"
31 +# Disable local CPU backend in LMCache
32 +os.environ["LMCACHE_LOCAL_CPU"] = "False"
33 +# Set local CPU memory buffer limit to 5.0 GB
34 +os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = "5.0"
35 +# Set the remote URL for LMCache server
36 +os.environ["LMCACHE_REMOTE_URL"] = f"lm://localhost:{port}"
37 +# Set the serializer/deserializer between vllm and LMCache server
38 +# `naive` indicates using raw bytes of the tensor without any compression
39 +os.environ["LMCACHE_REMOTE_SERDE"] = "naive"
40 +
41 +prompts = [
42 + "Hello, how are you?" * 1000,
43 +]
44 +
45 +
46 +def run_store(store_done, prompts):
47 + # We use GPU 0 for KV cache store process.
48 + os.environ["CUDA_VISIBLE_DEVICES"] = "0"
49 +
50 + sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
51 +
52 + ktc = KVTransferConfig.from_cli(
53 + '{"kv_connector":"LMCacheConnectorV1", "kv_role":"kv_both"}')
54 + # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
55 + # memory. Reduce the value if your GPU has less memory.
56 + llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
57 + kv_transfer_config=ktc,
58 + max_model_len=8000,
59 + gpu_memory_utilization=0.8,
60 + enforce_eager=True)
61 +
62 + outputs = llm.generate(prompts, sampling_params)
63 + for output in outputs:
64 + generated_text = output.outputs[0].text
65 + print(f"Generated text: {generated_text!r}")
66 + print("KV cache store is finished.")
67 + store_done.set()
68 +
69 + # Clean up lmcache backend
70 + LMCacheEngineBuilder.destroy(ENGINE_NAME)
71 +
72 +
73 +def run_retrieve(store_done, prompts, timeout=1):
74 + # We use GPU 1 for KV cache retrieve process.
75 + os.environ["CUDA_VISIBLE_DEVICES"] = "1"
76 +
77 + sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
78 +
79 + ktc = KVTransferConfig.from_cli(
80 + '{"kv_connector":"LMCacheConnectorV1", "kv_role":"kv_both"}')
81 + # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
82 + # of memory. Reduce the value if your GPU has less memory.
83 + llm = LLM(model="mistralai/Mistral-7B-Instruct-v0.2",
84 + kv_transfer_config=ktc,
85 + max_model_len=8000,
86 + gpu_memory_utilization=0.8,
87 + enforce_eager=True)
88 +
89 + print("Waiting for KV cache store to finish...")
90 + store_done.wait()
91 + time.sleep(timeout)
92 +
93 + outputs = llm.generate(prompts, sampling_params)
94 + for output in outputs:
95 + generated_text = output.outputs[0].text
96 + print(f"Generated text: {generated_text!r}")
97 +
98 + # Clean up lmcache backend
99 + LMCacheEngineBuilder.destroy(ENGINE_NAME)
100 +
101 +
102 +def run_lmcache_server(port):
103 + server_proc = subprocess.Popen([
104 + "python", "-m", "lmcache.experimental.server", "localhost",
105 + str(port)
106 + ])
107 + return server_proc
108 +
109 +
110 +def main():
111 + store_done = Event()
112 + store_process = Process(target=run_store, args=(store_done, prompts))
113 + retrieve_process = Process(target=run_retrieve, args=(store_done, prompts))
114 + lmcache_server_process = run_lmcache_server(port)
115 +
116 + # Start KV cache store process
117 + store_process.start()
118 +
119 + # Start KV cache retrieve process
120 + retrieve_process.start()
121 +
122 + # Clean up the processes
123 + store_process.join()
124 + retrieve_process.terminate()
125 + lmcache_server_process.terminate()
126 + lmcache_server_process.wait()
127 +
128 +
129 +if __name__ == "__main__":
130 + main()
vllm/distributed/kv_transfer/kv_connector/factory.py
+5
@@ -100,3 +100,8 @@ KVConnectorFactory.register_connector(
100 "SharedStorageConnector",
101 "vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector",
102 "SharedStorageConnector")
103 +
104 +KVConnectorFactory.register_connector(
105 + "LMCacheConnectorV1",
106 + "vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector",
107 + "LMCacheConnectorV1")
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py new
+131
@@ -0,0 +1,131 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import TYPE_CHECKING
3 +
4 +import torch
5 +from lmcache.integration.vllm.vllm_v1_adapter import LMCacheConnectorV1Impl
6 +
7 +from vllm.config import VllmConfig
8 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
9 + KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
10 +from vllm.logger import init_logger
11 +from vllm.v1.core.sched.output import SchedulerOutput
12 +
13 +if TYPE_CHECKING:
14 + from vllm.attention.backends.abstract import AttentionMetadata
15 + from vllm.forward_context import ForwardContext
16 + from vllm.v1.request import Request
17 +
18 +logger = init_logger(__name__)
19 +
20 +
21 +class LMCacheConnectorV1(KVConnectorBase_V1):
22 +
23 + def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
24 + super().__init__(vllm_config=vllm_config, role=role)
25 + self._lmcache_engine = LMCacheConnectorV1Impl(vllm_config, role, self)
26 +
27 + # ==============================
28 + # Worker-side methods
29 + # ==============================
30 + def start_load_kv(self, forward_context: "ForwardContext",
31 + **kwargs) -> None:
32 + """
33 + Start loading the KV cache from the connector to vLLM's paged
34 + KV buffer. This is called from the forward context before the
35 + forward pass to enable async loading during model execution.
36 +
37 + Args:
38 + forward_context (ForwardContext): the forward context.
39 + **kwargs: additional arguments for the load operation
40 +
41 + Note:
42 + The number of elements in kv_caches and layer_names should be
43 + the same.
44 +
45 + """
46 + self._lmcache_engine.start_load_kv(forward_context, **kwargs)
47 +
48 + def wait_for_layer_load(self, layer_name: str) -> None:
49 + """
50 + Block until the KV for a specific layer is loaded into vLLM's
51 + paged buffer. This is called from within attention layer to ensure
52 + async copying from start_load_kv is complete.
53 +
54 + This interface will be useful for layer-by-layer pipelining.
55 +
56 + Args:
57 + layer_name: the name of that layer
58 + """
59 + self._lmcache_engine.wait_for_layer_load(layer_name)
60 +
61 + def save_kv_layer(self, layer_name: str, kv_layer: torch.Tensor,
62 + attn_metadata: "AttentionMetadata", **kwargs) -> None:
63 + """
64 + Start saving the a layer of KV cache from vLLM's paged buffer
65 + to the connector. This is called from within attention layer to
66 + enable async copying during execution.
67 +
68 + Args:
69 + layer_name (str): the name of the layer.
70 + kv_layer (torch.Tensor): the paged KV buffer of the current
71 + layer in vLLM.
72 + attn_metadata (AttentionMetadata): the attention metadata.
73 + **kwargs: additional arguments for the save operation.
74 + """
75 + self._lmcache_engine.save_kv_layer(layer_name, kv_layer, attn_metadata,
76 + **kwargs)
77 +
78 + def wait_for_save(self):
79 + """
80 + Block until all the save operations is done. This is called
81 + as the forward context exits to ensure that the async saving
82 + from save_kv_layer is complete before finishing the forward.
83 +
84 + This prevents overwrites of paged KV buffer before saving done.
85 + """
86 + self._lmcache_engine.wait_for_save()
87 +
88 + # ==============================
89 + # Scheduler-side methods
90 + # ==============================
91 + def get_num_new_matched_tokens(
92 + self,
93 + request: "Request",
94 + num_computed_tokens: int,
95 + ) -> int:
96 + """
97 + Get number of new tokens that can be loaded from the
98 + external KV cache beyond the num_computed_tokens.
99 +
100 + Args:
101 + request (Request): the request object.
102 + num_computed_tokens (int): the number of locally
103 + computed tokens for this request
104 +
105 + Returns:
106 + the number of tokens that can be loaded from the
107 + external KV cache beyond what is already computed.
108 + """
109 + return self._lmcache_engine.get_num_new_matched_tokens(
110 + request, num_computed_tokens)
111 +
112 + def update_state_after_alloc(self, request: "Request",
113 + num_external_tokens: int):
114 + """
115 + Update KVConnector state after block allocation.
116 + """
117 + self._lmcache_engine.update_state_after_alloc(request,
118 + num_external_tokens)
119 +
120 + def build_connector_meta(
121 + self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata:
122 + """
123 + Build the connector metadata for this step.
124 +
125 + This function should NOT modify fields in the scheduler_output.
126 + Also, calling this function will reset the state of the connector.
127 +
128 + Args:
129 + scheduler_output (SchedulerOutput): the scheduler output object.
130 + """
131 + return self._lmcache_engine.build_connector_meta(scheduler_output)