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 4, 00:00 UTC
[Refactor] Nixl util using lazy init (#41392)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
May 9, 2026 at 17:46 UTC
f80aa53c9dc2273a19a6855092069db7e1306fff
4 files changed
+73
-43
vllm/distributed/eplb/eplb_communicator.py
+6
-7
@@ -19,6 +19,7 @@ from torch.distributed import (
19
batch_isend_irecv,
20
)
21
22
+import vllm.distributed.nixl_utils as nixl_utils
23
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
24
from vllm.distributed.device_communicators.pynccl_wrapper import (
25
ncclDataTypeEnum,
@@ -37,9 +38,7 @@ logger = init_logger(__name__)
38
39
def has_nixl() -> bool:
40
"""Whether the optional NIXL / RIXL package is available."""
40
- from vllm.distributed.nixl_utils import NixlWrapper
41
-
42
- return NixlWrapper is not None
41
+ return nixl_utils.NixlWrapper is not None
42
43
44
class EplbCommunicator(ABC):
@@ -233,10 +232,9 @@ class NixlEplbCommunicator(EplbCommunicator):
232
expert_weights: Sequence[torch.Tensor],
233
cuda_stream: torch.cuda.Stream | None = None,
234
) -> None:
236
- from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
237
-
235
assert expert_weights, "NixlEplbCommunicator requires non-empty expert_weights."
239
- if NixlWrapper is None:
236
+ nixl_wrapper_cls = nixl_utils.NixlWrapper
237
+ if nixl_wrapper_cls is None:
238
raise RuntimeError("NIXL/ RIXL is unavailable.")
239
self._cpu_group = cpu_group
240
self._cuda_stream = cuda_stream
@@ -254,12 +252,13 @@ class NixlEplbCommunicator(EplbCommunicator):
252
f"expected={self._device}, got={tensor.device}"
253
)
254
255
+ nixl_agent_config = nixl_utils.nixl_agent_config
256
config = (
257
nixl_agent_config(capture_telemetry=False)
258
if nixl_agent_config is not None
259
else None
260
)
262
- self._nixl_wrapper = NixlWrapper(self._make_agent_name(), config)
261
+ self._nixl_wrapper = nixl_wrapper_cls(self._make_agent_name(), config)
262
self._nixl_memory_type = "VRAM"
263
self._registered_desc: object | None = None
264
self._remote_agents: dict[int, str] = {}
vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py
+5
-3
@@ -4,7 +4,7 @@
4
5
import copy
6
from dataclasses import dataclass
7
-from typing import Any
7
+from typing import TYPE_CHECKING, Any
8
9
import numpy as np
10
@@ -15,9 +15,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
15
PromMetric,
16
PromMetricT,
17
)
18
-from vllm.distributed.nixl_utils import nixlXferTelemetry
18
from vllm.v1.metrics.utils import create_metric_per_engine
19
20
+if TYPE_CHECKING:
21
+ from vllm.distributed.nixl_utils import nixlXferTelemetry
22
+
23
24
@dataclass
25
class NixlKVConnectorStats(KVConnectorStats):
@@ -40,7 +42,7 @@ class NixlKVConnectorStats(KVConnectorStats):
42
"num_kv_expired_reqs": [],
43
}
44
43
- def record_transfer(self, res: nixlXferTelemetry):
45
+ def record_transfer(self, res: "nixlXferTelemetry"):
46
# Keep metrics units consistent with rest of the code: time us->s
47
self.data["transfer_duration"].append(res.xferDuration / 1e6)
48
self.data["post_duration"].append(res.postDuration / 1e6)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+3
-2
@@ -198,7 +198,8 @@ class NixlConnectorWorker:
198
engine_id: str,
199
kv_cache_config: "KVCacheConfig",
200
):
201
- if NixlWrapper is None:
201
+ nixl_wrapper_cls = NixlWrapper
202
+ if nixl_wrapper_cls is None:
203
logger.error("NIXL is not available")
204
raise RuntimeError("NIXL is not available")
205
logger.info("Initializing NIXL wrapper")
@@ -284,7 +285,7 @@ class NixlConnectorWorker:
285
else nixl_agent_config(num_threads=num_threads, capture_telemetry=True)
286
)
287
287
- self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()), config)
288
+ self.nixl_wrapper = nixl_wrapper_cls(str(uuid.uuid4()), config)
289
# Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}.
290
self._remote_agents: dict[EngineId, dict[int, str]] = defaultdict(dict)
291
vllm/distributed/nixl_utils.py
+59
-31
@@ -1,54 +1,82 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import importlib
5
import os
6
import sys
7
+from typing import Any
8
9
from vllm.logger import init_logger
10
from vllm.platforms import current_platform
11
12
logger = init_logger(__name__)
13
12
-if "UCX_RCACHE_MAX_UNRELEASED" not in os.environ:
14
+# declaration for static analyzers
15
+NixlWrapper: Any
16
+nixl_agent_config: Any
17
+nixlXferTelemetry: Any
18
+
19
+
20
+def _maybe_set_ucx_rcache_limit() -> None:
21
+ if "UCX_RCACHE_MAX_UNRELEASED" in os.environ:
22
+ return
23
+
24
if "nixl" in sys.modules or "rixl" in sys.modules:
25
logger.warning_once(
26
"NIXL was already imported, we can't reset "
27
"UCX_RCACHE_MAX_UNRELEASED. "
28
"Please set it to '1024' manually."
29
)
19
- else:
20
- logger.info_once(
21
- "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare "
22
- "memory leak in UCX when using NIXL."
23
- )
24
- os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024"
30
+ return
31
26
-try:
27
- if not current_platform.is_rocm():
28
- from nixl._api import nixl_agent as NixlWrapper
29
- else:
30
- from rixl._api import nixl_agent as NixlWrapper
32
+ logger.info_once(
33
+ "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare "
34
+ "memory leak in UCX when using NIXL."
35
+ )
36
+ os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024"
37
32
- logger.info_once("NIXL is available")
33
-except ImportError:
34
- logger.warning_once("NIXL is not available")
35
- NixlWrapper = None # type: ignore[assignment, misc]
38
37
-try:
38
- if not current_platform.is_rocm():
39
- from nixl._api import nixl_agent_config
40
- else:
41
- from rixl._api import nixl_agent_config
42
-except ImportError:
43
- nixl_agent_config = None # type: ignore[assignment]
44
- logger.warning_once("NIXL agent config is not available")
45
-
46
-try:
47
- if not current_platform.is_rocm():
48
- from nixl._bindings import nixlXferTelemetry
39
+def _get_nixl_module_name(name: str) -> str:
40
+ package_name = "rixl" if current_platform.is_rocm() else "nixl"
41
+ if name == "nixlXferTelemetry":
42
+ return f"{package_name}._bindings"
43
+ return f"{package_name}._api"
44
+
45
+
46
+def _load_nixl_attr(name: str) -> Any:
47
+ attr_name = {
48
+ "NixlWrapper": "nixl_agent",
49
+ "nixl_agent_config": "nixl_agent_config",
50
+ "nixlXferTelemetry": "nixlXferTelemetry",
51
+ }[name]
52
+
53
+ _maybe_set_ucx_rcache_limit()
54
+ try:
55
+ module = importlib.import_module(_get_nixl_module_name(name))
56
+ except ImportError:
57
+ if name == "NixlWrapper":
58
+ logger.warning_once("NIXL is not available")
59
+ elif name == "nixl_agent_config":
60
+ logger.warning_once("NIXL agent config is not available")
61
+ value = None
62
else:
50
- from rixl._bindings import nixlXferTelemetry
51
-except ImportError:
52
- nixlXferTelemetry = None # type: ignore[assignment, misc]
63
+ value = getattr(module, attr_name, None)
64
+ if name == "NixlWrapper":
65
+ if value is None:
66
+ logger.warning_once("NIXL is not available")
67
+ else:
68
+ logger.info_once("NIXL is available")
69
+ elif name == "nixl_agent_config" and value is None:
70
+ logger.warning_once("NIXL agent config is not available")
71
+
72
+ globals()[name] = value
73
+ return value
74
+
75
+
76
+def __getattr__(name: str) -> Any:
77
+ if name in __all__:
78
+ return _load_nixl_attr(name)
79
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
80
+
81
82
__all__ = ["NixlWrapper", "nixl_agent_config", "nixlXferTelemetry"]