?一、背景

在此前的內(nèi)容中，筆者曾介紹過 DeepSeek 的 DeepEP、字節(jié)跳動(dòng)的 Flux 和 Tilelink 等系統(tǒng)，這些系統(tǒng)在底層通信實(shí)現(xiàn)中均依賴于 NVIDIA 的 NVSHMEM 庫。事實(shí)上，字節(jié)跳動(dòng)后續(xù)的諸如 Comet、Triton-distributed，以及其他針對(duì)細(xì)粒度計(jì)算與通信重疊（Overlap）優(yōu)化的工作，也都廣泛使用了 NVSHMEM。

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

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

二、引言

2.1 DeepEP

DeepEP 是 DeepSeek 開源的專為 MoE 和專家并行（Expert Parallelism, EP）設(shè)計(jì)的通信庫。提供了一系列優(yōu)化的通信 Kernel，實(shí)現(xiàn)了以下能力：

• 高度優(yōu)化的 All2All 通信。
• 同時(shí)支持不同的通信類型：

節(jié)點(diǎn)內(nèi)（intra-node）：使用 NVLink + NVSwitch 通信。

節(jié)點(diǎn)間（inter-node）：使用 RDMA 通信。

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

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

2.2 NVSHMEM 簡介

NVSHMEM 是 NVIDIA 開發(fā)的一種并行編程接口，基于 OpenSHMEM 標(biāo)準(zhǔn)，專為 GPU 集群提供高效且可擴(kuò)展的通信。通過創(chuàng)建一個(gè)跨多個(gè) GPU 內(nèi)存的全局地址空間，實(shí)現(xiàn)細(xì)粒度的 GPU 發(fā)起的數(shù)據(jù)傳輸和同步操作，顯著減少了 CPU 的干預(yù)，從而降低同步開銷并提升性能。

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

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

2.3 IBRC & IBGDA

InfiniBand GPUDirect Async（IBGDA）是 NVSHMEM 中的一種新的通信方法，構(gòu)建在 GPUDirect Async 技術(shù)之上。IBGDA 在 NVSHMEM 2.6.0 中引入，并在 NVSHMEM 2.7.0 和 2.8.0 中得到顯著改進(jìn)。它可以使 GPU 繞過 CPU 進(jìn)行節(jié)點(diǎn)間 NVSHMEM 通信，而無需對(duì)現(xiàn)有應(yīng)用程序進(jìn)行修改，進(jìn)而使得 NVSHMEM 應(yīng)用程序的吞吐和擴(kuò)展得到顯著改進(jìn)。

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

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

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

三、NVSHMEM 關(guān)鍵概念

3.1 Processing Element（PE）

NVSHMEM 程序由多個(gè)進(jìn)程組成，每個(gè)進(jìn)程映射到一個(gè) GPU 上，稱為一個(gè)PE。所有 PE 運(yùn)行相同的程序，在啟動(dòng)時(shí)通過 nvshmem_init 或擴(kuò)展接口 nvshmemx_init_attr 集體初始化環(huán)境。

初始化過程中，每個(gè) PE 都會(huì)分配到唯一的 ID（可以使用 nvshmem_my_pe() 獲取本 PE 的 ID），并且同步總的 PE 數(shù)目（可以使用 nvshmem_n_pes() 獲?。?。假設(shè)有 2 臺(tái) 8 卡 H100 的機(jī)器，則啟動(dòng)后會(huì)有 16 個(gè) PE，對(duì)應(yīng)的 mype 為 [0, 1, 2, …, 15]，對(duì)應(yīng)的 npes 為 16。

PE 之間通過從 GPU 內(nèi)存中的對(duì)稱堆（symmetric heap）分配的對(duì)稱內(nèi)存進(jìn)行通信和數(shù)據(jù)共享。此類內(nèi)存需使用 CPU 端的 NVSHMEM 分配 API 進(jìn)行分配。采用其他任何方法分配的內(nèi)存均被視為分配 PE 的私有內(nèi)存，其他 PE 無法訪問。NVSHMEM 使用 CUDA 的統(tǒng)一虛擬地址 (UVA) 和 CUDA IPC 機(jī)制映射這些對(duì)稱緩沖區(qū)，以實(shí)現(xiàn) GPU 之間的高速訪問。在多 GPU 或分布式場景中，GPU 可直接發(fā)起通信請(qǐng)求，NVSHMEM 后端負(fù)責(zé)將這種 GPU 觸發(fā)的請(qǐng)求通過網(wǎng)絡(luò)傳輸?shù)侥繕?biāo) PE。

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

