?一、背景

在此前的內容中，筆者曾介紹過 DeepSeek 的 DeepEP、字節跳動的 Flux 和 Tilelink 等系統，這些系統在底層通信實現中均依賴于 NVIDIA 的 NVSHMEM 庫。事實上，字節跳動后續的諸如 Comet、Triton-distributed，以及其他針對細粒度計算與通信重疊（Overlap）優化的工作，也都廣泛使用了 NVSHMEM。

本文將深入剖析 NVSHMEM 的初始化流程及其核心概念，以便從開發者視角理解其機制，為后續的定制化改造和工程實踐打下基礎。?

也可以參考 NVSHMEM 的官方文檔：NVIDIA OpenSHMEM Library (NVSHMEM) Documentation [1]

二、引言

2.1 DeepEP

DeepEP 是 DeepSeek 開源的專為 MoE 和專家并行（Expert Parallelism, EP）設計的通信庫。提供了一系列優化的通信 Kernel，實現了以下能力：

• 高度優化的 All2All 通信。
• 同時支持不同的通信類型：

節點內（intra-node）：使用 NVLink + NVSwitch 通信。

節點間（inter-node）：使用 RDMA 通信。

• 針對不同場景的 Kernel：
• 常規（高吞吐） Kernel（Normal Kernel）：針對 Training 和 Inference Prefill。節點內 NVLink + 節點間 RDMA 通信。
• 低時延 Kernel（Low-Latency Kernel）：針對 Inference Decoding。使用純 RDMA 通信來最小化時延。
• 原生支持 FP8，減少數據傳輸需求，相比 FP16 通信量減半。
• 靈活的 GPU 資源（SM）控制，支持計算和通信的 Overlap。

代碼庫：DeepEP: an efficient expert-parallel communication library [2]

2.2 NVSHMEM 簡介

NVSHMEM 是 NVIDIA 開發的一種并行編程接口，基于 OpenSHMEM 標準，專為 GPU 集群提供高效且可擴展的通信。通過創建一個跨多個 GPU 內存的全局地址空間，實現細粒度的 GPU 發起的數據傳輸和同步操作，顯著減少了 CPU 的干預，從而降低同步開銷并提升性能。

NVSHMEM 通常被視為 MPI（Message-Passing Interface） 的替代方案，特別是在 GPU 集群通信中。與 MPI 不同，NVSHMEM 通過 GPU 發起操作減少了 CPU-GPU 同步開銷。雖然它可以與 MPI 結合使用（如在 Host 端通信中），但其核心優勢在于 GPU 間的直接通信。相比 NCCL，NVSHMEM 更專注于單邊通信，而 NCCL 則更適合集合通信，具體使用場景取決于應用程序的需求。

安裝指南和源代碼可以參考：NVSHMEM Installation Guide [3]

2.3 IBRC & IBGDA

InfiniBand GPUDirect Async（IBGDA）是 NVSHMEM 中的一種新的通信方法，構建在 GPUDirect Async 技術之上。IBGDA 在 NVSHMEM 2.6.0 中引入，并在 NVSHMEM 2.7.0 和 2.8.0 中得到顯著改進。它可以使 GPU 繞過 CPU 進行節點間 NVSHMEM 通信，而無需對現有應用程序進行修改，進而使得 NVSHMEM 應用程序的吞吐和擴展得到顯著改進。

如下圖所示，在引入 IBGDA 之前，NVSHMEM InfiniBand Reliable Connection (IBRC) 傳輸使用 CPU 上的代理線程來管理通信（PS：也就是數據傳輸不同用過主機內存，但是很多控制鏈路還在 CPU 上）。

與 IBRC 相比，IBGDA 利用 GPUDirect Async–Kernel-Initiated (GPUDirect Async–KI) ，使得 GPU SM 能夠直接與 NIC 交互。如下圖所示為其關鍵步驟，可以看出，IBGDA 將 CPU 上的通信控制鏈路移到 GPU 上，WQ 和 DBR 緩沖區也被移動到 GPU 內存。使用 IBGDA 時，GPU 和 NIC 直接交換通信所需信息，以提高 SM 訪問效率，同時通過 GPUDirect RDMA 保留 NIC 的訪問。

