A high-throughput and memory-efficient inference and serving engine for LLMs
[Mooncake] Use all HCAs on multi-NIC hosts instead of GPU-indexed RNIC selection (#43799)
Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Dao007forever committed
Jun 9, 2026 at 11:05 UTC
c1d754d681108096d8a125054f01a51afd60c618
3 files changed
+2
-163
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
-88
@@ -283,94 +283,6 @@ def test_get_configured_preferred_segment_rejects_empty_override():
283
rdma_utils.get_configured_preferred_segment({"preferred_segment": " "})
284
285
286
-def test_get_configured_worker_rnic_prefers_explicit_device_name(monkeypatch):
287
- store_config = worker.MooncakeStoreConfig(
288
- metadata_server="",
289
- local_buffer_size=1,
290
- protocol="rdma",
291
- device_name="rocep139s0",
292
- master_server_address="",
293
- )
294
-
295
- assert (
296
- rdma_utils.get_configured_worker_rnic(
297
- protocol=store_config.protocol,
298
- configured_device=store_config.device_name,
299
- )
300
- == "rocep139s0"
301
- )
302
-
303
-
304
-def test_get_configured_worker_rnic_selects_device_from_explicit_csv(monkeypatch):
305
- monkeypatch.setattr(
306
- rdma_utils,
307
- "get_current_physical_gpu_index",
308
- lambda: 1,
309
- )
310
- store_config = worker.MooncakeStoreConfig(
311
- metadata_server="",
312
- local_buffer_size=1,
313
- protocol="rdma",
314
- device_name="rocep139s0,rocep140s0",
315
- master_server_address="",
316
- )
317
-
318
- assert (
319
- rdma_utils.get_configured_worker_rnic(
320
- protocol=store_config.protocol,
321
- configured_device=store_config.device_name,
322
- )
323
- == "rocep140s0"
324
- )
325
-
326
-
327
-def test_get_configured_worker_rnic_warns_and_returns_empty_for_rdma_with_no_device(
328
- caplog, monkeypatch
329
-):
330
- """No device configured + protocol=rdma → emit a clear warning and return ""
331
- so the C++ side handles auto-selection. There is no Python-side fallback."""
332
- monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
333
- with caplog.at_level(logging.WARNING):
334
- result = rdma_utils.get_configured_worker_rnic(
335
- protocol="rdma",
336
- configured_device="",
337
- )
338
- assert result == ""
339
- warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
340
- assert any("No RDMA devices specified" in r.message for r in warnings), (
341
- f"expected fallback warning, got {[r.message for r in warnings]}"
342
- )
343
-
344
-
345
-def test_get_configured_worker_rnic_silent_for_tcp_with_no_device(caplog, monkeypatch):
346
- """protocol=tcp + no device → return "" silently (no RDMA, no warning)."""
347
- monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
348
- with caplog.at_level(logging.WARNING):
349
- result = rdma_utils.get_configured_worker_rnic(
350
- protocol="tcp",
351
- configured_device="",
352
- )
353
- assert result == ""
354
- warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
355
- assert not any("RDMA" in r.message for r in warnings), (
356
- "did not expect RDMA warning for tcp protocol, got "
357
- f"{[r.message for r in warnings]}"
358
- )
359
-
360
-
361
-def test_get_configured_worker_rnic_rejects_short_explicit_csv(monkeypatch):
362
- monkeypatch.setattr(
363
- rdma_utils,
364
- "get_current_physical_gpu_index",
365
- lambda: 2,
366
- )
367
- with pytest.raises(ValueError, match="does not cover local GPU 2"):
368
- rdma_utils.get_configured_worker_rnic(
369
- protocol="rdma",
370
- configured_device="rocep139s0,rocep140s0",
371
- )
372
-
373
-
286
class _ReplicaDesc:
287
def __init__(self, tier: str):
288
self.tier = tier
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/rdma_utils.py
-69
@@ -5,8 +5,6 @@
5
from collections.abc import Mapping
6
from typing import Any
7
8
-import torch
9
-
8
import vllm.envs as envs
9
from vllm.logger import init_logger
10
@@ -20,22 +18,6 @@ def normalize_string_override(value: Any) -> str | None:
18
return normalized or None
19
20
23
-def get_current_physical_gpu_index() -> int | None:
24
- try:
25
- from vllm.platforms import current_platform
26
- except ImportError:
27
- return None
28
-
29
- try:
30
- device_index = torch.accelerator.current_device_index()
31
- physical_device_id = current_platform.device_id_to_physical_device_id(
32
- device_index
33
- )
34
- return int(physical_device_id)
35
- except Exception:
36
- return None
37
-
38
-
21
def get_requester_local_hostname(local_ip: str) -> str:
22
override = normalize_string_override(envs.MOONCAKE_REQUESTER_LOCAL_HOSTNAME)
23
if override is not None:
@@ -62,54 +44,3 @@ def get_configured_preferred_segment(
44
)
45
return env_value
46
return None
65
-
66
-
67
-def _get_explicit_worker_rnic(device_list: str) -> str:
68
- entries = [entry.strip() for entry in device_list.split(",")]
69
- if any(not entry for entry in entries):
70
- raise ValueError(
71
- "Mooncake worker device_name contains an empty RDMA device entry"
72
- )
73
- if len(entries) == 1:
74
- return entries[0]
75
-
76
- gpu_index = get_current_physical_gpu_index()
77
- if gpu_index is None:
78
- raise RuntimeError(
79
- "Mooncake RDMA requester could not determine the local physical GPU index"
80
- )
81
- if gpu_index >= len(entries):
82
- raise ValueError(
83
- "Mooncake worker device list does not cover local GPU "
84
- f"{gpu_index}: {device_list}"
85
- )
86
- device_name = entries[gpu_index]
87
- logger.info(
88
- "Mooncake selected worker RNIC %s from explicit device list for local GPU %s",
89
- device_name,
90
- gpu_index,
91
- )
92
- return device_name
93
-
94
-
95
-def get_configured_worker_rnic(
96
- *,
97
- protocol: str,
98
- configured_device: str,
99
-) -> str:
100
- normalized_device = normalize_string_override(configured_device)
101
- if normalized_device is not None:
102
- return _get_explicit_worker_rnic(normalized_device)
103
-
104
- if protocol not in {"rdma", "efa"}:
105
- return ""
106
-
107
- logger.warning(
108
- "No RDMA devices specified for Mooncake backend (protocol=%s). "
109
- "Set 'device_name' in mooncake_config.json to a single RNIC name "
110
- "or a comma-separated CSV indexed by physical GPU; falling back to "
111
- "Mooncake's built-in auto-selection, which may converge on the same "
112
- "NIC across all DP ranks and saturate bandwidth.",
113
- protocol,
114
- )
115
- return ""
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+2
-6
@@ -141,7 +141,7 @@ class MooncakeStoreConfig:
141
)
142
143
@staticmethod
144
- def load_from_env() -> "MooncakeStoreConfig":
144
+ def load_from_config() -> "MooncakeStoreConfig":
145
config_path = os.getenv("MOONCAKE_CONFIG_PATH")
146
if not config_path:
147
raise ValueError(
@@ -983,16 +983,12 @@ class MooncakeStoreWorker:
983
)
984
985
# Initialize MooncakeDistributedStore with its own TransferEngine
986
- store_config = MooncakeStoreConfig.load_from_env()
986
+ store_config = MooncakeStoreConfig.load_from_config()
987
extra_config = (
988
vllm_config.kv_transfer_config.kv_connector_extra_config
989
if vllm_config.kv_transfer_config
990
else {}
991
)
992
- store_config.device_name = rdma_utils.get_configured_worker_rnic(
993
- protocol=store_config.protocol,
994
- configured_device=store_config.device_name,
995
- )
992
self.store = MooncakeDistributedStore()
993
local_ip = get_ip()
994
local_hostname = rdma_utils.get_requester_local_hostname(local_ip)