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 2, 20:45 UTC
[Misc][Refactor][ROCm] Convert MoRI-related envvars to extra config args (#43303)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Simon Danielsson committed
May 26, 2026 at 12:33 UTC
681d7dd38b596df4439301f8f47e60177792ebf3
5 files changed
+77
-47
tests/v1/kv_connector/unit/test_moriio_connector.py
+4
-12
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import importlib.util
4
-import os
4
import subprocess
5
import uuid
6
from unittest.mock import MagicMock, patch
@@ -202,6 +201,7 @@ def create_vllm_config(
201
enable_chunked_prefill: bool = True,
202
enable_permute_local_kv: bool = False,
203
role="kv_consumer",
204
+ read_mode: bool = False,
205
) -> VllmConfig:
206
"""Initialize VllmConfig for testing."""
207
scheduler_config = SchedulerConfig(
@@ -228,6 +228,7 @@ def create_vllm_config(
228
kv_connector="MoRIIOConnector",
229
kv_role=role,
230
enable_permute_local_kv=enable_permute_local_kv,
231
+ kv_connector_extra_config={"read_mode": read_mode},
232
)
233
return VllmConfig(
234
scheduler_config=scheduler_config,
@@ -238,15 +239,6 @@ def create_vllm_config(
239
)
240
241
241
-@pytest.fixture
242
-def moriio_read_mode():
243
- """Force the connector into read mode via env for tests."""
244
- os.environ["VLLM_MORIIO_CONNECTOR_READ_MODE"] = "True"
245
- yield
246
- # Cleanup after test
247
- os.environ.pop("VLLM_MORIIO_CONNECTOR_READ_MODE", None)
248
-
249
-
242
def test_write_mode_saves_local_block_ids():
243
"""Write mode records local block ids in MoRIIOConnectorMetadata.reqs_to_save."""
244
@@ -358,11 +350,11 @@ def test_write_mode_with_chunked_prefill_saves_local_block_ids():
350
assert block_id == block.block_id, f"{block_id} != {block.block_id}"
351
352
361
-def test_read_mode_loads_remote_block_ids(moriio_read_mode):
353
+def test_read_mode_loads_remote_block_ids():
354
"""Read mode loads remote block ids into local cache mapping."""
355
356
# Setup Scheduler and Request
365
- vllm_config = create_vllm_config(role="kv_consumer")
357
+ vllm_config = create_vllm_config(role="kv_consumer", read_mode=True)
358
scheduler = create_scheduler(vllm_config)
359
360
# 2 Full Blocks and 1 Half Block.
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+47
-4
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import contextlib
4
+import os
5
import threading
6
import time
7
from collections.abc import Iterator
@@ -12,8 +13,7 @@ import regex as re
13
import torch
14
import zmq
15
15
-from vllm import envs
16
-from vllm.config import VllmConfig
16
+from vllm.config import KVTransferConfig, VllmConfig
17
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
18
KVConnectorMetadata,
19
)
@@ -162,8 +162,10 @@ class TransferError(MoRIIOError):
162
pass
163
164
165
-def get_moriio_mode() -> MoRIIOMode:
166
- read_mode = envs.VLLM_MORIIO_CONNECTOR_READ_MODE
165
+def get_moriio_mode(kv_transfer_config: KVTransferConfig) -> MoRIIOMode:
166
+ read_mode = str(
167
+ kv_transfer_config.kv_connector_extra_config.get("read_mode", "false")
168
+ ).lower().strip() in ("true", "1")
169
logger.debug("MoRIIO Connector read_mode: %s", read_mode)
170
if read_mode:
171
return MoRIIOMode.READ
@@ -175,6 +177,26 @@ def get_port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int:
177
return (dp_rank) * tp_size + tp_rank
178
179
180
+_DEPRECATED_ENV_VARS: dict[str, str] = {
181
+ "VLLM_MORIIO_CONNECTOR_READ_MODE": "read_mode",
182
+ "VLLM_MORIIO_QP_PER_TRANSFER": "qp_per_transfer",
183
+ "VLLM_MORIIO_POST_BATCH_SIZE": "post_batch_size",
184
+ "VLLM_MORIIO_NUM_WORKERS": "num_workers",
185
+}
186
+
187
+
188
+def _warn_deprecated_env_vars() -> None:
189
+ for env_var, new_key in _DEPRECATED_ENV_VARS.items():
190
+ if env_var in os.environ:
191
+ logger.warning_once(
192
+ "The environment variable %s is deprecated and ignored. "
193
+ "Set %r inside kv_transfer_config.kv_connector_extra_config "
194
+ "instead.",
195
+ env_var,
196
+ new_key,
197
+ )
198
+
199
+
200
@dataclass
201
class MoRIIOConfig:
202
local_ip: str
@@ -189,6 +211,10 @@ class MoRIIOConfig:
211
dp_rank: int
212
dp_size: int
213
tp_size: int
214
+ read_mode: bool = False
215
+ qp_per_transfer: int = 1
216
+ post_batch_size: int = -1
217
+ num_workers: int = 1
218
backend: str = "rdma"
219
220
@classmethod
@@ -201,11 +227,24 @@ class MoRIIOConfig:
227
# notify_port -> For synchronizing stages between prefill and decode
228
# handshake_port -> For initial handshake between mori engine
229
230
+ # Optional tuning knobs
231
+ # read_mode -> If true, run the connector in READ mode (consumer
232
+ # pulls KV from producer) instead of the default
233
+ # WRITE mode.
234
+
235
+ # Knobs for RDMA transfers, ignored if on xgmi backend
236
+ # qp_per_transfer -> Number of RDMA Queue Pairs per KV transfer.
237
+ # post_batch_size -> Batch size for posting transfer work requests
238
+ # (-1 lets the MoRI backend choose).
239
+ # num_workers -> Number of background worker threads the MoRI
240
+ # engine uses for transfer processing.
241
+
242
# TODO : merge notify_port and handshake_port to simplify port management
243
# supports non-contiguous ports
244
assert vllm_config.kv_transfer_config is not None, (
245
"kv_transfer_config must be set for MoRIIOConnector"
246
)
247
+ _warn_deprecated_env_vars()
248
kv_transfer_config = vllm_config.kv_transfer_config
249
extra_config = kv_transfer_config.kv_connector_extra_config
250
tp_rank = get_tensor_model_parallel_rank()
@@ -234,6 +273,10 @@ class MoRIIOConfig:
273
dp_rank=dp_rank,
274
dp_size=dp_size,
275
tp_size=tp_size,
276
+ read_mode=get_moriio_mode(kv_transfer_config) == MoRIIOMode.READ,
277
+ qp_per_transfer=int(extra_config.get("qp_per_transfer", 1)),
278
+ post_batch_size=int(extra_config.get("post_batch_size", -1)),
279
+ num_workers=int(extra_config.get("num_workers", 1)),
280
backend=backend,
281
)
282
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+12
-4
@@ -108,7 +108,7 @@ class MoRIIOConnector(KVConnectorBase_V1):
108
+ ":"
109
+ str(self.kv_transfer_config.kv_connector_extra_config["handshake_port"])
110
)
111
- self.mode = get_moriio_mode()
111
+ self.mode = get_moriio_mode(self.kv_transfer_config)
112
if role == KVConnectorRole.SCHEDULER:
113
self.connector_scheduler: MoRIIOConnectorScheduler | None = (
114
MoRIIOConnectorScheduler(vllm_config, self.engine_id)
@@ -250,7 +250,7 @@ class MoRIIOConnectorScheduler:
250
self.kv_transfer_config = vllm_config.kv_transfer_config
251
self.block_size = vllm_config.cache_config.block_size
252
self.engine_id: EngineId = engine_id
253
- self.mode = get_moriio_mode()
253
+ self.mode = get_moriio_mode(self.kv_transfer_config)
254
self.host_ip = get_ip()
255
self.handshake_port = self.kv_transfer_config.kv_connector_extra_config[
256
"handshake_port"
@@ -615,8 +615,11 @@ class MoRIIOConnectorWorker:
615
"is installed and properly configured."
616
)
617
618
+ assert vllm_config.kv_transfer_config is not None
619
self.moriio_config = MoRIIOConfig.from_vllm_config(vllm_config)
619
- self.mode = get_moriio_mode()
620
+ self.mode = (
621
+ MoRIIOMode.READ if self.moriio_config.read_mode else MoRIIOMode.WRITE
622
+ )
623
624
logger.info("Initializing MoRIIO worker %s", engine_id)
625
@@ -700,7 +703,12 @@ class MoRIIOConnectorWorker:
703
if self.moriio_config.backend == "xgmi"
704
else BackendType.RDMA
705
)
703
- self.moriio_wrapper.set_backend_type(backend)
706
+ self.moriio_wrapper.set_backend_type(
707
+ backend,
708
+ qp_per_transfer=self.moriio_config.qp_per_transfer,
709
+ post_batch_size=self.moriio_config.post_batch_size,
710
+ num_workers=self.moriio_config.num_workers,
711
+ )
712
self.moriio_wrapper.notify_port = self.moriio_config.notify_port
713
self.local_kv_cache_metadata: list[bytes] = []
714
self.local_kv_cache_size: list[int] = []
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py
+14
-9
@@ -8,7 +8,6 @@ import msgpack
8
import torch
9
import zmq
10
11
-from vllm import envs
11
from vllm.logger import init_logger
12
from vllm.utils.network_utils import (
13
make_zmq_path,
@@ -16,7 +15,7 @@ from vllm.utils.network_utils import (
15
)
16
17
if TYPE_CHECKING:
19
- pass
18
+ from mori.io import BackendType
19
20
from queue import Empty, Queue
21
@@ -376,7 +375,13 @@ class MoRIIOWrapper:
375
)
376
self.moriio_engine = moriio_engine
377
379
- def set_backend_type(self, backend_type):
378
+ def set_backend_type(
379
+ self,
380
+ backend_type: "BackendType",
381
+ qp_per_transfer: int = 1,
382
+ post_batch_size: int = -1,
383
+ num_workers: int = 1,
384
+ ) -> None:
385
assert self.moriio_engine is not None, "MoRIIO engine must be set first"
386
if backend_type == BackendType.XGMI:
387
logger.info("Using MoRIIO backend: XGMI")
@@ -385,14 +390,14 @@ class MoRIIOWrapper:
390
logger.info(
391
"Using MoRIIO backend: RDMA "
392
"(qp_per_transfer=%d, post_batch_size=%d, num_workers=%d)",
388
- envs.VLLM_MORIIO_QP_PER_TRANSFER,
389
- envs.VLLM_MORIIO_POST_BATCH_SIZE,
390
- envs.VLLM_MORIIO_NUM_WORKERS,
393
+ qp_per_transfer,
394
+ post_batch_size,
395
+ num_workers,
396
)
397
rdma_cfg = RdmaBackendConfig(
393
- envs.VLLM_MORIIO_QP_PER_TRANSFER,
394
- envs.VLLM_MORIIO_POST_BATCH_SIZE,
395
- envs.VLLM_MORIIO_NUM_WORKERS,
398
+ qp_per_transfer,
399
+ post_batch_size,
400
+ num_workers,
401
PollCqMode.POLLING,
402
)
403
self.moriio_engine.create_backend(backend_type, rdma_cfg)
vllm/envs.py
-18
@@ -216,10 +216,6 @@ if TYPE_CHECKING:
216
VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None
217
VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB: int | None = None
218
VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB: int | None = None
219
- VLLM_MORIIO_CONNECTOR_READ_MODE: bool = False
220
- VLLM_MORIIO_QP_PER_TRANSFER: int = 1
221
- VLLM_MORIIO_POST_BATCH_SIZE: int = -1
222
- VLLM_MORIIO_NUM_WORKERS: int = 1
219
VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT: int = 480
220
VLLM_ENABLE_CUDAGRAPH_GC: bool = False
221
VLLM_LOOPBACK_IP: str = ""
@@ -1642,20 +1638,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
1638
"Use --linear-backend emulation.",
1639
lambda: bool(int(os.getenv("VLLM_USE_NVFP4_CT_EMULATIONS", "0"))),
1640
),
1645
- # Controls the read mode for the Mori-IO connector
1646
- "VLLM_MORIIO_CONNECTOR_READ_MODE": lambda: (
1647
- os.getenv("VLLM_MORIIO_CONNECTOR_READ_MODE", "False").lower() in ("true", "1")
1648
- ),
1649
- # Controls the QP (Queue Pair) per transfer configuration for the Mori-IO connector
1650
- "VLLM_MORIIO_QP_PER_TRANSFER": lambda: int(
1651
- os.getenv("VLLM_MORIIO_QP_PER_TRANSFER", "1")
1652
- ),
1653
- # Controls the post-processing batch size for the Mori-IO connector
1654
- "VLLM_MORIIO_POST_BATCH_SIZE": lambda: int(
1655
- os.getenv("VLLM_MORIIO_POST_BATCH_SIZE", "-1")
1656
- ),
1657
- # Controls the number of workers for Mori operations for the Mori-IO connector
1658
- "VLLM_MORIIO_NUM_WORKERS": lambda: int(os.getenv("VLLM_MORIIO_NUM_WORKERS", "1")),
1641
# Timeout (in seconds) for MooncakeConnector in PD disaggregated setup.
1642
"VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT": lambda: int(
1643
os.getenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "480")