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, 17:15 UTC
[Feature][Disaggregated] Support XpYd disaggregated prefill with MooncakeStore (#12957)
Signed-off-by: Shangming Cai <caishangming@linux.alibaba.com>
shangmingc committed
Mar 29, 2025 at 19:01 UTC
6fa7cd3dbcf3e78e36431ca31abd973e5617dd27
5 files changed
+906
-10
examples/online_serving/disagg_examples/disagg_proxy_demo.py
new
+450
@@ -0,0 +1,450 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+This file provides a disaggregated prefilling proxy demo to demonstrate an
4
+example usage of XpYd disaggregated prefilling.
5
+We can launch multiple vllm instances (2 for prefill and 2 for decode), and
6
+launch this proxy demo through:
7
+ python3 examples/online_serving/disagg_examples/disagg_proxy_demo.py \
8
+ --model $model_name \
9
+ --prefill localhost:8100 localhost:8101 \
10
+ --decode localhost:8200 localhost:8201 \
11
+ --port 8000
12
+
13
+Note: This demo will be removed once the PDController implemented in PR 15343
14
+(https://github.com/vllm-project/vllm/pull/15343) supports XpYd.
15
+"""
16
+import argparse
17
+import ipaddress
18
+import itertools
19
+import json
20
+import logging
21
+import os
22
+import sys
23
+from abc import ABC, abstractmethod
24
+from typing import Callable, Optional
25
+
26
+import aiohttp
27
+import requests
28
+import uvicorn
29
+from fastapi import (APIRouter, Depends, FastAPI, Header, HTTPException,
30
+ Request, status)
31
+from fastapi.responses import JSONResponse, StreamingResponse
32
+
33
+AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
34
+logger = logging.getLogger()
35
+logging.basicConfig(level=logging.INFO)
36
+
37
+
38
+class SchedulingPolicy(ABC):
39
+
40
+ @abstractmethod
41
+ def schedule(self, cycler: itertools.cycle):
42
+ raise NotImplementedError("Scheduling Proxy is not set.")
43
+
44
+
45
+class Proxy:
46
+
47
+ def __init__(
48
+ self,
49
+ prefill_instances: list[str],
50
+ decode_instances: list[str],
51
+ model: str,
52
+ scheduling_policy: SchedulingPolicy,
53
+ custom_create_completion: Optional[Callable[[Request],
54
+ StreamingResponse]] = None,
55
+ custom_create_chat_completion: Optional[Callable[
56
+ [Request], StreamingResponse]] = None,
57
+ ):
58
+ self.prefill_instances = prefill_instances
59
+ self.decode_instances = decode_instances
60
+ self.prefill_cycler = itertools.cycle(prefill_instances)
61
+ self.decode_cycler = itertools.cycle(decode_instances)
62
+ self.model = model
63
+ self.scheduling_policy = scheduling_policy
64
+ self.custom_create_completion = custom_create_completion
65
+ self.custom_create_chat_completion = custom_create_chat_completion
66
+ self.router = APIRouter()
67
+ self.setup_routes()
68
+
69
+ def setup_routes(self):
70
+ self.router.post(
71
+ "/v1/completions",
72
+ dependencies=[
73
+ Depends(self.validate_json_request)
74
+ ])(self.custom_create_completion if self.
75
+ custom_create_completion else self.create_completion)
76
+ self.router.post(
77
+ "/v1/chat/completions",
78
+ dependencies=[
79
+ Depends(self.validate_json_request)
80
+ ])(self.custom_create_chat_completion if self.
81
+ custom_create_chat_completion else self.create_chat_completion)
82
+ self.router.get("/status",
83
+ response_class=JSONResponse)(self.get_status)
84
+ self.router.post("/instances/add",
85
+ dependencies=[Depends(self.api_key_authenticate)
86
+ ])(self.add_instance_endpoint)
87
+
88
+ async def validate_json_request(self, raw_request: Request):
89
+ content_type = raw_request.headers.get("content-type", "").lower()
90
+ if content_type != "application/json":
91
+ raise HTTPException(
92
+ status_code=415,
93
+ detail=
94
+ "Unsupported Media Type: Only 'application/json' is allowed",
95
+ )
96
+
97
+ def api_key_authenticate(self, x_api_key: str = Header(...)):
98
+ expected_api_key = os.environ.get("ADMIN_API_KEY")
99
+ if not expected_api_key:
100
+ logger.error("ADMIN_API_KEY is not set in the environment.")
101
+ raise HTTPException(
102
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
103
+ detail="Server configuration error.",
104
+ )
105
+ if x_api_key != expected_api_key:
106
+ logger.warning("Unauthorized access attempt with API Key: %s",
107
+ x_api_key)
108
+ raise HTTPException(
109
+ status_code=status.HTTP_403_FORBIDDEN,
110
+ detail="Forbidden: Invalid API Key.",
111
+ )
112
+
113
+ async def validate_instance(self, instance: str) -> bool:
114
+ url = f"http://{instance}/v1/models"
115
+ try:
116
+ async with aiohttp.ClientSession(
117
+ timeout=AIOHTTP_TIMEOUT) as client:
118
+ logger.info("Verifying %s ...", instance)
119
+ async with client.get(url) as response:
120
+ if response.status == 200:
121
+ data = await response.json()
122
+ if "data" in data and len(data["data"]) > 0:
123
+ model_cur = data["data"][0].get("id", "")
124
+ if model_cur == self.model:
125
+ logger.info("Instance: %s could be added.",
126
+ instance)
127
+ return True
128
+ else:
129
+ logger.warning("Mismatch model %s : %s != %s",
130
+ instance, model_cur, self.model)
131
+ return False
132
+ else:
133
+ return False
134
+ else:
135
+ return False
136
+ except aiohttp.ClientError as e:
137
+ logger.error(str(e))
138
+ return False
139
+ except Exception as e:
140
+ logger.error(str(e))
141
+ return False
142
+
143
+ async def add_instance_endpoint(self, request: Request):
144
+ try:
145
+ data = await request.json()
146
+ logger.warning(str(data))
147
+ instance_type = data.get("type")
148
+ instance = data.get("instance")
149
+ if instance_type not in ["prefill", "decode"]:
150
+ raise HTTPException(status_code=400,
151
+ detail="Invalid instance type.")
152
+ if not instance or ":" not in instance:
153
+ raise HTTPException(status_code=400,
154
+ detail="Invalid instance format.")
155
+ host, port_str = instance.split(":")
156
+ try:
157
+ if host != "localhost":
158
+ ipaddress.ip_address(host)
159
+ port = int(port_str)
160
+ if not (0 < port < 65536):
161
+ raise HTTPException(status_code=400,
162
+ detail="Invalid port number.")
163
+ except Exception as e:
164
+ raise HTTPException(status_code=400,
165
+ detail="Invalid instance address.") from e
166
+
167
+ is_valid = await self.validate_instance(instance)
168
+ if not is_valid:
169
+ raise HTTPException(status_code=400,
170
+ detail="Instance validation failed.")
171
+
172
+ if instance_type == "prefill":
173
+ if instance not in self.prefill_instances:
174
+ self.prefill_instances.append(instance)
175
+ self.prefill_cycler = itertools.cycle(
176
+ self.prefill_instances)
177
+ else:
178
+ raise HTTPException(status_code=400,
179
+ detail="Instance already exists.")
180
+ else:
181
+ if instance not in self.decode_instances:
182
+ self.decode_instances.append(instance)
183
+ self.decode_cycler = itertools.cycle(self.decode_instances)
184
+ else:
185
+ raise HTTPException(status_code=400,
186
+ detail="Instance already exists.")
187
+
188
+ return JSONResponse(content={
189
+ "message":
190
+ f"Added {instance} to {instance_type}_instances."
191
+ })
192
+ except HTTPException as http_exc:
193
+ raise http_exc
194
+ except Exception as e:
195
+ logger.error("Error in add_instance_endpoint: %s", str(e))
196
+ raise HTTPException(status_code=500, detail=str(e)) from e
197
+
198
+ async def forward_request(self, url, data, use_chunked=True):
199
+ async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
200
+ headers = {
201
+ "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"
202
+ }
203
+ try:
204
+ async with session.post(url=url, json=data,
205
+ headers=headers) as response:
206
+ if 200 <= response.status < 300 or 400 <= response.status < 500: # noqa: E501
207
+ if use_chunked:
208
+ async for chunk_bytes in response.content.iter_chunked( # noqa: E501
209
+ 1024):
210
+ yield chunk_bytes
211
+ else:
212
+ content = await response.read()
213
+ yield content
214
+ else:
215
+ error_content = await response.text()
216
+ try:
217
+ error_content = json.loads(error_content)
218
+ except json.JSONDecodeError:
219
+ error_content = error_content
220
+ logger.error("Request failed with status %s: %s",
221
+ response.status, error_content)
222
+ raise HTTPException(
223
+ status_code=response.status,
224
+ detail=
225
+ f"Request failed with status {response.status}: "
226
+ f"{error_content}",
227
+ )
228
+ except aiohttp.ClientError as e:
229
+ logger.error("ClientError occurred: %s", str(e))
230
+ raise HTTPException(
231
+ status_code=502,
232
+ detail=
233
+ "Bad Gateway: Error communicating with upstream server.",
234
+ ) from e
235
+ except Exception as e:
236
+ logger.error("Unexpected error: %s", str(e))
237
+ raise HTTPException(status_code=500, detail=str(e)) from e
238
+
239
+ def schedule(self, cycler: itertools.cycle) -> str:
240
+ return self.scheduling_policy.schedule(cycler)
241
+
242
+ async def get_status(self):
243
+ status = {
244
+ "prefill_node_count": len(self.prefill_instances),
245
+ "decode_node_count": len(self.decode_instances),
246
+ "prefill_nodes": self.prefill_instances,
247
+ "decode_nodes": self.decode_instances,
248
+ }
249
+ return status
250
+
251
+ async def create_completion(self, raw_request: Request):
252
+ try:
253
+ request = await raw_request.json()
254
+
255
+ kv_prepare_request = request.copy()
256
+ kv_prepare_request["max_tokens"] = 1
257
+
258
+ prefill_instance = self.schedule(self.prefill_cycler)
259
+ try:
260
+ async for _ in self.forward_request(
261
+ f"http://{prefill_instance}/v1/completions",
262
+ kv_prepare_request):
263
+ continue
264
+ except HTTPException as http_exc:
265
+ self.remove_instance_endpoint("prefill", prefill_instance)
266
+ raise http_exc
267
+
268
+ # Perform kv recv and decoding stage
269
+ decode_instance = self.schedule(self.decode_cycler)
270
+
271
+ try:
272
+ generator = self.forward_request(
273
+ f"http://{decode_instance}/v1/completions", request)
274
+ except HTTPException as http_exc:
275
+ self.remove_instance_endpoint("decode", decode_instance)
276
+ raise http_exc
277
+ response = StreamingResponse(generator)
278
+ return response
279
+ except Exception:
280
+ import sys
281
+
282
+ exc_info = sys.exc_info()
283
+ print("Error occurred in disagg proxy server")
284
+ print(exc_info)
285
+
286
+ async def create_chat_completion(self, raw_request: Request):
287
+ try:
288
+ request = await raw_request.json()
289
+
290
+ # add params to request
291
+ kv_prepare_request = request.copy()
292
+ kv_prepare_request["max_tokens"] = 1
293
+
294
+ # prefill stage
295
+ prefill_instance = self.schedule(self.prefill_cycler)
296
+ try:
297
+ async for _ in self.forward_request(
298
+ f"http://{prefill_instance}/v1/chat/completions",
299
+ kv_prepare_request):
300
+ continue
301
+ except HTTPException as http_exc:
302
+ self.remove_instance_endpoint("prefill", prefill_instance)
303
+ raise http_exc
304
+ # Perform kv recv and decoding stage
305
+ decode_instance = self.schedule(self.decode_cycler)
306
+
307
+ try:
308
+ generator = self.forward_request(
309
+ "http://" + decode_instance + "/v1/chat/completions",
310
+ request)
311
+ except HTTPException as http_exc:
312
+ self.remove_instance_endpoint("decode", decode_instance)
313
+ raise http_exc
314
+ response = StreamingResponse(content=generator)
315
+ return response
316
+ except Exception:
317
+ exc_info = sys.exc_info()
318
+ error_messages = [str(e) for e in exc_info if e]
319
+ print("Error occurred in disagg proxy server")
320
+ print(error_messages)
321
+ return StreamingResponse(content=iter(error_messages),
322
+ media_type="text/event-stream")
323
+
324
+ def remove_instance_endpoint(self, instance_type, instance):
325
+ if (instance_type == "decode" and instance in self.decode_instances):
326
+ self.decode_instances.remove(instance)
327
+ self.decode_cycler = itertools.cycle(self.decode_instances)
328
+ if (instance_type == "prefill" and instance in self.decode_instances):
329
+ self.prefill_instances.remove(instance)
330
+ self.prefill_cycler = itertools.cycle(self.decode_instances)
331
+
332
+
333
+class RoundRobinSchedulingPolicy(SchedulingPolicy):
334
+
335
+ def __init__(self):
336
+ super().__init__()
337
+
338
+ def schedule(self, cycler: itertools.cycle) -> str:
339
+ return next(cycler)
340
+
341
+
342
+class ProxyServer:
343
+
344
+ def __init__(
345
+ self,
346
+ args: argparse.Namespace,
347
+ scheduling_policy: Optional[SchedulingPolicy] = None,
348
+ create_completion: Optional[Callable[[Request],
349
+ StreamingResponse]] = None,
350
+ create_chat_completion: Optional[Callable[[Request],
351
+ StreamingResponse]] = None,
352
+ ):
353
+ self.validate_parsed_serve_args(args)
354
+ self.port = args.port
355
+ self.proxy_instance = Proxy(
356
+ prefill_instances=[] if args.prefill is None else args.prefill,
357
+ decode_instances=[] if args.decode is None else args.decode,
358
+ model=args.model,
359
+ scheduling_policy=(scheduling_policy if scheduling_policy
360
+ is not None else RoundRobinSchedulingPolicy()),
361
+ custom_create_completion=create_completion,
362
+ custom_create_chat_completion=create_chat_completion,
363
+ )
364
+
365
+ def validate_parsed_serve_args(self, args: argparse.Namespace):
366
+ if not args.prefill:
367
+ raise ValueError("Please specify at least one prefill node.")
368
+ if not args.decode:
369
+ raise ValueError("Please specify at least one decode node.")
370
+ self.validate_instances(args.prefill)
371
+ self.validate_instances(args.decode)
372
+ self.verify_model_config(args.prefill, args.model)
373
+ self.verify_model_config(args.decode, args.model)
374
+
375
+ def validate_instances(self, instances: list):
376
+ for instance in instances:
377
+ if len(instance.split(":")) != 2:
378
+ raise ValueError(f"Invalid instance format: {instance}")
379
+ host, port = instance.split(":")
380
+ try:
381
+ if host != "localhost":
382
+ ipaddress.ip_address(host)
383
+ port = int(port)
384
+ if not (0 < port < 65536):
385
+ raise ValueError(
386
+ f"Invalid port number in instance: {instance}")
387
+ except Exception as e:
388
+ raise ValueError(
389
+ f"Invalid instance {instance}: {str(e)}") from e
390
+
391
+ def verify_model_config(self, instances: list, model: str) -> None:
392
+ model_suffix = model.split("/")[-1]
393
+ for instance in instances:
394
+ try:
395
+ response = requests.get(f"http://{instance}/v1/models")
396
+ if response.status_code == 200:
397
+ model_cur = response.json()["data"][0]["id"]
398
+ model_cur_suffix = model_cur.split("/")[-1]
399
+ if model_cur_suffix != model_suffix:
400
+ raise ValueError(
401
+ f"{instance} serves a different model: "
402
+ f"{model_cur} != {model}")
403
+ else:
404
+ raise ValueError(f"Cannot get model id from {instance}!")
405
+ except requests.RequestException as e:
406
+ raise ValueError(
407
+ f"Error communicating with {instance}: {str(e)}") from e
408
+
409
+ def run_server(self):
410
+ app = FastAPI()
411
+ app.include_router(self.proxy_instance.router)
412
+ config = uvicorn.Config(app, port=self.port, loop="uvloop")
413
+ server = uvicorn.Server(config)
414
+ server.run()
415
+
416
+
417
+if __name__ == "__main__":
418
+ # Todo: allow more config
419
+ parser = argparse.ArgumentParser("vLLM disaggregated proxy server.")
420
+ parser.add_argument("--model",
421
+ "-m",
422
+ type=str,
423
+ required=True,
424
+ help="Model name")
425
+
426
+ parser.add_argument(
427
+ "--prefill",
428
+ "-p",
429
+ type=str,
430
+ nargs="+",
431
+ help="List of prefill node URLs (host:port)",
432
+ )
433
+
434
+ parser.add_argument(
435
+ "--decode",
436
+ "-d",
437
+ type=str,
438
+ nargs="+",
439
+ help="List of decode node URLs (host:port)",
440
+ )
441
+
442
+ parser.add_argument(
443
+ "--port",
444
+ type=int,
445
+ default=8000,
446
+ help="Server port number",
447
+ )
448
+ args = parser.parse_args()
449
+ proxy_server = ProxyServer(args=args)
450
+ proxy_server.run_server()
vllm/distributed/kv_transfer/kv_connector/factory.py
+5
@@ -53,3 +53,8 @@ KVConnectorFactory.register_connector(
53
"LMCacheConnector",
54
"vllm.distributed.kv_transfer.kv_connector.lmcache_connector",
55
"LMCacheConnector")
56
+
57
+KVConnectorFactory.register_connector(
58
+ "MooncakeStoreConnector",
59
+ "vllm.distributed.kv_transfer.kv_connector.mooncake_store_connector",
60
+ "MooncakeStoreConnector")
\ No newline at end of file
vllm/distributed/kv_transfer/kv_connector/mooncake_store_connector.py
new
+216
@@ -0,0 +1,216 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+MooncakeStore Connector for Distributed Machine Learning Inference
4
+
5
+The MooncakeStoreConnector transfers KV caches between prefill vLLM workers
6
+(KV cache producer) and decode vLLM workers (KV cache consumer) using a
7
+database-style KVStore.
8
+"""
9
+import hashlib
10
+from typing import TYPE_CHECKING, List, Tuple, Union
11
+
12
+import torch
13
+
14
+from vllm import _custom_ops as ops
15
+from vllm.config import VllmConfig
16
+from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
17
+from vllm.logger import init_logger
18
+from vllm.sequence import IntermediateTensors
19
+
20
+if TYPE_CHECKING:
21
+ from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
22
+
23
+logger = init_logger(__name__)
24
+
25
+
26
+class MooncakeStoreConnector(KVConnectorBase):
27
+
28
+ def __init__(
29
+ self,
30
+ rank: int,
31
+ local_rank: int,
32
+ config: VllmConfig,
33
+ ):
34
+ self.config = config.kv_transfer_config
35
+ self.tp_size = config.parallel_config.tensor_parallel_size
36
+
37
+ self.local_tp_rank = local_rank
38
+
39
+ # Init kv_store
40
+ if self.config.kv_connector == "MooncakeStoreConnector":
41
+ # Check if MOONCAKE_CONFIG_PATH is set
42
+ import os
43
+ use_mooncake_store = os.getenv('MOONCAKE_CONFIG_PATH') is not None
44
+
45
+ if not use_mooncake_store:
46
+ raise ValueError(
47
+ "To use MooncakeStoreConnector, you need to pass the ENV: "
48
+ "'MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json'.")
49
+ else:
50
+ from vllm.distributed.kv_transfer.kv_lookup_buffer.mooncake_store import ( # noqa: E501
51
+ MooncakeStore)
52
+ logger.info(
53
+ "Initializing KVStoreConnector under kv_transfer_config %s",
54
+ self.config)
55
+ self.kv_store = MooncakeStore(config)
56
+ else:
57
+ logger.error("Can not find %s", self.config.kv_connector)
58
+
59
+ assert self.kv_store is not None
60
+
61
+ def close(self) -> None:
62
+ """Close the buffer and release resources.
63
+ This method is responsible for cleaning up resources related to the
64
+ connector when it is no longer needed.
65
+ Raises:
66
+ NotImplementedError: This method must be implemented in subclasses.
67
+ """
68
+ self.kv_store.close()
69
+
70
+ def send_kv_caches_and_hidden_states(
71
+ self,
72
+ model_executable: torch.nn.Module,
73
+ model_input: "ModelInputForGPUWithSamplingMetadata",
74
+ kv_caches: List[torch.Tensor],
75
+ hidden_or_intermediate_states: Union[torch.Tensor,
76
+ IntermediateTensors],
77
+ ) -> None:
78
+ input_tokens_tensor = model_input.input_tokens
79
+ seq_lens = model_input.attn_metadata.seq_lens
80
+ slot_mapping_flat = model_input.attn_metadata.slot_mapping.flatten()
81
+ start_layer = model_executable.model.start_layer
82
+ end_layer = model_executable.model.end_layer
83
+
84
+ model_config = model_executable.model.config
85
+ num_heads = int(model_config.num_key_value_heads / self.tp_size)
86
+ hidden_size = model_config.hidden_size
87
+ num_attention_heads = model_config.num_attention_heads
88
+ head_size = int(hidden_size / num_attention_heads)
89
+
90
+ for idx, slen in enumerate(seq_lens):
91
+ start_pos = sum(seq_lens[:idx])
92
+ end_pos = start_pos + slen
93
+
94
+ current_tokens = input_tokens_tensor[start_pos:end_pos]
95
+ store_key_prefix = self.tensor_hash(current_tokens)
96
+ keys, values = [], []
97
+
98
+ for layer_id in range(start_layer, end_layer):
99
+ kv_cache = kv_caches[layer_id - start_layer]
100
+
101
+ key_cache = kv_cache[0].reshape(-1, num_heads, head_size)
102
+ value_cache = kv_cache[1].reshape(-1, num_heads, head_size)
103
+
104
+ current_slot_mapping = slot_mapping_flat[start_pos:end_pos]
105
+
106
+ keys.append(key_cache[current_slot_mapping].unsqueeze(0))
107
+ values.append(value_cache[current_slot_mapping].unsqueeze(0))
108
+
109
+ keys = torch.cat(keys, dim=0)
110
+ values = torch.cat(values, dim=0)
111
+ kvcache_to_sent = torch.stack((keys, values), dim=0)
112
+ store_kvcache_key = f"{store_key_prefix}_{self.local_tp_rank}"
113
+ self.kv_store.put(store_kvcache_key, kvcache_to_sent)
114
+
115
+ hidden_key = f"{store_key_prefix}_hidden_{self.local_tp_rank}"
116
+ self.kv_store.put(hidden_key,
117
+ hidden_or_intermediate_states[start_pos:end_pos])
118
+
119
+ logger.debug("[rank%d]: KV send DONE.", torch.distributed.get_rank())
120
+
121
+ def recv_kv_caches_and_hidden_states(
122
+ self, model_executable: torch.nn.Module,
123
+ model_input: "ModelInputForGPUWithSamplingMetadata",
124
+ kv_caches: List[torch.Tensor]
125
+ ) -> Tuple[Union[torch.Tensor, IntermediateTensors], bool,
126
+ "ModelInputForGPUWithSamplingMetadata"]:
127
+ bypass_model_exec = True
128
+ input_tokens_tensor = model_input.input_tokens
129
+ seq_lens = model_input.attn_metadata.seq_lens
130
+ num_prefill_tokens = model_input.attn_metadata.num_prefill_tokens
131
+ slot_mapping = model_input.attn_metadata.slot_mapping.flatten()
132
+ start_layer = model_executable.model.start_layer
133
+ end_layer = model_executable.model.end_layer
134
+ hidden_or_intermediate_states_for_one_req = []
135
+
136
+ for idx, slen in enumerate(seq_lens):
137
+ start_pos = sum(seq_lens[:idx])
138
+ end_pos = start_pos + slen
139
+
140
+ if start_pos >= num_prefill_tokens:
141
+ # This can happen during inflight batching. See:
142
+ # vllm/worker/model_runner.py::_prepare_model_input_tensors:
143
+ # - input_tokens[:num_prefill_tokens] contains prefill tokens.
144
+ # - input_tokens[num_prefill_tokens:] contains decode tokens.
145
+ logger.warning("You should set --enable_chunked_prefill=False "
146
+ "and --max_num_batched_tokens "
147
+ "should be equal to max_seq_len_to_capture")
148
+ bypass_model_exec = False
149
+ assert start_pos == num_prefill_tokens
150
+ break
151
+
152
+ current_tokens = input_tokens_tensor[start_pos:end_pos]
153
+
154
+ # get roi for current seq
155
+ load_key_prefix = self.tensor_hash(current_tokens)
156
+ load_kvcache_key = f"{load_key_prefix}_{self.local_tp_rank}"
157
+ remote_kv = self.kv_store.get(load_kvcache_key)
158
+ hidden_key = f"{load_key_prefix}_hidden_{self.local_tp_rank}"
159
+ hidden = self.kv_store.get(hidden_key)
160
+
161
+ if remote_kv is None or hidden is None:
162
+ # didn't find any match.
163
+ bypass_model_exec = False
164
+ continue
165
+
166
+ num_computed_tokens = current_tokens.shape[0]
167
+
168
+ # update the end position based on how many tokens are cached.
169
+ end_pos = start_pos + num_computed_tokens
170
+
171
+ # call self.kv_store to get kv layer by layer
172
+ for layer_id in range(start_layer, end_layer):
173
+ layer = model_executable.model.layers[layer_id]
174
+ # get kvcache object
175
+ kv_cache = kv_caches[layer_id - start_layer]
176
+ key_cache, value_cache = kv_cache[0], kv_cache[1]
177
+ # get remote kvcache
178
+
179
+ remote_k, remote_v = remote_kv[0][layer_id], remote_kv[1][
180
+ layer_id]
181
+ # use ops.reshape_and_cache_flash to put kv into kvcache
182
+ ops.reshape_and_cache_flash(
183
+ remote_k.to(key_cache.device),
184
+ remote_v.to(value_cache.device),
185
+ key_cache,
186
+ value_cache,
187
+ slot_mapping[start_pos:end_pos],
188
+ layer.self_attn.attn.kv_cache_dtype,
189
+ layer.self_attn.attn._k_scale,
190
+ layer.self_attn.attn._v_scale,
191
+ )
192
+
193
+ hidden_or_intermediate_states_for_one_req.append(hidden)
194
+
195
+ if not bypass_model_exec:
196
+ logger.warning(
197
+ "[rank%d]: Failed to receive all KVs and hidden "
198
+ "states, redo model forwarding.", torch.distributed.get_rank())
199
+ hidden_or_intermediate_states = None
200
+
201
+ else:
202
+ logger.debug(
203
+ "[rank%d]: Successfully received all KVs and hidden "
204
+ "states, skip model forwarding.", torch.distributed.get_rank())
205
+ hidden_or_intermediate_states = torch.cat(
206
+ hidden_or_intermediate_states_for_one_req, dim=0)
207
+
208
+ return hidden_or_intermediate_states, bypass_model_exec, model_input
209
+
210
+ @staticmethod
211
+ def tensor_hash(tensor: torch.Tensor) -> int:
212
+ """Calculate the hash value of the tensor."""
213
+ tensor_bytes = tensor.clone().detach().cpu().numpy().tobytes()
214
+ hash_object = hashlib.blake2b(tensor_bytes)
215
+ hash_hex = hash_object.hexdigest()
216
+ return int(hash_hex[:16], 16)
vllm/distributed/kv_transfer/kv_lookup_buffer/base.py
+75
-10
@@ -1,11 +1,15 @@
1
# SPDX-License-Identifier: Apache-2.0
2
"""
3
-This file contains a new class `KVLookupBufferBase` that allows developers to
4
-think of KV cache operations as inserting new KV cache entries (`insert`)
5
-into the lookup buffer and querying existing KV caches (`drop_select`)
3
+This file contains a new class `KVLookupBufferBase` that allows developers to
4
+think of KV cache operations as inserting new KV cache entries (`insert`)
5
+into the lookup buffer and querying existing KV caches (`drop_select`)
6
from the lookup buffer.
7
8
-All distributed communications are abstracted behind this class.
8
+This file also contains a new class `KVStoreBufferBase` that allows developers
9
+to manage the KVCache buffer as a simple key-value storage buffer with basic
10
+put/get operations.
11
+
12
+These classes above are abstracted behind class `KVCacheBufferBase`.
13
"""
14
15
from abc import ABC, abstractmethod
@@ -14,9 +18,27 @@ from typing import List, Optional
18
import torch
19
20
17
-class KVLookupBufferBase(ABC):
21
+class KVCacheBufferBase(ABC):
22
+ """
23
+ Abstract base class for a KVCache buffer.
24
"""
19
- Abstract base class for a lookup buffer.
25
+
26
+ @abstractmethod
27
+ def close(self) -> None:
28
+ """Close the buffer and release resources.
29
+
30
+ This method is responsible for cleaning up resources related to the
31
+ KVCache buffer when it is no longer needed.
32
+
33
+ Raises:
34
+ NotImplementedError: This method must be implemented in subclasses.
35
+ """
36
+ raise NotImplementedError
37
+
38
+
39
+class KVLookupBufferBase(KVCacheBufferBase):
40
+ """
41
+ Abstract base class for a KVCache lookup buffer.
42
43
This class provides an abstraction for a key-value (KV) cache lookup buffer.
44
@@ -96,12 +118,55 @@ class KVLookupBufferBase(ABC):
118
"""
119
raise NotImplementedError
120
121
+
122
+class KVStoreBufferBase(KVCacheBufferBase):
123
+ """
124
+ Abstract base class for a KVCache storage buffer with key-value semantics.
125
+ This class provides a simple key-value storage buffer abstract with basic
126
+ put/get operations, which enables flexible KVCache transfer granular
127
+ control.
128
+
129
+ The functionality is similar to a distributed key-value store, where:
130
+ - Key: A unique string identifier for the cached entry
131
+ - Value:
132
+ - Tensor to be stored and retrieved
133
+ - None (indicating deletion or empty value)
134
+ """
135
+
136
+ @abstractmethod
137
+ def put(
138
+ self,
139
+ key: str,
140
+ value: Optional[torch.Tensor],
141
+ ) -> None:
142
+ """Store a key-value pair in the buffer.
143
+
144
+ Args:
145
+ key (str): Unique identifier for a tensor, this tensor could be the
146
+ key cache tensor, value cache tensor, or hidden state tensor
147
+ generated during model forwarding.
148
+
149
+ value (Optional[torch.Tensor]): Tensor to be stored.
150
+
151
+ Raises:
152
+ NotImplementedError: This method must be implemented in subclasses.
153
+ """
154
+ raise NotImplementedError
155
+
156
@abstractmethod
100
- def close(self) -> None:
101
- """Close the buffer and release resources.
157
+ def get(
158
+ self,
159
+ key: str,
160
+ ) -> Optional[torch.Tensor]:
161
+ """Retrieve a value from the buffer by key.
162
+
163
+ Args:
164
+ key (str): Unique identifier for a tensor, this tensor could be the
165
+ key cache tensor, value cache tensor, or hidden state tensor
166
+ generated during model forwarding.
167
103
- This method is responsible for cleaning up resources related to the
104
- lookup buffer when it is no longer needed.
168
+ Returns:
169
+ Optional[torch.Tensor]: Stored tensor if exists, None otherwise.
170
171
Raises:
172
NotImplementedError: This method must be implemented in subclasses.
vllm/distributed/kv_transfer/kv_lookup_buffer/mooncake_store.py
new
+160
@@ -0,0 +1,160 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+This file contains a new class `MooncakeStore` that allows developers to
4
+think of KV cache transfer operations as putting new KV cache entries
5
+into a remote KVStore-based lookup buffer and getting existing KV caches
6
+from this remote lookup buffer.
7
+"""
8
+import json
9
+import os
10
+from dataclasses import dataclass
11
+from typing import Optional
12
+
13
+import torch
14
+from safetensors.torch import load as safetensors_load
15
+from safetensors.torch import save as safetensors_save
16
+
17
+from vllm.config import VllmConfig
18
+from vllm.distributed.kv_transfer.kv_lookup_buffer.base import (
19
+ KVStoreBufferBase)
20
+from vllm.logger import init_logger
21
+
22
+DEFAULT_GLOBAL_SEGMENT_SIZE = 3355443200 # 3.125 GiB
23
+DEFAULT_LOCAL_BUFFER_SIZE = 1073741824 # 1.0 GiB
24
+
25
+logger = init_logger(__name__)
26
+
27
+
28
+@dataclass
29
+class MooncakeStoreConfig:
30
+ local_hostname: str
31
+ metadata_server: str
32
+ global_segment_size: int
33
+ local_buffer_size: int
34
+ protocol: str
35
+ device_name: str
36
+ master_server_address: str
37
+
38
+ @staticmethod
39
+ def from_file(file_path: str) -> 'MooncakeStoreConfig':
40
+ """Load the config from a JSON file."""
41
+ with open(file_path) as fin:
42
+ config = json.load(fin)
43
+ return MooncakeStoreConfig(
44
+ local_hostname=config.get("local_hostname"),
45
+ metadata_server=config.get("metadata_server"),
46
+ global_segment_size=config.get("global_segment_size",
47
+ DEFAULT_GLOBAL_SEGMENT_SIZE),
48
+ local_buffer_size=config.get("local_buffer_size",
49
+ DEFAULT_LOCAL_BUFFER_SIZE),
50
+ protocol=config.get("protocol", "tcp"),
51
+ device_name=config.get("device_name", ""),
52
+ master_server_address=config.get("master_server_address"),
53
+ )
54
+
55
+ @staticmethod
56
+ def load_from_env() -> 'MooncakeStoreConfig':
57
+ """Load config from a file specified in the environment variable."""
58
+ config_file_path = os.getenv('MOONCAKE_CONFIG_PATH')
59
+ if config_file_path is None:
60
+ raise ValueError(
61
+ "The environment variable 'MOONCAKE_CONFIG_PATH' is not set.")
62
+ return MooncakeStoreConfig.from_file(config_file_path)
63
+
64
+
65
+class MooncakeStore(KVStoreBufferBase):
66
+
67
+ def __init__(
68
+ self,
69
+ config: VllmConfig,
70
+ ):
71
+
72
+ try:
73
+ from mooncake_vllm_adaptor import MooncakeDistributedStore
74
+ except ImportError as e:
75
+ raise ImportError(
76
+ "Please install mooncake by following the instructions at "
77
+ "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md " # noqa: E501
78
+ "to run vLLM with MooncakeConnector.") from e
79
+
80
+ try:
81
+ self.store = MooncakeDistributedStore()
82
+ self.config = MooncakeStoreConfig.load_from_env()
83
+ logger.info("Mooncake Configuration loaded successfully.")
84
+
85
+ self.store.setup(self.config.local_hostname,
86
+ self.config.metadata_server,
87
+ self.config.global_segment_size,
88
+ self.config.local_buffer_size,
89
+ self.config.protocol, self.config.device_name,
90
+ self.config.master_server_address)
91
+
92
+ except ValueError as e:
93
+ logger.error("Configuration loading failed: %s", e)
94
+ raise
95
+ except Exception as exc:
96
+ logger.error(
97
+ "An error occurred while loading the configuration: %s", exc)
98
+ raise
99
+
100
+ def close(self):
101
+ # MooncakeDistributedStore will automatically call the destructor, so
102
+ # it is unnecessary to close it manually.
103
+ pass
104
+
105
+ def put(
106
+ self,
107
+ key: str,
108
+ value: Optional[torch.Tensor],
109
+ ) -> None:
110
+ # A message queue needs to be introduced before making it asynchronous.
111
+ if value is not None:
112
+ self._put_impl(key, value)
113
+
114
+ def get(
115
+ self,
116
+ key: str,
117
+ ) -> Optional[torch.Tensor]:
118
+ # A message queue needs to be introduced before making it asynchronous.
119
+ value = self._get_impl(key)
120
+ return value
121
+
122
+ def _put_impl(
123
+ self,
124
+ key: str,
125
+ value: torch.Tensor,
126
+ ) -> None:
127
+ """Put KVCache to Mooncake Store"""
128
+ device_id = value.device.index if value.device.type == 'cuda' else -1
129
+ device_tensor = torch.tensor(device_id, dtype=torch.int32)
130
+ value_bytes = safetensors_save({
131
+ "tensor": value,
132
+ "device_id": device_tensor
133
+ })
134
+ try:
135
+ self.store.put(key, value_bytes)
136
+ except TypeError as err:
137
+ logger.error("Failed to put value into Mooncake Store: %s", err)
138
+ raise TypeError("Mooncake Store Put Type Error.") from err
139
+
140
+ def _get_impl(
141
+ self,
142
+ key: str,
143
+ ) -> Optional[torch.Tensor]:
144
+ """Get KVCache from Mooncake Store"""
145
+ try:
146
+ data = self.store.get(key)
147
+ except TypeError as err:
148
+ logger.error("Failed to get value from Mooncake Store: %s", err)
149
+ raise TypeError("Mooncake Store Get Type Error.") from err
150
+
151
+ if data:
152
+ loaded_tensors = safetensors_load(data)
153
+ tensor = loaded_tensors["tensor"]
154
+ device_id_tensor = loaded_tensors["device_id"]
155
+ device_id = int(device_id_tensor.item())
156
+ device = torch.device(
157
+ 'cuda', device_id) if device_id >= 0 else torch.device('cpu')
158
+ return tensor.to(device)
159
+
160
+ return None