| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Mooncake data-plane engine and memory-registration ownership. |
| 4 | |
| 5 | The wrapper isolates optional-engine initialization, session addressing, |
| 6 | synchronous batched writes, and the lifetime of transient source registrations. |
| 7 | """ |
| 8 | |
| 9 | from __future__ import annotations |
| 10 | |
| 11 | import threading |
| 12 | from dataclasses import dataclass |
| 13 | |
| 14 | import torch |
| 15 | |
| 16 | from vllm.logger import init_logger |
| 17 | |
| 18 | logger = init_logger(__name__) |
| 19 | |
| 20 | _MOONCAKE_IMPORT_ERROR: ImportError | None = None |
| 21 | try: |
| 22 | from mooncake.engine import TransferEngine |
| 23 | except ImportError as error: |
| 24 | TransferEngine = None # type: ignore[misc, assignment] |
| 25 | _MOONCAKE_IMPORT_ERROR = error |
| 26 | |
| 27 | |
| 28 | def ensure_mooncake_available() -> None: |
| 29 | if _MOONCAKE_IMPORT_ERROR is not None: |
| 30 | raise ImportError( |
| 31 | "Install mooncake-transfer-engine to use ECMooncakeConnector." |
| 32 | ) from _MOONCAKE_IMPORT_ERROR |
| 33 | |
| 34 | |
| 35 | @dataclass |
| 36 | class _SourceRegistration: |
| 37 | """Retain one transient source range while batches reference it.""" |
| 38 | |
| 39 | tensor: torch.Tensor |
| 40 | nbytes: int |
| 41 | users: int = 1 |
| 42 | |
| 43 | |
| 44 | class MooncakeTransfer: |
| 45 | """Own a lazy Mooncake engine and transient memory registrations.""" |
| 46 | |
| 47 | def __init__(self, hostname: str, protocol: str) -> None: |
| 48 | self._hostname = hostname |
| 49 | self._protocol = protocol |
| 50 | self._engine: TransferEngine | None = None |
| 51 | self._engine_lock = threading.Lock() |
| 52 | self._source_registrations: dict[int, _SourceRegistration] = {} |
| 53 | self._pending_unregister: dict[int, torch.Tensor] = {} |
| 54 | self._registration_lock = threading.Lock() |
| 55 | self._closed = False |
| 56 | |
| 57 | def _ensure_engine(self) -> TransferEngine: |
| 58 | if self._engine is not None: |
| 59 | return self._engine |
| 60 | with self._engine_lock: |
| 61 | if self._engine is not None: |
| 62 | return self._engine |
| 63 | engine = TransferEngine() |
| 64 | ret = engine.initialize(self._hostname, "P2PHANDSHAKE", self._protocol, "") |
| 65 | if ret != 0: |
| 66 | raise RuntimeError("Mooncake TransferEngine initialization failed.") |
| 67 | self._engine = engine |
| 68 | logger.info( |
| 69 | "ECMooncakeConnector TransferEngine ready at %s:%d", |
| 70 | self._hostname, |
| 71 | engine.get_rpc_port(), |
| 72 | ) |
| 73 | return self._engine |
| 74 | |
| 75 | def ensure_ready(self) -> None: |
| 76 | self._ensure_engine() |
| 77 | |
| 78 | def local_session(self) -> str: |
| 79 | engine = self._ensure_engine() |
| 80 | return f"{self._hostname}:{engine.get_rpc_port()}" |
| 81 | |
| 82 | def register_memory(self, tensor: torch.Tensor) -> int: |
| 83 | return self._ensure_engine().batch_register_memory( |
| 84 | [tensor.data_ptr()], [tensor.nbytes] |
| 85 | ) |
| 86 | |
| 87 | def unregister_memory(self, tensor: torch.Tensor) -> bool: |
| 88 | engine = self._ensure_engine() |
| 89 | address = tensor.data_ptr() |
| 90 | ret = engine.unregister_memory(address) |
| 91 | with self._registration_lock: |
| 92 | if ret != 0: |
| 93 | logger.error( |
| 94 | "Mooncake EC memory unregistration failed for address %d: %d", |
| 95 | address, |
| 96 | ret, |
| 97 | ) |
| 98 | self._pending_unregister[address] = tensor |
| 99 | return False |
| 100 | self._pending_unregister.pop(address, None) |
| 101 | return True |
| 102 | |
| 103 | @staticmethod |
| 104 | def _source_range(tensor: torch.Tensor) -> tuple[int, int]: |
| 105 | # Encoder batches commonly split one storage into sibling tensor views. |
| 106 | # Registering each view's exact bytes avoids overlapping memory regions. |
| 107 | return tensor.data_ptr(), tensor.nbytes |
| 108 | |
| 109 | def acquire_sources(self, tensors: list[torch.Tensor]) -> list[int]: |
| 110 | ranges: dict[int, tuple[int, torch.Tensor]] = {} |
| 111 | for tensor in tensors: |
| 112 | address, nbytes = self._source_range(tensor) |
| 113 | ranges.setdefault(address, (nbytes, tensor)) |
| 114 | |
| 115 | engine = self._ensure_engine() |
| 116 | acquired: list[int] = [] |
| 117 | new_addresses: list[int] = [] |
| 118 | new_lengths: list[int] = [] |
| 119 | with self._registration_lock: |
| 120 | for address, (nbytes, tensor) in ranges.items(): |
| 121 | entry = self._source_registrations.get(address) |
| 122 | if entry is not None: |
| 123 | if entry.nbytes != nbytes: |
| 124 | raise RuntimeError( |
| 125 | "Mooncake EC source storage changed size while registered" |
| 126 | ) |
| 127 | entry.users += 1 |
| 128 | acquired.append(address) |
| 129 | continue |
| 130 | new_addresses.append(address) |
| 131 | new_lengths.append(nbytes) |
| 132 | self._source_registrations[address] = _SourceRegistration( |
| 133 | tensor=tensor, |
| 134 | nbytes=nbytes, |
| 135 | ) |
| 136 | acquired.append(address) |
| 137 | |
| 138 | if new_addresses: |
| 139 | ret = engine.batch_register_memory(new_addresses, new_lengths) |
| 140 | if ret != 0: |
| 141 | for address in acquired: |
| 142 | entry = self._source_registrations[address] |
| 143 | entry.users -= 1 |
| 144 | if entry.users == 0: |
| 145 | del self._source_registrations[address] |
| 146 | raise RuntimeError("Mooncake EC source registration failed") |
| 147 | return acquired |
| 148 | |
| 149 | def release_sources(self, addresses: list[int]) -> bool: |
| 150 | if not addresses: |
| 151 | return True |
| 152 | with self._registration_lock: |
| 153 | unused = [] |
| 154 | for address in addresses: |
| 155 | entry = self._source_registrations.get(address) |
| 156 | if entry is None: |
| 157 | continue |
| 158 | entry.users -= 1 |
| 159 | if entry.users == 0: |
| 160 | unused.append(address) |
| 161 | if not unused: |
| 162 | return True |
| 163 | ret = self._ensure_engine().batch_unregister_memory(unused) |
| 164 | if ret != 0: |
| 165 | logger.warning( |
| 166 | "Keeping %d EC source tensors registered after Mooncake " |
| 167 | "unregistration failure", |
| 168 | len(unused), |
| 169 | ) |
| 170 | return False |
| 171 | for address in unused: |
| 172 | del self._source_registrations[address] |
| 173 | self._pending_unregister.pop(address, None) |
| 174 | return True |
| 175 | |
| 176 | def write( |
| 177 | self, |
| 178 | session: str, |
| 179 | sources: list[int], |
| 180 | destinations: list[int], |
| 181 | lengths: list[int], |
| 182 | ) -> None: |
| 183 | """Write one synchronous batch, returning only at terminal status.""" |
| 184 | ret = self._ensure_engine().batch_transfer_sync_write( |
| 185 | session, sources, destinations, lengths |
| 186 | ) |
| 187 | if ret != 0: |
| 188 | raise RuntimeError( |
| 189 | f"Mooncake EC push to {session} failed with status {ret}" |
| 190 | ) |
| 191 | |
| 192 | def close(self) -> None: |
| 193 | if self._closed: |
| 194 | return |
| 195 | self._closed = True |
| 196 | engine = self._engine |
| 197 | if engine is None: |
| 198 | return |
| 199 | with self._registration_lock: |
| 200 | addresses = list(self._source_registrations) |
| 201 | addresses.extend(self._pending_unregister) |
| 202 | if not addresses: |
| 203 | return |
| 204 | ret = engine.batch_unregister_memory(list(dict.fromkeys(addresses))) |
| 205 | if ret != 0: |
| 206 | logger.error("Mooncake EC batch memory unregistration failed: %d", ret) |
| 207 | return |
| 208 | self._source_registrations.clear() |
| 209 | self._pending_unregister.clear() |