3.2 Team 管理

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

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

3.3 Transport 機(jī)制

NVSHMEM 支持多種傳輸通道來實(shí)現(xiàn) GPU 間通信。

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

NVSHMEM 對(duì) GPU 內(nèi)存的處理邏輯包括：如果目標(biāo) PE 在本節(jié)點(diǎn)，優(yōu)先使用 CUDA IPC 或 P2P 拷貝；如果在遠(yuǎn)端，則調(diào)用網(wǎng)絡(luò)傳輸接口，通過已經(jīng)注冊(cè)的 GPU 映射區(qū)完成遠(yuǎn)程寫入/讀取。所有 GPU 通信都在 NVSHMEM 的對(duì)稱地址空間模型下進(jìn)行：當(dāng)用戶調(diào)用 RMA 操作時(shí)，內(nèi)部先將對(duì)稱地址轉(zhuǎn)換為物理地址，然后通過上述通道發(fā)起數(shù)據(jù)傳輸。

默認(rèn)情況下，NVSHMEM 借助 CUDA 的 VMM 來映射 PE 間 GPU 對(duì)稱地址。用戶可通過 NVSHMEM_DISABLE_CUDA_VMM 禁用這一機(jī)制。對(duì)稱緩沖區(qū)會(huì)通過 CUDA IPC 注冊(cè)到各進(jìn)程地址空間，從而在同一節(jié)點(diǎn)上實(shí)現(xiàn) ZeroCopy。此外，NVSHMEM 通過 Mellanox OFED、nv_peer_mem 驅(qū)動(dòng)和 GDRCopy 庫支持 GPU 直連 RDMA，將 GPU 顯存注冊(cè)到網(wǎng)卡，可直接發(fā)起 RDMA 操作。也可以使用環(huán)境變量  NVSHMEM_DISABLE_GDRCOPY 來禁止在 IBRC 下使用 GDRCopy。

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

3.4 集合通信

在 NVSHMEM 程序中，PE 通過兩種方式進(jìn)行通信：

• 一種是點(diǎn)對(duì)點(diǎn)通信，需明確指定目標(biāo) PE 的編號(hào)。
• 另一種是集合通信，它們作用于一組 PE 之上。

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

如果未指定 Team 參數(shù)，則默認(rèn)作用于所有 PE。

3.5 Bootstrap

NVSHMEM 的 Bootstrap 模塊負(fù)責(zé)在多進(jìn)程環(huán)境中引導(dǎo)通信和 PE 管理。支持多種啟動(dòng)方式：默認(rèn)通過進(jìn)程管理器的 PMI 接口交互式啟動(dòng)，也可以直接利用現(xiàn)有的 MPI 通信域、OpenSHMEM 環(huán)境或者 Plugin 模式?？梢酝ㄟ^ NVSHMEM_BOOTSTRAP 環(huán)境變量來指定：

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

四、DeepEP 初始化

4.1 概覽

如下圖所示，在 DeepEP 中會(huì)調(diào)用 NVSHMEM 的如下幾個(gè)接口完成初始化（后續(xù)會(huì)詳細(xì)解釋這些函數(shù)的作用）：

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

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

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

PS：DeepEP 針對(duì)高吞吐模式的特殊邏輯導(dǎo)致其與 NVSHMEM 的拓?fù)溆成浞椒?nbsp;NVSHMEM_HCA_PE_MAPPING 不兼容，后續(xù)會(huì)具體介紹。

4.2 nvshmemx_set_attr_uniqueid_args

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

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

4.3 nvshmemx_init_attr

如下圖所示（nvshemem_src/src/include/host/nvshmemx_api.h），nvshmemx_init_attr 實(shí)際是調(diào)用了 nvshmemi_init_thread：

4.4 nvshmem_team_split_strided

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

4.5 nvshmem_barrier_all

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

4.6 nvshmem_my_pe

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

五、nvshmemi_init_thread