因此，IBGDA 非常適合控制鏈路開銷比較大的小消息的傳輸。?

三、NVSHMEM 關鍵概念

3.1 Processing Element（PE）

NVSHMEM 程序由多個進程組成，每個進程映射到一個 GPU 上，稱為一個PE。所有 PE 運行相同的程序，在啟動時通過 nvshmem_init 或擴展接口 nvshmemx_init_attr 集體初始化環境。

初始化過程中，每個 PE 都會分配到唯一的 ID（可以使用 nvshmem_my_pe() 獲取本 PE 的 ID），并且同步總的 PE 數目（可以使用 nvshmem_n_pes() 獲取）。假設有 2 臺 8 卡 H100 的機器，則啟動后會有 16 個 PE，對應的 mype 為 [0, 1, 2, …, 15]，對應的 npes 為 16。

PE 之間通過從 GPU 內存中的對稱堆（symmetric heap）分配的對稱內存進行通信和數據共享。此類內存需使用 CPU 端的 NVSHMEM 分配 API 進行分配。采用其他任何方法分配的內存均被視為分配 PE 的私有內存，其他 PE 無法訪問。NVSHMEM 使用 CUDA 的統一虛擬地址 (UVA) 和 CUDA IPC 機制映射這些對稱緩沖區，以實現 GPU 之間的高速訪問。在多 GPU 或分布式場景中，GPU 可直接發起通信請求，NVSHMEM 后端負責將這種 GPU 觸發的請求通過網絡傳輸到目標 PE。

需要注意，從 2.4.1 版開始，NVSHMEM 支持一個 GPU 上運行多個 PE 的情形（通過 CUDA 多進程服務 MPS 或時間共享方式），但這會對同步和集合通信的使用有額外限制。

3.2 Team 管理

NVSHMEM 支持在全局 PE 集合之外定義邏輯 Team（nvshmem_team_t），以在一部分 PE 上執行集合通信操作。默認情況下，NVSHMEM_TEAM_WORLD 表示包含所有 PE 的“全局 Team”，PE 在該 Team 中的編號即為 nvshmem_my_pe() 返回的值。

用戶可以通過“split”現有 Team 來創建新 Team（如 nvshmem_team_split_strided() 等），也可以使用擴展接口 nvshmemx_team_get_uniqueid 和 nvshmemx_team_init 任意指定新 Team 成員。后者需要選定一個 Root PE 來生成 Team 的唯一 ID，并將此 ID 分發給 Team 內所有 PE。創建 Team 時，每個 PE 提供在新 Team 中的編號（pe_idx_in_team），新 Team 成員按照該編號 0 到 N–1 重新編號。

3.3 Transport 機制

NVSHMEM 支持多種傳輸通道來實現 GPU 間通信。

• 對于節點內通信，NVSHMEM 主要利用 CUDA 級聯技術（比如 NVLink/PCIe 上的 P2P 訪問），使得一個節點上不同進程的 GPU 可以直接通過顯存共享或復制互相訪問。
• 對于節點間通信，NVSHMEM 默認使用 IBRC，并可選用其他協議如 UCX 或 Libfabric。具體選擇由環境變量 NVSHMEM_REMOTE_TRANSPORT 控制，允許設置為 "ibrc"（默認）、"ucx"、"libfabric"、"ibdevx" 等。

NVSHMEM 對 GPU 內存的處理邏輯包括：如果目標 PE 在本節點，優先使用 CUDA IPC 或 P2P 拷貝；如果在遠端，則調用網絡傳輸接口，通過已經注冊的 GPU 映射區完成遠程寫入/讀取。所有 GPU 通信都在 NVSHMEM 的對稱地址空間模型下進行：當用戶調用 RMA 操作時，內部先將對稱地址轉換為物理地址，然后通過上述通道發起數據傳輸。

