[Bugfix] Fix CustomAllreduce nvlink topology detection (#3974)

[Bugfix] Fix CustomAllreduce pcie nvlink topology detection (#3974) (#4159)

[Bugfix] Fix CustomAllreduce nvlink topology detection (#3974)
[Bugfix] Fix CustomAllreduce pcie nvlink topology detection (#3974) (#4159)
8f9c28fd · Adam Tilghman · GitHub · cd2f63fb · 8f9c28fd
Unverified Commit 8f9c28fd authored Apr 18, 2024 by Adam Tilghman Committed by GitHub Apr 18, 2024
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 2 deletions

vllm/distributed/device_communicators/custom_all_reduce.py vllm/distributed/device_communicators/custom_all_reduce.py +4 -2

No files found.
--- a/vllm/distributed/device_communicators/custom_all_reduce.py
+++ b/vllm/distributed/device_communicators/custom_all_reduce.py
@@ -145,8 +145,10 @@ def _is_full_nvlink(rank, world_size):
    for i in range(world_size):
        if i != rank:
            try:
-                link_state = pynvml.nvmlDeviceGetNvLinkState(handle, i)
+                peer_handle = pynvml.nvmlDeviceGetHandleByIndex(i)
-                if not link_state:
+                p2p_status = pynvml.nvmlDeviceGetP2PStatus(
+                    handle, peer_handle, pynvml.NVML_P2P_CAPS_INDEX_NVLINK)
+                if p2p_status != pynvml.NVML_P2P_STATUS_OK:
                    return False
            except pynvml.NVMLError as error:
                logger.info(