Prhub

#42083 [Feat] Add support for per GPU worker RDMA NIC selection

原始 PR 作者 rajkiranjoshi 合并时间 2026-05-29 03:45 文件变更 7 提交数 5 评论 23 代码增减 +365 / -0

执行摘要

支持 per-GPU worker RDMA NIC 选择

在 multi-NIC multi-GPU 节点上,RDMA 性能取决于 GPU-NIC PCIe 亲和性,但某些拓扑(如 Azure ND96isr_v5 的 flat PCIe)下网络库无法自动发现最优配对,导致所有 worker 竞争同一 NIC。PR body 指出:'In multi-NIC, multi-GPU nodes, RDMA performance depends on GPU-NIC PCIe affinity. When the topology is flat or not properly discoverable by networking libraries, all workers may default to the same NIC or incorrect NICs, leading to suboptimal RDMA performance.'

值得精读,尤其 PCI BDF 规范化和 sysfs 遍历的实现可供其他 RDMA 相关特性参考。设计决策(仅 NVML、去除 prefetch)体现了简化优先的务实思路。

讨论亮点

主要讨论聚焦于实现简化:

  • tlrmchlsmth 建议放弃非 NVML 路径和 nvidia-smi 子进程探测,仅保留 pynvml 实现,从而去除 prefetch 机制,使 PR 大幅简化(commit 1843c28)。
  • tlrmchlsmth 提议使用 envs.py 中注册的变量而非直接 os.environ.get,已被采用。
  • tlrmchlsmth 要求为 normalize_pci 添加单元测试,rajkiranjoshi 在 commit 7385778 中实现。
  • 关于 local_rank 传递的疑问,rajkiranjoshi 解释了 kwargs.get('local_rank', 0) 仅是 fallback,实际值在 spawn 时已正确设置,tlrmchlsmth 表示理解并同意简化调用。

实现拆解

  1. 环境变量注册 (vllm/envs.py):添加 VLLM_GPU_NIC_PCIE_MAPPINGVLLM_NIC_SELECTION_VARS,必须同时设置。
  2. 平台抽象扩展 (vllm/platforms/interface.py, vllm/platforms/cuda.py):在 Platform 接口中添加 get_all_gpu_pci_bus_ids() 抽象方法,NvmlCudaPlatform 通过 NVML 实现;非 NVML 平台抛出 NotImplementedError
  3. 核心映射与设备发现模块 (vllm/v1/executor/vllm_net_devices.py):实现 normalize_pci() 规范化 PCI BDF 地址,parse_gpu_nic_mapping() 解析映射字符串,rdma_name_for_nic_pci() 通过 /sys/class/infiniband 反向查找 RDMA 设备名,set_worker_net_device() 串联上述步骤为当前 worker 设置环境变量。
  4. 执行器集成 (multiproc_executor.py, uniproc_executor.py):在 worker 进程初始化前调用 set_worker_net_device(local_rank, vllm_config)
  5. 单元测试 (tests/v1/executor/test_vllm_net_devices.py):覆盖 normalize_pci 的多种格式和异常路径。
文件 模块 状态 重要度
vllm/v1/executor/vllm_net_devices.py 执行器 added 9.08
tests/v1/executor/test_vllm_net_devices.py 测试 added 7.07
vllm/platforms/cuda.py 平台层 modified 6.74
vllm/platforms/interface.py 平台层 modified 6.39
vllm/envs.py 配置 modified 5.31
vllm/v1/executor/multiproc_executor.py 执行器 modified 5.23
vllm/v1/executor/uniproc_executor.py 执行器 modified 4.83

关键符号

normalize_pci parse_gpu_nic_mapping rdma_name_for_nic_pci parse_nic_selection_vars set_worker_net_device get_all_gpu_pci_bus_ids

关键源码片段

tests/v1/executor/test_vllm_net_devices.py test-coverage

新增完整的 normalize_pci 单元测试,确保解析逻辑正确性。

import pytestfrom vllm.v1.executor.vllm_net_devices import normalize_pci# 测试完整的 domain:bus:dev.fn 格式
@pytest.mark.parametrize(
    "addr, expected",
    [
        ("0000:3f:00.0", (0, 0x3F, 0, 0)),
        ("0001:00:00.0", (1, 0, 0, 0)),
        ("00000001:00:00.0", (1, 0, 0, 0)),
        ("0000:0a:1f.7", (0, 0x0A, 0x1F, 7)),
    ],
)
def test_normalize_pci_full_domain(addr, expected):
    assert normalize_pci(addr) == expected# 测试简短的 bus:dev.fn 格式(domain 默认为 0)
