@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 842aba501d
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 12, 02:15 UTC

[P/D] Introduce Mooncake Transfer Engine as kv_connector (#24718)

Signed-off-by: Tianchen Ding <dtcccc@linux.alibaba.com> Signed-off-by: dtc <dtcccc@linux.alibaba.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com>

dtc committed Dec 4, 2025 at 17:51 UTC 842aba501d92ac77874d45612b7e3c6fed2ca243
6 files changed +1114 -125
docs/features/mooncake_connector_usage.md new
+58
@@ -0,0 +1,58 @@
1 +# MooncakeConnector Usage Guide
2 +
3 +## About Mooncake
4 +
5 +Mooncake aims to enhance the inference efficiency of large language models (LLMs), especially in slow object storage environments, by constructing a multi-level caching pool on high-speed interconnected DRAM/SSD resources. Compared to traditional caching systems, Mooncake utilizes (GPUDirect) RDMA technology to transfer data directly in a zero-copy manner, while maximizing the use of multi-NIC resources on a single machine.
6 +
7 +For more details about Mooncake, please refer to [Mooncake project](https://github.com/kvcache-ai/Mooncake) and [Mooncake documents](https://kvcache-ai.github.io/Mooncake/).
8 +
9 +## Prerequisites
10 +
11 +### Installation
12 +
13 +Install mooncake through pip: `uv pip install mooncake-transfer-engine`.
14 +
15 +Refer to [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for more installation instructions
16 +
17 +## Usage
18 +
19 +### Prefiller Node (192.168.0.2)
20 +
21 +```bash
22 +vllm serve Qwen/Qwen2.5-7B-Instruct --port 8010 --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer"}'
23 +```
24 +
25 +### Decoder Node (192.168.0.3)
26 +
27 +```bash
28 +vllm serve Qwen/Qwen2.5-7B-Instruct --port 8020 --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer"}'
29 +```
30 +
31 +### Proxy
32 +
33 +```bash
34 +python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --prefiller-host 192.168.0.2 --prefiller-port 8010 --decoder-host 192.168.0.3 --decoder-port 8020
35 +```
36 +
37 +> NOTE: The Mooncake Connector currently uses the proxy from nixl_integration. This will be replaced with a self-developed proxy in the future.
38 +
39 +Now you can send requests to the proxy server through port 8000.
40 +
41 +## Environment Variables
42 +
43 +- `VLLM_MOONCAKE_BOOTSTRAP_PORT`: Port for Mooncake bootstrap server
44 + - Default: 8998
45 + - Required only for prefiller instances
46 + - Each vLLM worker needs a unique port on its host; using the same port number across different hosts is fine
47 + - For TP/DP deployments, each worker's port on a node is computed as: base_port + dp_rank * tp_size + tp_rank
48 + - Used for the decoder notifying the prefiller
49 +
50 +- `VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing the prefiller’s KV cache for a particular request. (Optional)
51 + - Default: 480
52 + - If a request is aborted and the decoder has not yet notified the prefiller, the prefill instance will release its KV-cache blocks after this timeout to avoid holding them indefinitely.
53 +
54 +## KV Role Options
55 +
56 +- **kv_producer**: For prefiller instances that generate KV caches
57 +- **kv_consumer**: For decoder instances that consume KV caches from prefiller
58 +- **kv_both**: Enables symmetric functionality where the connector can act as both producer and consumer. This provides flexibility for experimental setups and scenarios where the role distinction is not predetermined.
vllm/distributed/kv_transfer/kv_connector/factory.py
+5
@@ -190,3 +190,8 @@ KVConnectorFactory.register_connector(
190 "vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector",
191 "DecodeBenchConnector",
192 )
193 +KVConnectorFactory.register_connector(
194 + "MooncakeConnector",
195 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake_connector",
196 + "MooncakeConnector",
197 +)
vllm/distributed/kv_transfer/kv_connector/utils.py
+124
@@ -4,10 +4,13 @@
4 KV cache helper for store.
5 """
6
7 +from dataclasses import dataclass
8 from typing import TYPE_CHECKING, Literal
9
10 import torch
11
12 +from vllm.attention.backends.abstract import AttentionBackend
13 +from vllm.attention.backends.registry import AttentionBackendEnum
14 from vllm.config import get_current_vllm_config
15 from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
16 from vllm.logger import init_logger
@@ -181,3 +184,124 @@ def copy_kv_blocks(
184 src_tensor = src_kv_caches[layer_name]
185 dst_tensor = dst_kv_caches[layer_name]
186 copy_fn(src_tensor, dst_tensor, src_indices, dst_indices)
187 +
188 +
189 +@dataclass
190 +class TpKVTopology:
191 + """
192 + Helper class for tensor parallel and KV topology information for
193 + mapping between local and remote TP workers.
194 + """
195 +
196 + tp_rank: int
197 + remote_tp_size: dict[str, int]
198 + is_mla: bool
199 + total_num_kv_heads: int
200 + attn_backend: type[AttentionBackend]
201 + engine_id: str
202 + remote_block_size: dict[str, int]
203 +
204 + def __post_init__(self):
205 + # Figure out whether the first dimension of the cache is K/V
206 + # or num_blocks. This is used to register the memory regions correctly.
207 + kv_cache_shape = self.attn_backend.get_kv_cache_shape(
208 + num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
209 + )
210 + # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D],
211 + # we just mock num_blocks to 1 for the dimension check below.
212 + self._is_kv_layout_blocks_first = (
213 + len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1
214 + )
215 +
216 + attn_backend = AttentionBackendEnum[self.attn_backend.get_name()]
217 + self._use_pallas = attn_backend == AttentionBackendEnum.PALLAS
218 +
219 + @property
220 + def is_kv_layout_blocks_first(self) -> bool:
221 + return self._is_kv_layout_blocks_first
222 +
223 + @property
224 + def split_k_and_v(self) -> bool:
225 + # Whether to register regions for K and V separately (when present).
226 + return not (self.is_mla or self._use_pallas or self.is_kv_layout_blocks_first)
227 +
228 + @property
229 + def tp_size(self) -> int:
230 + return self.remote_tp_size[self.engine_id]
231 +
232 + @property
233 + def block_size(self) -> int:
234 + return self.remote_block_size[self.engine_id]
235 +
236 + def tp_ratio(
237 + self,
238 + remote_tp_size: int,
239 + ) -> int:
240 + """
241 + Calculate the tensor parallel ratio between local and remote TP.
242 + We can think of it as the number of local TP workers-per-remote TP
243 + workers. Local workers will read from the same remote TP worker in
244 + groups of size `tp_ratio`.
245 + """
246 + assert self.tp_size % remote_tp_size == 0, (
247 + f"Local tensor parallel size {self.tp_size} is not divisible "
248 + f"by remote tensor parallel size {remote_tp_size}."
249 + )
250 + return self.tp_size // remote_tp_size
251 +
252 + def block_size_ratio(
253 + self,
254 + remote_block_size: int,
255 + ) -> float:
256 + """
257 + Calculate the block size ratio between local and remote TP.
258 + """
259 + assert self.block_size % remote_block_size == 0, (
260 + f"Local block size {self.block_size} is not divisible "
261 + f"by remote block size {remote_block_size} or vice versa."
262 + )
263 + return self.block_size // remote_block_size
264 +
265 + def tp_ratio_from_engine_id(
266 + self,
267 + remote_engine_id: str,
268 + ) -> int:
269 + remote_tp_size = self.remote_tp_size[remote_engine_id]
270 + return self.tp_ratio(remote_tp_size)
271 +
272 + def block_size_ratio_from_engine_id(
273 + self,
274 + remote_engine_id: str,
275 + ) -> float:
276 + remote_block_size = self.remote_block_size[remote_engine_id]
277 + return self.block_size_ratio(remote_block_size)
278 +
279 + def is_kv_replicated(self, engine_id: str) -> bool:
280 + """
281 + Whether the KV cache is replicated across TP workers due to the
282 + number of TP workers being greater than the number of KV heads.
283 + """
284 + tp_size = self.remote_tp_size[engine_id]
285 + return tp_size // self.total_num_kv_heads >= 1
286 +
287 + def replicates_kv_cache(self, remote_engine_id: str) -> bool:
288 + # MLA is always replicated as the hidden dim can't be split.
289 + return self.is_mla or self.is_kv_replicated(remote_engine_id)
290 +
291 + def get_target_remote_rank(
292 + self,
293 + remote_tp_size: int,
294 + ) -> int:
295 + """
296 + Get the remote TP rank (on P) that the current local TP rank
297 + (on D) will read from.
298 + """
299 + tp_ratio = self.tp_ratio(remote_tp_size)
300 + return self.tp_rank // tp_ratio
301 +
302 + def get_target_remote_rank_from_engine_id(
303 + self,
304 + remote_engine_id: str,
305 + ) -> int:
306 + remote_tp_size = self.remote_tp_size[remote_engine_id]
307 + return self.get_target_remote_rank(remote_tp_size)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake_connector.py new
+914
@@ -0,0 +1,914 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import asyncio
4 +import threading
5 +import time
6 +import uuid
7 +from collections import defaultdict
8 +from concurrent.futures import ThreadPoolExecutor
9 +from dataclasses import dataclass
10 +from typing import TYPE_CHECKING, Any, Optional
11 +
12 +import msgspec
13 +import numpy as np
14 +import torch
15 +import zmq
16 +import zmq.asyncio
17 +
18 +from vllm import envs
19 +from vllm.attention.backends.abstract import AttentionMetadata
20 +from vllm.attention.selector import get_attn_backend
21 +from vllm.config import VllmConfig
22 +from vllm.distributed.kv_transfer.kv_connector.utils import TpKVTopology
23 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
24 + KVConnectorBase_V1,
25 + KVConnectorMetadata,
26 + KVConnectorRole,
27 +)
28 +from vllm.distributed.parallel_state import (
29 + get_tensor_model_parallel_rank,
30 + get_tensor_model_parallel_world_size,
31 + get_tp_group,
32 +)
33 +from vllm.forward_context import ForwardContext
34 +from vllm.logger import init_logger
35 +from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket
36 +from vllm.v1.attention.backends.utils import get_kv_cache_layout
37 +from vllm.v1.core.sched.output import SchedulerOutput
38 +from vllm.v1.request import RequestStatus
39 +
40 +try:
41 + from mooncake.engine import TransferEngine
42 +except ImportError as e:
43 + raise ImportError(
44 + "Please install mooncake by following the instructions at "
45 + "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md " # noqa: E501
46 + "to run VLLM with MooncakeTransferEngine."
47 + ) from e
48 +
49 +if TYPE_CHECKING:
50 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
51 + from vllm.v1.kv_cache_interface import KVCacheConfig
52 + from vllm.v1.request import Request
53 +
54 +EngineId = str
55 +ReqId = str
56 +
57 +TRANS_DONE = b"trans_done"
58 +TRANS_ERROR = b"trans_error"
59 +
60 +logger = init_logger(__name__)
61 +
62 +
63 +class MooncakeAgentMetadata(
64 + msgspec.Struct,
65 + omit_defaults=True, # type: ignore[call-arg]
66 + # required for @cached_property.
67 + dict=True,
68 +):
69 + remote_hostname: str
70 + remote_port: int
71 + request_ids: list[ReqId]
72 + kv_caches_base_addr: list[int]
73 + block_ids: list[list[int]]
74 +
75 +
76 +@dataclass
77 +class RecvReqMeta:
78 + local_block_ids: list[int]
79 + remote_host: str
80 + remote_port: int
81 +
82 +
83 +@dataclass
84 +class SendBlockMeta:
85 + local_block_ids: list[int]
86 + ready: threading.Event
87 + expire_time: float = float("inf")
88 +
89 +
90 +@dataclass
91 +class SendReqMeta:
92 + reqs: dict[ReqId, SendBlockMeta]
93 + lock: threading.Lock
94 +
95 +
96 +@dataclass
97 +class FinishedSendReqSet:
98 + set: set[ReqId]
99 + lock: threading.Lock
100 +
101 +
102 +@dataclass
103 +class FinishedReceiveReqSet:
104 + set: set[ReqId]
105 + lock: asyncio.Lock
106 +
107 +
108 +class MooncakeConnectorMetadata(KVConnectorMetadata):
109 + def __init__(self):
110 + self.reqs_to_recv: dict[ReqId, RecvReqMeta] = {}
111 + self.reqs_to_send: dict[ReqId, list[int]] = {}
112 +
113 + def add_new_req(
114 + self,
115 + request_id: ReqId,
116 + local_block_ids: list[int],
117 + kv_transfer_params: dict[str, Any],
118 + load_remote_cache: bool = True,
119 + ):
120 + if load_remote_cache:
121 + self.reqs_to_recv[request_id] = RecvReqMeta(
122 + local_block_ids=local_block_ids,
123 + remote_host=kv_transfer_params["remote_host"],
124 + remote_port=kv_transfer_params["remote_port"],
125 + )
126 + else:
127 + self.reqs_to_send[request_id] = local_block_ids
128 +
129 +
130 +class MooncakeConnector(KVConnectorBase_V1):
131 + def __init__(
132 + self,
133 + vllm_config: VllmConfig,
134 + role: KVConnectorRole,
135 + kv_cache_config: Optional["KVCacheConfig"] = None,
136 + ):
137 + super().__init__(vllm_config, role, kv_cache_config)
138 +
139 + assert vllm_config.kv_transfer_config is not None
140 + assert vllm_config.kv_transfer_config.engine_id is not None
141 + self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
142 +
143 + if role == KVConnectorRole.SCHEDULER:
144 + self.connector_scheduler: MooncakeConnectorScheduler | None = (
145 + MooncakeConnectorScheduler(vllm_config, self.engine_id)
146 + )
147 + self.connector_worker: MooncakeConnectorWorker | None = None
148 + elif role == KVConnectorRole.WORKER:
149 + self.connector_scheduler = None
150 + self.connector_worker = MooncakeConnectorWorker(vllm_config, self.engine_id)
151 +
152 + ############################################################
153 + # Scheduler Side Methods
154 + ############################################################
155 +
156 + def get_num_new_matched_tokens(
157 + self, request: "Request", num_computed_tokens: int
158 + ) -> tuple[int, bool]:
159 + assert self.connector_scheduler is not None
160 + return self.connector_scheduler.get_num_new_matched_tokens(
161 + request, num_computed_tokens
162 + )
163 +
164 + def update_state_after_alloc(
165 + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
166 + ):
167 + assert self.connector_scheduler is not None
168 + return self.connector_scheduler.update_state_after_alloc(
169 + request, blocks, num_external_tokens
170 + )
171 +
172 + def build_connector_meta(
173 + self,
174 + scheduler_output: SchedulerOutput,
175 + ) -> KVConnectorMetadata:
176 + assert self.connector_scheduler is not None
177 + return self.connector_scheduler.build_connector_meta(scheduler_output)
178 +
179 + def request_finished(
180 + self,
181 + request: "Request",
182 + block_ids: list[int],
183 + ) -> tuple[bool, dict[str, Any] | None]:
184 + assert self.connector_scheduler is not None
185 + return self.connector_scheduler.request_finished(request, block_ids)
186 +
187 + ############################################################
188 + # Worker Side Methods
189 + ############################################################
190 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
191 + assert self.connector_worker is not None
192 + self.connector_worker.register_kv_caches(kv_caches)
193 +
194 + def get_finished(
195 + self, finished_req_ids: set[str]
196 + ) -> tuple[set[str] | None, set[str] | None]:
197 + """Get the finished recving and sending requests."""
198 + assert self.connector_worker is not None
199 + return self.connector_worker.get_finished()
200 +
201 + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
202 + assert self.connector_worker is not None
203 + assert isinstance(self._connector_metadata, MooncakeConnectorMetadata)
204 + self.connector_worker.start_load_kv(self._connector_metadata)
205 +
206 + def wait_for_layer_load(self, layer_name: str) -> None:
207 + """MooncakeConnector does not do layerwise saving."""
208 + pass
209 +
210 + def save_kv_layer(
211 + self,
212 + layer_name: str,
213 + kv_layer: torch.Tensor,
214 + attn_metadata: AttentionMetadata,
215 + **kwargs,
216 + ) -> None:
217 + """MooncakeConnector does not save explicitly."""
218 + pass
219 +
220 + def wait_for_save(self):
221 + pass
222 +
223 +
224 +class MooncakeConnectorScheduler:
225 + """Implementation of Scheduler side methods"""
226 +
227 + def __init__(self, vllm_config: VllmConfig, engine_id: str):
228 + self.vllm_config = vllm_config
229 + self.engine_id: EngineId = engine_id
230 + self.side_channel_host = get_ip()
231 + self.side_channel_port = get_mooncake_side_channel_port(vllm_config)
232 +
233 + assert vllm_config.kv_transfer_config
234 + self.kv_role = vllm_config.kv_transfer_config.kv_role
235 + logger.info("Initializing Mooncake Transfer Engine Scheduler %s", engine_id)
236 +
237 + # Requests that need to start recv/send.
238 + # New requests are added by update_state_after_alloc in
239 + # the scheduler. Used to make metadata passed to Worker.
240 + self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
241 + self._reqs_need_send: dict[ReqId, list[int]] = {}
242 +
243 + def get_num_new_matched_tokens(
244 + self, request: "Request", num_computed_tokens: int
245 + ) -> tuple[int, bool]:
246 + """
247 + For remote prefill, pull all prompt blocks from remote
248 + asynchronously relative to engine execution.
249 +
250 + Args:
251 + request (Request): the request object.
252 + num_computed_tokens (int): the number of locally
253 + computed tokens for this request
254 + Returns:
255 + * the number of tokens that can be loaded from the
256 + external KV cache beyond what is already computed.
257 + * true if the external KV cache tokens will be loaded
258 + asynchronously (between scheduler steps).
259 + """
260 +
261 + params = request.kv_transfer_params
262 + logger.debug(
263 + "MooncakeConnector get_num_new_matched_tokens: "
264 + "num_computed_tokens=%s, kv_transfer_params=%s",
265 + num_computed_tokens,
266 + params,
267 + )
268 +
269 + if params is not None and params.get("do_remote_prefill"):
270 + # Remote prefill: get all prompt blocks from remote.
271 + token_ids = request.prompt_token_ids or []
272 + count = len(token_ids) - num_computed_tokens
273 + if count > 0:
274 + return count, True
275 +
276 + # No remote prefill for this request.
277 + return 0, False
278 +
279 + def update_state_after_alloc(
280 + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
281 + ):
282 + params = request.kv_transfer_params
283 + logger.debug(
284 + "MooncakeConnector update_state_after_alloc: "
285 + "num_external_tokens=%s, kv_transfer_params=%s",
286 + num_external_tokens,
287 + params,
288 + )
289 +
290 + if not params:
291 + return
292 +
293 + if params.get("do_remote_prefill"):
294 + assert self.kv_role != "kv_producer"
295 + if all(p in params for p in ("remote_host", "remote_port")):
296 + # If remote_blocks and num_external_tokens = 0, we have
297 + # a full prefix cache hit on the D worker. We need to call
298 + # send_notif in _read_blocks to free the memory on the P.
299 + local_block_ids = (
300 + blocks.get_unhashed_block_ids() if num_external_tokens > 0 else []
301 + )
302 + # Get unhashed blocks to pull from remote.
303 + self._reqs_need_recv[request.request_id] = (request, local_block_ids)
304 + else:
305 + logger.warning(
306 + "Got invalid KVTransferParams: %s. This "
307 + "request will not utilize KVTransfer",
308 + params,
309 + )
310 + # Only trigger 1 KV transfer per request.
311 + params["do_remote_prefill"] = False
312 +
313 + elif params.get("do_remote_decode"):
314 + # Add an empty list to worker to create event.
315 + self._reqs_need_send[request.request_id] = []
316 +
317 + def build_connector_meta(
318 + self,
319 + scheduler_output: SchedulerOutput,
320 + ) -> KVConnectorMetadata:
321 + meta = MooncakeConnectorMetadata()
322 +
323 + # Loop through scheduled reqs and convert to RecvReqMeta.
324 + if self.kv_role != "kv_producer":
325 + for req_id, (req, block_ids) in self._reqs_need_recv.items():
326 + assert req.kv_transfer_params is not None
327 + meta.add_new_req(
328 + request_id=req_id,
329 + local_block_ids=block_ids,
330 + kv_transfer_params=req.kv_transfer_params,
331 + )
332 + self._reqs_need_recv.clear()
333 +
334 + if self.kv_role != "kv_consumer":
335 + for req_id, block_ids in self._reqs_need_send.items():
336 + meta.add_new_req(
337 + request_id=req_id,
338 + local_block_ids=block_ids,
339 + kv_transfer_params={},
340 + load_remote_cache=False,
341 + )
342 + self._reqs_need_send.clear()
343 +
344 + return meta
345 +
346 + def request_finished(
347 + self,
348 + request: "Request",
349 + block_ids: list[int],
350 + ) -> tuple[bool, dict[str, Any] | None]:
351 + """
352 + Once a request is finished, determine whether request blocks
353 + should be freed now or will be sent asynchronously and freed later.
354 + """
355 +
356 + params = request.kv_transfer_params
357 + logger.debug(
358 + "MooncakeConnector request_finished, request_status=%s, "
359 + "kv_transfer_params=%s",
360 + request.status,
361 + params,
362 + )
363 + if not params:
364 + return False, None
365 +
366 + if params.get("do_remote_prefill"):
367 + # If do_remote_prefill is still True when the request is finished,
368 + # update_state_after_alloc must not have been called (the request
369 + # must have been aborted before it was scheduled).
370 + # To avoid stranding the prefill blocks in the prefill instance,
371 + # we must add empty block_ids to _reqs_need_recv so that our
372 + # worker side will notify and free blocks in the prefill instance.
373 + assert self.kv_role != "kv_producer"
374 + self._reqs_need_recv[request.request_id] = (request, [])
375 + params["do_remote_prefill"] = False
376 + return False, None
377 +
378 + if (
379 + not params.get("do_remote_decode")
380 + or request.status != RequestStatus.FINISHED_LENGTH_CAPPED
381 + ):
382 + return False, None
383 +
384 + assert self.kv_role != "kv_consumer"
385 +
386 + # TODO: check whether block_ids actually ever be 0. If not we could
387 + # remove the conditional below
388 + delay_free_blocks = len(block_ids) > 0
389 +
390 + if delay_free_blocks:
391 + self._reqs_need_send[request.request_id] = block_ids
392 +
393 + return delay_free_blocks, dict(
394 + do_remote_prefill=True,
395 + do_remote_decode=False,
396 + remote_host=self.side_channel_host,
397 + remote_port=self.side_channel_port,
398 + )
399 +
400 +
401 +class MooncakeConnectorWorker:
402 + """Implementation of Worker side methods"""
403 +
404 + def __init__(self, vllm_config: VllmConfig, engine_id: str):
405 + logger.info("Initializing Mooncake Transfer Engine worker %s", engine_id)
406 +
407 + self.vllm_config = vllm_config
408 +
409 + self.engine = TransferEngine()
410 + self.hostname = get_ip()
411 + ret_value = self.engine.initialize(self.hostname, "P2PHANDSHAKE", "rdma", "")
412 + if ret_value != 0:
413 + raise RuntimeError("Mooncake Transfer Engine initialization failed.")
414 +
415 + self.rpc_port = self.engine.get_rpc_port()
416 +
417 + logger.debug(
418 + "Mooncake Transfer Engine initialized at %s:%d",
419 + self.hostname,
420 + self.rpc_port,
421 + )
422 +
423 + # Mooncake handshake port.
424 + self.side_channel_port: int = get_mooncake_side_channel_port(vllm_config)
425 +
426 + self.engine_id: EngineId = engine_id
427 + self.tp_rank = get_tensor_model_parallel_rank()
428 + self.world_size = get_tensor_model_parallel_world_size()
429 + self.tp_group = get_tp_group()
430 + self.num_blocks = 0
431 +
432 + assert vllm_config.kv_transfer_config
433 + self.kv_role = vllm_config.kv_transfer_config.kv_role
434 + self.num_workers = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
435 + "num_workers", 10
436 + )
437 +
438 + self.kv_caches_base_addr: list[int] = []
439 + self.device_kv_caches: dict[str, torch.Tensor] = {}
440 + self.reqs_need_send: SendReqMeta = SendReqMeta(reqs={}, lock=threading.Lock())
441 +
442 + # For kv_both, we will act both prefiller and decoder.
443 + if self.kv_role != "kv_consumer":
444 + # Background thread for sending kvcaches to D.
445 + self._mooncake_sender_t: threading.Thread | None = None
446 + # Background thread for processing new sending requests.
447 + self._sender_executor = ThreadPoolExecutor(
448 + max_workers=self.num_workers, thread_name_prefix="vllm-mooncake-sender"
449 + )
450 + logger.debug(
451 + "Mooncake Prefiller: use %d workers to send kvcaches", self.num_workers
452 + )
453 + if self.kv_role != "kv_producer":
454 + self.receiver_loop = asyncio.new_event_loop()
455 + self._mooncake_receiver_t = threading.Thread(
456 + target=self._receiver_loop, args=(self.receiver_loop,), daemon=True
457 + )
458 + self._mooncake_receiver_t.start()
459 + logger.debug("Mooncake Decoder: start receiver thread")
460 +
461 + self.finished_sending_reqs: FinishedSendReqSet = FinishedSendReqSet(
462 + set(), threading.Lock()
463 + )
464 + self.finished_recving_reqs: FinishedReceiveReqSet = FinishedReceiveReqSet(
465 + set(), asyncio.Lock()
466 + )
467 +
468 + self.block_size = vllm_config.cache_config.block_size
469 + self.model_config = vllm_config.model_config
470 + self.cache_config = vllm_config.cache_config
471 + self.use_mla = self.model_config.use_mla
472 +
473 + backend = get_attn_backend(
474 + self.model_config.get_head_size(),
475 + self.model_config.dtype,
476 + self.cache_config.cache_dtype,
477 + self.block_size,
478 + use_mla=self.use_mla,
479 + )
480 + self.backend_name = backend.get_name()
481 + self.kv_cache_layout = get_kv_cache_layout()
482 + logger.debug("Detected attention backend %s", self.backend_name)
483 + logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
484 +
485 + self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
486 + self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}
487 + self.kv_topo = TpKVTopology(
488 + tp_rank=self.tp_rank,
489 + engine_id=self.engine_id,
490 + remote_tp_size=self._tp_size, # shared state
491 + remote_block_size=self._block_size, # shared state
492 + is_mla=self.use_mla,
493 + total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
494 + attn_backend=backend,
495 + )
496 + self._use_pallas = self.kv_topo._use_pallas
497 +
498 + self.zmq_ctx = zmq.Context()
499 + self.async_zmq_ctx = zmq.asyncio.Context()
500 + self._encoder = msgspec.msgpack.Encoder()
501 + self._decoder = msgspec.msgpack.Decoder(MooncakeAgentMetadata)
502 +
503 + def __del__(self):
504 + self.shutdown()
505 +
506 + def shutdown(self):
507 + """Cleanup background threads on destruction."""
508 + self.zmq_ctx.term()
509 + self.async_zmq_ctx.term()
510 + if self.kv_role != "kv_consumer":
511 + self._sender_executor.shutdown(wait=False)
512 + if self._mooncake_sender_t:
513 + self._mooncake_sender_t.join()
514 + if self.kv_role != "kv_producer" and self.receiver_loop.is_running():
515 + self.receiver_loop.call_soon_threadsafe(self.receiver_loop.stop)
516 + self._mooncake_receiver_t.join()
517 +
518 + def _receiver_loop(self, loop: asyncio.AbstractEventLoop):
519 + asyncio.set_event_loop(loop)
520 + loop.run_forever()
521 +
522 + def _mooncake_sender(
523 + self, ready_event: threading.Event, base_port: int, tp_rank: int
524 + ):
525 + """
526 + Background thread that listens for Mooncake requests, dispatches them
527 + to a thread pool, and sends acknowledgments upon completion.
528 + """
529 +
530 + frontend_path = make_zmq_path("tcp", self.hostname, base_port + tp_rank)
531 + frontend = make_zmq_socket(self.zmq_ctx, frontend_path, zmq.ROUTER)
532 + logger.debug("Mooncake sender starting listening on path: %s", frontend_path)
533 +
534 + backend_path = make_zmq_path("inproc", str(uuid.uuid4()))
535 + backend = make_zmq_socket(self.zmq_ctx, backend_path, zmq.PULL)
536 +
537 + poller = zmq.Poller()
538 + poller.register(frontend, zmq.POLLIN)
539 + poller.register(backend, zmq.POLLIN)
540 +
541 + ready_event.set()
542 +
543 + try:
544 + while True:
545 + sockets = dict(poller.poll())
546 +
547 + if frontend in sockets:
548 + identity, _, metadata_bytes = frontend.recv_multipart()
549 + self._sender_executor.submit(
550 + self._sender_worker,
551 + identity,
552 + metadata_bytes,
553 + backend_path,
554 + )
555 +
556 + if backend in sockets:
557 + identity, status = backend.recv_multipart()
558 + frontend.send_multipart((identity, b"", status))
559 +
560 + except zmq.ContextTerminated:
561 + logger.debug("ZMQ context terminated, exiting Mooncake sender thread.")
562 + except Exception as e:
563 + logger.error("Error in Mooncake sender thread: %s. Exiting thread.", str(e))
564 + finally:
565 + frontend.close()
566 + backend.close()
567 +
568 + def _sender_worker(
569 + self, identity: bytes, metadata_bytes: bytes, worker_channel_path: str
570 + ):
571 + status = TRANS_ERROR
572 +
573 + try:
574 + metadata = self._decoder.decode(metadata_bytes)
575 + self.send_kv_to_decode(metadata)
576 + status = TRANS_DONE
577 + except Exception as e:
578 + logger.error("Error processing Mooncake handshake: %s", e)
579 + finally:
580 + pusher = make_zmq_socket(self.zmq_ctx, worker_channel_path, zmq.PUSH)
581 + try:
582 + pusher.send_multipart((identity, status))
583 + except zmq.ZMQError as e:
584 + logger.warning(
585 + "Internal error, maybe the server is shutting down. Error: %s",
586 + e,
587 + )
588 + finally:
589 + pusher.close()
590 +
591 + def send_kv_to_decode(self, meta: MooncakeAgentMetadata):
592 + send_reqs: list[tuple[ReqId, SendBlockMeta]] = []
593 + with self.reqs_need_send.lock:
594 + for req_id in meta.request_ids:
595 + send_meta = self.reqs_need_send.reqs.get(req_id)
596 + if send_meta is None:
597 + logger.warning("Request %s not found in reqs_need_send", req_id)
598 + return
599 + # Mark it as not expired. We will send it now.
600 + send_meta.expire_time = float("inf")
601 + send_reqs.append((req_id, send_meta))
602 +
603 + self._send_blocks(send_reqs, meta)
604 +
605 + with self.reqs_need_send.lock:
606 + for req_id in meta.request_ids:
607 + del self.reqs_need_send.reqs[req_id]
608 +
609 + with self.finished_sending_reqs.lock:
610 + self.finished_sending_reqs.set.update(meta.request_ids)
611 +
612 + def _send_blocks(
613 + self,
614 + send_reqs: list[tuple[ReqId, SendBlockMeta]],
615 + agent_meta: MooncakeAgentMetadata,
616 + ):
617 + src_ptrs = []
618 + dst_ptrs = []
619 + lengths = []
620 + local_base_addr = self.kv_caches_base_addr
621 + remote_base_addr = agent_meta.kv_caches_base_addr
622 + block_len = self.block_len
623 + remote_session = f"{agent_meta.remote_hostname}:{agent_meta.remote_port}"
624 +
625 + assert len(send_reqs) == len(agent_meta.block_ids)
626 + for (req_id, send_meta), remote_block_ids in zip(
627 + send_reqs, agent_meta.block_ids
628 + ):
629 + send_meta.ready.wait()
630 +
631 + num_remote_blocks = len(remote_block_ids)
632 + if num_remote_blocks == 0:
633 + continue
634 +
635 + local_block_ids = send_meta.local_block_ids
636 + # Partial prefix cache hit: just read uncomputed blocks.
637 + num_local_blocks = len(local_block_ids)
638 + assert num_local_blocks >= num_remote_blocks
639 + if num_local_blocks > num_remote_blocks:
640 + local_block_ids = local_block_ids[-num_remote_blocks:]
641 +
642 + # Group by indices
643 + group_local_block_ids, group_remote_block_ids = group_concurrent_contiguous(
644 + local_block_ids, remote_block_ids
645 + )
646 +
647 + for local_layer_addr, remote_layer_addr in zip(
648 + local_base_addr, remote_base_addr
649 + ):
650 + for group_local_block_id, group_remote_block_id in zip(
651 + group_local_block_ids, group_remote_block_ids
652 + ):
653 + src_ptrs.append(
654 + local_layer_addr + group_local_block_id[0] * block_len
655 + )
656 + dst_ptrs.append(
657 + remote_layer_addr + group_remote_block_id[0] * block_len
658 + )
659 + lengths.append(block_len * len(group_local_block_id))
660 +
661 + logger.debug(
662 + "Sending kv_caches for request %s (%d blocks) to %s",
663 + req_id,
664 + num_remote_blocks,
665 + remote_session,
666 + )
667 +
668 + start_time = time.perf_counter()
669 + ret_value = self.engine.batch_transfer_sync_write(
670 + remote_session, src_ptrs, dst_ptrs, lengths
671 + )
672 + if ret_value != 0:
673 + raise RuntimeError(f"Error in batch_transfer_sync_write: {ret_value}")
674 +
675 + logger.debug(
676 + "Sending to %s done, took %s",
677 + remote_session,
678 + time.perf_counter() - start_time,
679 + )
680 +
681 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
682 + """Register the KV Cache data in mooncake."""
683 +
684 + logger.info("Registering KV_Caches. use_mla: %s", self.use_mla)
685 +
686 + kv_data_ptrs = []
687 + kv_data_lens = []
688 + seen_base_addresses = []
689 +
690 + split_k_and_v = self.kv_topo.split_k_and_v
691 + tensor_size_bytes = None
692 + for layer_name, cache_or_caches in kv_caches.items():
693 + logger.debug(
694 + "registering layer %s with shape %s", layer_name, cache_or_caches.shape
695 + )
696 + cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
697 +
698 + for cache in cache_list:
699 + base_addr = cache.data_ptr()
700 + if base_addr in seen_base_addresses:
701 + continue
702 +
703 + seen_base_addresses.append(base_addr)
704 + curr_tensor_size_bytes = cache.nbytes
705 +
706 + if tensor_size_bytes is None:
707 + tensor_size_bytes = curr_tensor_size_bytes
708 + self.num_blocks = cache.shape[0]
709 +
710 + assert tensor_size_bytes == curr_tensor_size_bytes, (
711 + "All kv cache tensors must have the same size"
712 + )
713 + kernel_block_size = cache.shape[-2 if self.use_mla else -3]
714 + assert self.block_size == kernel_block_size
715 + kv_data_ptrs.append(base_addr)
716 + kv_data_lens.append(tensor_size_bytes)
717 +
718 + self.kv_caches_base_addr = seen_base_addresses
719 +
720 + ret_value = self.engine.batch_register_memory(kv_data_ptrs, kv_data_lens)
721 + if ret_value != 0:
722 + raise RuntimeError("Mooncake batch memory registration failed.")
723 +
724 + assert tensor_size_bytes is not None
725 + assert self.num_blocks != 0
726 + assert tensor_size_bytes % self.num_blocks == 0
727 + self.block_len = tensor_size_bytes // self.num_blocks
728 + self.device_kv_caches = kv_caches
729 + logger.debug(
730 + "registered num_blocks=%d block_len=%d", self.num_blocks, self.block_len
731 + )
732 +
733 + # No need to launch server for D node.
734 + if self.kv_role == "kv_consumer":
735 + return
736 +
737 + ready_event = threading.Event()
738 + self._mooncake_sender_t = threading.Thread(
739 + target=self._mooncake_sender,
740 + args=(ready_event, self.side_channel_port, self.tp_rank),
741 + daemon=True,
742 + name="mooncake_sender",
743 + )
744 + self._mooncake_sender_t.start()
745 + ready_event.wait() # Wait for listener ZMQ socket to be ready.
746 +
747 + async def fetch_finished_recving_reqs(self) -> set[ReqId]:
748 + async with self.finished_recving_reqs.lock:
749 + finished_recving_reqs = self.finished_recving_reqs.set
750 + self.finished_recving_reqs.set = set()
751 + return finished_recving_reqs
752 +
753 + def get_finished(self) -> tuple[set[str] | None, set[str] | None]:
754 + """
755 + Get requests that are done sending or recving on this specific worker.
756 + The scheduler process (via the MultiprocExecutor) will use this output
757 + to track which workers are done.
758 + """
759 + fut = None
760 + if self.kv_role != "kv_producer":
761 + fut = asyncio.run_coroutine_threadsafe(
762 + self.fetch_finished_recving_reqs(), self.receiver_loop
763 + )
764 +
765 + if self.kv_role != "kv_consumer":
766 + with self.finished_sending_reqs.lock:
767 + finished_sending_reqs = self.finished_sending_reqs.set
768 + self.finished_sending_reqs.set = set()
769 + else:
770 + finished_sending_reqs = set()
771 +
772 + finished_recving_reqs = fut.result() if fut else set()
773 +
774 + if finished_sending_reqs or finished_recving_reqs:
775 + logger.debug(
776 + "Rank %s, get_finished: %s requests done sending "
777 + "and %s requests done recving",
778 + self.tp_rank,
779 + len(finished_sending_reqs),
780 + len(finished_recving_reqs),
781 + )
782 +
783 + # Handle timeout to avoid stranding blocks on remote.
784 + now = time.perf_counter()
785 + with self.reqs_need_send.lock:
786 + expired_reqs = [
787 + req_id
788 + for req_id, send_meta in self.reqs_need_send.reqs.items()
789 + if send_meta.expire_time < now
790 + ]
791 + for req_id in expired_reqs:
792 + logger.warning(
793 + "Request %s timed out after %d seconds without "
794 + "being sent. Freeing its blocks on the producer side.",
795 + req_id,
796 + envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT,
797 + )
798 + del self.reqs_need_send.reqs[req_id]
799 + if expired_reqs:
800 + finished_sending_reqs.update(expired_reqs)
801 +
802 + return finished_sending_reqs or None, finished_recving_reqs or None
803 +
804 + async def receive_kv(self, path: str, req_blocks: list[tuple[str, list[int]]]):
805 + req_ids, block_ids = map(list, zip(*req_blocks))
806 + metadata = MooncakeAgentMetadata(
807 + remote_hostname=self.hostname,
808 + remote_port=self.rpc_port,
809 + request_ids=req_ids,
810 + kv_caches_base_addr=self.kv_caches_base_addr,
811 + block_ids=block_ids,
812 + )
813 +
814 + encoded_data = self._encoder.encode(metadata)
815 + logger.debug(
816 + "Size of encoded MooncakeAgentMetadata: %d bytes", len(encoded_data)
817 + )
818 + logger.debug("Sending kv transfer request for %s on path: %s", req_ids, path)
819 +
820 + # Send query for the request.
821 + sock: zmq.asyncio.Socket = make_zmq_socket(
822 + self.async_zmq_ctx, path, zmq.REQ, bind=False, linger=0
823 + )
824 + sock.setsockopt(zmq.RCVTIMEO, 60000)
825 + try:
826 + await sock.send(encoded_data)
827 + ret_msg = await sock.recv()
828 + if ret_msg != TRANS_DONE:
829 + logger.error(
830 + "Error happens during tranfering kvcache for %s, see logs in prefiller.", # noqa: E501
831 + req_ids,
832 + )
833 + return
834 + except zmq.ContextTerminated:
835 + logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.")
836 + except Exception as e:
837 + logger.error("MooncakeAgentMetadata transfer failed for %s: %s", req_ids, e)
838 + return
839 + finally:
840 + sock.close()
841 +
842 + async with self.finished_recving_reqs.lock:
843 + self.finished_recving_reqs.set.update(req_ids)
844 +
845 + logger.debug("pulling kv_caches for %s finished", req_ids)
846 +
847 + def group_kv_pull(self, metadata: MooncakeConnectorMetadata):
848 + kv_pulls = defaultdict(list)
849 + for req_id, meta in metadata.reqs_to_recv.items():
850 + logger.debug(
851 + "start_load_kv for request %s from remote engine. "
852 + "Num local_block_ids: %s.",
853 + req_id,
854 + len(meta.local_block_ids),
855 + )
856 + path = make_zmq_path(
857 + "tcp", meta.remote_host, meta.remote_port + self.tp_rank
858 + )
859 + kv_pulls[path].append((req_id, meta.local_block_ids))
860 +
861 + return kv_pulls
862 +
863 + def start_load_kv(self, metadata: MooncakeConnectorMetadata):
864 + if self.kv_role != "kv_producer":
865 + kv_pulls = self.group_kv_pull(metadata)
866 + for path, req_blocks in kv_pulls.items():
867 + asyncio.run_coroutine_threadsafe(
868 + self.receive_kv(path, req_blocks), self.receiver_loop
869 + )
870 +
871 + if self.kv_role != "kv_consumer":
872 + with self.reqs_need_send.lock:
873 + for req_id, block_ids in metadata.reqs_to_send.items():
874 + if block_ids:
875 + # Already gone through request_finished()
876 + send_meta = self.reqs_need_send.reqs[req_id]
877 + send_meta.local_block_ids = block_ids
878 + send_meta.ready.set()
879 + send_meta.expire_time = (
880 + time.perf_counter()
881 + + envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT
882 + )
883 + else:
884 + # From update_state_after_alloc(),
885 + # but not reach request_finished() yet
886 + self.reqs_need_send.reqs[req_id] = SendBlockMeta(
887 + local_block_ids=[], ready=threading.Event()
888 + )
889 +
890 +
891 +def group_concurrent_contiguous(
892 + src_indices: list[int], dst_indices: list[int]
893 +) -> tuple[list[list[int]], list[list[int]]]:
894 + """Vectorised NumPy implementation."""
895 + if len(src_indices) == 0:
896 + return [], []
897 +
898 + brk = np.where((np.diff(src_indices) != 1) | (np.diff(dst_indices) != 1))[0] + 1
899 + src_groups = np.split(src_indices, brk)
900 + dst_groups = np.split(dst_indices, brk)
901 +
902 + src_groups = [g.tolist() for g in src_groups]
903 + dst_groups = [g.tolist() for g in dst_groups]
904 +
905 + return src_groups, dst_groups
906 +
907 +
908 +def get_mooncake_side_channel_port(vllm_config: VllmConfig) -> int:
909 + # This logic is now centralized
910 + return (
911 + envs.VLLM_MOONCAKE_BOOTSTRAP_PORT
912 + + vllm_config.parallel_config.data_parallel_rank
913 + * vllm_config.parallel_config.tensor_parallel_size
914 + )
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+3 -125
@@ -20,10 +20,10 @@ import torch
20 import zmq
21
22 from vllm import envs
23 -from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata
24 -from vllm.attention.backends.registry import AttentionBackendEnum
23 +from vllm.attention.backends.abstract import AttentionMetadata
24 from vllm.attention.selector import get_attn_backend
25 from vllm.config import VllmConfig
26 +from vllm.distributed.kv_transfer.kv_connector.utils import TpKVTopology
27 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
28 CopyBlocksOp,
29 KVConnectorBase_V1,
@@ -668,128 +668,6 @@ class NixlConnectorScheduler:
668 class NixlConnectorWorker:
669 """Implementation of Worker side methods"""
670
671 - @dataclass
672 - class TpKVTopology:
673 - """
674 - Helper class for tensor parallel and KV topology information for
675 - mapping between local and remote TP workers.
676 - """
677 -
678 - tp_rank: int
679 - remote_tp_size: dict[EngineId, int]
680 - is_mla: bool
681 - total_num_kv_heads: int
682 - attn_backend: type[AttentionBackend]
683 - engine_id: EngineId
684 - remote_block_size: dict[EngineId, int]
685 -
686 - def __post_init__(self):
687 - # Figure out whether the first dimension of the cache is K/V
688 - # or num_blocks. This is used to register the memory regions correctly.
689 - kv_cache_shape = self.attn_backend.get_kv_cache_shape(
690 - num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
691 - )
692 - # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D],
693 - # we just mock num_blocks to 1 for the dimension check below.
694 - self._is_kv_layout_blocks_first = (
695 - len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1
696 - )
697 -
698 - attn_backend = AttentionBackendEnum[self.attn_backend.get_name()]
699 - self._use_pallas = attn_backend == AttentionBackendEnum.PALLAS
700 -
701 - @property
702 - def is_kv_layout_blocks_first(self) -> bool:
703 - return self._is_kv_layout_blocks_first
704 -
705 - @property
706 - def split_k_and_v(self) -> bool:
707 - # Whether to register regions for K and V separately (when present).
708 - return not (
709 - self.is_mla or self._use_pallas or self.is_kv_layout_blocks_first
710 - )
711 -
712 - @property
713 - def tp_size(self) -> int:
714 - return self.remote_tp_size[self.engine_id]
715 -
716 - @property
717 - def block_size(self) -> int:
718 - return self.remote_block_size[self.engine_id]
719 -
720 - def tp_ratio(
721 - self,
722 - remote_tp_size: int,
723 - ) -> int:
724 - """
725 - Calculate the tensor parallel ratio between local and remote TP.
726 - We can think of it as the number of local TP workers-per-remote TP
727 - workers. Local workers will read from the same remote TP worker in
728 - groups of size `tp_ratio`.
729 - """
730 - assert self.tp_size % remote_tp_size == 0, (
731 - f"Local tensor parallel size {self.tp_size} is not divisible "
732 - f"by remote tensor parallel size {remote_tp_size}."
733 - )
734 - return self.tp_size // remote_tp_size
735 -
736 - def block_size_ratio(
737 - self,
738 - remote_block_size: int,
739 - ) -> float:
740 - """
741 - Calculate the block size ratio between local and remote TP.
742 - """
743 - assert self.block_size % remote_block_size == 0, (
744 - f"Local block size {self.block_size} is not divisible "
745 - f"by remote block size {remote_block_size} or vice versa."
746 - )
747 - return self.block_size // remote_block_size
748 -
749 - def tp_ratio_from_engine_id(
750 - self,
751 - remote_engine_id: EngineId,
752 - ) -> int:
753 - remote_tp_size = self.remote_tp_size[remote_engine_id]
754 - return self.tp_ratio(remote_tp_size)
755 -
756 - def block_size_ratio_from_engine_id(
757 - self,
758 - remote_engine_id: EngineId,
759 - ) -> float:
760 - remote_block_size = self.remote_block_size[remote_engine_id]
761 - return self.block_size_ratio(remote_block_size)
762 -
763 - def is_kv_replicated(self, engine_id: EngineId) -> bool:
764 - """
765 - Whether the KV cache is replicated across TP workers due to the
766 - number of TP workers being greater than the number of KV heads.
767 - """
768 - tp_size = self.remote_tp_size[engine_id]
769 - return tp_size // self.total_num_kv_heads >= 1
770 -
771 - def replicates_kv_cache(self, remote_engine_id: EngineId) -> bool:
772 - # MLA is always replicated as the hidden dim can't be split.
773 - return self.is_mla or self.is_kv_replicated(remote_engine_id)
774 -
775 - def get_target_remote_rank(
776 - self,
777 - remote_tp_size: int,
778 - ) -> int:
779 - """
780 - Get the remote TP rank (on P) that the current local TP rank
781 - (on D) will read from.
782 - """
783 - tp_ratio = self.tp_ratio(remote_tp_size)
784 - return self.tp_rank // tp_ratio
785 -
786 - def get_target_remote_rank_from_engine_id(
787 - self,
788 - remote_engine_id: EngineId,
789 - ) -> int:
790 - remote_tp_size = self.remote_tp_size[remote_engine_id]
791 - return self.get_target_remote_rank(remote_tp_size)
792 -
671 def __init__(self, vllm_config: VllmConfig, engine_id: str):
672 if NixlWrapper is None:
673 logger.error("NIXL is not available")
@@ -958,7 +836,7 @@ class NixlConnectorWorker:
836 self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
837 self.xfer_stats = NixlKVConnectorStats()
838
961 - self.kv_topo = self.TpKVTopology(
839 + self.kv_topo = TpKVTopology(
840 tp_rank=self.tp_rank,
841 engine_id=self.engine_id,
842 remote_tp_size=self._tp_size, # shared state
vllm/envs.py
+10
@@ -175,6 +175,7 @@ if TYPE_CHECKING:
175 VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False
176 VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
177 VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600
178 + VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998
179 VLLM_ALL2ALL_BACKEND: Literal[
180 "naive",
181 "pplx",
@@ -197,6 +198,7 @@ if TYPE_CHECKING:
198 VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True
199 VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None
200 VLLM_NIXL_ABORT_REQUEST_TIMEOUT: int = 480
201 + VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT: int = 480
202 VLLM_USE_CUDNN_PREFILL: bool = False
203 VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL: bool = False
204 VLLM_ENABLE_CUDAGRAPH_GC: bool = False
@@ -1260,6 +1262,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1262 "VLLM_NIXL_SIDE_CHANNEL_PORT": lambda: int(
1263 os.getenv("VLLM_NIXL_SIDE_CHANNEL_PORT", "5600")
1264 ),
1265 + # Port used for Mooncake handshake between remote agents.
1266 + "VLLM_MOONCAKE_BOOTSTRAP_PORT": lambda: int(
1267 + os.getenv("VLLM_MOONCAKE_BOOTSTRAP_PORT", "8998")
1268 + ),
1269 # all2all backend for vllm's expert parallel communication
1270 # Available options:
1271 # - "naive": naive all2all implementation using broadcasts
@@ -1369,6 +1375,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1375 "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": lambda: int(
1376 os.getenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", "480")
1377 ),
1378 + # Timeout (in seconds) for MooncakeConnector in PD disaggregated setup.
1379 + "VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT": lambda: int(
1380 + os.getenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "480")
1381 + ),
1382 # Controls whether or not to use cudnn prefill
1383 "VLLM_USE_CUDNN_PREFILL": lambda: bool(
1384 int(os.getenv("VLLM_USE_CUDNN_PREFILL", "0"))