默認情況下，NVSHMEM 借助 CUDA 的 VMM 來映射 PE 間 GPU 對稱地址。用戶可通過 NVSHMEM_DISABLE_CUDA_VMM 禁用這一機制。對稱緩沖區會通過 CUDA IPC 注冊到各進程地址空間，從而在同一節點上實現 ZeroCopy。此外，NVSHMEM 通過 Mellanox OFED、nv_peer_mem 驅動和 GDRCopy 庫支持 GPU 直連 RDMA，將 GPU 顯存注冊到網卡，可直接發起 RDMA 操作。也可以使用環境變量  NVSHMEM_DISABLE_GDRCOPY 來禁止在 IBRC 下使用 GDRCopy。

對于 GPU 和 NIC 的綁定，NVSHMEM 提供自動映射和手動配置的方式。默認情況下，各 PE 會分配到“最近的” NIC，但可通過設置 NVSHMEM_ENABLE_NIC_PE_MAPPING=1 使系統按照輪詢或用戶指定的方式綁定。進一步可用 NVSHMEM_HCA_LIST 和 NVSHMEM_HCA_PE_MAPPING 明確指定使用哪些 HCA 設備 Port，以及每個 Port 分配給哪些 PE。所有這些機制確保 GPU 的對稱內存可以高效地通過 NIC 通信。

3.4 集合通信

在 NVSHMEM 程序中，PE 通過兩種方式進行通信：

• 一種是點對點通信，需明確指定目標 PE 的編號。
• 另一種是集合通信，它們作用于一組 PE 之上。

基于 Team 的集合通信，通過 Team 句柄參數確定參與通信的 PE，并利用 Team 對象封裝的資源進行操作。

如果未指定 Team 參數，則默認作用于所有 PE。

3.5 Bootstrap

NVSHMEM 的 Bootstrap 模塊負責在多進程環境中引導通信和 PE 管理。支持多種啟動方式：默認通過進程管理器的 PMI 接口交互式啟動，也可以直接利用現有的 MPI 通信域、OpenSHMEM 環境或者 Plugin 模式。可以通過 NVSHMEM_BOOTSTRAP 環境變量來指定：

• PMI：可以選擇不同的版本，比如 PMI-1、PMI-2 或 PMIx。
• MPI：用戶可用 nvshmemx_init_attr 指定 NVSHMEMX_INIT_WITH_MPI_COMM 標志以傳入一個已有的 MPI_Comm 作為全局 Team。
• OpenSHMEM：同樣可以用 NVSHMEMX_INIT_WITH_SHMEM 指示在 OpenSHMEM 程序內啟動。
• UID：在沒有 MPI/SHMEM 的場景下，還可以采用網絡唯一 ID (UID) 方式，調用 nvshmemx_get_uniqueid 在一個 PE 上生成 UID，并通過用戶定義的機制（比如環境變量 NVSHMEM_BOOTSTRAP_UID_SESSION_ID）分發給各 PE，然后各 PE 在 nvshmemx_set_attr_uniqueid_args 中設置該 UID、自己的 PE 編號和總 PE 數，再集體調用 nvshmemx_init_attr。這種 UID 模式利用 TCP 套接字自動交換初始化信息，無需依賴 MPI。也可以像 NCCL 一樣配合 NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME 和 NVSHMEM_BOOTSTRAP_UID_SOCK_FAMILY 使用。

四、DeepEP 初始化

4.1 概覽

如下圖所示，在 DeepEP 中會調用 NVSHMEM 的如下幾個接口完成初始化（后續會詳細解釋這些函數的作用）：

• nvshmemx_set_attr_uniqueid_args()。
• nvshmemx_init_attr()。
• nvshmem_team_split_strided()。
• nvshmem_barrier_all()。
• nvshmem_my_pe()。

