@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4505849b30
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 10, 10:15 UTC

[ROCm][PD] add moriio kv connector. (#29304)

Signed-off-by: inkcherry <mingzhi.liu@amd.com>

inkcherry committed Jan 9, 2026 at 22:01 UTC 4505849b309b5f665a16778a4f0070d73c79e06d
10 files changed +3369 -3
docker/Dockerfile.rocm_base
+18 -1
@@ -11,6 +11,8 @@ ARG FA_BRANCH="0e60e394"
11 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
12 ARG AITER_BRANCH="6af8b687"
13 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
14 +ARG MORI_BRANCH="2d02c6a9"
15 +ARG MORI_REPO="https://github.com/ROCm/mori.git"
16
17 #TODO: When patch has been upstreamed, switch to the main repo/branch
18 # ARG RIXL_BRANCH="<TODO>"
@@ -31,6 +33,7 @@ ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
33 ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151
34 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
35 ENV AITER_ROCM_ARCH=gfx942;gfx950
36 +ENV MORI_GPU_ARCHS=gfx942;gfx950
37
38 # Required for RCCL in ROCm7.1
39 ENV HSA_NO_SCRATCH_RECLAIM=1
@@ -44,7 +47,7 @@ ENV DEBIAN_FRONTEND=noninteractive
47
48 # Install Python and other dependencies
49 RUN apt-get update -y \
47 - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 \
50 + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev \
51 && for i in 1 2 3; do \
52 add-apt-repository -y ppa:deadsnakes/ppa && break || \
53 { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
@@ -86,6 +89,18 @@ RUN cd /opt/rocm/share/amd_smi \
89 && pip wheel . --wheel-dir=dist
90 RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install
91
92 +FROM base AS build_mori
93 +ARG MORI_BRANCH
94 +ARG MORI_REPO
95 +RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
96 + pip install /install/*.whl
97 +RUN git clone ${MORI_REPO}
98 +RUN cd mori \
99 + && git checkout ${MORI_BRANCH} \
100 + && git submodule update --init --recursive \
101 + && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl
102 +RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install
103 +
104
105 ###
106 ### Pytorch build
@@ -253,6 +268,8 @@ RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
268 cp /install/*.whl /app/debs
269 RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
270 cp /install/*.whl /app/debs
271 +RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
272 + cp /install/*.whl /app/debs
273 RUN --mount=type=bind,from=build_rixl,src=/app/install/,target=/install \
274 cp /install/*.whl /app/debs
275
docs/getting_started/installation/gpu.rocm.inc.md
+17 -2
@@ -98,9 +98,24 @@ Currently, there are no pre-built ROCm wheels.
98 !!! note
99 - You will need to config the `$AITER_BRANCH_OR_COMMIT` for your purpose.
100 - The validated `$AITER_BRANCH_OR_COMMIT` can be found in the [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base).
101 -
101
103 -4. Build vLLM. For example, vLLM on ROCM 7.0 can be built with the following steps:
102 +
103 +4. If you want to use MORI for EP or PD disaggregation, you can install [MORI](https://github.com/ROCm/mori) using the following steps:
104 +
105 + ```bash
106 + git clone https://github.com/ROCm/mori.git
107 + cd mori
108 + git checkout $MORI_BRANCH_OR_COMMIT
109 + git submodule sync; git submodule update --init --recursive
110 + MORI_GPU_ARCHS="gfx942;gfx950" python3 install .
111 + ```
112 +
113 + !!! note
114 + - You will need to config the `$MORI_BRANCH_OR_COMMIT` for your purpose.
115 + - The validated `$MORI_BRANCH_OR_COMMIT` can be found in the [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base).
116 +
117 +
118 +5. Build vLLM. For example, vLLM on ROCM 7.0 can be built with the following steps:
119
120 ???+ console "Commands"
121
examples/online_serving/disaggregated_serving/moriio_toy_proxy_server.py new
+320
@@ -0,0 +1,320 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import asyncio
4 +import copy
5 +import logging
6 +import os
7 +import re
8 +import socket
9 +import threading
10 +import uuid
11 +
12 +import aiohttp
13 +import msgpack
14 +import zmq
15 +from quart import Quart, make_response, request
16 +
17 +logger = logging.getLogger(__name__)
18 +logger.setLevel(logging.DEBUG)
19 +prefill_instances: list[dict] = []
20 +decode_instances: list[dict] = []
21 +request_nums = 0
22 +app = Quart(__name__)
23 +
24 +IP_PORT_PATTERN = re.compile(r"//(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}):(\d+)")
25 +
26 +
27 +TRANSFER_TYPE = None
28 +
29 +
30 +def _append_whole_dict_unique(target_list, data_dict):
31 + new_filtered = {k: v for k, v in data_dict.items() if k != "index"}
32 + for existed in target_list:
33 + existed_filtered = {k: v for k, v in existed.items() if k != "index"}
34 + if existed_filtered == new_filtered:
35 + return False
36 + print("!!APPEND!!", data_dict)
37 + target_list.append(data_dict)
38 + transfer_mode = data_dict.get("transfer_mode", "unknown")
39 + global TRANSFER_TYPE
40 +
41 + if TRANSFER_TYPE is None:
42 + TRANSFER_TYPE = transfer_mode
43 + logger.info("SET TRANSFER TYPE TO %s", TRANSFER_TYPE)
44 + elif transfer_mode != TRANSFER_TYPE:
45 + raise ValueError(f"mismatched transfer mode {TRANSFER_TYPE} vs {transfer_mode}")
46 +
47 + return True
48 +
49 +
50 +_list_lock = threading.RLock()
51 +
52 +
53 +def _listen_for_register(hostname, port):
54 + context = zmq.Context()
55 + router_socket = context.socket(zmq.ROUTER)
56 + router_socket.bind(f"tcp://{hostname}:{port}")
57 + poller = zmq.Poller()
58 + poller.register(router_socket, zmq.POLLIN)
59 + global prefill_instances
60 + global decode_instances
61 +
62 + while True:
63 + socks = dict(poller.poll())
64 + if router_socket in socks:
65 + remote_addr, msg = router_socket.recv_multipart()
66 + data = msgpack.loads(msg)
67 + if data["type"] == "HELLO":
68 + pass
69 + elif (
70 + data["type"] == "register"
71 + and data["role"] == "P"
72 + and data["request_address"] not in prefill_instances
73 + ):
74 + with _list_lock:
75 + _append_whole_dict_unique(prefill_instances, data)
76 +
77 + elif (
78 + data["type"] == "register"
79 + and data["role"] == "D"
80 + and data["request_address"] not in decode_instances
81 + ):
82 + with _list_lock:
83 + _append_whole_dict_unique(decode_instances, data)
84 +
85 +
86 +def start_service_discovery(hostname, port):
87 + if not hostname:
88 + hostname = socket.gethostname()
89 + if port == 0:
90 + raise ValueError("Port cannot be 0")
91 +
92 + _listener_thread = threading.Thread(
93 + target=_listen_for_register, args=(hostname, port), daemon=True
94 + )
95 + _listener_thread.start()
96 + return _listener_thread
97 +
98 +
99 +async def send_request_to_prefill(
100 + endpoint, req_data, request_id, d_endpoint, dip, dport, selected_prefill_dp_rank
101 +):
102 + req_data_copy = req_data
103 +
104 + req_data_copy["kv_transfer_params"].update(
105 + {
106 + "do_remote_decode": True,
107 + "do_remote_prefill": False,
108 + "remote_handshake_port": d_endpoint["handshake_port"],
109 + "remote_notify_port": d_endpoint["notify_port"],
110 + "remote_engine_id": None,
111 + "remote_block_ids": None,
112 + "remote_host": dip,
113 + "remote_port": dport,
114 + }
115 + )
116 + req_data_copy["stream"] = False
117 + req_data_copy["max_tokens"] = 1
118 + if "max_completion_tokens" in req_data_copy:
119 + req_data_copy["max_completion_tokens"] = 1
120 + if "stream_options" in req_data_copy:
121 + del req_data_copy["stream_options"]
122 + async with aiohttp.ClientSession(
123 + timeout=aiohttp.ClientTimeout(total=6 * 6000 * 6000)
124 + ) as session:
125 + headers = {
126 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
127 + "X-Request-Id": request_id,
128 + }
129 + if selected_prefill_dp_rank is not None:
130 + headers["X-data-parallel-rank"] = str(selected_prefill_dp_rank)
131 + async with session.post(
132 + url=endpoint, json=req_data_copy, headers=headers
133 + ) as response:
134 + if response.status == 200:
135 + return await response.json()
136 +
137 + else:
138 + raise RuntimeError(
139 + "send_request_to_prefill response.status != 200response.status = ",
140 + response.status,
141 + )
142 +
143 +
144 +async def start_decode_request(endpoint, req_data, request_id):
145 + session = aiohttp.ClientSession(
146 + timeout=aiohttp.ClientTimeout(total=6 * 6000 * 6000)
147 + )
148 + headers = {
149 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
150 + "X-Request-Id": request_id,
151 + }
152 + response = await session.post(url=endpoint, json=req_data, headers=headers)
153 + return session, response
154 +
155 +
156 +async def stream_decode_response(session, response, request_id):
157 + try:
158 + if response.status == 200:
159 + async for chunk_bytes in response.content.iter_chunked(1024):
160 + yield chunk_bytes
161 + else:
162 + raise RuntimeError(
163 + f"decode response.status != 200, status = {response.status}"
164 + )
165 + finally:
166 + await session.close()
167 +
168 +
169 +async def send_request_to_decode(endpoint, req_data, request_id):
170 + async with aiohttp.ClientSession(
171 + timeout=aiohttp.ClientTimeout(total=6 * 6000 * 6000)
172 + ) as session:
173 + headers = {
174 + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
175 + "X-Request-Id": request_id,
176 + }
177 + async with session.post(
178 + url=endpoint, json=req_data, headers=headers
179 + ) as response:
180 + if response.status == 200:
181 + async for chunk_bytes in response.content.iter_chunked(1024):
182 + yield chunk_bytes
183 + else:
184 + raise RuntimeError(
185 + "send_request_to_decode response.status != 200,response.statuus = ",
186 + response.status,
187 + )
188 +
189 +
190 +def example_round_robin_dp_loader(request_number, dp_size):
191 + return request_nums % dp_size
192 +
193 +
194 +@app.route("/v1/completions", methods=["POST"])
195 +@app.route("/v1/chat/completions", methods=["POST"])
196 +async def handle_request():
197 + try:
198 + with _list_lock:
199 + global request_nums
200 + request_nums += 1
201 +
202 + def extract_ip_port_fast(url):
203 + match = IP_PORT_PATTERN.search(url)
204 + if not match:
205 + raise ValueError(f"Invalid URL format: {url}")
206 + return match.groups()
207 +
208 + req_data = await request.get_json()
209 + request_id = str(uuid.uuid4())
210 +
211 + prefill_instance_endpoint = None
212 + decode_instance_endpoint = None
213 + error_msg = (
214 + "Service Unavailable: No prefill or decode instances are registered."
215 + )
216 + if not prefill_instances or not decode_instances:
217 + return await make_response(
218 + (
219 + error_msg,
220 + 503,
221 + )
222 + )
223 + pid = request_nums % len(prefill_instances)
224 + did = request_nums % len(decode_instances)
225 + prefill_instance_endpoint = prefill_instances[pid]
226 + decode_instance_endpoint = decode_instances[did]
227 +
228 + selected_prefill_dp_rank = None
229 + if prefill_instance_endpoint["dp_size"] > 1:
230 + selected_prefill_dp_rank = example_round_robin_dp_loader(
231 + request_nums // len(prefill_instance_endpoint),
232 + prefill_instance_endpoint["dp_size"],
233 + )
234 +
235 + dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])
236 + ip, port = extract_ip_port_fast(prefill_instance_endpoint["request_address"])
237 +
238 + req_data_to_prefill = copy.deepcopy(req_data)
239 + req_data_to_prefill["kv_transfer_params"] = {}
240 + req_data["kv_transfer_params"] = {}
241 + req_data_to_prefill["kv_transfer_params"]["remote_dp_size"] = (
242 + decode_instance_endpoint["dp_size"]
243 + )
244 + req_data_to_prefill["kv_transfer_params"]["remote_tp_size"] = (
245 + decode_instance_endpoint["tp_size"]
246 + )
247 +
248 + send_prefill_task = asyncio.create_task(
249 + send_request_to_prefill(
250 + prefill_instance_endpoint["request_address"],
251 + req_data_to_prefill,
252 + request_id,
253 + decode_instance_endpoint,
254 + dip,
255 + dport,
256 + selected_prefill_dp_rank,
257 + )
258 + )
259 + ip, port = extract_ip_port_fast(prefill_instance_endpoint["request_address"])
260 +
261 + req_data["max_tokens"] -= 1
262 +
263 + req_data["kv_transfer_params"] = {
264 + "do_remote_decode": False,
265 + "do_remote_prefill": True,
266 + "remote_handshake_port": prefill_instance_endpoint["handshake_port"],
267 + "remote_notify_port": prefill_instance_endpoint["notify_port"],
268 + "remote_engine_id": None,
269 + "remote_block_ids": None,
270 + "remote_host": ip,
271 + "remote_port": port,
272 + }
273 + if TRANSFER_TYPE == "READ":
274 + # In read mode, prefill and decode are executed serially.
275 + prefill_response = await send_prefill_task
276 + req_data["kv_transfer_params"]["remote_engine_id"] = prefill_response[
277 + "kv_transfer_params"
278 + ]["remote_engine_id"]
279 + req_data["kv_transfer_params"]["remote_block_ids"] = prefill_response[
280 + "kv_transfer_params"
281 + ]["remote_block_ids"]
282 +
283 + req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[
284 + "dp_size"
285 + ]
286 + req_data["kv_transfer_params"]["remote_tp_size"] = prefill_instance_endpoint[
287 + "tp_size"
288 + ]
289 +
290 + if selected_prefill_dp_rank is not None:
291 + req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank
292 +
293 + decode_request_task = asyncio.create_task(
294 + start_decode_request(
295 + decode_instance_endpoint["request_address"], req_data, request_id
296 + )
297 + )
298 +
299 + session, decode_response = await decode_request_task
300 + stream_generator = stream_decode_response(session, decode_response, request_id)
301 + response = await make_response(stream_generator)
302 + return response
303 + except Exception as e:
304 + logger.exception("An error occurred while handling the request: %s", e)
305 + return await make_response(
306 + (
307 + f"Internal Server Error: {e!s}",
308 + 500,
309 + )
310 + )
311 +
312 +
313 +if __name__ == "__main__":
314 + t = start_service_discovery("0.0.0.0", 36367)
315 + app.debug = True
316 + app.config["BODY_TIMEOUT"] = 360000
317 + app.config["RESPONSE_TIMEOUT"] = 360000
318 +
319 + app.run(host="0.0.0.0", port=10001)
320 + t.join()
tests/v1/kv_connector/unit/test_moriio_connector.py new
+545
@@ -0,0 +1,545 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import importlib.util
4 +import os
5 +from unittest.mock import MagicMock, patch
6 +
7 +import msgspec
8 +import pytest
9 +import torch
10 +import zmq
11 +
12 +from tests.conftest import _find_free_port
13 +from vllm.config import (
14 + CacheConfig,
15 + DeviceConfig,
16 + KVTransferConfig,
17 + ModelConfig,
18 + SchedulerConfig,
19 + VllmConfig,
20 +)
21 +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
22 + MoRIIOAgentMetadata,
23 + MoRIIOConnectorMetadata,
24 + MoRIIOConstants,
25 + zmq_ctx,
26 +)
27 +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
28 + KVConnectorRole,
29 + MoRIIOConnector,
30 + MoRIIOConnectorWorker,
31 +)
32 +from vllm.platforms import current_platform
33 +from vllm.utils.network_utils import (
34 + get_ip,
35 + make_zmq_path,
36 +)
37 +
38 +from .utils import create_request, create_scheduler
39 +
40 +aiter_available = importlib.util.find_spec("aiter") is not None
41 +mori_available = importlib.util.find_spec("mori") is not None
42 +pytestmark = pytest.mark.skipif(
43 + not (current_platform.is_rocm() and mori_available),
44 + reason="MoRIIOs are only available on ROCm with aiter package installed",
45 +)
46 +
47 +
48 +@pytest.fixture
49 +def mock_parallel_groups():
50 + """Mock tensor/data parallel group functions for single-rank tests."""
51 + mock_group = MagicMock()
52 + mock_group.rank = 0
53 + mock_group.local_rank = 0
54 + mock_group.world_size = 1
55 +
56 + with (
57 + patch.multiple(
58 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common",
59 + get_tensor_model_parallel_rank=MagicMock(return_value=0),
60 + get_tensor_model_parallel_world_size=MagicMock(return_value=0),
61 + ),
62 + patch.multiple(
63 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector",
64 + get_tensor_model_parallel_world_size=MagicMock(return_value=0),
65 + get_world_group=MagicMock(return_value=mock_group),
66 + get_tp_group=MagicMock(return_value=mock_group),
67 + ),
68 + ):
69 + yield mock_group
70 +
71 +
72 +def _setup_kv_transfer_request(request, remote_host="127.0.0.1", fake_port=4789):
73 + """Setup KV transfer parameters for a request."""
74 + request.kv_transfer_params.update(
75 + {
76 + "remote_notify_port": fake_port,
77 + "remote_block_ids": None,
78 + "remote_host": remote_host,
79 + "remote_port": fake_port,
80 + "remote_handshake_port": fake_port,
81 + "remote_engine_id": "test_engine",
82 + }
83 + )
84 + return request
85 +
86 +
87 +class FakeMorIIOWrapper:
88 + # A fake MoRIIOWrapper for testing purposes
89 + def __init__(self, *args, **kwargs):
90 + pass
91 +
92 + def set_moriio_engine(self, moriio_engine):
93 + pass
94 +
95 + def set_backend_type(self, backend_type):
96 + pass
97 +
98 + def get_agent_metadata(self):
99 + pass
100 +
101 + def register_remote_engine(self, remote_packed_engine_metadata):
102 + pass
103 +
104 + def register_local_tensor(self, tensor: torch.Tensor):
105 + pass
106 +
107 + def get_unpack_memory_metadata(self, packed_memory_metadata):
108 + pass
109 +
110 + def build_session(self, local_memory_metadata, remote_memory_metadata):
111 + pass
112 +
113 + def read_remote_data(
114 + self, transfer_size_byte, local_offset=0, remote_offset=0, session=None
115 + ):
116 + pass
117 +
118 + def write_remote_data(
119 + self, transfer_size_byte, local_offset=0, remote_offset=0, session=None
120 + ):
121 + pass
122 +
123 + def write_remote_data_single(
124 + self, transfer_size_byte, local_offset=0, remote_offset=0, sess_idx=0
125 + ):
126 + pass
127 +
128 + def waiting_for_transfer_complete(self):
129 + pass
130 +
131 + def async_wait_reqid(self):
132 + pass
133 +
134 + def _handle_message(self, msg: bytes):
135 + pass
136 +
137 + def _handle_structured_message(self, data: dict):
138 + pass
139 +
140 + def _handle_completion_message(self, msg: str):
141 + pass
142 +
143 + def send_notify(self, req_ids, remote_ip, remote_port):
144 + pass
145 +
146 + def pop_finished_req_ids(self):
147 + pass
148 +
149 + def pop_finished_write_req_ids(self):
150 + pass
151 +
152 + def shutdown(self):
153 + pass
154 +
155 +
156 +class FakeMorIIOConnectorWorker(MoRIIOConnectorWorker):
157 + # Define a fake remote engine id for testing
158 + REMOTE_ENGINE_ID = "remote_engine"
159 +
160 + def __init__(
161 + self, *args, hand_shake_latency: float = 1.8, kv_cache_layout="HND", **kwargs
162 + ):
163 + super().__init__(*args, **kwargs)
164 +
165 +
166 +def create_vllm_config(
167 + model: str = "facebook/opt-125m",
168 + max_num_seqs: int = 16,
169 + max_num_batched_tokens: int = 64,
170 + block_size: int = 16,
171 + max_model_len: int = 10000,
172 + enable_chunked_prefill: bool = True,
173 + enable_permute_local_kv: bool = False,
174 + role="kv_consumer",
175 +) -> VllmConfig:
176 + """Initialize VllmConfig for testing."""
177 + scheduler_config = SchedulerConfig(
178 + max_num_seqs=max_num_seqs,
179 + max_num_batched_tokens=max_num_batched_tokens,
180 + max_model_len=max_model_len,
181 + enable_chunked_prefill=enable_chunked_prefill,
182 + is_encoder_decoder=False,
183 + )
184 + model_config = ModelConfig(
185 + model=model,
186 + trust_remote_code=True,
187 + dtype="bfloat16",
188 + seed=42,
189 + )
190 + # Cache config, optionally force APC
191 + cache_config = CacheConfig(
192 + block_size=block_size,
193 + gpu_memory_utilization=0.9,
194 + swap_space=0,
195 + cache_dtype="auto",
196 + enable_prefix_caching=True,
197 + )
198 + kv_transfer_config = KVTransferConfig(
199 + kv_connector="MoRIIOConnector",
200 + kv_role=role,
201 + enable_permute_local_kv=enable_permute_local_kv,
202 + )
203 + return VllmConfig(
204 + scheduler_config=scheduler_config,
205 + model_config=model_config,
206 + cache_config=cache_config,
207 + kv_transfer_config=kv_transfer_config,
208 + device_config=DeviceConfig("cpu"),
209 + )
210 +
211 +
212 +@pytest.fixture
213 +def moriio_read_mode():
214 + """Force the connector into read mode via env for tests."""
215 + os.environ["VLLM_MORIIO_CONNECTOR_READ_MODE"] = "True"
216 + yield
217 + # Cleanup after test
218 + os.environ.pop("VLLM_MORIIO_CONNECTOR_READ_MODE", None)
219 +
220 +
221 +def test_write_mode_saves_local_block_ids():
222 + """Write mode records local block ids in MoRIIOConnectorMetadata.reqs_to_save."""
223 +
224 + # Setup Scheduler and Request
225 + vllm_config = create_vllm_config(role="kv_producer")
226 + scheduler = create_scheduler(vllm_config)
227 +
228 + # 2 Full Blocks and 1 Half Block.
229 + BLOCK_SIZE = vllm_config.cache_config.block_size
230 + NUM_EXTERNAL_FULL_BLOCKS = 2
231 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
232 +
233 + request = create_request(
234 + request_id=1,
235 + block_size=BLOCK_SIZE,
236 + num_tokens=NUM_TOKENS,
237 + do_remote_decode=True,
238 + do_remote_prefill=False,
239 + )
240 + request_id = request.request_id
241 +
242 + scheduler.add_request(request)
243 +
244 + # Fake Config
245 + request = _setup_kv_transfer_request(request)
246 +
247 + # Remote Prefill, triggers MoRIIOConnectorMetadata.
248 + scheduler_output = scheduler.schedule()
249 + kv_connector_metadata = scheduler_output.kv_connector_metadata
250 + assert kv_connector_metadata is not None, "kv_connector_metadata is None"
251 + assert isinstance(kv_connector_metadata, MoRIIOConnectorMetadata)
252 +
253 + assert len(kv_connector_metadata.reqs_to_save) == 1, (
254 + "Unexpected number of reqs_to_save"
255 + )
256 + assert len(kv_connector_metadata.reqs_to_recv) == 0, (
257 + "Unexpected number of reqs_to_recv"
258 + )
259 + assert len(kv_connector_metadata.reqs_to_send) == 0, (
260 + "Unexpected number of reqs_to_send"
261 + )
262 + assert request_id in kv_connector_metadata.reqs_to_save, (
263 + "Request ID not in reqs_to_save"
264 + )
265 + req_meta = kv_connector_metadata.reqs_to_save[request_id]
266 +
267 + for block_id, block in zip(
268 + req_meta.local_block_ids,
269 + scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
270 + request_id
271 + ],
272 + ):
273 + assert block_id == block.block_id, f"{block_id} != {block.block_id}"
274 +
275 +
276 +def test_write_mode_with_chunked_prefill_saves_local_block_ids():
277 + """Write mode with chunked prefill still records correct local block ids."""
278 + # Setup Scheduler and Request
279 + MAX_NUM_BATCHED_TOKENS = 64
280 + NUM_TOKENS = MAX_NUM_BATCHED_TOKENS * 2 + MAX_NUM_BATCHED_TOKENS // 2
281 +
282 + vllm_config = create_vllm_config(
283 + max_num_batched_tokens=MAX_NUM_BATCHED_TOKENS, role="kv_producer"
284 + )
285 + BLOCK_SIZE = vllm_config.cache_config.block_size
286 +
287 + scheduler = create_scheduler(vllm_config)
288 +
289 + # 2 Full Blocks and 1 Half Block.
290 +
291 + request = create_request(
292 + request_id=1,
293 + block_size=BLOCK_SIZE,
294 + num_tokens=NUM_TOKENS,
295 + do_remote_decode=True,
296 + do_remote_prefill=False,
297 + )
298 + request_id = request.request_id
299 +
300 + scheduler.add_request(request)
301 +
302 + # Fake Config
303 + request = _setup_kv_transfer_request(request)
304 +
305 + # Remote Prefill with chunked prefill, triggers multiple schedules.
306 + expected_counts = [(0, 0, 0), (0, 0, 0), (1, 0, 0)]
307 + kv_connector_metadata = None
308 + for _, (expected_save, expected_recv, expected_send) in enumerate(expected_counts):
309 + scheduler_output = scheduler.schedule()
310 + kv_connector_metadata = scheduler_output.kv_connector_metadata
311 +
312 + assert len(kv_connector_metadata.reqs_to_save) == expected_save
313 + assert len(kv_connector_metadata.reqs_to_recv) == expected_recv
314 + assert len(kv_connector_metadata.reqs_to_send) == expected_send
315 + assert kv_connector_metadata is not None, "kv_connector_metadata is None"
316 + assert request_id in kv_connector_metadata.reqs_to_save, (
317 + "Request ID not in reqs_to_save"
318 + )
319 + req_meta = kv_connector_metadata.reqs_to_save[request_id]
320 +
321 + for block_id, block in zip(
322 + req_meta.local_block_ids,
323 + scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
324 + request_id
325 + ],
326 + ):
327 + assert block_id == block.block_id, f"{block_id} != {block.block_id}"
328 +
329 +
330 +def test_read_mode_loads_remote_block_ids(moriio_read_mode):
331 + """Read mode loads remote block ids into local cache mapping."""
332 +
333 + # Setup Scheduler and Request
334 + vllm_config = create_vllm_config(role="kv_consumer")
335 + scheduler = create_scheduler(vllm_config)
336 +
337 + # 2 Full Blocks and 1 Half Block.
338 + BLOCK_SIZE = vllm_config.cache_config.block_size
339 + NUM_EXTERNAL_FULL_BLOCKS = 2
340 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
341 +
342 + request = create_request(
343 + request_id=1,
344 + block_size=BLOCK_SIZE,
345 + num_tokens=NUM_TOKENS,
346 + do_remote_decode=False,
347 + do_remote_prefill=True,
348 + )
349 + request_id = request.request_id
350 +
351 + scheduler.add_request(request)
352 + block_list = scheduler.kv_cache_manager.coordinator.single_type_managers[
353 + 0
354 + ].req_to_blocks[request_id]
355 +
356 + request = _setup_kv_transfer_request(request)
357 +
358 + # Set remote block ids to be fetched.
359 + request.kv_transfer_params["remote_block_ids"] = block_list
360 +
361 + # Remote Prefill, triggers MorIIOConnectorMetadata.
362 +
363 + scheduler_output = scheduler.schedule()
364 + kv_connector_metadata = scheduler_output.kv_connector_metadata
365 + assert kv_connector_metadata is not None, "kv_connector_metadata is None"
366 + assert isinstance(kv_connector_metadata, MoRIIOConnectorMetadata), (
367 + "kv_connector_metadata is not MoRIIOConnectorMetadata"
368 + )
369 + assert len(kv_connector_metadata.reqs_to_save) == 0, (
370 + "Unexpected number of reqs_to_save"
371 + )
372 + assert len(kv_connector_metadata.reqs_to_recv) == 1, (
373 + "Unexpected number of reqs_to_recv"
374 + )
375 + assert len(kv_connector_metadata.reqs_to_send) == 0, (
376 + "Unexpected number of reqs_to_send"
377 + )
378 + assert request_id in kv_connector_metadata.reqs_to_recv, (
379 + "Request ID not in reqs_to_recv"
380 + )
381 + req_meta = kv_connector_metadata.reqs_to_recv[request_id]
382 +
383 + for block_id, block in zip(
384 + req_meta.local_block_ids,
385 + scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
386 + request_id
387 + ],
388 + ):
389 + assert block_id == block.block_id, f"{block_id} != {block.block_id}"
390 +
391 +
392 +@pytest.mark.skipif(
393 + not aiter_available, reason="Requires aiter package for ROCm FlashAttention backend"
394 +)
395 +def test_register_kv_caches(mock_parallel_groups):
396 + """Test that MoRIIOConnector.register_kv_caches correctly registers kv caches."""
397 + ROLE = "kv_consumer"
398 + IP = get_ip()
399 + vllm_config = create_vllm_config(role=ROLE)
400 + DEFAULT_PORT = 6301
401 + TP_RANK = 0
402 + DP_RANK = 0
403 + from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
404 +
405 + backend_cls = AiterFlashAttentionBackend
406 +
407 + # Create test kv cache tensors using proper backend shape
408 + kv_cache_shape = backend_cls.get_kv_cache_shape(
409 + num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
410 + )
411 + shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
412 + unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
413 + kv_caches = {
414 + "layer0": shared_tensor,
415 + "layer1": unique_tensor,
416 + "layer2": shared_tensor,
417 + }
418 +
419 + with (
420 + patch(
421 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector.threading.Event"
422 + ),
423 + patch(
424 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector.threading.Thread"
425 + ),
426 + ):
427 + # Create connector
428 + vllm_config.kv_transfer_config.kv_connector_extra_config.update(
429 + {
430 + "proxy_ip": "127.0.0.1",
431 + "proxy_ping_port": 12345,
432 + "http_port": 12346,
433 + }
434 + )
435 +
436 + connector = MoRIIOConnector(vllm_config, KVConnectorRole.WORKER)
437 + connector.connector_worker = FakeMorIIOConnectorWorker(
438 + vllm_config, connector.engine_id, hand_shake_latency=0
439 + )
440 +
441 + from mori.io import (
442 + MemoryDesc,
443 + )
444 +
445 + # Execute register_kv_caches
446 + connector.register_kv_caches(kv_caches)
447 +
448 + # Verify that the MemoryDesc stored in layer_name_to_local_kv_cache_metadata
449 + assert (
450 + shared_tensor.data_ptr()
451 + == MemoryDesc.unpack(
452 + connector.connector_worker.layer_name_to_local_kv_cache_metadata[
453 + "layer0"
454 + ][0]
455 + ).data
456 + )
457 + assert (
458 + unique_tensor.data_ptr()
459 + == MemoryDesc.unpack(
460 + connector.connector_worker.layer_name_to_local_kv_cache_metadata[
461 + "layer1"
462 + ][0]
463 + ).data
464 + )
465 + assert (
466 + shared_tensor.data_ptr()
467 + == MemoryDesc.unpack(
468 + connector.connector_worker.layer_name_to_local_kv_cache_metadata[
469 + "layer2"
470 + ][0]
471 + ).data
472 + )
473 +
474 + # Verify engine keys
475 + expected_engine_key = f"{ROLE[3:]}:{IP}:{DEFAULT_PORT}:tp{TP_RANK}:dp{DP_RANK}"
476 + assert (
477 + MemoryDesc.unpack(
478 + connector.connector_worker.layer_name_to_local_kv_cache_metadata[
479 + "layer0"
480 + ][0]
481 + ).engine_key
482 + == expected_engine_key
483 + )
484 +
485 +
486 +@pytest.mark.skipif(
487 + not aiter_available, reason="Requires aiter package for ROCm FlashAttention backend"
488 +)
489 +def test_moriio_handshake_returns_metadata(mock_parallel_groups):
490 + """MoRIIO handshake socket returns valid agent metadata over ZMQ."""
491 +
492 + ROLE = "kv_consumer"
493 + vllm_config = create_vllm_config(role=ROLE)
494 + from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
495 +
496 + backend_cls = AiterFlashAttentionBackend
497 +
498 + # Create test kv cache tensors using proper backend shape
499 + kv_cache_shape = backend_cls.get_kv_cache_shape(
500 + num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
501 + )
502 + shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
503 + unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
504 + kv_caches = {
505 + "layer0": shared_tensor,
506 + "layer1": unique_tensor,
507 + "layer2": shared_tensor,
508 + }
509 +
510 + with (
511 + patch(
512 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine.MoRIIOWrapper",
513 + FakeMorIIOWrapper,
514 + ),
515 + ):
516 + handshake_port = _find_free_port()
517 + # Create connector
518 + vllm_config.kv_transfer_config.kv_connector_extra_config.update(
519 + {
520 + "proxy_ip": "127.0.0.1",
521 + "proxy_ping_port": 12345,
522 + "http_port": 12346,
523 + "handshake_port": handshake_port,
524 + }
525 + )
526 + connector = MoRIIOConnector(vllm_config, KVConnectorRole.WORKER)
527 +
528 + # Execute register_kv_caches
529 + connector.register_kv_caches(kv_caches)
530 +
531 + # Connect to handshake socket and request metadata
532 + path = make_zmq_path("tcp", "127.0.0.1", handshake_port)
533 + with zmq_ctx(zmq.DEALER, path) as sock:
534 + sock.send(MoRIIOConstants.GET_META_MSG)
535 + received_frame = sock.recv_multipart()
536 +
537 + if len(received_frame) != 2 or received_frame[0] != b"":
538 + raise ValueError(f"Unexpected frame! {received_frame = }")
539 +
540 + metadata_bytes = received_frame[1]
541 + decoder = msgspec.msgpack.Decoder(MoRIIOAgentMetadata)
542 + metadata = decoder.decode(metadata_bytes)
543 + assert isinstance(metadata, MoRIIOAgentMetadata), (
544 + "Decoded metadata is not MoRIIOAgentMetadata"
545 + )
vllm/distributed/kv_transfer/kv_connector/factory.py
+6
@@ -179,6 +179,12 @@ KVConnectorFactory.register_connector(
179 "MultiConnector",
180 )
181
182 +KVConnectorFactory.register_connector(
183 + "MoRIIOConnector",
184 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector",
185 + "MoRIIOConnector",
186 +)
187 +
188 KVConnectorFactory.register_connector(
189 "OffloadingConnector",
190 "vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector",
vllm/distributed/kv_transfer/kv_connector/v1/moriio/__init__.py
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py new
+321
@@ -0,0 +1,321 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import contextlib
4 +import threading
5 +import time
6 +from collections.abc import Iterator
7 +from dataclasses import dataclass
8 +from typing import TYPE_CHECKING, Any, Optional
9 +
10 +import msgspec
11 +import torch
12 +import zmq
13 +
14 +from vllm import envs
15 +from vllm.config import VllmConfig
16 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
17 + KVConnectorMetadata,
18 +)
19 +from vllm.distributed.parallel_state import (
20 + get_tensor_model_parallel_rank,
21 + get_tensor_model_parallel_world_size,
22 +)
23 +from vllm.logger import init_logger
24 +from vllm.utils.network_utils import (
25 + get_ip,
26 + get_open_port,
27 + make_zmq_socket,
28 +)
29 +
30 +if TYPE_CHECKING:
31 + pass
32 +
33 +from dataclasses import field
34 +from enum import Enum
35 +
36 +logger = init_logger(__name__)
37 +
38 +
39 +Transfer = tuple[int, float]
40 +EngineId = str
41 +ReqId = str
42 +
43 +
44 +@dataclass
45 +class WriteTask:
46 + request_id: str
47 + dst_engine_id: str
48 + local_block_ids: list[int]
49 + remote_block_ids_hint: list[int] | None
50 + layer_name: str
51 + event: torch.cuda.Event
52 + remote_notify_port: int
53 + remote_ip: str
54 + enqueue_time: float = field(default_factory=time.perf_counter)
55 + retried: int = 0
56 +
57 +
58 +@dataclass
59 +class LayerTransferPlan:
60 + """Plan for transferring a single layer."""
61 +
62 + request_id: str
63 + layer_name: str
64 + sess_idx: int
65 + transfer_local_offsets: list[int]
66 + transfer_remote_offsets: list[int]
67 + transfer_sizes: list[int]
68 + use_batch: bool = True
69 +
70 +
71 +@dataclass
72 +class RemoteAllocInfo:
73 + """Information about remote block allocation."""
74 +
75 + block_ids: list[int]
76 + writes_done: int = 0
77 + decode_dp_rank: int = 0
78 + transfer_offset: tuple[list[int], list[int], list[int]] | None = None
79 +
80 +
81 +class ROLE(Enum):
82 + PRODUCER = "producer"
83 + CONSUMER = "consumer"
84 + NOTINIT = "notinit"
85 +
86 +
87 +class MoRIIOAgentMetadata(
88 + msgspec.Struct,
89 + omit_defaults=True, # type: ignore[call-arg]
90 + # required for @cached_property.d
91 + dict=True,
92 +):
93 + engine_id: str
94 + agent_metadata: bytes
95 + kv_caches_base_addr: list[int]
96 + num_blocks: int
97 + block_len: int
98 + attn_backend_name: str
99 +
100 +
101 +class RoleManager:
102 + """Manages role state across the connector."""
103 +
104 + _instance: Optional["RoleManager"] = None
105 + _lock = threading.Lock()
106 +
107 + def __init__(self) -> None:
108 + self._role: ROLE = ROLE.NOTINIT
109 +
110 + @classmethod
111 + def get_instance(cls) -> "RoleManager":
112 + if cls._instance is None:
113 + with cls._lock:
114 + if cls._instance is None:
115 + cls._instance = cls()
116 + return cls._instance
117 +
118 + def set_role(self, role: ROLE) -> None:
119 + """Set the current role."""
120 + with self._lock:
121 + self._role = role
122 +
123 + def get_role(self) -> ROLE:
124 + """Get the current role."""
125 + return self._role
126 +
127 +
128 +def set_role(role: ROLE):
129 + """Set the global role."""
130 + RoleManager.get_instance().set_role(role)
131 +
132 +
133 +def get_role() -> ROLE:
134 + """Get the global role."""
135 + return RoleManager.get_instance().get_role()
136 +
137 +
138 +class MoRIIOMode(Enum):
139 + READ = "read"
140 + WRITE = "write"
141 +
142 +
143 +class MoRIIOError(Exception):
144 + """Base exception for MoRIIO operations."""
145 +
146 + pass
147 +
148 +
149 +class HandshakeError(MoRIIOError):
150 + """Exception raised when handshake fails."""
151 +
152 + pass
153 +
154 +
155 +class TransferError(MoRIIOError):
156 + """Exception raised when transfer fails."""
157 +
158 + pass
159 +
160 +
161 +def get_moriio_mode() -> MoRIIOMode:
162 + read_mode = envs.VLLM_MORIIO_CONNECTOR_READ_MODE
163 + logger.debug("MoRIIO Connector read_mode: %s", read_mode)
164 + if read_mode:
165 + return MoRIIOMode.READ
166 + else:
167 + return MoRIIOMode.WRITE
168 +
169 +
170 +def get_port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int:
171 + return (dp_rank) * tp_size + tp_rank
172 +
173 +
174 +@dataclass
175 +class MoRIIOConfig:
176 + local_ip: str
177 + local_kv_port: int
178 + proxy_ip: str
179 + local_ping_port: int
180 + proxy_ping_port: int
181 + http_port: int
182 + handshake_port: int
183 + notify_port: int
184 + tp_rank: int
185 + dp_rank: int
186 + dp_size: int
187 + tp_size: int
188 +
189 + @classmethod
190 + def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig":
191 + # Port Configuration:
192 + # local_ping_port -> Outgoing heartbeat to proxy
193 + # proxy_ping_port -> Remote proxy's heartbeat ingress port
194 + # http_port -> Instance's HTTP service endpoint
195 + # local_kv_port -> service port for mori engine
196 + # notify_port -> For synchronizing stages between prefill and decode
197 + # handshake_port -> For initial handshake between mori engine
198 +
199 + # TODO : merge notify_port and handshake_port to simplify port management
200 + # supports non-contiguous ports
201 + assert vllm_config.kv_transfer_config is not None, (
202 + "kv_transfer_config must be set for MoRIIOConnector"
203 + )
204 + kv_transfer_config = vllm_config.kv_transfer_config
205 + extra_config = kv_transfer_config.kv_connector_extra_config
206 + tp_rank = get_tensor_model_parallel_rank()
207 + dp_rank = vllm_config.parallel_config.data_parallel_rank
208 + base_notify_port = int(extra_config["notify_port"])
209 + dp_size = vllm_config.parallel_config.data_parallel_size
210 + tp_size = get_tensor_model_parallel_world_size()
211 + port_offset = get_port_offset(dp_rank, tp_rank)
212 +
213 + return cls(
214 + local_ip=get_ip(),
215 + local_kv_port=get_open_port(),
216 + proxy_ip=extra_config["proxy_ip"],
217 + local_ping_port=get_open_port(),
218 + proxy_ping_port=int(extra_config["proxy_ping_port"]),
219 + http_port=int(extra_config["http_port"]),
220 + handshake_port=int(extra_config["handshake_port"]),
221 + notify_port=base_notify_port + port_offset,
222 + tp_rank=tp_rank,
223 + dp_rank=dp_rank,
224 + dp_size=dp_size,
225 + tp_size=tp_size,
226 + )
227 +
228 +
229 +class MoRIIOConstants:
230 + """Constants for MoRIIO connector."""
231 +
232 + # ZMQ message types
233 + GET_META_MSG = b"get_meta_msg"
234 + POP_DONE_RECV = b"pop_done_recv"
235 + OVER = b"OVER"
236 + COMPLETION_PREFIX = "cmpl"
237 +
238 + PING_INTERVAL = 5
239 + MAX_PING_RETRIES = 100
240 + DEFAULT_HANDSHAKE_PORT = "6301"
241 + DEFAULT_NOTIFY_PORT = "61005"
242 +
243 + VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT = 3600
244 +
245 +
246 +@dataclass
247 +class ReqMeta:
248 + """Metadata for a single request."""
249 +
250 + local_block_ids: list[int]
251 + remote_block_ids: list[int]
252 + remote_host: str
253 + remote_port: int
254 + remote_handshake_port: int
255 + remote_notify_port: int
256 + remote_engine_id: str
257 + tp_size: int
258 + remote_dp_size: int
259 +
260 +
261 +class MoRIIOConnectorMetadata(KVConnectorMetadata):
262 + def __init__(self):
263 + self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
264 + self.reqs_to_save: dict[ReqId, ReqMeta] = {}
265 + self.reqs_to_send: dict[ReqId, float] = {}
266 +
267 + def __repr__(self):
268 + return_str = ""
269 + for req_id, req_meta in self.reqs_to_recv.items():
270 + return_str += (
271 + f"{req_id = },{req_meta.local_block_ids = },"
272 + f"{req_meta.remote_host = },{req_meta.remote_port = }"
273 + f"{req_meta.remote_engine_id = },{req_meta.tp_size = }"
274 + )
275 + return_str = f"MoRIIOConnectorMetadata:reqs_to_recv:{return_str},"
276 +
277 + for req_id, expiry in self.reqs_to_send.items():
278 + return_str += f"{req_id = },{expiry = }"
279 + return_str = f"MoRIIOConnectorMetadata:reqs_to_send:{return_str},"
280 + return return_str
281 +
282 + def add_new_req(
283 + self,
284 + request_id: ReqId,
285 + local_block_ids: list[int],
286 + kv_transfer_params: dict[str, Any],
287 + write_mode=False,
288 + ):
289 + _req = ReqMeta(
290 + local_block_ids=local_block_ids,
291 + remote_block_ids=kv_transfer_params["remote_block_ids"],
292 + remote_engine_id=kv_transfer_params["remote_engine_id"],
293 + remote_host=kv_transfer_params["remote_host"],
294 + remote_port=kv_transfer_params["remote_port"],
295 + remote_handshake_port=kv_transfer_params["remote_handshake_port"],
296 + remote_notify_port=kv_transfer_params["remote_notify_port"],
297 + tp_size=kv_transfer_params.get("tp_size", 1),
298 + remote_dp_size=kv_transfer_params.get("remote_dp_size", 1),
299 + )
300 + if write_mode:
301 + self.reqs_to_save[request_id] = _req
302 + else:
303 + self.reqs_to_recv[request_id] = _req
304 +
305 +
306 +@contextlib.contextmanager
307 +def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
308 + """Context manager for a ZMQ socket"""
309 +
310 + if socket_type not in (zmq.ROUTER, zmq.REQ, zmq.DEALER):
311 + raise ValueError(f"Unexpected socket type: {socket_type}")
312 +
313 + ctx: zmq.Context | None = None
314 + try:
315 + ctx = zmq.Context() # type: ignore[attr-defined]
316 + yield make_zmq_socket(
317 + ctx=ctx, path=addr, socket_type=socket_type, bind=socket_type == zmq.ROUTER
318 + )
319 + finally:
320 + if ctx is not None:
321 + ctx.destroy(linger=0)
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py new
+1515
@@ -0,0 +1,1515 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import logging
4 +import math
5 +import queue
6 +import threading
7 +import time
8 +from collections import defaultdict
9 +from concurrent.futures import Future, ThreadPoolExecutor
10 +from typing import TYPE_CHECKING, Any, Optional
11 +
12 +import msgpack
13 +import msgspec
14 +import numpy as np
15 +import torch
16 +import zmq
17 +
18 +from vllm.attention.selector import get_attn_backend
19 +from vllm.config import VllmConfig
20 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
21 + KVConnectorBase_V1,
22 + KVConnectorMetadata,
23 + KVConnectorRole,
24 +)
25 +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
26 + ROLE,
27 + EngineId,
28 + HandshakeError,
29 + MoRIIOAgentMetadata,
30 + MoRIIOConfig,
31 + MoRIIOConnectorMetadata,
32 + MoRIIOConstants,
33 + MoRIIOMode,
34 + ReqId,
35 + ReqMeta,
36 + WriteTask,
37 + get_moriio_mode,
38 + get_port_offset,
39 + get_role,
40 + set_role,
41 + zmq_ctx,
42 +)
43 +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import (
44 + MoRIIOWrapper,
45 + MoRIIOWriter,
46 +)
47 +from vllm.distributed.parallel_state import (
48 + get_tensor_model_parallel_world_size,
49 + get_tp_group,
50 + get_world_group,
51 +)
52 +from vllm.forward_context import ForwardContext
53 +from vllm.logger import init_logger
54 +from vllm.utils.network_utils import (
55 + get_ip,
56 + make_zmq_path,
57 + make_zmq_socket,
58 +)
59 +from vllm.v1.core.sched.output import SchedulerOutput
60 +from vllm.v1.request import RequestStatus
61 +
62 +if TYPE_CHECKING:
63 + from vllm.attention.backends.abstract import AttentionMetadata
64 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
65 + from vllm.v1.kv_cache_interface import KVCacheConfig
66 + from vllm.v1.request import Request
67 +
68 +logger = init_logger(__name__)
69 +
70 +try:
71 + from mori.io import (
72 + BackendType,
73 + IOEngine,
74 + IOEngineConfig,
75 + )
76 +
77 + logger.info("MoRIIO is available")
78 + MoRIIO_enabled = True
79 +except ImportError:
80 + logger.error("MoRIIO is not available")
81 + MoRIIO_enabled = False
82 +
83 +
84 +def is_moriio_available() -> bool:
85 + return MoRIIO_enabled
86 +
87 +
88 +class MoRIIOConnector(KVConnectorBase_V1):
89 + def __init__(
90 + self,
91 + vllm_config: VllmConfig,
92 + role: KVConnectorRole,
93 + kv_cache_config: Optional["KVCacheConfig"] = None,
94 + ):
95 + super().__init__(vllm_config, role)
96 + assert vllm_config.kv_transfer_config is not None, (
97 + "kv_transfer_config must be set for MoRIIOConnector"
98 + )
99 +
100 + self.kv_transfer_config = vllm_config.kv_transfer_config
101 + self._set_port_defaults(vllm_config)
102 +
103 + self.engine_id = (
104 + str(get_ip())
105 + + ":"
106 + + str(self.kv_transfer_config.kv_connector_extra_config["handshake_port"])
107 + )
108 + self.mode = get_moriio_mode()
109 + if role == KVConnectorRole.SCHEDULER:
110 + self.connector_scheduler: MoRIIOConnectorScheduler | None = (
111 + MoRIIOConnectorScheduler(vllm_config, self.engine_id)
112 + )
113 + self.connector_worker: MoRIIOConnectorWorker | None = None
114 + elif role == KVConnectorRole.WORKER:
115 + self.connector_scheduler = None
116 + self.connector_worker = MoRIIOConnectorWorker(vllm_config, self.engine_id)
117 + logger.info(
118 + "Initialized MoRIIO Connector,engine_id:%s,role: %s",
119 + self.engine_id,
120 + role.value,
121 + )
122 +
123 + ############################################################
124 + # Scheduler Side Methods
125 + ############################################################
126 +
127 + def _set_port_defaults(self, vllm_config: VllmConfig):
128 + assert vllm_config.kv_transfer_config is not None, (
129 + "kv_transfer_config must be set for MoRIIOConnector"
130 + )
131 + kv_transfer_config = vllm_config.kv_transfer_config
132 + extra_config = kv_transfer_config.kv_connector_extra_config
133 +
134 + if "handshake_port" not in extra_config or not extra_config["handshake_port"]:
135 + extra_config["handshake_port"] = MoRIIOConstants.DEFAULT_HANDSHAKE_PORT
136 +
137 + if "notify_port" not in extra_config or not extra_config["notify_port"]:
138 + extra_config["notify_port"] = MoRIIOConstants.DEFAULT_NOTIFY_PORT
139 +
140 + def get_num_new_matched_tokens(
141 + self, request: "Request", num_computed_tokens: int
142 + ) -> tuple[int, bool]:
143 + assert self.connector_scheduler is not None
144 + return self.connector_scheduler.get_num_new_matched_tokens(
145 + request, num_computed_tokens
146 + )
147 +
148 + def update_state_after_alloc(
149 + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
150 + ):
151 + assert self.connector_scheduler is not None
152 + return self.connector_scheduler.update_state_after_alloc(
153 + request, blocks, num_external_tokens, self.connector_worker
154 + )
155 +
156 + def build_connector_meta(
157 + self,
158 + scheduler_output: SchedulerOutput,
159 + ) -> KVConnectorMetadata:
160 + assert self.connector_scheduler is not None
161 + return self.connector_scheduler.build_connector_meta(scheduler_output)
162 +
163 + def request_finished(
164 + self,
165 + request: "Request",
166 + block_ids: list[int],
167 + ) -> tuple[bool, dict[str, Any] | None]:
168 + assert self.connector_scheduler is not None
169 + return self.connector_scheduler.request_finished(request, block_ids)
170 +
171 + ############################################################
172 + # Worker Side Methods
173 + ############################################################
174 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
175 + assert self.connector_worker is not None
176 + self.connector_worker.register_kv_caches(kv_caches)
177 +
178 + def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
179 + """Get the finished recving and sending requests."""
180 + assert self.connector_worker is not None
181 + return self.connector_worker.get_finished()
182 +
183 + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
184 + assert self.connector_worker is not None
185 + if self.mode == MoRIIOMode.WRITE and get_role() == ROLE.CONSUMER:
186 + self.connector_worker.moriio_wrapper.async_wait_reqid()
187 +
188 + assert isinstance(self._connector_metadata, MoRIIOConnectorMetadata)
189 + self.connector_worker.start_load_kv(self._connector_metadata)
190 +
191 + def wait_for_layer_load(self, layer_name: str) -> None:
192 + pass
193 +
194 + def save_kv_layer(
195 + self,
196 + layer_name: str,
197 + kv_layer: torch.Tensor,
198 + attn_metadata: "AttentionMetadata",
199 + **kwargs,
200 + ) -> None:
201 + # Only producer/prefill saves KV Cache
202 + if get_role() == ROLE.CONSUMER:
203 + return
204 + assert self.connector_worker is not None, (
205 + "save_kv_layer called on scheduler role"
206 + )
207 +
208 + assert isinstance(self._connector_metadata, MoRIIOConnectorMetadata), (
209 + "Connector metadata not initialized yet"
210 + )
211 + self.connector_worker.save_kv_layer(
212 + self._connector_metadata, layer_name, kv_layer, attn_metadata, **kwargs
213 + )
214 +
215 + return None
216 +
217 + def wait_for_save(self):
218 + pass
219 +
220 + def shutdown(self):
221 + if self.connector_worker is not None:
222 + self.connector_worker.shutdown()
223 + if self.connector_scheduler is not None:
224 + self.connector_scheduler.shutdown()
225 +
226 + def has_connector_metadata(self) -> bool:
227 + """Check whether the connector metadata is currently set.
228 +
229 + Returns:
230 + bool: True if connector metadata exists, False otherwise.
231 + """
232 + try:
233 + return self._connector_metadata is not None
234 + except AttributeError:
235 + return False
236 +
237 +
238 +class MoRIIOConnectorScheduler:
239 + """Implementation of Scheduler side methods"""
240 +
241 + def __init__(self, vllm_config: VllmConfig, engine_id: str):
242 + self.vllm_config = vllm_config
243 +
244 + assert vllm_config.kv_transfer_config is not None, (
245 + "kv_transfer_config must be set for MoRIIOConnector"
246 + )
247 + self.kv_transfer_config = vllm_config.kv_transfer_config
248 + self.block_size = vllm_config.cache_config.block_size
249 + self.engine_id: EngineId = engine_id
250 + self.mode = get_moriio_mode()
251 + self.host_ip = get_ip()
252 + self.handshake_port = self.kv_transfer_config.kv_connector_extra_config[
253 + "handshake_port"
254 + ]
255 + logger.info("Initializing MoRIIO Scheduler engine_id = %s", engine_id)
256 +
257 + self.side_notify_port = self.kv_transfer_config.kv_connector_extra_config[
258 + "notify_port"
259 + ]
260 + self.tp_size = self.vllm_config.parallel_config.tensor_parallel_size
261 + self.dp_rank = self.vllm_config.parallel_config.data_parallel_rank
262 + self.is_producer = self.kv_transfer_config.kv_role == "kv_producer"
263 + # Requests that need to start recv/send.
264 + # New requests are added by update_state_after_alloc in
265 + # the scheduler. Used to make metadata passed to Worker.
266 + self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
267 + self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}
268 +
269 + # For chunked prefill, we perform layer-wise access within the final chunk.
270 + # TODO: Perform transfer at end chunk.
271 + self._reqs_need_pending_save: dict[ReqId, tuple[Request, list[int]]] = {}
272 +
273 + if self.is_producer:
274 + set_role(ROLE.PRODUCER)
275 + else:
276 + set_role(ROLE.CONSUMER)
277 + # Reqs to send and their expiration time
278 + self._reqs_need_send: dict[ReqId, float] = {}
279 + self.paths: dict[str, zmq.Socket] = {}
280 +
281 + def get_num_new_matched_tokens(
282 + self,
283 + request: "Request",
284 + num_computed_tokens: int,
285 + ) -> tuple[int, bool]:
286 + """
287 + For remote prefill, pull all prompt blocks from remote
288 + asynchronously relative to engine execution.
289 +
290 + Args:
291 + request (Request): the request object.
292 + num_computed_tokens (int): the number of locally
293 + computed tokens for this request
294 + Returns:
295 + * the number of tokens that can be loaded from the
296 + external KV cache beyond what is already computed.
297 + * true if the external KV cache tokens will be loaded
298 + asynchronously (between scheduler steps).
299 + """
300 + if self.is_producer:
301 + return 0, False
302 +
303 + token_ids = request.prompt_token_ids or []
304 + if self.mode == MoRIIOMode.WRITE:
305 + # MoriiO in write mode, no remote prefill
306 +
307 + return len(token_ids) - num_computed_tokens, True
308 +
309 + return len(token_ids) - 1 - num_computed_tokens, False
310 +
311 + def send_notify_block(
312 + self, req_id: str, block_notify_list: list[int], host=None, port=None
313 + ):
314 + path = make_zmq_path("tcp", host, port)
315 + if path not in self.paths:
316 + ctx = zmq.Context.instance()
317 + sock = make_zmq_socket(
318 + ctx=ctx, path=path, socket_type=zmq.DEALER, bind=False
319 + )
320 + self.paths[path] = sock
321 +
322 + data = {
323 + "req_id": req_id,
324 + "block_notify_list": block_notify_list or [],
325 + "decode_rank": self.dp_rank,
326 + "type": "remote_blocks",
327 + }
328 + serialized_data = msgpack.dumps(data)
329 + self.paths[path].send(serialized_data)
330 +
331 + def update_state_after_alloc(
332 + self,
333 + request: "Request",
334 + blocks: "KVCacheBlocks",
335 + num_external_tokens: int,
336 + connector_worker: Optional["MoRIIOConnectorWorker"] = None,
337 + ):
338 + params = request.kv_transfer_params
339 + if not params:
340 + return
341 + if params.get("do_remote_decode"):
342 + local_block_ids = blocks.get_block_ids()[0]
343 + self._reqs_need_save[request.request_id] = (request, local_block_ids)
344 +
345 + if params is not None and params.get("do_remote_prefill"):
346 + if self.mode == MoRIIOMode.READ:
347 + if remote_block_ids := params.get("remote_block_ids"):
348 + if all(
349 + p in params
350 + for p in ("remote_engine_id", "remote_host", "remote_port")
351 + ):
352 + # If remote_blocks and num_external_tokens = 0, we
353 + # a full prefix cache hit on the D worker. We need to call
354 + # send_notif in _read_blocks to free the memory on the P.
355 +
356 + # Get unhashed blocks to pull from remote.
357 + local_block_ids = blocks.get_block_ids()[0]
358 + assert len(local_block_ids) <= len(remote_block_ids)
359 + if len(local_block_ids) == len(remote_block_ids):
360 + pass
361 + else:
362 + local_block_ids = remote_block_ids[-len(local_block_ids) :]
363 +
364 + self._reqs_need_recv[request.request_id] = (
365 + request,
366 + local_block_ids,
367 + )
368 + else:
369 + logger.warning(
370 + "Got invalid KVTransferParams: %s. This "
371 + "request will not utilize KVTransfer",
372 + params,
373 + )
374 +
375 + else:
376 + assert request.kv_transfer_params is not None, (
377 + "kv_transfer_params should not be None"
378 + )
379 +
380 + remote_dp_rank = request.kv_transfer_params.get("remote_dp_rank", 0)
381 +
382 + for tp_index in range(self.tp_size):
383 + target_port = request.kv_transfer_params[
384 + "remote_notify_port"
385 + ] + get_port_offset(remote_dp_rank, tp_index)
386 +
387 + self.send_notify_block(
388 + req_id=request.request_id,
389 + block_notify_list=blocks.get_block_ids()[0],
390 + host=params.get("remote_host"),
391 + port=target_port,
392 + )
393 +
394 + # Only trigger 1 KV transfer per request.
395 +
396 + params["do_remote_prefill"] = False
397 +
398 + def build_connector_meta(
399 + self,
400 + scheduler_output: SchedulerOutput,
401 + ) -> KVConnectorMetadata:
402 + meta = MoRIIOConnectorMetadata()
403 +
404 + if self.mode == MoRIIOMode.WRITE:
405 + # when async_load_kv finished,
406 + # new reqs will be added to scheduler_output.scheduled_new_reqs
407 +
408 + if get_role() == ROLE.CONSUMER:
409 + for new_req in scheduler_output.scheduled_new_reqs:
410 + red_id = new_req.req_id
411 + local_block_ids = list(new_req.block_ids)[0]
412 + assert new_req.sampling_params is not None, (
413 + f"sampling_params is None for req {new_req.req_id}"
414 + )
415 + assert hasattr(new_req.sampling_params, "extra_args"), (
416 + f"sampling_params missing extra_args for req {new_req.req_id}"
417 + )
418 + kv_transfer_params = (
419 + new_req.sampling_params.extra_args.get("kv_transfer_params", {})
420 + if new_req.sampling_params.extra_args
421 + else {}
422 + )
423 + meta.add_new_req(
424 + red_id,
425 + local_block_ids,
426 + kv_transfer_params,
427 + )
428 + if get_role() == ROLE.PRODUCER:
429 + # This is the logic for checking against chunked prefill.
430 + # When the last chunk is identified,
431 + # It places the request metadata into the saving queue.
432 +
433 + for i, req_id in enumerate(
434 + scheduler_output.scheduled_cached_reqs.req_ids
435 + ):
436 + new_block_ids = (
437 + scheduler_output.scheduled_cached_reqs.new_block_ids[i]
438 + )
439 +
440 + if new_block_ids is not None:
441 + block_ids = new_block_ids[0]
442 + # TODO : hybrid attn, etc
443 + req, existing_blocks = self._reqs_need_pending_save[req_id]
444 + updated_blocks = list(existing_blocks) + (block_ids)
445 + self._reqs_need_pending_save[req_id] = (req, updated_blocks)
446 + if (
447 + len(self._reqs_need_pending_save[req_id][1])
448 + * self.block_size
449 + >= req.num_prompt_tokens
450 + ):
451 + meta.add_new_req(
452 + request_id=req_id,
453 + local_block_ids=self._reqs_need_pending_save[req_id][1],
454 + kv_transfer_params=req.kv_transfer_params or {},
455 + write_mode=True,
456 + )
457 + del self._reqs_need_pending_save[req_id]
458 +
459 + # Loop through scheduled reqs and convert to ReqMeta.
460 + for req_id, (req, block_ids) in self._reqs_need_recv.items():
461 + assert req.kv_transfer_params is not None
462 + meta.add_new_req(
463 + request_id=req_id,
464 + local_block_ids=block_ids,
465 + kv_transfer_params=req.kv_transfer_params,
466 + )
467 +
468 + for req_id, (req, block_ids) in self._reqs_need_save.items():
469 + assert req.kv_transfer_params is not None
470 + if req.num_prompt_tokens > len(block_ids) * self.block_size:
471 + # not last chunk prefill
472 + self._reqs_need_pending_save[req_id] = (req, block_ids)
473 + continue
474 + meta.add_new_req(
475 + request_id=req_id,
476 + local_block_ids=block_ids,
477 + kv_transfer_params=req.kv_transfer_params,
478 + write_mode=True,
479 + )
480 + # Clear the list once workers start the transfers
481 +
482 + meta.reqs_to_send = self._reqs_need_send
483 +
484 + self._reqs_need_recv.clear()
485 + self._reqs_need_save.clear()
486 + self._reqs_need_send = {}
487 +
488 + return meta
489 +
490 + def shutdown(self):
491 + for path, sock in self.paths.items():
492 + try:
493 + sock.close(linger=0)
494 + logger.debug("Closed ZMQ socket for path: %s", path)
495 + except Exception as e:
496 + logger.warning("Error closing ZMQ socket for path %s: %s", path, e)
497 + self.paths.clear()
498 +
499 + def request_finished(
500 + self,
501 + request: "Request",
502 + block_ids: list[int],
503 + ) -> tuple[bool, dict[str, Any] | None]:
504 + """
505 + Once a request is finished, determine whether request blocks
506 + should be freed now or will be sent asynchronously and freed later.
507 + """
508 +
509 + params = request.kv_transfer_params
510 + logger.debug(
511 + "MoriioConnector request_finished, request_status=%s, "
512 + "kv_transfer_params=%s",
513 + request.status,
514 + params,
515 + )
516 + if not params:
517 + return False, None
518 +
519 + if params.get("do_remote_prefill"):
520 + # If do_remote_prefill is still True when the request is finished,
521 + # update_state_after_alloc must not have been called (the request
522 + # must have been aborted before it was scheduled).
523 + # To avoid stranding the prefill blocks in the prefill instance,
524 + # we must add empty block_ids to _reqs_need_recv so that our
525 + # worker side will notify and free blocks in the prefill instance.
526 + self._reqs_need_recv[request.request_id] = (request, [])
527 + params["do_remote_prefill"] = False
528 + return False, None
529 +
530 + if (
531 + not params.get("do_remote_decode")
532 + or request.status != RequestStatus.FINISHED_LENGTH_CAPPED
533 + ):
534 + return False, None
535 +
536 + # computed_block_ids = block_ids if all_full else block_ids[:-1]
537 + computed_block_ids = block_ids
538 + # If prompt < block_size, no xfer so free blocks immediately.
539 + delay_free_blocks = len(computed_block_ids) > 0
540 +
541 + if delay_free_blocks:
542 + # Prefill request on remote. It will be read from D upon completion
543 + self._reqs_need_send[request.request_id] = (
544 + time.perf_counter()
545 + + MoRIIOConstants.VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT
546 + )
547 +
548 + # If we execute in P-D serial mode, no notification port is needed.
549 + return delay_free_blocks, dict(
550 + do_remote_prefill=True,
551 + do_remote_decode=False,
552 + remote_block_ids=computed_block_ids,
553 + remote_engine_id=self.engine_id,
554 + remote_host=self.host_ip,
555 + remote_port=self.handshake_port,
556 + tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
557 + )
558 +
559 +
560 +class MoRIIOConnectorWorker:
561 + """Implementation of Worker side methods"""
562 +
563 + def __init__(self, vllm_config: VllmConfig, engine_id: str):
564 + if not is_moriio_available():
565 + raise RuntimeError(
566 + "MoRIIO is not available. Please ensure the 'mori' package "
567 + "is installed and properly configured."
568 + )
569 +
570 + self.moriio_config = MoRIIOConfig.from_vllm_config(vllm_config)
571 + self.mode = get_moriio_mode()
572 +
573 + logger.info("Initializing MoRIIO worker %s", engine_id)
574 +
575 + logging.getLogger("aiter").disabled = True
576 +
577 + # Config.
578 + self.vllm_config = vllm_config
579 + assert vllm_config.kv_transfer_config is not None, (
580 + "kv_transfer_config must be set for MoRIIOConnector"
581 + )
582 + self.kv_transfer_config = vllm_config.kv_transfer_config
583 + self.is_producer = self.kv_transfer_config.is_kv_producer
584 +
585 + if self.is_producer:
586 + set_role(ROLE.PRODUCER)
587 + else:
588 + set_role(ROLE.CONSUMER)
589 + # mori engine
590 + self._rank = get_world_group().rank
591 + self._local_rank = get_world_group().local_rank
592 + self.tp_rank = self.moriio_config.tp_rank
593 + self.dp_rank = self.moriio_config.dp_rank
594 +
595 + self.local_ip = self.moriio_config.local_ip
596 + self.local_kv_port = self.moriio_config.local_kv_port
597 + self.proxy_ip = self.moriio_config.proxy_ip
598 + self.local_ping_port = self.moriio_config.local_ping_port
599 + self.proxy_ping_port = self.moriio_config.proxy_ping_port
600 + self.http_port = self.moriio_config.http_port
601 + self.handshake_port = self.moriio_config.handshake_port
602 + self.notify_port = self.moriio_config.notify_port
603 +
604 + self.zmq_context = zmq.Context()
605 + self.metadata_address = (
606 + f"{self.moriio_config.local_ip}:{self.moriio_config.local_ping_port}"
607 + )
608 + self.request_address = (
609 + f"{self.moriio_config.local_ip}:{self.moriio_config.http_port}"
610 + )
611 +
612 + self.moriio_engine = None
613 + self._handle_request_thread = None
614 + self._ping_thread = None
615 + self._writer = MoRIIOWriter(self)
616 +
617 + role = "producer" if self.is_producer else "consumer"
618 + engine_suffix = (
619 + f"{self.moriio_config.local_ip}:{self.moriio_config.handshake_port}:"
620 + f"tp{self.tp_rank}:dp{self.dp_rank}"
621 + )
622 + self.moriio_engine = IOEngine(
623 + f"{role}:{engine_suffix}",
624 + IOEngineConfig(
625 + self.moriio_config.local_ip, self.moriio_config.local_kv_port
626 + ),
627 + )
628 + logger.debug(
629 + "build MORI IOEngine %s (ip=%s port=%s)",
630 + f"{role}:{engine_suffix}",
631 + self.moriio_config.local_ip,
632 + self.moriio_config.local_kv_port,
633 + )
634 +
635 + if self._rank == 0 and self.moriio_config.proxy_ip:
636 + self._ping_thread = threading.Thread(
637 + target=self._ping, args=(self.zmq_context,), daemon=True
638 + )
639 + self._ping_thread.start()
640 +
641 + logger.info(
642 + "Initializing MoRIIO Engine, engine = %s, role = %s",
643 + self.moriio_engine,
644 + "producer" if self.is_producer else "consumer",
645 + )
646 +
647 + # Agent.
648 + self.moriio_wrapper = MoRIIOWrapper(tp_rank=self.tp_rank, dp_rank=self.dp_rank)
649 + self.moriio_wrapper.set_moriio_engine(self.moriio_engine)
650 + self.moriio_wrapper.set_backend_type(BackendType.RDMA)
651 + self.moriio_wrapper.notify_port = self.moriio_config.notify_port
652 + self.local_kv_cache_metadata: list[bytes] = []
653 + self.local_kv_cache_size: list[int] = []
654 + self.layer_name_to_local_kv_cache_metadata: dict[str, list[bytes]] = {}
655 +
656 + self.remote_kv_cache_metadata: list[bytes] = []
657 + self.remote_kv_cache_size: list[int] = []
658 + self.layer_name_to_remote_kv_cache_metadata: dict[str, dict[str, list[Any]]] = (
659 + dict()
660 + )
661 + self.remote_moriio_metadata: dict[EngineId, MoRIIOAgentMetadata] = {}
662 + self.slot_size_bytes = 0
663 +
664 + self.load_ready_flag: dict[str, bool] = {}
665 + self.write_ready_flags: dict[str, bool] = {}
666 + self.kv_cache_shape = None
667 + self.block_shape = None
668 + self.kv_element_size = 0
669 +
670 + # Map of engine_id -> {agent_name0, agent_name1..}.
671 + self._remote_agents: dict[EngineId, set[str]] = {}
672 +
673 + self.side_channel_port: int = (
674 + self.moriio_config.handshake_port
675 + + get_port_offset(self.dp_rank, self.tp_rank)
676 + )
677 + self.engine_id: EngineId = engine_id
678 +
679 + self.world_size = get_tensor_model_parallel_world_size()
680 + self.tp_group = get_tp_group()
681 +
682 + # KV Caches and moriio tracking data.
683 + self.kv_caches: dict[str, torch.Tensor] = {}
684 +
685 + # Map of engine_id -> kv_caches_base_addr. For TP case, each local
686 + # rank will still only pull from a single remote TP worker.
687 + self.kv_caches_base_addr: dict[EngineId, list[int]] = {}
688 +
689 + # Number of MoRIIO regions. Currently one region per cache
690 + # (so 1 per layer for MLA, otherwise 2 per layer)
691 + self.num_regions = 0
692 + self.num_layers = 0
693 +
694 + # Map of engine_id -> num_blocks. All ranks in the same deployment will
695 + # have the same number of blocks.
696 + self.dst_num_blocks: dict[EngineId, int] = {}
697 + # In progress transfers.
698 + self._recving_transfers: defaultdict[ReqId, list] = defaultdict(list)
699 + self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str]] = {}
700 +
701 + # Track the expiration time of requests that are waiting to be sent.
702 + self._reqs_to_send: dict[ReqId, float] = {}
703 +
704 + # Background thread for handling new handshake requests.
705 + self._moriio_handshake_listener_t: threading.Thread | None = None
706 + # Background thread for initializing new MoRIIO handshakes.
707 + self._handshake_initiation_executor = ThreadPoolExecutor(
708 + # MoRIIO is not guaranteed to be thread-safe, limit 1 worker.
709 + max_workers=1,
710 + thread_name_prefix="vllm-moriio-handshake-initiator",
711 + )
712 + self._ready_requests = queue.Queue[tuple[ReqId, ReqMeta]]()
713 + self._handshake_futures: dict[EngineId, Future[set[str]]] = {}
714 + # Protects _handshake_futures and _remote_agents.
715 + self._handshake_lock = threading.RLock()
716 +
717 + self.block_size = vllm_config.cache_config.block_size
718 + self.model_config = vllm_config.model_config
719 + self.cache_config = vllm_config.cache_config
720 +
721 + self.block_window_per_layer: list[int | None] = []
722 + self.use_mla = self.model_config.use_mla
723 + self.built_session = False
724 + self.built_write_session: defaultdict[str, list] = defaultdict(list)
725 + backend = get_attn_backend(
726 + self.model_config.get_head_size(),
727 + self.model_config.dtype,
728 + self.cache_config.cache_dtype,
729 + self.block_size,
730 + use_mla=self.use_mla,
731 + )
732 +
733 + # TODO: consider the integration of flashinfer or other backends.
734 + self.backend_name = backend.get_name()
735 + logger.debug("Detected attention backend %s", self.backend_name)
736 +
737 + def schedule_write_blocks(
738 + self,
739 + request_id: str,
740 + dst_engine_id: str,
741 + local_block_ids: list[int],
742 + remote_block_ids: list[int] | None,
743 + layer_name: str,
744 + kv_layer: torch.Tensor,
745 + remote_notify_port: int,
746 + remote_ip: str,
747 + ) -> None:
748 + """Schedule a block write operation.
749 +
750 + Args:
751 + request_id: Unique identifier for the request
752 + dst_engine_id: Destination engine ID
753 + local_block_ids: Local block IDs to transfer
754 + remote_block_ids: Hint for remote block IDs
755 + layer_name: Name of the layer
756 + kv_layer: KV cache tensor
757 + remote_notify_port: Port for completion notification
758 + remote_ip: IP address of remote node
759 + """
760 +
761 + # synchronization to prevent dirty reads between
762 + # transfer and attention operations
763 + # we can consider removing this synchronization after ibgda is enabled.
764 + # when mori-io supports ibgda functionality
765 +
766 + stream = torch.cuda.current_stream()
767 + event = torch.cuda.Event()
768 + event.record(stream)
769 +
770 + task = WriteTask(
771 + request_id=request_id,
772 + dst_engine_id=dst_engine_id,
773 + local_block_ids=local_block_ids,
774 + remote_block_ids_hint=remote_block_ids,
775 + layer_name=layer_name,
776 + event=event,
777 + remote_notify_port=remote_notify_port,
778 + remote_ip=remote_ip,
779 + )
780 + self._writer.schedule_write(task)
781 +
782 + def _get_built_session(self, remote_engine_id):
783 + if remote_engine_id not in self.built_write_session:
784 + cur_remote_engine_sessions = []
785 + for ln, local_meta in self.layer_name_to_local_kv_cache_metadata.items():
786 + unpacked_local_memory_meta = (
787 + self.moriio_wrapper.get_unpack_memory_metadata(local_meta[0])
788 + )
789 + unpacked_remote_memory_meta = (
790 + self.moriio_wrapper.get_unpack_memory_metadata(
791 + self.layer_name_to_remote_kv_cache_metadata[remote_engine_id][
792 + ln
793 + ][0]
794 + )
795 + )
796 + cur_remote_engine_sessions.append(
797 + self.moriio_wrapper.build_session(
798 + unpacked_local_memory_meta, unpacked_remote_memory_meta
799 + )
800 + )
801 + self.built_write_session[remote_engine_id] = cur_remote_engine_sessions
802 + return self.built_write_session[remote_engine_id], self.remote_moriio_metadata[
803 + remote_engine_id
804 + ]
805 +
806 + def _ping(self, zmq_context):
807 + http_request_address = f"http://{self.request_address}/v1/completions"
808 + role = "P" if self.is_producer else "D"
809 +
810 + retry_count = 0
811 + index = 1
812 + with zmq_context.socket(zmq.DEALER) as sock:
813 + sock.connect(f"tcp://{self.proxy_ip}:{self.proxy_ping_port}")
814 +
815 + while True:
816 + try:
817 + data = {
818 + "type": "register",
819 + "role": role,
820 + "index": str(index),
821 + "request_address": http_request_address,
822 + "handshake_port": self.handshake_port,
823 + "notify_port": self.notify_port,
824 + "dp_size": self.moriio_config.dp_size,
825 + "tp_size": self.moriio_config.tp_size,
826 + "transfer_mode": self.mode.name,
827 + }
828 +
829 + sock.send(msgpack.dumps(data))
830 + # logger.debug(f"Successfully sent ping message #{index}")
831 + retry_count = 0
832 +
833 + except ConnectionRefusedError:
834 + logger.info(
835 + "Connection refused: %s:%s -> %s:%s",
836 + self.local_ip,
837 + self.local_ping_port,
838 + self.proxy_ip,
839 + self.proxy_ping_port,
840 + )
841 + retry_count += 1
842 +
843 + except OSError as e:
844 + logger.info("OS error when sending ping: %s", e)
845 + retry_count += 1
846 +
847 + except Exception as e:
848 + logger.info("Unexpected error when sending ping: %s", e)
849 + retry_count += 1
850 + if retry_count >= MoRIIOConstants.MAX_PING_RETRIES:
851 + logger.error(
852 + "Max retries (%s) exceeded. Stopping ping loop.",
853 + MoRIIOConstants.MAX_PING_RETRIES,
854 + )
855 + raise RuntimeError(
856 + f"Ping failed after {retry_count} retries"
857 + ) from e
858 +
859 + finally:
860 + time.sleep(MoRIIOConstants.PING_INTERVAL)
861 + index += 1
862 +
863 + def shutdown(self):
864 + if hasattr(self, "moriio_wrapper") and self.moriio_wrapper:
865 + self.moriio_wrapper.shutdown()
866 +
867 + if hasattr(self, "_handshake_initiation_executor"):
868 + self._handshake_initiation_executor.shutdown(wait=False)
869 +
870 + if (
871 + hasattr(self, "_moriio_handshake_listener_t")
872 + and self._moriio_handshake_listener_t
873 + ):
874 + self._moriio_handshake_listener_t.join(timeout=0)
875 +
876 + if hasattr(self, "zmq_context") and self.zmq_context:
877 + self.zmq_context.destroy(linger=0)
878 + self.zmq_context = None
879 +
880 + def __del__(self):
881 + self.shutdown()
882 +
883 + @staticmethod
884 + def _moriio_handshake_listener(
885 + metadata: MoRIIOAgentMetadata,
886 + ready_event: threading.Event,
887 + base_port: int,
888 + tp_rank: int,
889 + dp_rank: int,
890 + layer_name_to_local_kv_cache_metadata: dict,
891 + ):
892 + """Background thread for getting new MoRIIO handshakes."""
893 +
894 + encoder = msgspec.msgpack.Encoder()
895 + encoded_data = encoder.encode(metadata)
896 + size_in_bytes = len(encoded_data)
897 + logger.debug(
898 + "Size of encoded MoRIIOAgentMetadata: %s bytes", str(size_in_bytes)
899 + )
900 +
901 + # Listen for new requests for metadata.
902 + host = "*"
903 +
904 + path = make_zmq_path("tcp", host, base_port)
905 + logger.debug("mori handshake starting listening on path: %s", path)
906 +
907 + with zmq_ctx(zmq.ROUTER, path) as sock:
908 + ready_event.set()
909 + while True:
910 + identity, msg = sock.recv_multipart()
911 + if (
912 + msg != MoRIIOConstants.GET_META_MSG
913 + and msg != MoRIIOConstants.POP_DONE_RECV
914 + ):
915 + logger.error("Connection listener got unexpected message")
916 + raise HandshakeError("handshake failed, unexpected msg type")
917 + elif msg == MoRIIOConstants.GET_META_MSG:
918 + sock.send_multipart(
919 + (identity, b"", encoded_data)
920 + ) # send local mori io engine meta data
921 + logger.debug("MoRIIO handshake listener sent metadata")
922 + # now we send tensor meta data for each block
923 + buf = msgpack.dumps(layer_name_to_local_kv_cache_metadata)
924 + sock.send_multipart((identity, b"", buf))
925 + elif msg == MoRIIOConstants.POP_DONE_RECV:
926 + _, req_id = sock.recv_multipart()
927 + logger.debug(
928 + "MoRIIO handshake listener received done recv for req",
929 + req_id.decode(),
930 + )
931 +
932 + def _moriio_handshake(
933 + self,
934 + host: str,
935 + port: int,
936 + remote_tp_size: int,
937 + expected_engine_id: str,
938 + remote_dp_rank: int = 0,
939 + ) -> set[str]:
940 + """Do a MoRIIO handshake with a remote instance."""
941 +
942 + start_time = time.perf_counter()
943 +
944 + # NOTE(rob): we need each rank to have a unique port. This is
945 + # a hack to keep us moving. We will switch when moving to etcd
946 + # or where we have a single ZMQ socket in the scheduler.
947 +
948 + port_offset = get_port_offset(remote_dp_rank, self.tp_rank)
949 + path = make_zmq_path("tcp", host, port + port_offset)
950 + logger.debug("handshake Querying metadata on path: %s", path)
951 +
952 + # Send query for the request.
953 + with zmq_ctx(zmq.DEALER, path) as sock:
954 + logger.debug("prepare send msg INSTAZNCE: %s", path)
955 + sock.send(MoRIIOConstants.GET_META_MSG)
956 + received_frame = sock.recv_multipart()
957 + if len(received_frame) != 2 or received_frame[0] != b"":
958 + raise HandshakeError(f"Unexpected frame! {received_frame = }")
959 +
960 + metadata_bytes = received_frame[1]
961 + decoder = msgspec.msgpack.Decoder(MoRIIOAgentMetadata)
962 + metadata = decoder.decode(metadata_bytes)
963 + got_metadata_time = time.perf_counter()
964 + logger.info(
965 + "MoRIIO handshake: get metadata took: %s",
966 + got_metadata_time - start_time,
967 + )
968 +
969 + self.moriio_wrapper.remote_engine_ip = host
970 + remote_agent_name = self.moriio_wrapper.register_remote_engine(
971 + metadata.agent_metadata
972 + )
973 +
974 + logger.debug(
975 + "MoRIIO handshake: registered"
976 + "remote agent %s for engine ID %s, path = %s",
977 + remote_agent_name,
978 + expected_engine_id,
979 + path,
980 + )
981 +
982 + if len(self.local_kv_cache_metadata) > 0:
983 + logger.warning(
984 + "len(self.local_kv_cache_metadata) = %s,"
985 + "maybe you didnt clear this buffer correctly",
986 + len(self.local_kv_cache_metadata),
987 + )
988 + self.local_kv_cache_metadata = []
989 + if len(self.remote_kv_cache_metadata) > 0:
990 + logger.warning(
991 + "len(self.remote_kv_cache_metadata) = %s,"
992 + "maybe you didnt clear this buffer correctly",
993 + len(self.remote_kv_cache_metadata),
994 + )
995 + self.remote_kv_cache_metadata = []
996 +
997 + received_frame = sock.recv_multipart()
998 + if len(received_frame) != 2 or received_frame[0] != b"":
999 + raise HandshakeError(f"unexpected frame! {received_frame = }")
1000 + buf = received_frame[1]
1001 + self.layer_name_to_remote_kv_cache_metadata[expected_engine_id] = (
1002 + msgpack.loads(buf)
1003 + )
1004 + self.remote_moriio_metadata[expected_engine_id] = metadata
1005 + setup_agent_time = time.perf_counter()
1006 + logger.debug(
1007 + "MoRIIO handshake: add agent took: %s",
1008 + setup_agent_time - got_metadata_time,
1009 + )
1010 +
1011 + return {remote_agent_name}
1012 +
1013 + def _background_moriio_handshake(
1014 + self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta
1015 + ):
1016 + # Do MoRIIO handshake in background and add to _ready_requests when done.
1017 + fut = None
1018 + if remote_engine_id is not None:
1019 + fut = self._handshake_futures.get(remote_engine_id)
1020 + if fut is None:
1021 + host = meta.remote_host
1022 + port = int(meta.remote_handshake_port)
1023 + tp_size = int(meta.tp_size)
1024 + remote_dp_size = int(meta.remote_dp_size)
1025 +
1026 + def request_ready(_f: Future[Any], entry=(req_id, meta)):
1027 + logger.info("MoRIIO handshake done for request %s", req_id)
1028 + self._ready_requests.put(entry)
1029 + self.load_ready_flag[remote_engine_id] = True
1030 + self.write_ready_flags[remote_engine_id] = True
1031 +
1032 + fut_list = []
1033 +
1034 + # In dp(prefill)<->dp(decode) communication, we require an all-to-all handshake.
1035 +
1036 + for cur_dp_rank in range(remote_dp_size):
1037 + dp_engine_id = self.get_engine_name_with_dp(remote_engine_id, cur_dp_rank)
1038 + future = self._handshake_initiation_executor.submit(
1039 + self._moriio_handshake, host, port, tp_size, dp_engine_id, cur_dp_rank
1040 + )
1041 + fut_list.append(future)
1042 +
1043 + def done_callback(f: Future[set[str]], eid=dp_engine_id):
1044 + with self._handshake_lock:
1045 + self._handshake_futures.pop(eid, None)
1046 + try:
1047 + self._remote_agents[eid] = f.result()
1048 + except Exception:
1049 + logger.exception("Handshake with %s failed", eid)
1050 +
1051 + future.add_done_callback(done_callback)
1052 + self._handshake_futures[dp_engine_id] = future
1053 +
1054 + # fut = fut_list
1055 + def wait_all_dp():
1056 + for future in fut_list:
1057 + future.result()
1058 + return True
1059 +
1060 + all_done_future = self._handshake_initiation_executor.submit(wait_all_dp)
1061 + all_done_future.add_done_callback(request_ready)
1062 +
1063 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
1064 + """Register the KV Cache data in moriio."""
1065 +
1066 + _, first_kv_cache = next(iter(kv_caches.items()))
1067 + kv_elem_size = first_kv_cache.element_size()
1068 +
1069 + use_mla = len(first_kv_cache.shape) == 3
1070 + assert use_mla == self.use_mla
1071 +
1072 + if use_mla:
1073 + # MLA case.
1074 + self.num_blocks = first_kv_cache.shape[0]
1075 + block_rank = 2 # [block_size, latent_dim]
1076 + block_shape = first_kv_cache.shape[-block_rank:]
1077 + block_size, kv_latent_dim = block_shape
1078 + self.slot_size_bytes = kv_elem_size * kv_latent_dim
1079 + else:
1080 + # [2 (k and v), num_blocks, ...]
1081 + self.num_blocks = first_kv_cache.shape[1]
1082 + block_rank = 3 # [block_size, kv_heads, head_dim]
1083 + block_shape = first_kv_cache.shape[-block_rank:]
1084 + block_size, n_kv_heads, head_dim = block_shape[-3:]
1085 + # head size in bytes.
1086 + self.slot_size_bytes = (
1087 + kv_elem_size * n_kv_heads * head_dim
1088 + ) # 1 token 1 layer size , slot size
1089 + assert block_size == self.block_size
1090 + # TODO(tms): self.block_len needs to be per-layer for sliding window,
1091 + # hybrid attn, etc
1092 + # block size in bytes
1093 + self.block_len = kv_elem_size * math.prod(block_shape)
1094 + self.kv_cache_shape = first_kv_cache.shape
1095 + self.block_shape = block_shape
1096 + self.kv_element_size = kv_elem_size
1097 +
1098 + self.dst_num_blocks[self.engine_id] = self.num_blocks
1099 + self.kv_caches = kv_caches # layer name to kv cache
1100 + kv_caches_base_addr = []
1101 + caches_data = []
1102 +
1103 + for cache_or_caches in kv_caches.values():
1104 + cache_list = [cache_or_caches] if use_mla else cache_or_caches
1105 + for cache in cache_list:
1106 + base_addr = cache.data_ptr()
1107 + region_len = self.num_blocks * self.block_len
1108 + caches_data.append((base_addr, region_len, cache.device.index, ""))
1109 + kv_caches_base_addr.append(base_addr)
1110 +
1111 + for layer_name, kv_cache in kv_caches.items():
1112 + if layer_name not in self.layer_name_to_local_kv_cache_metadata:
1113 + self.layer_name_to_local_kv_cache_metadata[layer_name] = []
1114 +
1115 + moriio_mem_metadata = self.moriio_wrapper.register_local_tensor(kv_cache)
1116 + self.layer_name_to_local_kv_cache_metadata[layer_name].append(
1117 + moriio_mem_metadata
1118 + )
1119 +
1120 + self.local_kv_cache_size.append(cache.nelement() * cache.element_size())
1121 +
1122 + self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr
1123 + self.num_regions = len(caches_data)
1124 + self.num_layers = len(self.kv_caches.keys())
1125 +
1126 + # Optimization for models with local attention (Llama 4)
1127 + if self.vllm_config.model_config.hf_config.model_type == "llama4":
1128 + from transformers import Llama4TextConfig
1129 +
1130 + assert isinstance(
1131 + self.vllm_config.model_config.hf_text_config, Llama4TextConfig
1132 + )
1133 + llama4_config = self.vllm_config.model_config.hf_text_config
1134 + no_rope_layers = llama4_config.no_rope_layers
1135 + chunk_size = llama4_config.attention_chunk_size
1136 + chunk_block_size = math.ceil(chunk_size / self.block_size)
1137 + for layer_idx in range(self.num_layers):
1138 + # no_rope_layers[layer_idx] == 0 means NoPE (global)
1139 + # Any other value means RoPE (local chunked)
1140 + is_local_attention = no_rope_layers[layer_idx] != 0
1141 + block_window = chunk_block_size if is_local_attention else None
1142 + self.block_window_per_layer.append(block_window)
1143 + logger.debug(
1144 + "Llama 4 block window per layer mapping: %s",
1145 + self.block_window_per_layer,
1146 + )
1147 + assert len(self.block_window_per_layer) == self.num_layers
1148 +
1149 + metadata = MoRIIOAgentMetadata(
1150 + engine_id=self.engine_id,
1151 + agent_metadata=self.moriio_wrapper.get_agent_metadata(),
1152 + kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id],
1153 + num_blocks=self.num_blocks,
1154 + block_len=self.block_len,
1155 + attn_backend_name=self.backend_name,
1156 + )
1157 + ready_event = threading.Event()
1158 + self._moriio_handshake_listener_t = threading.Thread(
1159 + target=self._moriio_handshake_listener,
1160 + args=(
1161 + metadata,
1162 + ready_event,
1163 + self.side_channel_port,
1164 + self.tp_rank,
1165 + self.dp_rank,
1166 + self.layer_name_to_local_kv_cache_metadata,
1167 + ),
1168 + daemon=True,
1169 + name="moriio_handshake_listener",
1170 + )
1171 + self._moriio_handshake_listener_t.start()
1172 + ready_event.wait() # Wait for listener ZMQ socket to be ready.
1173 + self.moriio_wrapper.async_wait_reqid()
1174 +
1175 + def get_finished(self) -> tuple[set[str], set[str]]:
1176 + """
1177 + Get requests that are done sending or recving on this specific worker.
1178 + The scheduler process (via the MultiprocExecutor) will use this output
1179 + to track which workers are done.
1180 + """
1181 +
1182 + done_sending, done_recving = set(), set()
1183 +
1184 + if self.is_producer:
1185 + done_sending = self.moriio_wrapper.pop_finished_req_ids()
1186 +
1187 + else:
1188 + if self.mode == MoRIIOMode.WRITE:
1189 + done_recving = self.moriio_wrapper.pop_finished_write_req_ids()
1190 + else:
1191 + done_recving = self._pop_done_transfers()
1192 +
1193 + return done_sending, done_recving
1194 +
1195 + def _pop_done_transfers(self) -> set[str]:
1196 + done_req_ids: set[str] = set()
1197 + with self.moriio_wrapper.lock:
1198 + to_remove = []
1199 + for req_id, status_list in self._recving_transfers.items():
1200 + if status_list[-1].Succeeded():
1201 + done_req_ids.add(req_id)
1202 +
1203 + self.moriio_wrapper.send_notify(
1204 + req_id,
1205 + self._recving_transfers_callback_addr[req_id][0],
1206 + self._recving_transfers_callback_addr[req_id][1],
1207 + )
1208 + to_remove.append(req_id)
1209 + for req_id in to_remove:
1210 + del self._recving_transfers[req_id]
1211 + del self._recving_transfers_callback_addr[req_id]
1212 +
1213 + return done_req_ids
1214 +
1215 + def save_kv_layer(
1216 + self,
1217 + metadata: MoRIIOConnectorMetadata,
1218 + layer_name: str,
1219 + kv_layer: torch.Tensor,
1220 + attn_metadata: "AttentionMetadata",
1221 + **kwargs,
1222 + ):
1223 + if not self.is_producer:
1224 + return
1225 + if self.mode == MoRIIOMode.READ:
1226 + return
1227 + remote_engine_id = None
1228 +
1229 + for req_id, meta in metadata.reqs_to_save.items():
1230 + # we only need to check if dp0 in rank
1231 + remote_engine_id = (
1232 + str(meta.remote_host) + ":" + str(meta.remote_handshake_port)
1233 + )
1234 +
1235 + meta.remote_engine_id = remote_engine_id
1236 +
1237 + dp0_remote_engine_id = self.get_engine_name_with_dp(remote_engine_id, 0)
1238 + if dp0_remote_engine_id not in self._remote_agents:
1239 + # Initiate handshake with remote engine to exchange metadata.
1240 + with self._handshake_lock:
1241 + if remote_engine_id not in self._remote_agents:
1242 + self._background_moriio_handshake(
1243 + req_id, remote_engine_id, meta
1244 + )
1245 +
1246 + continue
1247 + self._write_blocks_for_req(req_id, meta, layer_name, kv_layer)
1248 +
1249 + while True:
1250 + if (
1251 + self._ready_requests.empty()
1252 + and remote_engine_id not in self.write_ready_flags
1253 + ):
1254 + continue
1255 + elif not self._ready_requests.empty() and (
1256 + remote_engine_id in self.write_ready_flags
1257 + ):
1258 + self._write_blocks_for_req(
1259 + *self._ready_requests.get_nowait(), layer_name, kv_layer
1260 + )
1261 + break
1262 + else:
1263 + break
1264 +
1265 + def get_engine_name_with_dp(self, engine_name, dp_rank):
1266 + return f"{engine_name}_dp{dp_rank}"
1267 +
1268 + def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
1269 + """
1270 + Start loading by triggering non-blocking moriio_xfer.
1271 + We check for these trnxs to complete in each step().
1272 + """
1273 + if self.is_producer:
1274 + self.moriio_wrapper.async_wait_reqid()
1275 + return
1276 + if self.mode == MoRIIOMode.WRITE:
1277 + return
1278 +
1279 + wait_handshake_readd_req = False
1280 + remote_engine_id = None
1281 +
1282 + for req_id, meta in metadata.reqs_to_recv.items():
1283 + remote_engine_id = (
1284 + str(meta.remote_host) + ":" + str(meta.remote_handshake_port)
1285 + )
1286 + meta.remote_engine_id = remote_engine_id
1287 + dp0_remote_engine_id = self.get_engine_name_with_dp(remote_engine_id, 0)
1288 + if dp0_remote_engine_id not in self._remote_agents:
1289 + # Initiate handshake with remote engine to exchange metadata.
1290 + with self._handshake_lock:
1291 + if remote_engine_id not in self._remote_agents:
1292 + self._background_moriio_handshake(
1293 + req_id, remote_engine_id, meta
1294 + )
1295 + wait_handshake_readd_req = True
1296 +
1297 + continue
1298 +
1299 + # Handshake already completed, start async read xfer.
1300 + self._read_blocks_for_req(req_id, meta)
1301 + # Start transfers for requests whose handshakes have now finished.
1302 +
1303 + while True:
1304 + if (
1305 + self._ready_requests.empty()
1306 + and remote_engine_id not in self.load_ready_flag
1307 + and wait_handshake_readd_req
1308 + ):
1309 + continue
1310 + elif (
1311 + not self._ready_requests.empty()
1312 + and remote_engine_id in self.load_ready_flag
1313 + ):
1314 + self._read_blocks_for_req(*self._ready_requests.get_nowait())
1315 + break
1316 + else:
1317 + break
1318 +
1319 + self._reqs_to_send.update(metadata.reqs_to_send)
1320 +
1321 + def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
1322 + logger.debug(
1323 + "Remote agent %s available, calling _read_blocks for req %s",
1324 + meta.remote_engine_id,
1325 + req_id,
1326 + )
1327 + self._read_blocks(
1328 + request_id=req_id,
1329 + dst_engine_id=meta.remote_engine_id,
1330 + local_block_ids=meta.local_block_ids,
1331 + remote_block_ids=meta.remote_block_ids,
1332 + remote_host=meta.remote_host,
1333 + remote_notify_port=meta.remote_notify_port,
1334 + )
1335 +
1336 + def _write_blocks_for_req(self, req_id: str, meta: ReqMeta, layer_name, kv_layer):
1337 + self.schedule_write_blocks(
1338 + request_id=req_id,
1339 + dst_engine_id=meta.remote_engine_id,
1340 + local_block_ids=meta.local_block_ids,
1341 + remote_block_ids=meta.remote_block_ids,
1342 + layer_name=layer_name,
1343 + kv_layer=kv_layer,
1344 + remote_notify_port=meta.remote_notify_port,
1345 + remote_ip=meta.remote_host,
1346 + )
1347 +
1348 + def _is_last_layer(self, layer_name):
1349 + return layer_name == list(self.kv_caches.keys())[-1]
1350 +
1351 + def merge_contiguous_blocks(
1352 + self,
1353 + offsets_local: list[int],
1354 + offsets_remote: list[int],
1355 + sizes: list[int],
1356 + assume_sorted: bool = False,
1357 + ) -> tuple[list[int], list[int], list[int]]:
1358 + n = len(offsets_local)
1359 + if n == 0:
1360 + return [], [], []
1361 + if not (n == len(offsets_remote) == len(sizes)):
1362 + raise ValueError("Input list lengths mismatch")
1363 + local_arr = np.fromiter(offsets_local, dtype=np.int64, count=n)
1364 + remote_arr = np.fromiter(offsets_remote, dtype=np.int64, count=n)
1365 + sizes_arr = np.fromiter(sizes, dtype=np.int64, count=n)
1366 +
1367 + if assume_sorted:
1368 + local_sorted = local_arr
1369 + remote_sorted = remote_arr
1370 + sizes_sorted = sizes_arr
1371 + else:
1372 + if np.all(local_arr[:-1] <= local_arr[1:]):
1373 + local_sorted = local_arr
1374 + remote_sorted = remote_arr
1375 + sizes_sorted = sizes_arr
1376 + else:
1377 + sort_idx = np.argsort(local_arr, kind="stable")
1378 + local_sorted = local_arr[sort_idx]
1379 + remote_sorted = remote_arr[sort_idx]
1380 + sizes_sorted = sizes_arr[sort_idx]
1381 +
1382 + if n == 1:
1383 + return (
1384 + [int(local_sorted[0])],
1385 + [int(remote_sorted[0])],
1386 + [int(sizes_sorted[0])],
1387 + )
1388 +
1389 + diff_local = local_sorted[1:] - local_sorted[:-1]
1390 + diff_remote = remote_sorted[1:] - remote_sorted[:-1]
1391 + prev_size = sizes_sorted[:-1]
1392 +
1393 + contiguous = (diff_local == prev_size) & (diff_remote == prev_size)
1394 +
1395 + if not contiguous.any():
1396 + return local_sorted.tolist(), remote_sorted.tolist(), sizes_sorted.tolist()
1397 +
1398 + if contiguous.all():
1399 + total_size = int(sizes_sorted.sum())
1400 + return [int(local_sorted[0])], [int(remote_sorted[0])], [total_size]
1401 +
1402 + break_positions = np.flatnonzero(~contiguous) + 1
1403 + segment_starts = np.concatenate(([0], break_positions))
1404 + segment_ends = np.concatenate((break_positions, [n]))
1405 +
1406 + seg_count = len(segment_starts)
1407 + merged_local = [0] * seg_count
1408 + merged_remote = [0] * seg_count
1409 + merged_sizes = [0] * seg_count
1410 +
1411 + for si in range(seg_count):
1412 + s = segment_starts[si]
1413 + e = segment_ends[si]
1414 + merged_local[si] = int(local_sorted[s])
1415 + merged_remote[si] = int(remote_sorted[s])
1416 +
1417 + merged_sizes[si] = int(
1418 + local_sorted[e - 1] + sizes_sorted[e - 1] - local_sorted[s]
1419 + )
1420 +
1421 + return merged_local, merged_remote, merged_sizes
1422 +
1423 + def _compute_block_transfer_offsets(
1424 + self,
1425 + layer_name: str,
1426 + local_block_ids: list[int],
1427 + remote_block_ids: list[int],
1428 + remote_moriio_meta: MoRIIOAgentMetadata,
1429 + ) -> tuple[list[int], list[int], list[int]]:
1430 + """Compute transfer offsets for block data.
1431 +
1432 + Args:
1433 + layer_name: Name of the layer to transfer
1434 + local_block_ids: IDs of local blocks
1435 + remote_block_ids: IDs of remote blocks
1436 + remote_moriio_meta: Metadata of the remote MoRIIO agent
1437 + Returns:
1438 + Tuple of (local_offsets, remote_offsets, transfer_sizes)
1439 + """
1440 + assert self.kv_cache_shape is not None, "KV caches shape not initialized"
1441 + is_mla = len(self.kv_cache_shape) == 3
1442 + stride = self.kv_caches[layer_name].stride()
1443 + sz = self.kv_caches[layer_name].element_size()
1444 + if is_mla:
1445 + blknum, blksize, hs = self.kv_cache_shape
1446 + hn = 1
1447 + block_stride = stride[0]
1448 + else:
1449 + _, blknum, blksize, hn, hs = self.kv_cache_shape
1450 + local_ktov_stride = stride[0]
1451 + block_stride = stride[1]
1452 + remote_ktov_stride = block_stride * remote_moriio_meta.num_blocks
1453 +
1454 + transfer_size_byte = blksize * hn * hs * sz
1455 + per_block = 1 if is_mla else 2
1456 + total = len(local_block_ids) * per_block
1457 + offset_local = [0] * total
1458 + offset_remote = [0] * total
1459 + sizes = [transfer_size_byte] * total
1460 +
1461 + w = 0
1462 + for i, lb in enumerate(local_block_ids):
1463 + rb = remote_block_ids[i]
1464 + # K
1465 + offset_local[w] = sz * (lb * block_stride)
1466 + offset_remote[w] = sz * (rb * block_stride)
1467 + w += 1
1468 + if not is_mla:
1469 + # V
1470 + # Handle num_block variations originating from PD (different kv strides)
1471 + # TODO: address block_sz differences in heterogeneous TP scenarios
1472 + # In MLA, we don't need to consider these two cases.
1473 + offset_local[w] = sz * (1 * local_ktov_stride + lb * block_stride)
1474 + offset_remote[w] = sz * (1 * remote_ktov_stride + rb * block_stride)
1475 + w += 1
1476 +
1477 + merged_l, merged_r, merged_s = self.merge_contiguous_blocks(
1478 + offset_local, offset_remote, sizes, assume_sorted=False
1479 + )
1480 + return merged_l, merged_r, merged_s
1481 +
1482 + def _read_blocks(
1483 + self,
1484 + local_block_ids: list[int],
1485 + remote_block_ids: list[int],
1486 + dst_engine_id: str,
1487 + request_id: str,
1488 + remote_host: str,
1489 + remote_notify_port: int,
1490 + ) -> None:
1491 + if self.mode == MoRIIOMode.WRITE:
1492 + return
1493 +
1494 + dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0)
1495 + sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id)
1496 +
1497 + first_layer = list(self.layer_name_to_local_kv_cache_metadata.keys())[0]
1498 + offs = self._compute_block_transfer_offsets(
1499 + first_layer, local_block_ids, remote_block_ids, remote_moriio_meta
1500 + )
1501 +
1502 + for layer_name in self.layer_name_to_local_kv_cache_metadata:
1503 + sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index(
1504 + layer_name
1505 + )
1506 + # TODO : apply multi-session batch-read when moriio support it
1507 + transfer_status = self.moriio_wrapper.read_remote_data(
1508 + offs[2], offs[0], offs[1], sessions[sess_idx]
1509 + )
1510 + with self.moriio_wrapper.lock:
1511 + self._recving_transfers[request_id].append(transfer_status)
1512 + self._recving_transfers_callback_addr[request_id] = (
1513 + remote_host,
1514 + str(remote_notify_port + self.tp_rank),
1515 + )
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py new
+609
@@ -0,0 +1,609 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import threading
4 +from typing import TYPE_CHECKING, Any, Optional
5 +from weakref import ref as weakref_ref
6 +
7 +import msgpack
8 +import torch
9 +import zmq
10 +
11 +from vllm import envs
12 +from vllm.logger import init_logger
13 +from vllm.utils.network_utils import (
14 + make_zmq_path,
15 + make_zmq_socket,
16 +)
17 +
18 +if TYPE_CHECKING:
19 + pass
20 +
21 +from queue import Empty, Queue
22 +
23 +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
24 + ROLE,
25 + HandshakeError,
26 + LayerTransferPlan,
27 + MoRIIOAgentMetadata,
28 + MoRIIOConstants,
29 + MoRIIOError,
30 + RemoteAllocInfo,
31 + TransferError,
32 + WriteTask,
33 + get_port_offset,
34 + get_role,
35 + zmq_ctx,
36 +)
37 +
38 +if TYPE_CHECKING:
39 + from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
40 + MoRIIOConnectorWorker,
41 + )
42 +
43 +logger = init_logger(__name__)
44 +try:
45 + from mori.io import (
46 + EngineDesc,
47 + IOEngine,
48 + MemoryDesc,
49 + PollCqMode,
50 + RdmaBackendConfig,
51 + )
52 +
53 + logger.info("MoRIIO is available")
54 +except ImportError:
55 + logger.error("MoRIIO is not available")
56 +
57 +
58 +"""Write task execution logic for MoRIIO connector."""
59 +
60 +
61 +class MoRIIOWriter:
62 + """Handles write operations for KV cache transfers.
63 + Implements distributed KV cache transfer using the MoRIIO library
64 + for RDMA-based communication between prefill and decode instances."""
65 +
66 + def __init__(self, worker: "MoRIIOConnectorWorker"):
67 + """Initialize the writer.
68 +
69 + Args:
70 + worker: Reference to the parent worker
71 + """
72 + self._worker_ref: weakref_ref[MoRIIOConnectorWorker] = weakref_ref(worker)
73 + self._write_task_q: Queue[WriteTask] = Queue()
74 + self._write_worker_started = False
75 + self._write_worker_lock = threading.Lock()
76 + self._deferred_tasks: list[WriteTask] = []
77 +
78 + @property
79 + def worker(self) -> "MoRIIOConnectorWorker":
80 + """Get the worker instance.
81 +
82 + Returns:
83 + The parent worker instance
84 +
85 + Raises:
86 + RuntimeError: If worker has been garbage collected
87 + """
88 + worker = self._worker_ref()
89 + if worker is None:
90 + raise RuntimeError("Parent worker has been garbage collected")
91 + return worker
92 +
93 + def ensure_worker_started(self) -> None:
94 + """Ensure the background write worker is running."""
95 + if self._write_worker_started:
96 + return
97 + self._write_worker_started = True
98 + with self._write_worker_lock:
99 + thread = threading.Thread(
100 + target=self._write_worker_loop, daemon=True, name="moriio-write-worker"
101 + )
102 + thread.start()
103 + logger.info("Started MoRIIO write worker thread")
104 +
105 + def schedule_write(self, task: WriteTask) -> None:
106 + """Schedule a write task.
107 +
108 + Args:
109 + task: The write task to schedule
110 + """
111 + self.ensure_worker_started()
112 + self._write_task_q.put(task)
113 +
114 + def _write_worker_loop(self) -> None:
115 + """Main loop for the write worker thread."""
116 +
117 + while True:
118 + # Process deferred tasks first
119 + self._process_deferred_tasks()
120 +
121 + # Get new task
122 + try:
123 + task = self._write_task_q.get(timeout=0.01)
124 + except Empty:
125 + continue
126 +
127 + # Check if remote blocks are ready
128 + if not self._is_remote_ready(task):
129 + # task.retry_count += 1
130 + self._deferred_tasks.append(task)
131 + # logger.debug(
132 + # "Deferred task for request %s (retry %d)",
133 + # task.request_id, task.retry_count
134 + # )
135 + continue
136 +
137 + # Execute the task
138 +
139 + self._execute_write_task(task)
140 +
141 + def _process_deferred_tasks(self) -> None:
142 + """Process tasks that were previously deferred."""
143 + if not self._deferred_tasks:
144 + return
145 +
146 + still_deferred: list[WriteTask] = []
147 + for task in self._deferred_tasks:
148 + if self._is_remote_ready(task):
149 + self._execute_write_task(task)
150 + else:
151 + still_deferred.append(task)
152 +
153 + self._deferred_tasks = still_deferred
154 +
155 + def _is_remote_ready(self, task: WriteTask) -> bool:
156 + """Check if remote blocks are allocated for this task.
157 +
158 + Args:
159 + task: The write task
160 +
161 + Returns:
162 + True if remote blocks are ready
163 + """
164 + return (
165 + task.request_id in self.worker.moriio_wrapper.done_remote_allocate_req_dict
166 + )
167 +
168 + def _get_remote_alloc_info(self, request_id: str) -> RemoteAllocInfo:
169 + """Get remote allocation info for a request.
170 +
171 + Args:
172 + request_id: The request ID
173 +
174 + Returns:
175 + Remote allocation information
176 +
177 + Raises:
178 + KeyError: If allocation info is missing
179 + """
180 + try:
181 + return self.worker.moriio_wrapper.done_remote_allocate_req_dict[request_id]
182 + except KeyError as e:
183 + raise KeyError(
184 + f"Remote allocation info missing for request {request_id}"
185 + ) from e
186 +
187 + def _execute_write_task(self, task: WriteTask) -> None:
188 + """Execute a single write task.
189 +
190 + Args:
191 + task: The write task to execute
192 +
193 + """
194 + # Get remote allocation info
195 + request_info = self._get_remote_alloc_info(task.request_id)
196 +
197 + if request_info.block_ids is None:
198 + logger.debug("Request %s remote block IDs not ready", task.request_id)
199 + return
200 +
201 + # Wait for CUDA event
202 + # The attention computation of the current layer cannot
203 + # overlap with the kv transfer task,
204 + # otherwise it will cause precision issues.
205 + # This event is used to synchronize the kv transfer and computation tasks.
206 + task.event.synchronize()
207 +
208 + # Update engine ID with DP rank
209 + task.dst_engine_id = self.worker.get_engine_name_with_dp(
210 + task.dst_engine_id, request_info.decode_dp_rank
211 + )
212 +
213 + # Get or create sessions
214 + sessions, remote_moriio_meta = self.worker._get_built_session(
215 + task.dst_engine_id
216 + )
217 +
218 + # Prepare transfer plan
219 + plan = self._prepare_transfer_plan(task, request_info, remote_moriio_meta)
220 +
221 + # Execute transfer
222 + self._do_layer_write(plan, sessions)
223 +
224 + # Finalize if all layers complete
225 + self._finalize_if_complete(task, request_info)
226 +
227 + def _prepare_transfer_plan(
228 + self,
229 + task: WriteTask,
230 + request_info: RemoteAllocInfo,
231 + remote_moriio_meta: MoRIIOAgentMetadata,
232 + ) -> LayerTransferPlan:
233 + """Prepare the transfer plan for a layer.
234 +
235 + Args:
236 + task: The write task
237 + request_info: Remote allocation information
238 +
239 + Returns:
240 + The transfer plan
241 + """
242 + # Compute offsets if not cached
243 + if request_info.transfer_offset is None:
244 + offsets = self.worker._compute_block_transfer_offsets(
245 + task.layer_name,
246 + task.local_block_ids,
247 + request_info.block_ids,
248 + remote_moriio_meta,
249 + )
250 + request_info.transfer_offset = offsets
251 +
252 + # Get session index
253 + layer_names = list(self.worker.layer_name_to_local_kv_cache_metadata.keys())
254 + sess_idx = layer_names.index(task.layer_name)
255 +
256 + local_off, remote_off, sizes = request_info.transfer_offset
257 +
258 + return LayerTransferPlan(
259 + request_id=task.request_id,
260 + layer_name=task.layer_name,
261 + sess_idx=sess_idx,
262 + transfer_local_offsets=local_off,
263 + transfer_remote_offsets=remote_off,
264 + transfer_sizes=sizes,
265 + use_batch=True,
266 + )
267 +
268 + def _do_layer_write(self, plan: LayerTransferPlan, sessions: list) -> None:
269 + """Perform the actual layer write.
270 +
271 + Args:
272 + plan: The transfer plan
273 + sessions: List of transfer sessions
274 + """
275 + if plan.use_batch:
276 + self.worker.moriio_wrapper.write_remote_data(
277 + plan.transfer_sizes,
278 + plan.transfer_local_offsets,
279 + plan.transfer_remote_offsets,
280 + sessions[plan.sess_idx],
281 + )
282 + else:
283 + for i in range(len(plan.transfer_local_offsets)):
284 + self.worker.moriio_wrapper.write_remote_data_single(
285 + plan.transfer_sizes[i],
286 + plan.transfer_local_offsets[i],
287 + plan.transfer_remote_offsets[i],
288 + plan.sess_idx,
289 + )
290 +
291 + def _finalize_if_complete(
292 + self, task: WriteTask, request_info: RemoteAllocInfo
293 + ) -> None:
294 + """Finalize transfer if all layers are complete.
295 +
296 + Args:
297 + task: The write task
298 + request_info: Remote allocation information
299 + """
300 + request_info.writes_done += 1
301 +
302 + if request_info.writes_done >= self.worker.num_layers:
303 + # Wait for transfer to complete
304 + self.worker.moriio_wrapper.waiting_for_transfer_complete()
305 +
306 + remote_port = task.remote_notify_port + get_port_offset(
307 + request_info.decode_dp_rank, self.worker.tp_rank
308 + )
309 + # Consider using RDMA immediate data in decode side
310 + # to eliminate the need for this notification.
311 + # Consider including the first gen token from prefill in the notification
312 +
313 + # Send completion notification
314 + self.worker.moriio_wrapper.send_notify(
315 + task.request_id, task.remote_ip, remote_port
316 + )
317 + # mark request as done, then we can free the blocks
318 + with self.worker.moriio_wrapper.lock:
319 + self.worker.moriio_wrapper.done_req_ids.append(task.request_id)
320 + del self.worker.moriio_wrapper.done_remote_allocate_req_dict[
321 + task.request_id
322 + ]
323 + logger.debug(
324 + "Completed transfer for request %s, notified port %d",
325 + task.request_id,
326 + remote_port,
327 + )
328 +
329 +
330 +class MoRIIOWrapper:
331 + """Wrapper for MoRIIO engine operations.
332 +
333 + Handles both producer and consumer roles for KV cache transfers.
334 +
335 + Args:
336 + moriio_engine: MoRIIO engine instance
337 + tp_rank: Tensor parallel rank
338 + dp_rank: Data parallel rank
339 + """
340 +
341 + def __init__(
342 + self,
343 + moriio_engine: Optional["IOEngine"] = None,
344 + tp_rank: int = 0,
345 + dp_rank: int = 0,
346 + ):
347 + self.tp_rank = tp_rank
348 + self.dp_rank = dp_rank
349 + self.moriio_engine = moriio_engine
350 + self.remote_memory_metadata = None
351 + self.local_memory_registered = False
352 + self.local_memory_metadata = None
353 + self.transfer_status: list[Any] = []
354 + self.remote_engine_ip: str | None = None
355 + self.notify_port: int | None = None
356 + self.lock = threading.Lock()
357 + self.done_req_ids: list[str] = []
358 + self.done_remote_allocate_req_dict: dict[str, RemoteAllocInfo] = {}
359 + self.done_write_cache_req_ids: list[str] = []
360 + self.notify_thread: threading.Thread | None = None
361 + self.sessions: list[IOEngine.Session] = []
362 + self.paths: dict[str, zmq.Socket] = {}
363 +
364 + def set_moriio_engine(self, moriio_engine):
365 + assert moriio_engine is not None, (
366 + "You Cannot pass None engine to MoRIIOWrapper!"
367 + )
368 + self.moriio_engine = moriio_engine
369 +
370 + def set_backend_type(self, backend_type):
371 + assert self.moriio_engine is not None, "MoRIIO engine must be set first"
372 + qp_per_transfer = envs.VLLM_MORIIO_QP_PER_TRANSFER
373 + post_batch_size = envs.VLLM_MORIIO_POST_BATCH_SIZE
374 + num_worker_threads = envs.VLLM_MORIIO_NUM_WORKERS
375 + poll_mode = PollCqMode.POLLING
376 + rdma_cfg = RdmaBackendConfig(
377 + qp_per_transfer,
378 + post_batch_size,
379 + num_worker_threads,
380 + poll_mode,
381 + )
382 + self.moriio_engine.create_backend(backend_type, rdma_cfg)
383 +
384 + def get_agent_metadata(self):
385 + assert self.moriio_engine is not None, "MoRIIO engine must be set first"
386 + engine_metadata = self.moriio_engine.get_engine_desc()
387 + engine_metadata_packed = engine_metadata.pack()
388 + return engine_metadata_packed
389 +
390 + def register_remote_engine(self, remote_packed_engine_metadata):
391 + assert self.moriio_engine is not None, "MoRIIO engine must be set first"
392 + consumer_engine_metadata = EngineDesc.unpack(remote_packed_engine_metadata)
393 + self.moriio_engine.register_remote_engine(consumer_engine_metadata)
394 + return consumer_engine_metadata.key
395 +
396 + def register_local_tensor(self, tensor: torch.Tensor):
397 + assert self.moriio_engine is not None, "MoRIIO engine must be set first"
398 + try:
399 + self.local_memory_metadata = self.moriio_engine.register_torch_tensor(
400 + tensor
401 + )
402 + assert self.local_memory_metadata is not None, (
403 + "register_torch_tensor returned None"
404 + )
405 + local_memory_metadata_packed = self.local_memory_metadata.pack()
406 + except Exception as e:
407 + raise MoRIIOError(f"Failed to register local memory: {e}") from e
408 + self.local_memory_registered = True
409 + return local_memory_metadata_packed
410 +
411 + def get_unpack_memory_metadata(self, packed_memory_metadata):
412 + return MemoryDesc.unpack(packed_memory_metadata)
413 +
414 + def build_session(self, local_memory_metadata, remote_memory_metadata):
415 + assert self.moriio_engine is not None, "MoRIIO engine must be set first"
416 + return self.moriio_engine.create_session(
417 + local_memory_metadata, remote_memory_metadata
418 + )
419 +
420 + def read_remote_data(
421 + self, transfer_size_byte, local_offset=0, remote_offset=0, session=None
422 + ):
423 + assert self.local_memory_registered, "You have not register local memory data!"
424 + assert self.moriio_engine is not None, "MoRIIO engine must be set first"
425 + transfer_status = session.batch_read(
426 + local_offset,
427 + remote_offset,
428 + transfer_size_byte,
429 + self.moriio_engine.allocate_transfer_uid(),
430 + )
431 +
432 + return transfer_status
433 +
434 + def write_remote_data(
435 + self, transfer_size_byte, local_offset=0, remote_offset=0, session=None
436 + ):
437 + assert self.local_memory_registered, "You have not register local memory data!"
438 + assert self.moriio_engine is not None, "MoRIIO engine must be set first"
439 + write_uid = self.moriio_engine.allocate_transfer_uid()
440 +
441 + transfer_status = session.batch_write(
442 + local_offset, remote_offset, transfer_size_byte, write_uid
443 + )
444 + with self.lock:
445 + self.transfer_status.append(transfer_status)
446 +
447 + def write_remote_data_single(
448 + self, transfer_size_byte, local_offset=0, remote_offset=0, sess_idx=0
449 + ):
450 + assert self.local_memory_registered, "You have not register local memory data!"
451 + assert self.moriio_engine is not None, "MoRIIO engine must be set first"
452 + transfer_status = self.sessions[sess_idx].write(
453 + local_offset,
454 + remote_offset,
455 + transfer_size_byte,
456 + self.moriio_engine.allocate_transfer_uid(),
457 + )
458 + with self.lock:
459 + self.transfer_status.append(transfer_status)
460 +
461 + def waiting_for_transfer_complete(self):
462 + if not self.transfer_status:
463 + return
464 +
465 + transfers_to_wait = []
466 + with self.lock:
467 + transfers_to_wait = self.transfer_status[:]
468 + self.transfer_status.clear()
469 +
470 + for status in transfers_to_wait:
471 + try:
472 + status.Wait()
473 + if not status.Succeeded():
474 + logger.error(
475 + "Transfer failed: %s, Code: %s", status.Message(), status.Code()
476 + )
477 + raise TransferError("MoRIIO transfer failed!")
478 + except Exception as e:
479 + logger.error("Transfer %s failed: %s", status, e)
480 + raise
481 +
482 + def async_wait_reqid(self):
483 + assert self.notify_port is not None, "Notify port cannot be None"
484 +
485 + if self.notify_thread is not None:
486 + return
487 +
488 + def _async_wait():
489 + host = "*"
490 + path = make_zmq_path("tcp", host, self.notify_port)
491 + logger.info("Node starting to listen notify from path = %s", path)
492 +
493 + with zmq_ctx(zmq.ROUTER, path) as sock:
494 + while True:
495 + try:
496 + identity, msg = sock.recv_multipart()
497 + self._handle_message(msg)
498 + except Exception as e:
499 + logger.error("Error processing message: %s", e)
500 + raise HandshakeError(f"Error processing message: {e}") from e
501 +
502 + self.notify_thread = threading.Thread(
503 + target=_async_wait, daemon=True, name="moriio-notify-listener"
504 + )
505 + self.notify_thread.start()
506 +
507 + def _handle_message(self, msg: bytes):
508 + """Handles incoming messages from remote nodes."""
509 + # Handles incoming remote messages:
510 + # Prefill Role:
511 + # [write] mode: receives block information (allocation)
512 + # [read] mode: receives block release messages from decode side
513 + # Decode Role:
514 + # [write] mode: receives KV cache write completion notifications
515 + handled = False
516 + try:
517 + data = msgpack.loads(msg)
518 + if isinstance(data, dict) and "req_id" in data:
519 + self._handle_structured_message(data)
520 +
521 + return
522 + except (msgpack.exceptions.ExtraData, msgpack.exceptions.UnpackException):
523 + logger.debug("Failed to decode msgpack message, will try as string")
524 + pass
525 +
526 + try:
527 + msg_str = msg.decode("UTF-8")
528 + if msg_str.startswith(MoRIIOConstants.COMPLETION_PREFIX):
529 + self._handle_completion_message(msg_str)
530 + handled = True
531 + except UnicodeDecodeError:
532 + logger.warning("Received non-UTF8 message: %s", msg_str)
533 + if not handled:
534 + raise MoRIIOError(f"Unhandled message format: {msg_str}")
535 +
536 + def _handle_structured_message(self, data: dict):
537 + assert get_role() == ROLE.PRODUCER, "Only prefill can get block messages"
538 + req_id = data["req_id"]
539 + block_notify_list = data.get("block_notify_list", [])
540 + decode_dp_rank = data.get("decode_rank", 0)
541 + assert len(block_notify_list) > 0, (
542 + "block_notify_list cannot be empty in remote allocate message"
543 + )
544 +
545 + with self.lock:
546 + self.done_remote_allocate_req_dict[req_id] = RemoteAllocInfo(
547 + block_ids=block_notify_list, decode_dp_rank=decode_dp_rank
548 + )
549 +
550 + def _handle_completion_message(self, msg: str):
551 + with self.lock:
552 + if get_role() == ROLE.PRODUCER:
553 + self.done_req_ids.append(msg)
554 + else:
555 + self.done_write_cache_req_ids.append(msg)
556 +
557 + def send_notify(self, req_ids, remote_ip, remote_port):
558 + if not remote_ip or not remote_port:
559 + logger.warning("Missing remote_ip or remote_port for notification")
560 + return
561 +
562 + path = make_zmq_path("tcp", remote_ip, remote_port)
563 +
564 + if path not in self.paths:
565 + ctx = zmq.Context.instance()
566 + sock = make_zmq_socket(
567 + ctx=ctx, path=path, socket_type=zmq.DEALER, bind=False
568 + )
569 + self.paths[path] = sock
570 +
571 + req_list = req_ids if isinstance(req_ids, list) else [req_ids]
572 +
573 + sock = self.paths[path]
574 + try:
575 + for req_id in req_list:
576 + if not isinstance(req_id, str):
577 + logger.warning(
578 + "Invalid req_id type: %s, expected str", type(req_id)
579 + )
580 + continue
581 + sock.send(req_id.encode("utf-8"))
582 + except Exception as e:
583 + logger.error("Failed to send notification to %s: %s", path, e)
584 + self.paths.pop(path, None)
585 + raise
586 +
587 + def pop_finished_req_ids(self):
588 + # producer invocation: get the set of completed requests at the decode
589 + with self.lock:
590 + done_send = set(self.done_req_ids)
591 + self.done_req_ids = []
592 + return done_send
593 +
594 + def pop_finished_write_req_ids(self):
595 + # Call the consumer in write mode to get the collection after write completion
596 + with self.lock:
597 + done_write_cache = set(self.done_write_cache_req_ids)
598 + self.done_write_cache_req_ids = []
599 + return done_write_cache
600 +
601 + def shutdown(self):
602 + logger.debug("Closing MoRIIOWrapper and cleaning up ZMQ sockets")
603 + for path, sock in self.paths.items():
604 + try:
605 + sock.close(linger=0)
606 + logger.debug("Closed ZMQ socket for path: %s", path)
607 + except Exception as e:
608 + logger.warning("Error closing ZMQ socket for path %s: %s", path, e)
609 + self.paths.clear()
vllm/envs.py
+18
@@ -204,6 +204,10 @@ if TYPE_CHECKING:
204 VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True
205 VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None
206 VLLM_NIXL_ABORT_REQUEST_TIMEOUT: int = 480
207 + VLLM_MORIIO_CONNECTOR_READ_MODE: bool = False
208 + VLLM_MORIIO_QP_PER_TRANSFER: int = 1
209 + VLLM_MORIIO_POST_BATCH_SIZE: int = -1
210 + VLLM_MORIIO_NUM_WORKERS: int = 1
211 VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT: int = 480
212 VLLM_USE_CUDNN_PREFILL: bool = False
213 VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL: bool = False
@@ -1383,6 +1387,20 @@ environment_variables: dict[str, Callable[[], Any]] = {
1387 "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": lambda: int(
1388 os.getenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", "480")
1389 ),
1390 + # Controls the read mode for the Mori-IO connector
1391 + "VLLM_MORIIO_CONNECTOR_READ_MODE": lambda: (
1392 + os.getenv("VLLM_MORIIO_CONNECTOR_READ_MODE", "False").lower() in ("true", "1")
1393 + ),
1394 + # Controls the QP (Queue Pair) per transfer configuration for the Mori-IO connector
1395 + "VLLM_MORIIO_QP_PER_TRANSFER": lambda: int(
1396 + os.getenv("VLLM_MORIIO_QP_PER_TRANSFER", "1")
1397 + ),
1398 + # Controls the post-processing batch size for the Mori-IO connector
1399 + "VLLM_MORIIO_POST_BATCH_SIZE": lambda: int(
1400 + os.getenv("VLLM_MORIIO_POST_BATCH_SIZE", "-1")
1401 + ),
1402 + # Controls the number of workers for Mori operations for the Mori-IO connector
1403 + "VLLM_MORIIO_NUM_WORKERS": lambda: int(os.getenv("VLLM_MORIIO_NUM_WORKERS", "1")),
1404 # Timeout (in seconds) for MooncakeConnector in PD disaggregated setup.
1405 "VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT": lambda: int(
1406 os.getenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "480")