5.1 概覽

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

nvshmemid_hostlib_init_attr：負(fù)責(zé)完成 NVSHMEM Host 端初始化流程的主控函數(shù)。

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

_nvshmemi_init_device_only_state：進(jìn)一步設(shè)置 Device 端的集合通信等內(nèi)部狀態(tài)。主要是調(diào)用 nvshmemi_setup_collective_launch。

cudaGetDevice：調(diào)用 cudaGetDevice 獲取當(dāng)前 GPU 的 device id，存入 

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

5.2 nvshmemid_hostlib_init_attr

包括版本兼容性檢查、bootstrap 啟動(dòng)、全局和本地狀態(tài)初始化、設(shè)備狀態(tài)注冊(cè)、以及多種初始化的管理。包括以下幾個(gè)關(guān)鍵步驟。

5.2.1 IBGDA 狀態(tài)初始化

如果啟用 IBGDA，則提前設(shè)置一下狀態(tài)信息：

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

5.2.2 nvshmemi_options_init 環(huán)境變量初始化

如下圖所示：該函數(shù)的作用是批量初始化 NVSHMEM 的所有環(huán)境變量配置項(xiàng)，并將結(jié)果寫入 options 結(jié)構(gòu)體。

5.2.3 nvshmemi_bootstrap_preinit 預(yù)初始化

如下圖所示（位于 nvshmem_src/src/host/init/init.cu），實(shí)際調(diào)用 bootstrap_preinit：

bootstrap_preinit 核心就是通過 dlopen 提前加載動(dòng)態(tài)庫，當(dāng)然，這里主要是 UID 模式的動(dòng)態(tài)庫：

5.2.4 nvshmemi_bootstrap 初始化

如下圖所示（位于 nvshmem_src/src/host/init/init.cu），nvshmemi_bootstrap 負(fù)責(zé)完成 NVSHMEM 啟動(dòng)階段的“進(jìn)程組通信環(huán)境”初始化。通過底層 bootstrap 機(jī)制（如 MPI、UID、SHMEM 等）建立全局進(jìn)程組信息，確定每個(gè) PE 的全局編號(hào)、節(jié)點(diǎn)內(nèi)編號(hào)、節(jié)點(diǎn)內(nèi) PE 數(shù)量等，并進(jìn)行一致性檢查，為后續(xù)通信和資源分配打下基礎(chǔ)。

• bootstrap_set_bootattr：根據(jù) flags 和 attr 構(gòu)造 bootstrap 所需的屬性結(jié)構(gòu)體，會(huì)針對(duì) MPI、SHMEM、UID、PMI、PLUGIN 等不同 mode 對(duì)應(yīng)處理。
• bootstrap_init：傳入 flags、屬性和 nvshmemi_boot_handle，由底層 Plugin/機(jī)制完成通信環(huán)境初始化（根據(jù) MPI、SHMEM、UID、PMI、PLUGIN 等不同 mode 調(diào)用不同的初始化 bootstrap_loader_init），填充 nvshmemi_boot_handle（包括 allgather、barrier、pg_rank、pg_size 等回調(diào)和參數(shù)）。
• 計(jì)算全局和節(jié)點(diǎn)內(nèi)編號(hào)：

通過 nvshmemi_boot_handle.pg_rank 和 pg_size 獲取本進(jìn)程全局編號(hào)和總進(jìn)程數(shù)。

通過 getHostHash 獲取本節(jié)點(diǎn)唯一標(biāo)識(shí)（host hash）。

使用 allgather 收集所有進(jìn)程的 host hash，統(tǒng)計(jì)每個(gè)節(jié)點(diǎn)上的 PE 數(shù)量（npes_node）和本 PE 在節(jié)點(diǎn)內(nèi)的編號(hào)（mype_node）。

將 mype_node 和 npes_node 寫入 nvshmemi_boot_handle。

• 檢查所有節(jié)點(diǎn) PE 數(shù)一致性：遍歷所有 host hash，確保每個(gè)節(jié)點(diǎn)上的 PE 數(shù)量一致，否則報(bào)錯(cuò)退出。
• 設(shè)置 PE 分布類型：檢查 PE 分布是否為 block 或 round robin，設(shè)置 nvshmemi_pe_dist，用于后續(xù)通信優(yōu)化。