而 DeepEP 中在 Buffer 的 sync 中會調用該 init() 接口，并且調用之前會完成 rank 和 num_ranks 的計算。這里會針對上述的“高吞吐”和“低時延”模式進行不同的處理。假設 2 臺機器，各 8 個 GPU，則與 PyTorch 對應的 rank 為 [0, 15]，num_ranks 為 16。

• 高吞吐模式：將上述 rank 轉為 rdma_rank（[0, 1]） 和 nvl_rank；num_ranks 轉為 num_rdma_ranks（為 2，其中 NUM_MAX_NVL_PEERS  為 8）。并且使用 rdma_rank 和 num_rdma_ranks 初始化。（PS：節點內使用 NVLink）
• 低時延模式：直接使用 rank 和 num_ranks。

PS：DeepEP 針對高吞吐模式的特殊邏輯導致其與 NVSHMEM 的拓撲映射方法 NVSHMEM_HCA_PE_MAPPING 不兼容，后續會具體介紹。

4.2 nvshmemx_set_attr_uniqueid_args

如下圖所示（nvshemem_src/src/host/init/init.cu），nvshmemx_set_attr_uniqueid_args 其實就是設置了 id、myrank 和 nranks：

PS：如果是高吞吐模式，這里的 myrank 和 nranks 已經不再等價于 PyTorch 里的 rank 和 num_ranks。

4.3 nvshmemx_init_attr

如下圖所示（nvshemem_src/src/include/host/nvshmemx_api.h），nvshmemx_init_attr 實際是調用了 nvshmemi_init_thread：

4.4 nvshmem_team_split_strided

如果是 low_latency_mode，并且 num_ranks > NUM_MAX_NVL_PEERS(8)，也就是多機時才會執行 nvshmem_team_split_strided，將 NVSHMEM 的 team 劃分為多組 sub-RDMA team，這樣可以實現每組內部的高效通信，同時保證整體的可擴展性和性能。

4.5 nvshmem_barrier_all

如下圖所示（nvshmem_src/src/host/coll/barrier/barrier.cpp），使用 nvshmem_barrier_all 進行全局同步。

4.6 nvshmem_my_pe

如下圖所示（nvshmem_src/src/host/init/query_host.cpp），用 nvshmem_my_pe 返回當前的 pe： 

五、nvshmemi_init_thread

5.1 概覽

nvshmemi_init_thread（位于 nvshmem_src\src\device\init\init_device.cu） 是 NVSHMEM 初始化流程中的關鍵函數，主要作用是初始化運行環境，確保后續在 GPU 上能夠安全、高效地進行 NVSHMEM 通信和同步操作。其主要功能和步驟如下：

nvshmemid_hostlib_init_attr：負責完成 NVSHMEM Host 端初始化流程的主控函數。

nvshmemid_init_status：Host 端初始化狀態檢查，如果狀態大于 NVSHMEM_STATUS_IS_BOOTSTRAPPED，說明 Host 端初始化已完成，可以安全地初始化 Device 端狀態。

_nvshmemi_init_device_only_state：進一步設置 Device 端的集合通信等內部狀態。主要是調用 nvshmemi_setup_collective_launch。

cudaGetDevice：調用 cudaGetDevice 獲取當前 GPU 的 device id，存入 

nvshmemi_device_only_state.cuda_device_id，便于后續 Device 端操作定位到正確的 GPU。

5.2 nvshmemid_hostlib_init_attr

包括版本兼容性檢查、bootstrap 啟動、全局和本地狀態初始化、設備狀態注冊、以及多種初始化的管理。包括以下幾個關鍵步驟。

5.2.1 IBGDA 狀態初始化

如果啟用 IBGDA，則提前設置一下狀態信息：

如下圖所示（nvshmem_src/src/host/init/init.cu）：

5.2.2 nvshmemi_options_init 環境變量初始化

如下圖所示：該函數的作用是批量初始化 NVSHMEM 的所有環境變量配置項，并將結果寫入 options 結構體。