@pytest.mark.parametrize(
    "addr, expected",
    [
        ("01:00.0", (0, 1, 0, 0)),
        ("3f:00.0", (0, 0x3F, 0, 0)),
        ("ff:1f.7", (0, 0xFF, 0x1F, 7)),
    ],
)
def test_normalize_pci_short_form(addr, expected):
    assert normalize_pci(addr) == expected# 验证大小写不敏感
def test_normalize_pci_case_insensitive():
    assert normalize_pci("0A:1F.7") == normalize_pci("0a:1f.7")# 验证去除首尾空白和 0x 前缀
def test_normalize_pci_strips_whitespace():
    assert normalize_pci("  0001:00:00.0  ") == (1, 0, 0, 0)def test_normalize_pci_strips_0x_prefix():
    assert normalize_pci("0x0001:00:00.0") == (1, 0, 0, 0)# 测试各种异常输入
def test_normalize_pci_missing_function_raises():
    with pytest.raises(ValueError, match="missing function suffix"):
        normalize_pci("0001:00:00")def test_normalize_pci_invalid_function_char_raises():
    with pytest.raises(ValueError, match="invalid PCI function"):
        normalize_pci("0001:00:00.z")def test_normalize_pci_too_many_segments_raises():
    with pytest.raises(ValueError, match="invalid PCI BDF"):
        normalize_pci("a:b:c:d.0")def test_normalize_pci_bus_out_of_range_raises():
    with pytest.raises(ValueError, match="out of range"):
        normalize_pci("0000:1ff:00.0")def test_normalize_pci_device_out_of_range_raises():
    with pytest.raises(ValueError, match="out of range"):
        normalize_pci("0000:00:20.0")def test_normalize_pci_empty_string_raises():
    with pytest.raises(ValueError):
        normalize_pci("")

评论区精华

简化:仅支持 NVML 路径 设计

tlrmchlsmth 提出 PR 可以通过仅支持 NVML 路径大幅简化,去掉 nvidia-smi 子进程和 prefetch 优化 : "This PR would be a lot simpler if we only supported this for the NVML path. We could get rid of the prefetching, and querying nvidia-smi completely."

结论:rajkiranjoshi 采纳建议,移除了非 NVML 支持和 prefetch 机制,使实现更加简洁。 · 已解决

使用 envs.py 注册的变量 style

tlrmchlsmth 建议使用 envs.py 中定义的变量而非直接 os.environ.get: "Instead of calling os.environ.get, we should use the variable created in envs.py"

结论:rajkiranjoshi 在后续 commit 中改为引用 vllm.envs 中的变量。 · 已解决

worker 间 local_rank 传递 正确性

tlrmchlsmth 质疑 set_worker_net_device 为什么只对 local_rank=0 设置 : "Why are we setting this for local rank 0? What about the other local ranks? Should we be doing this in the worker directly?" rajkiranjoshi 解释 kwargs.get 的默认值 0 只是 fallback,实际每个 worker 都携带了各自的 local_rank。

结论:tlrmchlsmth 表示理解并同意保留当前调用方式。 · 已解决

添加 normalize_pci 单元测试 测试

tlrmchlsmth 在 review 中要求为 normalize_pci 添加单元测试 : "this deserves a unit test"

结论:rajkiranjoshi 在 commit 7385778 中新增了完整的单元测试套件,覆盖正常和异常路径。 · 已解决

风险与影响

新模块 vllm_net_devices.py 依赖 /sys/class/infiniband,仅 Linux 有效。NVML 路径被设为唯一实现,非 NVML 平台(如 Jetson)或 AMD GPU 将抛出 NotImplementedError,需要后续扩展。环境变量配置错误(如 PCI BDF 格式不匹配)会抛出 ValueError 并终止 worker 初始化。在多 DeepEP 场景中已验证性能提升,但其他 RDMA 库(如 UCCL)的实际测试尚未覆盖。

用户侧需要手动设置两个环境变量,但 PR body 提供了详细示例,降低了使用门槛。系统侧正确配置后可显著降低 RDMA 传输延迟,提升跨节点通信效率。团队侧新增一个独立模块,后续维护成本低,且库无关设计易于扩展支持新通信库。

仅 NVML 平台支持 依赖 /sys/class/infiniband(仅 Linux) 环境变量必须同时设置 缺少非 NVML 平台的实现在未来可能影响扩展 深层性能测试仅覆盖 UCX 和 NVSHMEM

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论