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, 14:45 UTC
[Misc] Update usage with mooncake lib for kv transfer (#16523)
Signed-off-by: Shangming Cai <caishangming@linux.alibaba.com>
shangmingc committed
Apr 14, 2025 at 19:31 UTC
1dd23386ecab7b7c50ea61b8ff37ca14d2dbc0f7
2 files changed
+11
-11
vllm/distributed/kv_transfer/kv_lookup_buffer/mooncake_store.py
+1
-1
@@ -70,7 +70,7 @@ class MooncakeStore(KVStoreBufferBase):
70
):
71
72
try:
73
- from mooncake_vllm_adaptor import MooncakeDistributedStore
73
+ from mooncake.store import MooncakeDistributedStore
74
except ImportError as e:
75
raise ImportError(
76
"Please install mooncake by following the instructions at "
vllm/distributed/kv_transfer/kv_pipe/mooncake_pipe.py
+10
-10
@@ -57,14 +57,14 @@ class MooncakeTransferEngine:
57
58
def __init__(self, kv_rank: int, local_rank: int):
59
try:
60
- import mooncake_vllm_adaptor as mva
60
+ from mooncake.engine import TransferEngine
61
except ImportError as e:
62
raise ImportError(
63
"Please install mooncake by following the instructions at "
64
"https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md " # noqa: E501
65
"to run vLLM with MooncakeConnector.") from e
66
67
- self.engine = mva.mooncake_vllm_adaptor()
67
+ self.engine = TransferEngine()
68
self.local_rank = local_rank
69
70
try:
@@ -140,12 +140,12 @@ class MooncakeTransferEngine:
140
"Mooncake Configuration error. `metadata_backend`"
141
f" should be one of {supported_backend}.")
142
143
- self.engine.initializeExt(local_hostname, metadata_server,
144
- protocol, device_name, metadata_backend)
143
+ self.engine.initialize_ext(local_hostname, metadata_server,
144
+ protocol, device_name, metadata_backend)
145
146
def allocate_managed_buffer(self, length: int) -> int:
147
"""Allocate a managed buffer of the specified length."""
148
- ret = self.engine.allocateManagedBuffer(length)
148
+ ret = self.engine.allocate_managed_buffer(length)
149
if ret <= 0:
150
logger.error("Allocation Return Error")
151
raise Exception("Allocation Return Error")
@@ -153,13 +153,13 @@ class MooncakeTransferEngine:
153
154
def free_managed_buffer(self, buffer: int, length: int) -> int:
155
"""Free a previously allocated managed buffer."""
156
- return self.engine.freeManagedBuffer(buffer, length)
156
+ return self.engine.free_managed_buffer(buffer, length)
157
158
def transfer_sync(self, buffer: int, peer_buffer_address: int,
159
length: int) -> int:
160
"""Synchronously transfer data to the specified address."""
161
- ret = self.engine.transferSync(self.remote_url, buffer,
162
- peer_buffer_address, length)
161
+ ret = self.engine.transfer_sync_read(self.remote_url, buffer,
162
+ peer_buffer_address, length)
163
if ret < 0:
164
logger.error("Transfer Return Error")
165
raise Exception("Transfer Return Error")
@@ -168,11 +168,11 @@ class MooncakeTransferEngine:
168
def write_bytes_to_buffer(self, buffer: int, user_data: bytes,
169
length: int) -> int:
170
"""Write bytes to the allocated buffer."""
171
- return self.engine.writeBytesToBuffer(buffer, user_data, length)
171
+ return self.engine.write_bytes_to_buffer(buffer, user_data, length)
172
173
def read_bytes_from_buffer(self, buffer: int, length: int) -> bytes:
174
"""Read bytes from the allocated buffer."""
175
- return self.engine.readBytesFromBuffer(buffer, length)
175
+ return self.engine.read_bytes_from_buffer(buffer, length)
176
177
def wait_for_ack(self, src_ptr: int, length: int) -> None:
178
"""Asynchronously wait for ACK from the receiver."""