5.2.5 nvshmemi_try_common_init 初始化

實(shí)際上是調(diào)用 nvshmemi_common_init，這是最核心的初始化操作，后面會(huì)詳細(xì)介紹：

5.3 nvshmemi_setup_collective_launch

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

• 檢查當(dāng)前 GPU 的 SM 數(shù)量。
• 檢查當(dāng)前 GPU 是否支持 Cooperative Launch。
• 檢查 CUDA Stream 優(yōu)先級(jí)范圍。
• 創(chuàng)建一個(gè)最高優(yōu)先級(jí)的非阻塞 CUDA Stream。
• 創(chuàng)建兩個(gè) CUDA Event，用于同步 Collective  Launch 的起止。

六、nvshmemi_common_init

6.1 概覽

nvshmemi_common_init 的主要作用是完成 NVSHMEM 運(yùn)行時(shí)的核心初始化工作。負(fù)責(zé)初始化和配置 NVSHMEM 的所有關(guān)鍵組件，使得后續(xù)的通信和內(nèi)存管理功能可以正常工作。其主要任務(wù)包括：

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

6.2 nvshmemi_coll_common_cpu_init

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

• nvshmemi_coll_common_cpu_read_env：將環(huán)境變量和配置項(xiàng)寫入全局結(jié)構(gòu)體（如 barrier、reduce、broadcast 等算法和閾值參數(shù)）。
• NCCL 相關(guān)支持邏輯（編譯時(shí)需要加上 NVSHMEM_USE_NCCL）

默認(rèn)嘗試啟用 NCCL，如果設(shè)置了 DISABLE_NCCL 選項(xiàng)，則禁用 NCCL，直接返回。

NCCL 不支持 MPG，如果設(shè)置 MPG 要禁用 NCCL。

加載 NCCL 動(dòng)態(tài)庫（“l(fā)ibnccl.so.2”）。

檢查 NCCL 版本和兼容性。

通過 dlsym 加載 NCCL 所需的符號(hào)，并填充到 nccl_ftable。

6.3 nvshmemi_transport_init

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

• 分配 transport 結(jié)構(gòu)體數(shù)組：

若 state->transports 為空，則分配一個(gè)大小為 NVSHMEM_TRANSPORT_COUNT（6） 的數(shù)組，每個(gè)元素對(duì)應(yīng)一種 Transport。

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

PS：IBRC、UCX、IBDEVX、Libfabric 只會(huì)選擇其中的一個(gè)，主要是因?yàn)轫樞驒z查：代碼會(huì)按照 IBRC -> UCX -> IBDEVX -> LIBFABRIC 的順序，逐一檢查環(huán)境變量，只要檢查到一個(gè)就會(huì)跳轉(zhuǎn)到 transport_init，從而跳過后續(xù)的幾個(gè)。

6.4 nvshmemi_build_transport_map

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

分配映射表內(nèi)存：

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

檢查每個(gè) PE 的可達(dá)性：

• 外層循環(huán)遍歷所有 PE（i），即本地 PE 需要與哪些遠(yuǎn)端 PE 通信。
• 內(nèi)層循環(huán)遍歷所有已初始化的 transport（j），判斷本地通過 transport j 能否訪問到 PE i。如果能訪問，則在 local_map[i] 的相應(yīng) bit 位置標(biāo)記，并更新 state->transport_bitmap。同時(shí)記錄每個(gè) transport 對(duì)每個(gè) PE 的能力（state->transports[j]->cap[i]）。

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

6.5 nvshmemi_setup_connections

6.5.1 概覽

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

遍歷所有已初始化的 Transport：

• 只處理 state->transport_bitmap 標(biāo)記為可用的 Transport。

計(jì)算每個(gè) PE 可用的設(shè)備數(shù)：

• tcurr->n_devices 是當(dāng)前 Host 上的 NIC 數(shù)，state->npes_node 是當(dāng)前 Host 上的 PE 數(shù)。
• 平均分配，并且保證每個(gè) PE 至少有 1 個(gè) NIC。假設(shè) 8 個(gè) GPU，2 個(gè) NIC，則會(huì)有 4 個(gè) GPU 共享 1 個(gè) NIC。