5.2.3 nvshmemi_bootstrap_preinit 預初始化

如下圖所示（位于 nvshmem_src/src/host/init/init.cu），實際調用 bootstrap_preinit：

bootstrap_preinit 核心就是通過 dlopen 提前加載動態庫，當然，這里主要是 UID 模式的動態庫：

5.2.4 nvshmemi_bootstrap 初始化

如下圖所示（位于 nvshmem_src/src/host/init/init.cu），nvshmemi_bootstrap 負責完成 NVSHMEM 啟動階段的“進程組通信環境”初始化。通過底層 bootstrap 機制（如 MPI、UID、SHMEM 等）建立全局進程組信息，確定每個 PE 的全局編號、節點內編號、節點內 PE 數量等，并進行一致性檢查，為后續通信和資源分配打下基礎。

• bootstrap_set_bootattr：根據 flags 和 attr 構造 bootstrap 所需的屬性結構體，會針對 MPI、SHMEM、UID、PMI、PLUGIN 等不同 mode 對應處理。
• bootstrap_init：傳入 flags、屬性和 nvshmemi_boot_handle，由底層 Plugin/機制完成通信環境初始化（根據 MPI、SHMEM、UID、PMI、PLUGIN 等不同 mode 調用不同的初始化 bootstrap_loader_init），填充 nvshmemi_boot_handle（包括 allgather、barrier、pg_rank、pg_size 等回調和參數）。
• 計算全局和節點內編號：

通過 nvshmemi_boot_handle.pg_rank 和 pg_size 獲取本進程全局編號和總進程數。

通過 getHostHash 獲取本節點唯一標識（host hash）。

使用 allgather 收集所有進程的 host hash，統計每個節點上的 PE 數量（npes_node）和本 PE 在節點內的編號（mype_node）。

將 mype_node 和 npes_node 寫入 nvshmemi_boot_handle。

• 檢查所有節點 PE 數一致性：遍歷所有 host hash，確保每個節點上的 PE 數量一致，否則報錯退出。
• 設置 PE 分布類型：檢查 PE 分布是否為 block 或 round robin，設置 nvshmemi_pe_dist，用于后續通信優化。

5.2.5 nvshmemi_try_common_init 初始化

實際上是調用 nvshmemi_common_init，這是最核心的初始化操作，后面會詳細介紹：

5.3 nvshmemi_setup_collective_launch

nvshmemi_setup_collective_launch 是 NVSHMEM  Device 端集合通信相關的初始化函數。它的主要作用是為后續在 GPU 上安全、高效地發起通信 kernel 啟動做準備，確保相關 CUDA 資源和屬性已正確配置。主要包括以下幾個部分：

• 檢查當前 GPU 的 SM 數量。
• 檢查當前 GPU 是否支持 Cooperative Launch。
• 檢查 CUDA Stream 優先級范圍。
• 創建一個最高優先級的非阻塞 CUDA Stream。
• 創建兩個 CUDA Event，用于同步 Collective  Launch 的起止。

六、nvshmemi_common_init

6.1 概覽

nvshmemi_common_init 的主要作用是完成 NVSHMEM 運行時的核心初始化工作。負責初始化和配置 NVSHMEM 的所有關鍵組件，使得后續的通信和內存管理功能可以正常工作。其主要任務包括：

