Skip to content

vllm.utils.nccl

Functions:

find_nccl_include_paths()

Return possible include paths containing nccl.h.

Considers VLLM_NCCL_INCLUDE_PATH and the nvidia-nccl-cuXX package.

Source code in vllm/utils/nccl.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def find_nccl_include_paths() -> list[str] | None:
    """Return possible include paths containing `nccl.h`.

    Considers `VLLM_NCCL_INCLUDE_PATH` and the `nvidia-nccl-cuXX` package.
    """
    paths: list[str] = []
    inc = envs.VLLM_NCCL_INCLUDE_PATH
    if inc and os.path.isdir(inc):
        paths.append(inc)

    try:
        spec = importlib.util.find_spec("nvidia.nccl")
        if spec and (locs := getattr(spec, "submodule_search_locations", None)):
            for loc in locs:
                inc_dir = os.path.join(loc, "include")
                if os.path.exists(os.path.join(inc_dir, "nccl.h")):
                    paths.append(inc_dir)
    except Exception as e:
        logger.debug("Failed to find nccl include path from nvidia.nccl package: %s", e)

    seen: set[str] = set()
    out: list[str] = []
    for p in paths:
        if p and p not in seen:
            out.append(p)
            seen.add(p)
    return out or None

find_nccl_library()

Return NCCL/RCCL shared library name to load.

Uses VLLM_NCCL_SO_PATH if set; otherwise chooses by torch backend.

Source code in vllm/utils/nccl.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def find_nccl_library() -> str:
    """Return NCCL/RCCL shared library name to load.

    Uses `VLLM_NCCL_SO_PATH` if set; otherwise chooses by torch backend.
    """
    so_file = envs.VLLM_NCCL_SO_PATH
    if so_file:
        logger.info(
            "Found nccl from environment variable VLLM_NCCL_SO_PATH=%s", so_file
        )
    else:
        if torch.version.cuda is not None:
            so_file = "libnccl.so.2"
        elif torch.version.hip is not None:
            so_file = "librccl.so.1"
        else:
            raise ValueError("NCCL only supports CUDA and ROCm backends.")
        logger.debug_once("Found nccl from library %s", so_file)
    return so_file

find_nccl_library_paths()

Return possible library paths containing libnccl.so.

Looks inside the nvidia-nccl-cuXX pip package.

Source code in vllm/utils/nccl.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def find_nccl_library_paths() -> list[str] | None:
    """Return possible library paths containing `libnccl.so`.

    Looks inside the `nvidia-nccl-cuXX` pip package.
    """
    paths: list[str] = []
    try:
        spec = importlib.util.find_spec("nvidia.nccl")
        if spec and (locs := getattr(spec, "submodule_search_locations", None)):
            for loc in locs:
                lib_dir = os.path.join(loc, "lib")
                if os.path.isdir(lib_dir):
                    paths.append(lib_dir)
    except Exception as e:
        logger.debug("Failed to find nccl library path from nvidia.nccl package: %s", e)
    return paths or None

query_nccl_gin_type(group, *, railed=False)

Return the full or railed GIN type, or None on query failure.

Source code in vllm/utils/nccl.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def query_nccl_gin_type(
    group: torch.distributed.ProcessGroup, *, railed: bool = False
) -> int | None:
    """Return the full or railed GIN type, or ``None`` on query failure."""
    from vllm.distributed.device_communicators.pynccl_wrapper import (
        NCCL_COMM_PROPERTIES_LAYOUT_VERSION,
        NCCLLibrary,
        ncclCommProperties,
    )

    try:
        backend = group._get_backend(torch.device("cuda"))
        # GIN is a property of this initialized communicator, not just the
        # NCCL version. ncclCommQueryProperties requires its ncclComm_t.
        comm_ptr = backend._comm_ptr()
        if comm_ptr == 0:
            return None
    except Exception:
        logger.warning(
            "Failed to extract NCCL comm pointer from process group",
            exc_info=True,
        )
        return None

    try:
        nccl = NCCLLibrary()
        query_fn = nccl._funcs.get("ncclCommQueryProperties")
        if query_fn is None:
            return None

        props = ncclCommProperties()
        ctypes.memset(ctypes.addressof(props), 0, ctypes.sizeof(props))
        props.size = ctypes.sizeof(props)
        props.magic = 0xCAFEBEEF
        props.version = min(
            nccl.ncclGetRawVersion(), NCCL_COMM_PROPERTIES_LAYOUT_VERSION
        )
        result = query_fn(ctypes.c_void_p(comm_ptr), ctypes.byref(props))
    except Exception:
        logger.warning("Failed to query NCCL communicator properties", exc_info=True)
        return None

    if result != 0:
        logger.warning("ncclCommQueryProperties returned error %d", result)
        return None
    return props.railedGinType if railed else props.ginType