選擇本 PE 要使用的 NIC：

• 如果 Transport 只有一個(gè) NIC，則直接使用。
• 如果用戶打開了 ENABLE_NIC_PE_MAPPING 環(huán)境變量，則使用輪詢分配方式。
• 否則，調(diào)用 nvshmemi_get_devices_by_distance，根據(jù) PCIe 拓?fù)渚嚯x等信息，智能選擇最優(yōu)的設(shè)備分配給本 PE。

檢查設(shè)備分配有效性：

建立 Endpoint 連接：

調(diào)用 transport 的 connect_endpoints 回調(diào)，傳入選中的設(shè)備列表，建立與其他 PE 的通信端點(diǎn)。

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

更新設(shè)備狀態(tài)。

6.5.2 拓?fù)溆成?/h4>

當(dāng)用戶想要通過 ENABLE_NIC_PE_MAPPING 進(jìn)行 PE（實(shí)際也就是 GPU）與 NIC 的映射時(shí)，可以通過以下兩個(gè)環(huán)境變量實(shí)現(xiàn)：

NVSHMEM_HCA_LIST：直接提供 NIC name 和 port 索引的列表（PS：內(nèi)部會(huì)進(jìn)行排序）。

• “^mlx5_0”：排除 mlx5_0 NIC。
• “mlx5_1:1,mlx5_2:2”：如果有 8 個(gè) GPU，則 4 個(gè) GPU 對(duì)應(yīng) mlx5_1 的 Port 1，4 個(gè) GPU 對(duì)應(yīng) mlx5_2 的 Port 2。

NVSHMEM_HCA_PE_MAPPING：和上述的 NVSHMEM_HCA_LIST 類似，只不過可以進(jìn)一步設(shè)置 PE 的數(shù)量（PS：內(nèi)部也會(huì)進(jìn)行排序）。

• “mlx5_0:1:3,mlx5_0:2:5”：還是 8 個(gè) GPU，則 3 個(gè) GPU 對(duì)應(yīng) mlx5_1 的 Port 1，5 個(gè) GPU 對(duì)應(yīng) mlx5_0 的 Port 2。

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

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

由于高吞吐模式時(shí)這里的 mype_node 實(shí)際為 rdma_rank，一個(gè)節(jié)點(diǎn)上所有 PE 的 mype_node 相同，實(shí)際都是 0，導(dǎo)致這里實(shí)際上這里所有 PE 選擇了相同的 NIC，沒有達(dá)到 Mapping 的目的，反而導(dǎo)致了 NIC 熱點(diǎn)。

6.5.3 智能拓?fù)涓兄?/h4>

主要是調(diào)用 nvshmemi_get_devices_by_distance 為本節(jié)點(diǎn)上的每個(gè) PE 分配最優(yōu)的 NIC，以實(shí)現(xiàn)高效的 GPU-NIC 通信。分配策略基于 PCIe 拓?fù)渚嚯x，優(yōu)先選擇距離 GPU 最近的 NIC，并在多 GPU/多 NIC 情況下盡量負(fù)載均衡（位于 nvshmem_src/nvshmem/src/host/topo/topo.cpp）。整體思路與 NCCL 中類似，這里主要包括幾個(gè)步驟：

收集 GPU 和 NIC 路徑：

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

計(jì)算所有 GPU-NIC 距離：

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

第一輪分配（優(yōu)先分配最近的 NIC）：

• 遍歷 pe_dev_pairs（距離優(yōu)先），為每個(gè)本地 PE 依次分配最優(yōu)的 NIC。
• 如果當(dāng)前距離比已分配的更優(yōu)，則分配該 NIC，并記錄距離和使用計(jì)數(shù)。
• 如果遇到更差的距離，后續(xù)的 NIC 都不再考慮，標(biāo)記為無更優(yōu)分配（-2）。
• 直到所有本地 PE 的分配完成。

第二輪分配（負(fù)載均衡）：

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

輸出本 PE 的分配結(jié)果：

• 將本 PE 分配到的 NIC 索引寫入 device_arr。
• 統(tǒng)計(jì)本 PE 實(shí)際分配到的 NIC 數(shù)量。

七、參考鏈接

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??

本文轉(zhuǎn)載自???AI閑談?????，作者：AI閑談