• CUDA 相關初始化：加載 CUDA 符號表、初始化 CUDA 庫、查詢 CUDA 驅動版本，決定對稱堆類型（如 SYSMEM/VIDMEM），并獲取 CUDA 上下文和流優先級等。
• Rail 優化設置：如果啟動軌道優化，則檢查是否滿足軌道優化的支持，并完成相關設置。
• 對稱堆初始化：根據配置和硬件能力，初始化對稱堆，為后續的通信和內存操作分配統一的內存空間。
• 設備狀態注冊與更新：注冊和更新設備狀態指針，確保主機和設備之間的狀態同步。
• 通信與拓撲初始化：初始化 transport，構建 transport 映射表，建立與其他 PE 的連接。
• CUDA 句柄和 Event 初始化：為每個 peer 分配 CUDA stream 和 event，設置相關句柄。
• 集合通信和 Team 初始化：初始化 CPU 端的集合通信，設置 NVSHMEM 的 Team 結構，支持多種 Team 劃分和同步。
• 每個 GPU 的多 PE（MPG）支持：檢測和配置多 PE 共享同一 GPU 的場景，包括 MPS 相關的共享內存和事件管理。
• 性能與兼容性檢查：檢測如 NVLS、64-bit stream memops 等高級特性支持，并給出性能建議或警告。
• 全局同步：在關鍵階段通過 barrier 保證所有 PE 的同步，確保初始化過程一致。
• 最終狀態標記：標記設備狀態為已初始化，確保后續 API 調用安全。

6.2 nvshmemi_coll_common_cpu_init

初始化 NVSHMEM Host 端的集合通信相關環境參數，并根據環境和配置決定是否啟用 NCCL 作為底層通信庫，如下圖所示（位于 nvshmem_src\src\host\coll\cpu_coll.cpp）。

• nvshmemi_coll_common_cpu_read_env：將環境變量和配置項寫入全局結構體（如 barrier、reduce、broadcast 等算法和閾值參數）。
• NCCL 相關支持邏輯（編譯時需要加上 NVSHMEM_USE_NCCL）

默認嘗試啟用 NCCL，如果設置了 DISABLE_NCCL 選項，則禁用 NCCL，直接返回。

NCCL 不支持 MPG，如果設置 MPG 要禁用 NCCL。

加載 NCCL 動態庫（“libnccl.so.2”）。

檢查 NCCL 版本和兼容性。

通過 dlsym 加載 NCCL 所需的符號，并填充到 nccl_ftable。

6.3 nvshmemi_transport_init

初始化 NVSHMEM 支持的所有 Host 端 Transport 模塊，包括 P2P、本地和遠程網絡通信插件（如 IBRC、UCX、Libfabric、IBDEVX、IBGDA 等），為后續 PE 間數據傳輸和同步提供底層支撐（位于 nvshmem_src\src\host\transport\transport.cpp）。包括如下關鍵步驟：

• 分配 transport 結構體數組：

若 state->transports 為空，則分配一個大小為 NVSHMEM_TRANSPORT_COUNT（6） 的數組，每個元素對應一種 Transport。

• 初始化 P2P 通信（僅用于同一節點不同 GPU）：
• 如果沒有禁用 P2P，則初始化本地內存緩存，并調用 nvshmemt_p2p_init 初始化 P2P 通信。
• 初始化成功則填充相關字段（如 boot_handle、heap_base、cap、index、granularity 等），并將其加入 transports 列表。
• 初始化遠程通信插件（用于不同節點）：
• 支持多種插件（IBRC、UCX、IBDEVX、Libfabric），通過宏控制編譯。
• 判斷環境變量 REMOTE_TRANSPORT，決定是否跳過某個插件。
• 若選擇某插件，拼接動態庫（如 nvshmem_transport_ucx.so.1），然后進行相應初始化。
• 動態加載并初始化遠程通信插件：
• 使用 dlopen 加載對應的動態庫，并使用 dlsym 獲取初始化函數指針（nvshmemt_init）。
• 初始化本地緩存，調用插件的初始化函數，填充 transport 結構體相關字段（heap_base、cap、index、my_pe、n_pes、cache_handle 等）。
• IBGDA 特殊支持（需要打開編譯選項 NVSHMEM_IBRC_SUPPORT）：
• 如果環境變量啟用 IBGDA，則單獨加載并初始化 IBGDA 插件，流程與上面類似。
• 初始化成功后，調用 nvshmemi_ibgda_get_device_state 獲取設備狀態，并設置全局標志。
• 檢查至少有一個 transport 初始化成功。
• 記錄已初始化的 transport 數量：
• 將成功初始化的 transport 數量寫入 state->num_initialized_transports。

