Allow NVSHMEM PE to NIC mapping to be initialized by DeepEP rank

[smarterclayton]

The nvshmemi_get_devices_by_distance default initialization method in NVSHMEM does not work optimally for GPU configurations where 2 GPU and 2 RDMA NIC share a PCIe bus, such as the x86 based GCP A3 Ultra H200 and A4 B200 instance types: https://cloud.google.com/compute/docs/gpus/gpu-network-bandwidth#h200-gpus. GPU0 and GPU1 (on two independent processes) observe NIC0 and NIC1 on the same PCIe switch as being equidistant and the default configuration for DeepEP + NVSHMEM results in both GPUs being assigned NIC0 from nvshmemi_get_devices_by_distance, halving the observed bandwidth for RDMA in test_internode.py and in vLLM wide-EP (because 4 of 8 NIC are enabled).

The alternative is a static mapping between GPU host index (PE) and NIC index (HCA), but the NVSHMEMX_INIT_WITH_UNIQUEID initialization method bypasses setting mype_node and npes_node.

The nvshmemi_boot_handle.pg_rank for this initialization method is always 0 and the nvshmem_boot_handle.pg_size is always 2, and as a result mype_node and npes_node are set to 0 and 2 respectively for all PEs initializing. This prevents NVSHMEM_ENABLE_NIC_PE_MAPPING=1 from selecting from a static list of devices by mype_node / local rank in transport.cpp#nvshmemi_setup_connections:

selected_devices[0] =
  nvshmemi_state->mype_node % (tcurr->n_devices > 0
    ? tcurr->n_devices : 1);

To allow static assignment, we introduce a new DEEP_EP_DEVICE_TO_HCA_MAPPING environment variable during deep_ep.Buffer initialization that accepts <cuda_device_id>:<HCA_name>:<HCA_port> and uses torch.cuda.current_device() to set NVSHMEM_HCA_LIST to the appropriate <HCA_name>:<HCA_port>, or error if no such device was listed.

We are proposing the change to DeepEP because the choice of initialization method determines how HCA to GPU binding can be achieved across multiple existing NVSHMEM versions. Because vLLM 0.11.0 uses CUDA_VISIBLE_DEVICES to associate GPU devices to each rank, we reverse map the visible devices to the current cuda device.

On GCP we would propose this configuration for all full host workloads on A3U and A4 instance types (H200/B200 + x86 with shared PCIe hubs per 2 NIC / 2 GPU):

DEEP_EP_DEVICE_TO_HCA_MAPPING=0:mlx5_0:1,1:mlx5_1:1,2:mlx5_2:1,3:mlx5_3:1,4:mlx5_4:1,5:mlx5_5:1,6:mlx5_6:1,7:mlx5_7:1

Tested with NVSHMEM 3.3.20 and vLLM 0.11.0 and vLLM main @ 2025/10/23.

Co-Authored-By: Keon Jang [email protected]

[sphish]

Why not just set the NVSHMEM environment variables directly?

[smarterclayton]

They have to be set per rank, so we would need to change all callers of deep_ep.Buffer to set those environment variables before invoking, including test_low_latency and test_internode (which we use to validate network performance of the supporting RDMA clusters). Since we expected all consumers of DeepEP to need this mapping on our H200 and B200 VM families, we wanted to push the fix as low in the software stack as possible (so below SGLang, vLLM, NeMO, etc) so that those frameworks don't have to change how they call DeepEP. We considered changing NVSHMEM but that would likely still require a change to Buffer initialization to pass local rank.

[sphish]

NVSHMEM stores local rank information, for example in mype_node in the code.
I believe the issue you’re experiencing is quite common — applications using NVSHMEM outside of DeepEP are likely to encounter it as well. It would be best to submit a PR to NVSHMEM to resolve this problem.

[sphish]

@seth-howell Could you help have a look at this？

[keonjang-google]

mype_node is 0 for all the PEs because DeepEP initializes NVSHMEM separately for each rank IIUC.