[PR #466] Allow NVSHMEM PE to NIC mapping to be initialized by DeepEP rank #586

New issue

Open

opened 2026-10-27 21:04:18 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:04:18 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/466
Author: @smarterclayton
Created: 10/24/2025
Status: 🔄 Open

Base: main ← Head: nic_pe_alignment

📝 Commits (2)

55d97a6 Allow NVSHMEM PE to NIC to be initialized by rank
8a07e7e If CUDA_VISIBLE_DEVICES is set, map cuda.current_device()

📊 Changes

1 file changed (+32 additions, -0 deletions)

View changed files

📝 deep_ep/buffer.py (+32 -0)

📄 Description

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 keonjang@google.com

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/466 **Author:** [@smarterclayton](https://github.com/smarterclayton) **Created:** 10/24/2025 **Status:** 🔄 Open **Base:** `main` ← **Head:** `nic_pe_alignment` --- ### 📝 Commits (2) - [`55d97a6`](https://github.com/deepseek-ai/DeepEP/commit/55d97a68bd9edee3beb1eb9635673a630c6ce5cc) Allow NVSHMEM PE to NIC to be initialized by rank - [`8a07e7e`](https://github.com/deepseek-ai/DeepEP/commit/8a07e7e7ed57cfc63005f670d64d410cb30b061b) If CUDA_VISIBLE_DEVICES is set, map cuda.current_device() ### 📊 Changes **1 file changed** (+32 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `deep_ep/buffer.py` (+32 -0) </details> ### 📄 Description 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 <keonjang@google.com> --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-27 21:04:18 +00:00

bsakaguchi referenced this issue

2026-10-27 21:04:51 +00:00

[PR #586] docs: improve documentation with examples and fix typos #646