PS：IBRC、UCX、IBDEVX、Libfabric 只會選擇其中的一個，主要是因為順序檢查：代碼會按照 IBRC -> UCX -> IBDEVX -> LIBFABRIC 的順序，逐一檢查環境變量，只要檢查到一個就會跳轉到 transport_init，從而跳過后續的幾個。

6.4 nvshmemi_build_transport_map

為每個 PE 建立一張“可達性”與“可用 Transport”的映射表。檢測本地 PE 與所有其他 PE 之間，哪些 Transport 能夠訪問對方，并將結果記錄下來，最終形成全局的 Transport 映射表（state->transport_map），為后續高效通信做準備（位于 nvshmem_src/nvshmem/src/host/topo/topo.cpp）。主要流程包括：

分配映射表內存：

• state->transport_map：大小為 npes × npes，存儲所有 PE 間的 Transport 可達性信息。
• local_map：大小為 npes，存儲本地 PE 到所有其他 PE 的可達性信息。
• state->transport_bitmap：用來記錄本 PE 能用到的所有 Transport 類型（每一位代表一種 Transport）。

檢查每個 PE 的可達性：

• 外層循環遍歷所有 PE（i），即本地 PE 需要與哪些遠端 PE 通信。
• 內層循環遍歷所有已初始化的 transport（j），判斷本地通過 transport j 能否訪問到 PE i。如果能訪問，則在 local_map[i] 的相應 bit 位置標記，并更新 state->transport_bitmap。同時記錄每個 transport 對每個 PE 的能力（state->transports[j]->cap[i]）。

匯總所有 PE 的可達性信息：通過 allgather，將每個 PE 的 local_map 匯總到全局的 state->transport_map，這樣每個進程都能知道所有 PE 之間的可達性和可用 transport。

6.5 nvshmemi_setup_connections

6.5.1 概覽

為每個可用的 transport 在本 PE 上選擇合適的底層設備（如 NIC 等），并建立與其他 PE 的通信端點（endpoint）連接。這是 NVSHMEM 初始化流程中，真正“連通”各個進程間網絡通信的關鍵步驟（位于 nvshmem_src/src/host/transport/transport.cpp）。關鍵步驟如下圖所示：

遍歷所有已初始化的 Transport：

• 只處理 state->transport_bitmap 標記為可用的 Transport。

計算每個 PE 可用的設備數：

• tcurr->n_devices 是當前 Host 上的 NIC 數，state->npes_node 是當前 Host 上的 PE 數。
• 平均分配，并且保證每個 PE 至少有 1 個 NIC。假設 8 個 GPU，2 個 NIC，則會有 4 個 GPU 共享 1 個 NIC。

選擇本 PE 要使用的 NIC：

• 如果 Transport 只有一個 NIC，則直接使用。
• 如果用戶打開了 ENABLE_NIC_PE_MAPPING 環境變量，則使用輪詢分配方式。
• 否則，調用 nvshmemi_get_devices_by_distance，根據 PCIe 拓撲距離等信息，智能選擇最優的設備分配給本 PE。

檢查設備分配有效性：

建立 Endpoint 連接：

調用 transport 的 connect_endpoints 回調，傳入選中的設備列表，建立與其他 PE 的通信端點。

之后進行 barrier，同步所有進程，確保連接建立一致。

更新設備狀態。

6.5.2 拓撲映射

當用戶想要通過 ENABLE_NIC_PE_MAPPING 進行 PE（實際也就是 GPU）與 NIC 的映射時，可以通過以下兩個環境變量實現：

NVSHMEM_HCA_LIST：直接提供 NIC name 和 port 索引的列表（PS：內部會進行排序）。

• “^mlx5_0”：排除 mlx5_0 NIC。
• “mlx5_1:1,mlx5_2:2”：如果有 8 個 GPU，則 4 個 GPU 對應 mlx5_1 的 Port 1，4 個 GPU 對應 mlx5_2 的 Port 2。

NVSHMEM_HCA_PE_MAPPING：和上述的 NVSHMEM_HCA_LIST 類似，只不過可以進一步設置 PE 的數量（PS：內部也會進行排序）。

• “mlx5_0:1:3,mlx5_0:2:5”：還是 8 個 GPU，則 3 個 GPU 對應 mlx5_1 的 Port 1，5 個 GPU 對應 mlx5_0 的 Port 2。

我們前面提到，DeepEP 的高吞吐模式與 ENABLE_NIC_PE_MAPPING 不兼容也正是這里的問題。

假設還是 8 個 GPU，設置 NVSHMEM_HCA_PE_MAPPING 為 “mlx5_0:1:2,mlx5_0:2:2,mlx5_1:1:2,mlx5_1:2:2”，預期是 PE0 和 PE1 使用 mlx5_0 的 Port 1， PE6 和 PE7 使用 mlx5_1 的 Port 2。

由于高吞吐模式時這里的 mype_node 實際為 rdma_rank，一個節點上所有 PE 的 mype_node 相同，實際都是 0，導致這里實際上這里所有 PE 選擇了相同的 NIC，沒有達到 Mapping 的目的，反而導致了 NIC 熱點。

6.5.3 智能拓撲感知

主要是調用 nvshmemi_get_devices_by_distance 為本節點上的每個 PE 分配最優的 NIC，以實現高效的 GPU-NIC 通信。分配策略基于 PCIe 拓撲距離，優先選擇距離 GPU 最近的 NIC，并在多 GPU/多 NIC 情況下盡量負載均衡（位于 nvshmem_src/nvshmem/src/host/topo/topo.cpp）。整體思路與 NCCL 中類似，這里主要包括幾個步驟：

收集 GPU 和 NIC 路徑：

• 獲取本 PE 的 GPU PCIe Bus ID。
• 通過 allgather 收集所有 PE 的 GPU Bus ID。
• 遍歷所有 PE，篩選出本節點的 PE（hostHash 相同），并獲取其 GPU 的 PCIe 路徑，填入 cuda_device_paths。
• 填充所有 NIC 的 PCIe 路徑。

計算所有 GPU-NIC 距離：

• 針對每個 PE 的 GPU，都會通過 get_pci_distance 計算每個 NIC 與 GPU 的距離，并構造一個三元組 {pe_id, dev_id, distance}，插入 pe_dev_pairs，按距離升序排列（距離越小越優）。
• PCIe 距離的優先級為 PIX > PXB > PHB > NODE > SYS，和 NCCL 一樣。

第一輪分配（優先分配最近的 NIC）：

• 遍歷 pe_dev_pairs（距離優先），為每個本地 PE 依次分配最優的 NIC。
• 如果當前距離比已分配的更優，則分配該 NIC，并記錄距離和使用計數。
• 如果遇到更差的距離，后續的 NIC 都不再考慮，標記為無更優分配（-2）。
• 直到所有本地 PE 的分配完成。

第二輪分配（負載均衡）：

• 檢查是否有 NIC 被多個 PE 分配（nic_density >= 2）。
• 嘗試為這些 PE 找到同等距離但負載更低的 NIC，減少單個 NIC 的壓力。
• 只在不會降低距離優先級的前提下進行重新分配。

輸出本 PE 的分配結果：

• 將本 PE 分配到的 NIC 索引寫入 device_arr。
• 統計本 PE 實際分配到的 NIC 數量。

七、參考鏈接

1. ??https://docs.nvidia.com/nvshmem/api/index.html??
2. ??https://github.com/deepseek-ai/DeepEP??
3. ??https://docs.nvidia.com/nvshmem/release-notes-install-guide/install-guide/abstract.html??

本文轉載自???AI閑談?????，作者：AI閑談