CHAPTER 01

第 1 章:執行與現象:從一個 AllReduce 開始看外部行為

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 1 章 / 共 25 章

在深入任何核心程式碼之前,我們先把 NCCL 跑起來,觀察它對外暴露的行為。這一章不讀核心,只做一件事:建立一個可驗證的參照系——任何後續的內部機制分析,最終都要能解釋這裡看到的外部行為。

1.1 從建置入口看 NCCL 的工程結構

直覺模型

建置系統就像一棟大樓的施工圖紙:它不決定樓裡住誰,但決定了有哪些房間、門朝哪開。如果建置入口混亂,你連「跑起來」這第一步都邁不出去。NCCL 同時提供 Makefile 和 CMake 兩套建置入口,理解它們的差異,是理解這個專案工程組織的第一步。

兩套建置入口的結構

頂層Makefile是一個極薄的排程層,它本身不編譯任何原始檔,而是把工作轉發給各個子目錄的 Makefile。

📎 Makefile:44-45定義了src.%模式規則,把src.build、src.install等目標轉發給src/Makefile:

code
src.%:
	${MAKE} -C src $* BUILDDIR=${ABSBUILDDIR}

📎 Makefile:47-48定義了examples目標,它依賴src.build,然後進入docs/examples目錄建置範例:

code
examples: src.build
	${MAKE} -C docs/examples NCCL_HOME=${ABSBUILDDIR}

注意這裡的依賴關係:範例的建置依賴src.build先完成,因為範例需要連結 NCCL 函式庫,而NCCL_HOME環境變數把建置產物目錄傳給範例的 Makefile。這就是「先有函式庫,再有範例」的建置順序約束。

📎 Makefile:29列出了所有可清理的目標集合:

code
TARGETS := src pkg nccl4py ir

📎 Makefile:30用 GNU Make 的替換引用語法${TARGETS:%=%.clean}把src pkg nccl4py ir展開成src.clean pkg.clean nccl4py.clean ir.clean,一次性定義所有清理目標。這是 Makefile 裡常見的「用資料驅動規則」技巧——新增一個模組只需往TARGETS裡加一個詞。

CMake 入口:版本號從哪來

CMake 入口比 Makefile 複雜得多,因為它要處理跨平台、CUDA 版本探測、架構選擇等。我們只關注與「跑起來」直接相關的部分。

📎 CMakeLists.txt:5-11展示了版本號的來源——它不是硬編碼在 CMakeLists.txt 裡,而是從makefiles/version.mk讀取後用正則提取:

cmake
file(READ ${CMAKE_SOURCE_DIR}/makefiles/version.mk VERSION_CONTENT)
string(REGEX REPLACE ".*NCCL_MAJOR[ ]*:=[ ]*([0-9]+).*" "\\1" NCCL_MAJOR "${VERSION_CONTENT}")
...
math(EXPR NCCL_VERSION_CODE "(${NCCL_MAJOR} * 10000) + (${NCCL_MINOR} * 100) + ${NCCL_PATCH}")
〔設計推斷與架構權衡〕

把版本號集中放在version.mk裡,讓 Makefile 和 CMake 兩套建置系統共享同一個版本源,避免「兩套建置系統版本號不一致」這個經典工程陷阱。NCCL_VERSION_CODE的計算公式MAJOR*10000 + MINOR*100 + PATCH與標頭檔裡的NCCL_VERSION巨集保持一致。

📎 CMakeLists.txt:14-20把這些版本號透過add_compile_definitions注入到所有 C++ 原始碼檔案:

cmake
add_compile_definitions(
    NCCL_USE_CMAKE
    NCCL_MAJOR=${NCCL_MAJOR}
    NCCL_MINOR=${NCCL_MINOR}
    NCCL_PATCH=${NCCL_PATCH}
    NCCL_VERSION_CODE=${NCCL_VERSION_CODE}
)

📎 CMakeLists.txt:24-25宣告了專案語言為 CUDA、CXX、C:

cmake
project(NCCL VERSION ${NCCL_MAJOR}.${NCCL_MINOR}.${NCCL_PATCH}
        LANGUAGES CUDA CXX C)

CUDA 架構選擇:為什麼預設值這麼複雜

📎 CMakeLists.txt:140-171是一大段根據 CUDA 版本決定CMAKE_CUDA_ARCHITECTURES的邏輯。以 CUDA 12.8 及以上為例:

cmake
elseif(${CUDA_MAJOR} EQUAL 12)
    if(${CUDA_MINOR} LESS 8)
        set(CMAKE_CUDA_ARCHITECTURES "50;60;61;70;80;90")
    else()
        set(CMAKE_CUDA_ARCHITECTURES "50;60;61;70;80;90;100;120")
    endif()
〔設計推斷與架構權衡〕

這段邏輯的設計動機是:新架構(如 100、120)的 PTX 只有較新的 CUDA 工具鏈才認識,如果對舊 CUDA 強行指定新架構,編譯會直接失敗。所以預設架構列表必須隨 CUDA 版本動態調整。對讀者而言,這意味著:如果你不顯式設定CMAKE_CUDA_ARCHITECTURES,編譯產物會包含一長串架構的 fatbin,編譯時間會顯著變長。生產環境通常顯式指定目標架構來加速建置。

建置流程決策圖

下面這張圖展示了從執行make到產出可執行範例的完整決策路徑:

mermaid
flowchart TD
    start["执行 make 或 make examples"] --> check_ir{"EMIT_LLVM_IR 或<br/>NCCL_EMIT_LTO_IR 非 0?"}
    check_ir -->|是| add_ir["IR_GOALS 加入 llvm_ir/ltoir<br/>default 依赖 ir-emit"]
    check_ir -->|否| only_src["default 仅依赖 src.build"]
    add_ir --> src_build["make -C src build<br/>BUILDDIR=build"]
    only_src --> src_build
    src_build --> build_ok{"src.build 成功?"}
    build_ok -->|否| fail["构建失败,终止"]
    build_ok -->|是| is_examples{"目标是 examples?"}
    is_examples -->|是| ex_build["make -C docs/examples<br/>NCCL_HOME=build"]
    is_examples -->|否| done["产出 libnccl.so"]
    ex_build --> ex_ok{"示例链接成功?"}
    ex_ok -->|否| fail
    ex_ok -->|是| runnable["产出可执行示例"]

這張圖的關鍵分支在於IR_GOALS是否非空——它決定了預設建置是否額外觸發 LLVM IR 生成。對只想「跑起來」的讀者,保持EMIT_LLVM_IR=0即可走最短路徑。

1.2 最小可執行程序的前置條件

直覺模型

寫一個 NCCL 程式,就像組織一場多方電話會議。你需要先確認:有幾個人參加(裝置數)、每個人是誰(rank)、用什麼線路通話(stream)。缺任何一樣,會議都開不起來。這一節我們透過01_communicators範例,看清楚這三個前置條件在程式碼裡長什麼樣。

資料結構:三個陣列承載全部狀態

📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:88-92定義了範例的核心變數:

c
int num_gpus;                 // Number of available CUDA devices
ncclComm_t *comms = NULL;     // Array of NCCL communicators (one per GPU)
cudaStream_t *streams = NULL; // Array of CUDA streams (one per GPU)
int *devices = NULL;          // Array of device IDs to use

這裡體現了 NCCL 單行程多卡程式設計模型的核心:每個 GPU 一個通訊域、一個 stream、一個裝置號。三個陣列的長度都是num_gpus,下標i對應第i個 GPU。

ncclComm_t在標頭檔裡被定義為不透明指標。📎 src/nccl.h.in:36給出了它的真實型別:

c
typedef struct ncclComm* ncclComm_t;
〔設計推斷與架構權衡〕

「不透明指標」(opaque pointer)是 C 語言裡實現資訊隱藏的經典手法:標頭檔只暴露struct ncclComm*這個指標型別,使用者程式碼無法存取結構體內部的欄位,所有操作必須透過 API 函式完成。這樣 NCCL 就能在不破壞 ABI 的前提下自由修改ncclComm的內部佈局。對小白讀者,可以理解為「你拿到的是一個黑盒句柄,只能透過官方介面操作它」。

Step-by-Step:從裝置探測到通訊域建立

第一步:探測裝置數。 📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:96-104呼叫cudaGetDeviceCount並檢查是否為 0:

c
CUDACHECK(cudaGetDeviceCount(&num_gpus));

if (num_gpus == 0) {
    fprintf(stderr, "ERROR: No CUDA devices found on this system\n");
    ...
    return 1;
}

這一步在幹什麼:向 CUDA 執行時詢問「這台機器上有幾張 GPU」。如果回傳 0,說明沒有可用裝置,程式直接退出——這是最前置的守衛條件。

第二步:分配宿主記憶體並填充裝置列表。 📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:114-121分配三個陣列並檢查分配是否成功:

c
devices = (int *)malloc(num_gpus * sizeof(int));
comms = (ncclComm_t *)malloc(num_gpus * sizeof(ncclComm_t));
streams = (cudaStream_t *)malloc(num_gpus * sizeof(cudaStream_t));

if (!devices || !comms || !streams) {
    fprintf(stderr, "ERROR: Failed to allocate memory for device arrays\n");
    return 1;
}

📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:126-136用迴圈填充devices[i] = i,並列印每個裝置的屬性:

c
for (int i = 0; i < num_gpus; i++) {
    devices[i] = i; // Use device i for communicator i
    cudaDeviceProp prop;
    CUDACHECK(cudaGetDeviceProperties(&prop, devices[i]));
    printf("  GPU %d: %s (CUDA Device %d)\n", i, prop.name, devices[i]);
    ...
}

第三步:為每個 GPU 建立 stream。 📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:140-145是關鍵:

c
for (int i = 0; i < num_gpus; i++) {
    CUDACHECK(cudaSetDevice(devices[i]));
    CUDACHECK(cudaStreamCreate(&streams[i]));
}

注意cudaSetDevice必須在cudaStreamCreate之前呼叫。這是 CUDA 程式設計的基本規則:stream 屬於當前活躍裝置,如果不先切換裝置,stream 會建立在錯誤的 GPU 上。這是新手最容易踩的坑之一。

第四步:建立通訊域。 📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:169是整個範例的核心呼叫:

c
NCCLCHECK(ncclCommInitAll(comms, num_gpus, devices));

ncclCommInitAll是單行程多卡場景的便捷入口。標頭檔📎 src/nccl.h.in:301-301給出了它的契約:

c
/* Creates a clique of communicators (single process version).
 * This is a convenience function to create a single-process communicator clique.
 * Returns an array of ndev newly initialized communicators in comm.
 * comm should be pre-allocated with size at least ndev*sizeof(ncclComm_t).
 * If devlist is NULL, the first ndev CUDA devices are used.
 * Order of devlist defines user-order of processors within the communicator. */
ncclResult_t  ncclCommInitAll(ncclComm_t* comm, int ndev, const int* devlist);

三個參數的含義:comm是預分配的通訊域陣列,ndev是裝置數,devlist是裝置號列表(傳 NULL 則用前ndev個裝置)。呼叫回傳後,comms[i]就是第i個裝置的通訊域,其 rank 為i。

第五步:驗證通訊域屬性。 📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:185-189用三個查詢 API 驗證:

c
NCCLCHECK(ncclCommUserRank(comms[i], &rank));
NCCLCHECK(ncclCommCount(comms[i], &size));
NCCLCHECK(ncclCommCuDevice(comms[i], &device));

這三個 API 在標頭檔裡的定義分別是📎 src/nccl.h.in:396、📎 src/nccl.h.in:400、📎 src/nccl.h.in:404。它們分別回答三個問題:我是誰(rank)、一共有幾個人(size)、我在哪張卡上(device)。

通訊域建立流程時序圖

mermaid
sequenceDiagram
    participant App as 应用主线程
    participant CUDA as CUDA Runtime
    participant NCCL as NCCL 库
    App->>CUDA: cudaGetDeviceCount(&num_gpus)
    CUDA-->>App: num_gpus = N
    loop i in 0..N-1
        App->>CUDA: cudaSetDevice(devices[i])
        App->>CUDA: cudaStreamCreate(&streams[i])
        CUDA-->>App: streams[i]
    end
    App->>NCCL: ncclCommInitAll(comms, N, devices)
    Note over NCCL: 内部为每个设备建立通信域<br/>分配 rank 0..N-1
    NCCL-->>App: comms[0..N-1]
    loop i in 0..N-1
        App->>NCCL: ncclCommUserRank(comms[i], &rank)
        NCCL-->>App: rank = i
        App->>NCCL: ncclCommCount(comms[i], &size)
        NCCL-->>App: size = N
    end

這張時序圖揭示了關鍵點:ncclCommInitAll是一個同步阻塞呼叫,它內部會完成所有裝置間的協調,回傳時所有通訊域都已就緒。

設計思考:為什麼需要 ncclCommInitAll

〔設計推斷與架構權衡〕

多進程場景下,每個進程只管理一張 GPU,用ncclCommInitRank各自初始化即可。但單進程多卡場景下,如果讓使用者手動為每張卡呼叫ncclCommInitRank,就必須處理「多個 rank 之間的同步」——而單進程裡只有一個執行緒,無法同時推進多個 rank 的初始化,會死鎖。ncclCommInitAll把這種協調封裝在函式庫內部,用內部機制(通常是多執行緒或狀態機)完成所有 rank 的同步初始化,對使用者暴露成一個簡單的同步呼叫。這就是「便捷函式」存在的根本原因。

1.3 一次 AllReduce 的完整外部行為

直覺模型

AllReduce 是集合通訊裡最常用的操作:每個參與者貢獻一份資料,所有人拿到所有資料的總和。就像小組作業算總分——每個人報上自己的分數,最後每個人手裡都有一份全班總分。這一節我們追蹤03_collectives/01_allreduce範例,看一次 AllReduce 從呼叫到結果驗證的完整外部行為。

資料結構:資料緩衝區與初始化

📎 docs/examples/03_collectives/01_allreduce/c/main.cc:59-63定義了核心變數:

c
int num_gpus = 0;
ncclComm_t *comms;
cudaStream_t *streams;
float **sendbuff;
float **recvbuff;

注意sendbuff和recvbuff是float**——指向指標陣列的指標。每個sendbuff[i]是第i張 GPU 上的裝置記憶體位址。

📎 docs/examples/03_collectives/01_allreduce/c/main.cc:99定義了資料規模:

c
const size_t size = 32 * 1024 * 1024; // 32M floats for demonstration

32M 個 float,每個 4 位元組,即 128 MB 的發送緩衝和 128 MB 的接收緩衝,每張卡各一份。

📎 docs/examples/03_collectives/01_allreduce/c/main.cc:101-120是每個裝置的初始化迴圈:

c
for (int i = 0; i < num_gpus; i++) {
    CUDACHECK(cudaSetDevice(i));
    CUDACHECK(cudaStreamCreate(&streams[i]));
    CUDACHECK(cudaMalloc((void **)&sendbuff[i], size * sizeof(float)));
    CUDACHECK(cudaMalloc((void **)&recvbuff[i], size * sizeof(float)));
    CUDACHECK(cudaMemset(sendbuff[i], 0, size * sizeof(float)));
    float rank_value = (float)i;
    CUDACHECK(cudaMemcpy(sendbuff[i], &rank_value, sizeof(float),
                         cudaMemcpyHostToDevice));
    printf("  Device %d initialized with data value %d\n", i, i);
}

這段程式碼的巧妙之處:先把整個發送緩衝區清零,然後只把第一個元素設為i(該裝置的 rank 值)。這樣 AllReduce 求和後,第一個元素的結果就是0 + 1 + 2 + ... + (num_gpus-1),而其餘元素都是 0。驗證時只需檢查第一個元素,就能確認 AllReduce 是否正確。

Step-by-Step:AllReduce 呼叫與驗證

第一步:Group 包裹。 📎 docs/examples/03_collectives/01_allreduce/c/main.cc:130-136是核心呼叫:

c
NCCLCHECK(ncclGroupStart());
for (int i = 0; i < num_gpus; i++) {
    NCCLCHECK(ncclAllReduce(sendbuff[i], recvbuff[i], size, ncclFloat, ncclSum,
                            comms[i], streams[i]));
}
NCCLCHECK(ncclGroupEnd());

這裡有一個極其重要的細節:註解📎 docs/examples/03_collectives/01_allreduce/c/main.cc:128-129明確說明:

c
// NOTE: ncclGroupStart and ncclGroupEnd are essential to avoid
// deadlock when using ncclCommInitAll and multiple communication calls.

為什麼必須用 Group?標頭檔📎 src/nccl.h.in:844-864給出了解釋:

c
/* Group semantics
 *
 * When managing multiple GPUs from a single thread, and since NCCL collective
 * calls may perform inter-CPU synchronization, we need to "group" calls for
 * different ranks/devices into a single call.
 * ...
 * Both collective communication and ncclCommInitRank can be used in conjunction
 * of ncclGroupStart/ncclGroupEnd, but not together.
 */
〔設計推斷與架構權衡〕

核心矛盾在於:集合通訊需要所有 rank 同時參與,但單執行緒裡你只能一個一個呼叫ncclAllReduce。如果第一個ncclAllReduce呼叫就阻塞等待其他 rank,而其他 rank 的呼叫還沒發出,就會死鎖。Group 機制的作用是:ncclGroupStart之後的所有呼叫只做「登記」,不實際啟動;ncclGroupEnd時才把所有登記的操作一起提交,讓它們能並發推進。這就像點外送時先把所有菜加進購物車,最後一起結算,而不是一道菜一道菜地下單。

第二步:同步 stream。 📎 docs/examples/03_collectives/01_allreduce/c/main.cc:139-142:

c
for (int i = 0; i < num_gpus; i++) {
    CUDACHECK(cudaSetDevice(i));
    CUDACHECK(cudaStreamSynchronize(streams[i]));
}

標頭檔📎 src/nccl.h.in:854-856強調:ncclGroupEnd只保證操作被入隊到 stream,不保證操作完成。所以必須顯式同步 stream,才能安全讀取結果。

第三步:驗證結果。 📎 docs/examples/03_collectives/01_allreduce/c/main.cc:152-169:

c
float expected = (float)(num_gpus * (num_gpus - 1) / 2);
...
for (int i = 0; i < num_gpus; i++) {
    float result;
    CUDACHECK(cudaSetDevice(i));
    CUDACHECK(cudaMemcpy(&result, recvbuff[i], sizeof(float),
                         cudaMemcpyDeviceToHost));
    if (result != expected) {
        printf("  Device %d received incorrect result: %.0f (expected %.0f)\n", i,
               result, expected);
        success = false;
    } else {
        printf("  Device %d correctly received sum: %.0f\n", i, result);
    }
}

期望值是等差數列求和0 + 1 + ... + (N-1) = N*(N-1)/2。每張卡都應該收到相同的值——這正是 AllReduce 的定義。

AllReduce 資料流圖

mermaid
flowchart LR
    subgraph dev0["GPU 0 (rank 0)"]
        s0["sendbuff[0]<br/>首元素=0"]
        r0["recvbuff[0]"]
    end
    subgraph dev1["GPU 1 (rank 1)"]
        s1["sendbuff[1]<br/>首元素=1"]
        r1["recvbuff[1]"]
    end
    subgraph dev2["GPU 2 (rank 2)"]
        s2["sendbuff[2]<br/>首元素=2"]
        r2["recvbuff[2]"]
    end
    s0 -->|ncclAllReduce<br/>ncclFloat ncclSum| reduce["归约求和<br/>0+1+2=3"]
    s1 -->|ncclAllReduce<br/>ncclFloat ncclSum| reduce
    s2 -->|ncclAllReduce<br/>ncclFloat ncclSum| reduce
    reduce -->|广播结果| r0
    reduce -->|广播结果| r1
    reduce -->|广播结果| r2

這張圖展示了 AllReduce 的兩個階段:先歸約(reduce),再廣播(broadcast)。每個 rank 的recvbuff最終都得到相同的結果。

設計思考:為什麼用 Group 而不是逐個呼叫

〔設計推斷與架構權衡〕

如果去掉ncclGroupStart/ncclGroupEnd,程式碼會變成:

c
for (int i = 0; i < num_gpus; i++) {
    ncclAllReduce(sendbuff[i], recvbuff[i], size, ncclFloat, ncclSum,
                  comms[i], streams[i]);
}

在單執行緒裡,第一次迭代呼叫ncclAllReduce時,NCCL 需要等待所有 rank 都發起 AllReduce 才能推進。但其他 rank 的呼叫還在迴圈裡沒執行到,於是第一次呼叫永遠等不到其他 rank,死鎖。Group 機制把「發起」和「執行」分離,讓所有 rank 的呼叫先全部登記,再一起執行,從根本上避免了單執行緒死鎖。

1.4 通訊域的生命週期與資源清理

直覺模型

通訊域就像一場會議。開會前要簽到(初始化),開完會要散會(銷毀)。如果散會順序不對——比如人還沒走就把會議室鎖了——就會出問題。這一節我們看 NCCL 通訊域的銷毀順序,以及為什麼這個順序不能顛倒。

銷毀的兩個階段:Finalize 與 Destroy

📎 docs/examples/03_collectives/01_allreduce/c/main.cc:176-183展示了標準的銷毀流程:

c
NCCLCHECK(ncclGroupStart());
for (int i = 0; i < num_gpus; i++) {
    NCCLCHECK(ncclCommFinalize(comms[i]));
}
NCCLCHECK(ncclGroupEnd());
for (int i = 0; i < num_gpus; i++) {
    NCCLCHECK(ncclCommDestroy(comms[i]));
}

標頭檔📎 src/nccl.h.in:309-309解釋了ncclCommFinalize的語義:

c
/* Finalize a communicator. ncclCommFinalize flushes all issued communications,
 * and marks communicator state as ncclInProgress. The state will change to ncclSuccess
 * when the communicator is globally quiescent and related resources are freed; then,
 * calling ncclCommDestroy can locally free the rest of the resources (e.g. communicator
 * itself) without blocking. */
ncclResult_t  ncclCommFinalize(ncclComm_t comm);

📎 src/nccl.h.in:313-313解釋了ncclCommDestroy:

c
/* Frees local resources associated with communicator object. */
ncclResult_t  ncclCommDestroy(ncclComm_t comm);
〔設計推斷與架構權衡〕

為什麼銷毀要分兩步?ncclCommFinalize是全域操作——它需要所有 rank 都參與,確保沒有在途的通訊。ncclCommDestroy是本地操作——它只釋放本進程的資源,不阻塞。這個設計讓「等待所有 rank 靜默」和「釋放本地資源」解耦:前者可能耗時較長(要等網路對端),後者是純本地操作。如果只有一個ncclCommDestroy,它就必須同時承擔這兩個職責,要麼阻塞太久,要麼無法保證全域靜默。

銷毀順序的完整鏈條

📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:221-249展示了完整的清理順序,註解📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:218-219強調:

c
// IMPORTANT: Proper cleanup is critical for NCCL applications
// Resources must be cleaned up in the correct order to avoid issues

順序是:

1. 同步所有 stream(📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:224-227)

2. Finalize + Destroy 通訊域(📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:233-240)

3. 銷毀 CUDA stream(📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:246-249)

4. 釋放宿主記憶體(📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:253-255)

通訊域狀態機

ncclCommFinalize的文件明確提到了狀態轉換,這符合狀態機的准入條件:

mermaid
stateDiagram-v2
    [*] --> Active : ncclCommInitAll() 成功
    Active --> InProgress : ncclCommFinalize()<br/>刷新在途通信
    InProgress --> Quiescent : 全局静默<br/>相关资源释放
    Quiescent --> Destroyed : ncclCommDestroy()<br/>释放本地资源
    Destroyed --> [*]
    Active --> Aborted : ncclCommAbort()<br/>中止在途操作
    Aborted --> [*]

這個狀態機的關鍵轉換是InProgress -> Quiescent:它由「全域靜默」這個事件觸發,而不是由某個函式呼叫直接觸發。這意味著ncclCommFinalize返回後,通訊域可能還處於InProgress狀態,需要輪詢ncclCommGetAsyncError才能知道何時進入Quiescent。

設計思考:為什麼銷毀順序不能顛倒

〔設計推斷與架構權衡〕

如果先銷毀 CUDA stream 再銷毀通訊域,會出什麼問題?通訊域內部可能持有對 stream 的參考(比如用於非同步操作的完成通知)。如果 stream 先被銷毀,通訊域在 Finalize 時存取已銷毀的 stream,會導致未定義行為。同理,如果先釋放宿主記憶體(comms陣列)再銷毀通訊域,ncclCommDestroy就拿到了野指標。這就是為什麼順序必須是「先同步、再銷毀通訊域、再銷毀 stream、最後釋放宿主記憶體」——依賴關係決定了銷毀順序必須與建立順序相反。

1.5 生產避坑指南

坑一:忘記 Group 導致死鎖

這是新手最常踩的坑。在單行程多卡場景下,如果直接迴圈呼叫ncclAllReduce而不加 Group,程式會在第一次呼叫時死鎖。症狀是:程式卡住不動,CPU 佔用率接近 0,沒有任何輸出。

排查方法:用gdbattach 到行程,看堆疊是否停在 NCCL 內部的等待邏輯上。如果是,檢查是否遺漏了ncclGroupStart/ncclGroupEnd。

坑二:忘記同步 stream 就讀結果

📎 src/nccl.h.in:854-856明確說明ncclGroupEnd只保證入列,不保證完成。如果省略📎 docs/examples/03_collectives/01_allreduce/c/main.cc:139-142的 stream 同步,直接讀取recvbuff,會讀到未完成的資料。

症狀是:結果時對時錯,或者讀到全 0。這是因為cudaMemcpy預設是同步的,但它同步的是當前 stream,而 AllReduce 可能在其他 stream 上執行。排查方法:在讀取結果前加cudaStreamSynchronize,如果問題消失,就是這個坑。

坑三:銷毀順序錯誤導致段錯誤

如果在ncclCommDestroy之前就cudaFree了sendbuff/recvbuff,通訊域在 Finalize 時可能還在存取這些緩衝區,導致段錯誤或資料損壞。

症狀是:程式在退出階段崩潰,或者偶發地讀到垃圾資料。排查方法:檢查清理程式碼的順序,確保通訊域銷毀在所有 CUDA 資源釋放之前。

坑四:裝置號與 rank 混淆

📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:198-200有一個驗證:

c
if (device != devices[i]) {
    printf(" [WARNING: Expected device %d]", devices[i]);
}
〔設計推斷與架構權衡〕

rank 和 device 是兩個不同的概念。rank 是通訊域內的邏輯編號(0 到 nRanks-1),device 是實體 GPU 編號。在ncclCommInitAll的預設用法裡,devices[i] = i,所以 rank 和 device 恰好相等。但如果傳入自訂的devlist(比如{2, 0, 1}),rank 0 就對應 device 2。混淆這兩個概念會導致資料發到錯誤的 GPU 上。

本章小結

本章我們完成了三件事:

1. 建置入口:理解了 Makefile 的轉發機制和 CMake 的版本號來源、CUDA 架構選擇邏輯。關鍵結論是make examples會先建置函式庫再建置範例,NCCL_HOME把建置產物目錄傳給範例。

2. 最小可執行程式的三要素:裝置數(cudaGetDeviceCount)、rank(由ncclCommInitAll自動分配)、stream(每個 GPU 一個)。ncclCommInitAll是單行程多卡的便捷入口,它把多 rank 同步初始化封裝在函式庫內部。

3. 一次 AllReduce 的完整外部行為:從ncclGroupStart包裹多個ncclAllReduce呼叫,到ncclGroupEnd提交,再到cudaStreamSynchronize等待完成,最後驗證結果。Group 機制是單執行緒多卡場景避免死鎖的關鍵。

4. 通訊域生命週期:ncclCommFinalize(全域靜默)+ncclCommDestroy(本地釋放)的兩階段銷毀,以及「先同步、再銷毀通訊域、再銷毀 stream、最後釋放宿主記憶體」的順序約束。

本章思考與自測

Q1: 如果把📎 docs/examples/03_collectives/01_allreduce/c/main.cc:130-136的 ncclGroupStart/ncclGroupEnd 去掉,改成直接迴圈呼叫 ncclAllReduce,在單行程多卡場景下會發生什麼?為什麼?

參考解析:會發生死鎖。標頭檔📎 src/nccl.h.in:844-864解釋了原因:集合通訊呼叫可能執行 inter-CPU 同步,需要所有 rank 同時參與。在單執行緒裡,第一次迴圈迭代呼叫ncclAllReduce(comms[0], ...)時,NCCL 需要等待其他 rank 也發起 AllReduce 才能推進。但其他 rank 的呼叫還在迴圈裡沒執行到(因為當前執行緒被阻塞在第一次呼叫上),於是第一次呼叫永遠等不到其他 rank,死鎖。

Group 機制的作用是把「發起」和「執行」分離:ncclGroupStart之後的所有呼叫只做登記,ncclGroupEnd時才把所有登記的操作一起提交,讓它們能並行推進。這從根本上避免了單執行緒死鎖。

驗證方法:去掉 Group 後執行程式,用gdbattach 看堆疊,會停在 NCCL 內部的等待邏輯上,CPU 佔用率接近 0。

Q2: 📎 docs/examples/03_collectives/01_allreduce/c/main.cc:139-142的 cudaStreamSynchronize 能否用 cudaDeviceSynchronize 替代?兩者在語意上有什麼區別?在什麼場景下這個替代會出問題?

參考解析:可以用cudaDeviceSynchronize替代,但語意不同。cudaStreamSynchronize(streams[i])只等待指定 stream 上的操作完成;cudaDeviceSynchronize等待當前裝置上所有stream 的操作完成。

在單行程多卡場景下,cudaDeviceSynchronize只同步當前裝置(由cudaSetDevice決定),所以需要配合cudaSetDevice(i)迴圈使用。如果省略cudaSetDevice,cudaDeviceSynchronize只會同步預設裝置(通常是 device 0),其他裝置的 AllReduce 可能還沒完成。

標頭檔📎 src/nccl.h.in:854-856強調ncclGroupEnd只保證入列不保證完成,所以同步是必須的。用cudaStreamSynchronize更精確,因為它只等待相關 stream,不會誤等無關操作。用cudaDeviceSynchronize的問題是:如果裝置上有其他無關的長時間執行 kernel,會被誤等,降低效能。

Q3: 📎 docs/examples/01_communicators/01_multiple_devices_single_process/c/main.cc:233-240的銷毀順序是「先 Finalize 所有通訊域,再 Destroy 所有通訊域」。如果改成「對每個通訊域先 Finalize 再 Destroy」(即在一個迴圈裡完成兩個操作),會有什麼問題?

參考解析:會破壞 Group 語意。當前的寫法是:

c
ncclGroupStart();
for (i) ncclCommFinalize(comms[i]);
ncclGroupEnd();
for (i) ncclCommDestroy(comms[i]);

ncclCommFinalize被 Group 包裹,意味著所有通訊域的 Finalize 會一起提交,能並行推進。如果改成:

c
for (i) {
    ncclCommFinalize(comms[i]);
    ncclCommDestroy(comms[i]);
}

第一次迭代的ncclCommFinalize(comms[0])會阻塞等待所有 rank 靜默,但其他通訊域的 Finalize 還沒發起,導致死鎖——這與 Q1 的死鎖是同一類問題。

另外,標頭檔📎 src/nccl.h.in:309-309說明ncclCommFinalize返回時通訊域可能還處於ncclInProgress狀態,需要等待全域靜默才能進入ncclSuccess。如果緊接著就ncclCommDestroy,可能在通訊域還沒完全靜默時就釋放本地資源,導致未定義行為。正確做法是 Finalize 後輪詢ncclCommGetAsyncError確認狀態,再 Destroy。

這些外部行為構成了後續所有原始碼分析的參照系。第 2 章我們將建立核心心智模型:通訊域、通道、演算法、協定、傳輸層這五件套,看看 NCCL 內部是如何組織這些概念的。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 02

第 2 章:核心抽象模型:通訊算子、拓撲、演算法、協定與傳輸層

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 2 章 / 共 25 章

上一章我們讓 NCCL 跑了起來,觀察了 ncclCommInitRank、ncclAllReduce、ncclCommDestroy 三個 API 的外部行為。但外部行為只是冰山一角——當 ncclAllReduce 返回時,GPU 上到底發生了什麼?資料走了哪條路?為什麼同樣的 AllReduce 在不同機器上效能差異巨大?要回答這些問題,必須先建立 NCCL 的公共詞彙表。本章將逐一拆解五個核心抽象:通訊域(ncclComm)、通道(channel)、演算法(algorithm)、協定(protocol)、傳輸層(transport)。這五個概念貫穿全書,後續每一章的分析都會用到它們。理解它們之間的關係,就理解了 NCCL 的骨架。

2.1 通訊域 ncclComm:一個行程的通訊上下文

直覺模型

把ncclComm想像成一個「群聊」:每個行程加入群聊後拿到一個群 ID,之後所有訊息都在這個群裡發。群裡有幾個人(nRanks)、我是誰(rank)、走什麼線路(channels)、用什麼規則(config),全都記在這個群聊物件裡。

如果沒有ncclComm,NCCL 就不知道「誰和誰通訊」「資料發到哪裡去」——每次呼叫 API 都得重新協商 rank 列表、重建連線,開銷無法承受。

資料結構與記憶體佈局

ncclComm是整個 NCCL 最核心的結構體,定義在src/include/comm.h中。它極其龐大(近 300 行),我們按功能分組來看關鍵欄位。

身分標識與生命週期哨兵

📎 src/include/comm.h:576-580定義了startMagic,📎 src/include/comm.h:879-881定義了endMagic。這兩個欄位不是安全金鑰,而是記憶體越界檢測哨兵。在📎 src/include/comm.h:883-885處有兩個static_assert:

c
static_assert(offsetof(struct ncclComm, startMagic) == 0, "startMagic must be the first field of ncclComm");
static_assert(offsetof(struct ncclComm, endMagic) == sizeof(struct ncclComm) - sizeof(uint64_t),
              "endMagic must be the last field of ncclComm");
〔設計推斷與架構權衡〕

這兩個斷言在編譯期強制startMagic位於結構體首位址、endMagic位於末尾。執行時可以透過檢查這兩個魔數是否被篡改,快速判斷ncclComm指標是否有效——這在多執行緒環境下排查「野指標存取已銷毀通訊域」類 bug 時非常有用。

Rank 與拓撲資訊

📎 src/include/comm.h:628-629定義了rank和nRanks——我在通訊域中的編號和總參與者數。📎 src/include/comm.h:644-652定義了節點相關欄位:node(我所在節點編號)、nNodes(總節點數)、localRank(節點內編號)、localRanks(節點內 GPU 數),以及三張映射表rankToNode、rankToLocalRank、localRankToRank。

〔設計推斷與架構權衡〕

這三張映射表是拓撲感知演算法的基礎。比如 Ring 演算法需要知道「我的下一個 rank 是否在同一節點內」來決定走 NVLink 還是網路。如果沒有這些映射表,每次演算法選擇都要重新查詢拓撲圖,開銷巨大。

通道與緩衝區

📎 src/include/comm.h:593-593定義了channels[MAXCHANNELS]——這是通信域內所有通道的陣列。📎 src/include/comm.h:674-676定義了通道數量:nChannels(連接通道數)、collChannels(集合通信入隊通道數)、nvlsChannels(NVLS 通道數)。

📎 src/include/comm.h:691-693定義了緩衝區大小:buffSizes[NCCL_NUM_PROTOCOLS](每種協議的緩衝區大小)、p2pChunkSize(P2P 塊大小)、nvlsChunkSize(NVLS 塊大小)。

〔設計推斷與架構權衡〕

buffSizes陣列的索引就是協議枚舉值(LL/LL128/Simple),這意味著每種協議有獨立的緩衝區大小配置。LL 協議需要小緩衝區以降低延遲,Simple 協議需要大緩衝區以提高頻寬——這個陣列讓兩種需求共存。

工作佇列與 FIFO

📎 src/include/comm.h:719-728定義了工作 FIFO 相關欄位:workFifoBytes(FIFO 大小,2 的冪)、workFifoBuf(主機側 FIFO 緩衝區)、workFifoBufDev(裝置側 FIFO 緩衝區)、workFifoProduced(已生產位元組數)、workFifoConsumed(已消費位元組數)。

〔設計推斷與架構權衡〕

這是一個典型的生產者-消費者環形緩衝區。主機側(生產者)把工作描述寫入 FIFO,GPU kernel(消費者)讀取並執行。workFifoBytes必須是 2 的冪,這樣可以用位元遮罩代替取模運算,加速索引計算。

行程內同步屏障

📎 src/include/comm.h:731-731定義了行程內多通信域同步機制:

c
struct ncclComm* intraComm0; // leader of intra-process comms (self possible)
struct ncclComm* intraNext; // next of intra-process comms, intraComm0 is head
int intraRank;
int intraRanks;
uint32_t intraBarrierPhase;
char intraPad1[64 - sizeof(uint64_t)];
uint64_t intraBarrierCounter; // only used if this is intraComm0
char intraPad2[64 - sizeof(uint64_t)];
uint64_t intraBarrierGate; // only used if this is intraComm0

注意intraPad1和intraPad2的大小是64 - sizeof(uint64_t),即 56 位元組。加上前面的uint64_t欄位,每個欄位組恰好佔 64 位元組——這是一個快取行(Cache Line)。

〔設計推斷與架構權衡〕

這是典型的快取行填充(Cache Line Padding)技術。intraBarrierCounter和intraBarrierGate會被多個執行緒高頻讀寫,如果它們共享同一個快取行,會導致偽共享(False Sharing):一個執行緒修改intraBarrierCounter會使另一個執行緒的intraBarrierGate快取失效,造成效能急劇下降。用 56 位元組填充把它們隔開到不同快取行,是高效能並發程式設計的標準手法。

非同步錯誤狀態

📎 src/include/comm.h:705-705定義了asyncResult——這個欄位記錄通信域的非同步操作狀態。上一章我們提到ncclCommFinalize返回時通信域可能還處於ncclInProgress狀態,就是透過這個欄位追蹤的。

場景驅動 Walkthrough:從 ncclCommInitRank 到結構體填充

當使用者呼叫ncclCommInitRank(&comm, nranks, commId, rank)時,NCCL 內部會分配一個ncclComm結構體並逐欄位填充。我們跟隨這個流程看關鍵欄位如何被設定:

第一步:分配與清零

NCCL 使用ncclCalloc分配ncclComm,確保所有欄位初始為 0。此時startMagic和endMagic被設定為NCCL_MAGIC(📎 src/include/comm.h:563-569定義為0x0280028002800280,註解說 "Nickel atomic number is 28")。

第二步:填充身份資訊

rank、nRanks、cudaDev從參數和 CUDA API 取得。commHash由ncclCommId雜湊得到,用於後續網路通信中的一致性校驗。

第三步:建構拓撲圖

NCCL 呼叫拓撲探測模組枚舉所有 GPU、網卡、PCI 交換器,建構topo欄位(📎 src/include/comm.h:595-595)。這個拓撲圖決定了後續演算法選擇和路徑規劃。

第四步:初始化通道

channels[MAXCHANNELS]陣列被逐個初始化。每個通道的id被設定為陣列索引,peers和devPeers指標被分配。

第五步:建立傳輸連接

根據拓撲圖,NCCL 為每對 rank 選擇傳輸層(P2P/SHM/NET),呼叫對應的setup和connect回呼。連接資訊存儲在channels[i].peers[j]中。

第六步:設定魔數

最後,endMagic被設定為NCCL_MAGIC,標記結構體初始化完成。

設計思考與生產踩坑

為什麼ncclComm這麼大?

〔設計推斷與架構權衡〕

ncclComm包含近 300 個欄位,因為它承載了一個通信域的全部狀態。NCCL 的設計哲學是「一次初始化,多次複用」——初始化時把所有可能用到的資訊都算好存下來,執行時直接查表,避免重複計算。代價是記憶體佔用較大(每個通信域約幾 KB),但相比 GPU 顯存和網路頻寬,這點記憶體微不足道。

踩坑場景一:多執行緒共享通信域

〔設計推斷與架構權衡〕

ncclComm不是執行緒安全的。如果兩個執行緒同時對同一個ncclComm呼叫ncclAllReduce,workFifoProduced等欄位會競爭,導致資料損壞。 正確做法是每個執行緒使用獨立的通信域,或者用外部鎖串行化呼叫。

踩坑場景二:銷毀後存取

ncclCommDestroy釋放結構體記憶體後,如果還有執行緒持有指標並存取,會讀到已釋放記憶體。startMagic和endMagic可以幫助檢測這種情況——如果魔數不匹配,說明指標已失效。

踩坑場景三:快取行偽共享

在多行程場景下(每個行程一個 rank),intraBarrierCounter和intraBarrierGate的填充尤為重要。如果省略填充,多個行程的屏障操作會互相干擾,導致同步延遲從奈秒級上升到微秒級。

2.2 通道 channel:把一次通信切成多條流水線

直覺模型

搬家時不止開一條傳送帶,而是同時開好幾條,每條負責一部分箱子,整體搬得更快。channel就是 NCCL 的「傳送帶」——把一次集合通訊的資料切分成多份,每條通道獨立搬運一份,並行推進以提高頻寬利用率。

如果沒有 channel,所有資料只能走一條路徑,GPU 之間的多條物理鏈路(多張網卡、多組 NVLink)無法同時利用,頻寬利用率會大幅下降。

資料結構與記憶體佈局

ncclChannel定義在📎 src/include/comm.h:169-191:

c
struct ncclChannel {
  struct ncclChannelPeer** peers;
  struct ncclDevChannelPeer** devPeers;
  /* devPeer pointer array used for host side access */
  struct ncclDevChannelPeer** devPeersHostPtr;
  struct ncclRing ring;
  int* devRingUserRanks;
  struct ncclTree tree;

  struct ncclTree collnetChain;
  struct ncclDirect collnetDirect;

  struct ncclNvls nvls;

  int id; // index of this channel
  uint32_t workFifoProduced; // +1 successor of last used work fifo byte

  /* comm split sharable resources */
  struct ncclChannelPeer* collnetPeers;
  struct ncclDevChannelPeer* collnetDevPeers;
  struct ncclChannelPeer* nvlsPeers;
  struct ncclDevChannelPeer* nvlsDevPeers;
};

關鍵欄位解析

  • peers / devPeers:指向該通道內所有 rank 的連接資訊。peers是主機側視圖,devPeers是裝置側視圖(GPU kernel 直接存取)。
  • ring:Ring 演算法的拓撲描述——每個 rank 的前驅和後繼。
  • tree:Tree 演算法的拓撲描述——父節點和子節點列表。
  • collnetChain / collnetDirect:CollNet 演算法的兩種變體拓撲。
  • nvls:NVLink SHARP 的拓撲描述。
  • id:通道索引,從 0 到nChannels-1。
  • workFifoProduced:該通道的工作 FIFO 生產指標。
〔設計推斷與架構權衡〕

注意ring、tree、collnetChain、collnetDirect、nvls這五個欄位是並列的——同一個通道可以同時持有多種演算法的拓撲描述。執行時根據演算法選擇決定使用哪個欄位。這種設計讓演算法切換不需要重建通道,只需切換讀取的欄位。

通道數量計算

通道數量在ncclComm中定義(📎 src/include/comm.h:674-676):

c
int nChannels; // connection nChannels
int collChannels; // enqueue nChannels
int nvlsChannels; // enqueue nChannels
〔設計推斷與架構權衡〕

nChannels是實際建立的連接數,collChannels是集合通訊入隊時使用的通道數,nvlsChannels是 NVLS 專用通道數。三者可能不同——比如某些通道只用於 P2P 不用於集合通訊。

P2P 通道排程

📎 src/include/channel.h:21-33定義了ncclP2pChannelBaseForRound函式,用於計算 P2P 通訊中每個 round 使用的通道基址:

c
inline uint8_t ncclP2pChannelBaseForRound(struct ncclComm* comm, int p2pRound) {
  int base;
  if (comm->nNodes > 1) {
    int localSize = comm->p2pSchedGroupSize;
    int groupDelta = p2pRound / localSize;
    int localDelta = p2pRound % localSize;
    base = groupDelta * divUp(localSize, NCCL_MAX_DEV_WORK_P2P_PER_BATCH);
    base += localDelta / NCCL_MAX_DEV_WORK_P2P_PER_BATCH;
  } else {
    base = p2pRound;
  }
  return reverseBits(base, log2Up(comm->p2pnChannels));
}
〔設計推斷與架構權衡〕

這個函式的邏輯是:多節點場景下,P2P 通訊按「組」排程,每組內的 rank 使用相鄰通道;單節點場景下,每個 round 直接映射到一個通道。reverseBits是位反轉操作,用於打散通道分配,避免熱點集中。

場景驅動 Walkthrough:一次 AllReduce 如何分配通道

假設 8 個 rank、4 個通道,執行一次 AllReduce。資料被切成 4 份,每份由一個通道負責。

第一步:演算法選擇

NCCL 的 tuning 模組根據訊息大小和拓撲選擇演算法(比如 Ring)和協定(比如 Simple)。

第二步:通道分配

ncclTaskColl結構體(📎 src/include/comm.h:212-273)被建立,其中nChannels欄位被設定為 4(📎 src/include/comm.h:254-254)。channelLo和channelHi欄位(📎 src/include/comm.h:256-257)標記該任務使用的通道範圍。

第三步:資料切分

每個通道負責count / nChannels個元素。通道 0 處理第 0 到 count/4-1 個元素,通道 1 處理第 count/4 到 count/2-1 個元素,以此類推。

第四步:並行執行

4 個通道的 GPU kernel 同時啟動,各自在自己的資料切片上執行 Ring AllReduce。由於通道之間沒有資料依賴,可以完全並行。

第五步:結果合併

所有通道完成後,每個 rank 的 recv buffer 中就是完整的 AllReduce 結果。

並發控制與硬體互動

通道與 GPU 資源的映射

〔設計推斷與架構權衡〕

每個通道通常綁定到一個獨立的 CUDA stream 或 GPU 硬體佇列。這樣不同通道的 kernel 可以在 GPU 上並發執行,充分利用 SM(串流多處理器)資源。

通道與網路裝置的映射

在多網卡場景下,不同通道可以綁定到不同網卡。比如 4 個通道、2 張網卡,通道 0 和 1 走網卡 A,通道 2 和 3 走網卡 B。這樣兩張網卡的頻寬都能被利用。

通道數量的選擇

〔設計推斷與架構權衡〕

通道數量不是越多越好。通道數增加會帶來:

  • 更多 kernel 啟動開銷
  • 更多連接建立開銷
  • 更複雜的同步

NCCL 的 tuning 模組會根據訊息大小自動選擇最優通道數。小訊息用少量通道(減少開銷),大訊息用多通道(提高頻寬)。

生產避坑指南

踩坑場景一:通道數配置不當

〔設計推斷與架構權衡〕

如果手動設定NCCL_NCHANNELS過大,小訊息場景下 kernel 啟動開銷會超過收益,效能反而下降。 建議讓 NCCL 自動選擇,除非有明確的調優需求。

踩坑場景二:通道與拓撲不匹配

〔設計推斷與架構權衡〕

如果通道數超過物理鏈路數,部分通道會共享鏈路,無法實現真正的並行。 比如 2 張網卡配 8 個通道,實際只有 2 個通道能同時傳輸,其餘 6 個在排隊。

踩坑場景三:P2P 通道衝突

ncclP2pChannelBaseForRound的reverseBits操作如果實作有誤,會導致多個 round 映射到同一通道,造成串行化。📎 src/include/channel.h:32-32的reverseBits(base, log2Up(comm->p2pnChannels))確保通道分配均勻。

2.3 演算法 algorithm:Tree/Ring/CollNet/NVLS/PAT 的拓撲組織

直覺模型

從北京到上海可以坐高鐵、飛機或自駕,每種方式適合不同的距離和人數。NCCL 的演算法就是這些「出行方式」——Ring 適合大消息的穩定頻寬,Tree 適合小消息的低延遲,CollNet 利用網卡卸載,NVLS 利用 NVLink SHARP 硬體加速,PAT 是 NVLS 的平行化變體。

如果沒有演算法選擇,NCCL 只能用一種固定模式通訊,無法適應不同消息大小和拓撲結構,效能會大打折扣。

資料結構與記憶體佈局

Ring 演算法

Ring 演算法的核心是ncclRing結構體(在src/include/comm.h中透過channels[i].ring引用)。📎 src/include/collectives.h:81-116定義了RingAlgorithm基類:

c
class RingAlgorithm {
protected:
  int refCount;
  int nRanks;
  int nStepsPerLoop;
  int chunkSteps;
  int sliceSteps;
  ssize_t sliceSize;
  ssize_t loopSize;
  ssize_t channelSize;
  uint8_t* sendbuff;
  uint8_t* recvbuff;
  void* sendMhandle;
  void* recvMhandle;
  void* srecvMhandle;

public:
  virtual void getNextSendAddr(int curStep, uint8_t** sendbuffOut, size_t* sizeOut, void** mhandleOut) = 0;
  virtual void getNextRecvAddr(int curStep, uint8_t** recvbuffOut, size_t* sizeOut, void** mhandleOut) = 0;
  int incRefCount() {
    return (int)COMPILER_ATOMIC_ADD_FETCH(&refCount, 1, std::memory_order_relaxed);
  }
  int decRefCount() {
    return (int)COMPILER_ATOMIC_SUB_FETCH(&refCount, 1, std::memory_order_release);
  }
  RingAlgorithm() {
    refCount = 0;
  }
  virtual ~RingAlgorithm() {};
};

關鍵欄位解析

  • refCount:引用計數,用於 proxy 執行緒和 GPU kernel 共享演算法物件。
  • nRanks:環上節點數。
  • nStepsPerLoop:每輪迴圈的步數。AllReduce 是2*(nRanks-1)*chunkSteps(📎 src/include/collectives.h:218-218)。
  • chunkSteps / sliceSteps:塊步數和切片步數,控制流水線粒度。
  • sliceSize / loopSize / channelSize:切片大小、迴圈大小、通道大小。
  • sendbuff / recvbuff:發送和接收緩衝區指標。
  • sendMhandle / recvMhandle / srecvMhandle:記憶體句柄,用於網路註冊。

引用計數的原子操作

📎 src/include/collectives.h:106-108展示了incRefCount和decRefCount:

c
int incRefCount() {
  return (int)COMPILER_ATOMIC_ADD_FETCH(&refCount, 1, std::memory_order_relaxed);
}
int decRefCount() {
  return (int)COMPILER_ATOMIC_SUB_FETCH(&refCount, 1, std::memory_order_release);
}
〔設計推斷與架構權衡〕

incRefCount使用memory_order_relaxed——增加引用計數不需要同步,只要保證原子性即可。decRefCount使用memory_order_release——減少引用計數時,需要確保之前的寫操作對其他執行緒可見(因為可能觸發物件銷毀)。

RingARAlgorithm:AllReduce 的 Ring 實作

📎 src/include/collectives.h:118-234定義了RingARAlgorithm,繼承自RingAlgorithm。核心方法是getNextSendAddr和getNextRecvAddr。

📎 src/include/collectives.h:126-167的getNextSendAddr邏輯:

c
void getNextSendAddr(int curStep, uint8_t** sendbuffOut, size_t* sizeOut, void** mhandleOut) {
  int curLoop = curStep / nStepsPerLoop;
  int curLoopStage = (curStep % nStepsPerLoop) / chunkSteps;
  int chunkStage = curLoopStage % nRanks;
  int sliceStage = (curStep % chunkSteps) / sliceSteps;
  ssize_t elemOffset = curLoop * loopSize;
  ssize_t remSize = channelSize - elemOffset;
  // ... 计算 chunkOffset, sliceOffset, curSliceSize ...
  if (remSize < loopSize) {
    curChunkSize = alignUp(divUp(remSize / elemSize, nRanks), 16 / elemSize) * elemSize;
  } else {
    curChunkSize = chunkSize;
  }
  chunkId = (ringIndex + nRanks - 1 - chunkStage) % nRanks;
  chunkOffset = chunkId * curChunkSize;
  nelem = std::min(remSize - chunkOffset, curChunkSize);
  curSliceSize = std::max(divUp(nelem / elemSize, 16 * slicePerChunk) * 16, sliceSize / elemSize / 32) * elemSize;
  sliceOffset = sliceStage * curSliceSize;
  // ... 设置 sendbuffOut, sizeOut, mhandleOut ...
}
〔設計推斷與架構權衡〕

這段程式碼的核心是位址計算:給定當前步數curStep,計算出應該發送哪個資料塊的哪個切片。chunkId的計算(ringIndex + nRanks - 1 - chunkStage) % nRanks實作了環上的反向傳播——每個 rank 從前驅接收資料,處理後發送給後繼。

PAT 演算法

PAT(Parallel Aggregated Tree)是 NVLS 的平行化變體。📎 src/include/collectives.h:416-423定義了ncclPatStep:

c
struct ncclPatStep {
  int recvDim, sendDim, recvOffset, sendOffset, stepOffset, postRecv, postSend, nelem, last, flags;
  // PAT algo computation thread step number; -1 while the slot is free.
  int step;
  // This PAT group's offset within the shared NVLS slot.
  int nvlsOffset;
  size_t inpIx, outIx;
};

📎 src/include/collectives.h:425-435定義了ncclPatPeer:

c
struct ncclPatPeer {
  uint64_t step;
  struct ncclConnInfo* conn;
  struct ncclConnFifo* connFifo;
  void* buff;
  uint64_t* headPtr;
  uint64_t* tailPtr;
  uint64_t stepCache;
  long long int accSize;
  int connStepSize;
};
〔設計推斷與架構權衡〕

PAT 演算法的核心思想是聚合多個小步驟為一個大步驟,減少同步開銷。ncclPatStep描述一個聚合步驟的收發維度、偏移量、元素數等資訊。ncclPatPeer描述一個對等節點的連接狀態和緩衝區指標。

場景驅動 Walkthrough:Ring AllReduce 的步驟演化

假設 4 個 rank(0, 1, 2, 3),每個 rank 有 4 個元素,執行 Ring AllReduce。

Reduce-Scatter 階段

  • 步驟 0:rank 0 發送元素 0 給 rank 1,rank 1 發送元素 1 給 rank 2,rank 2 發送元素 2 給 rank 3,rank 3 發送元素 3 給 rank 0。
  • 步驟 1:每個 rank 將收到的元素與本地對應元素相加,然後發送給下一個 rank。
  • 步驟 2:繼續累加和傳遞。
  • 步驟 3:此時每個 rank 擁有一個完整的歸約結果(rank 0 有元素 3 的結果,rank 1 有元素 0 的結果,等等)。

AllGather 階段

  • 步驟 4-6:每個 rank 將自己擁有的歸約結果沿環傳播,最終所有 rank 擁有完整結果。

📎 src/include/collectives.h:218-218的nStepsPerLoop = 2 * (nRanks - 1) * chunkSteps正好對應這個流程:Reduce-Scatter 需要(nRanks-1)*chunkSteps步,AllGather 也需要(nRanks-1)*chunkSteps步,總共2*(nRanks-1)*chunkSteps步。

設計思考與生產踩坑

為什麼 Ring 和 Tree 並存?

〔設計推斷與架構權衡〕

Ring 演算法的頻寬利用率高(每條鏈路都在傳輸),但延遲隨 rank 數線性增長。Tree 演算法的延遲是對數級的,但頻寬利用率低(只有部分鏈路在工作)。NCCL 根據消息大小自動選擇:小消息用 Tree(延遲敏感),大消息用 Ring(頻寬敏感)。

踩坑場景一:演算法選擇錯誤

〔設計推斷與架構權衡〕

如果手動強制使用 Ring 處理小消息,延遲會顯著增加。 建議讓 tuning 模組自動選擇,除非有明確的效能分析資料支持手動干預。

踩坑場景二:NVLS 硬體不支援

NVLS 需要特定的硬體支援(NVLink SHARP)。如果硬體不支援但程式碼強制使用 NVLS,會回退到 Ring 或 Tree,但可能伴隨效能抖動。📎 src/include/comm.h:755-755的nvlsSupport欄位標記硬體是否支援 NVLS。

踩坑場景三:PAT 演算法的聚合因子配置

PAT 演算法的aggFactor決定了聚合多少個步驟。📎 src/include/collectives.h:537-560展示了aggFactor的計算邏輯:

c
aggFactor = 1;
size_t channelSize = end - offset;
while (stepSize / (channelSize * sizeof(T) * aggFactor) >= 2 && aggFactor < nranks / 2) {
  aggFactor *= 2;
  aggDelta /= 2;
}
postFreq = aggFactor;
if (postFreq < parallelFactor) parallelFactor = postFreq;
int d = stepDepth;
while (d > 1 && aggFactor < nranks / 2) {
  d /= 2;
  aggFactor *= 2;
  aggDelta /= 2;
}
〔設計推斷與架構權衡〕

aggFactor過小會導致同步開銷大,過大則會導致流水線氣泡。NCCL 根據stepSize、channelSize、nranks自動計算最佳值。

2.4 協定 protocol:LL/LL128/Simple 三種資料搬運策略

直覺模型

寄快遞可以選「同城閃送」「次日達」或「普通快遞」,速度和成本不同。NCCL 的協議就是這些「寄法」——LL(Low Latency)適合小訊息的低延遲傳輸,LL128 適合中等訊息的 128 位元組對齊傳輸,Simple 適合大訊息的高頻寬傳輸。

如果沒有協議選擇,NCCL 只能用一種固定策略搬運資料,無法在延遲和頻寬之間取得平衡。

資料結構與記憶體佈局

協議列舉

📎 src/include/comm.h:55-57定義了協議相關的執行緒閾值:

c
#define NCCL_LL_THREAD_THRESHOLD 8
#define NCCL_LL128_THREAD_THRESHOLD 8
#define NCCL_SIMPLE_THREAD_THRESHOLD 64
〔設計推斷與架構權衡〕

這些閾值決定了每種協議使用多少個執行緒。LL 和 LL128 用 8 個執行緒(低延遲,少量執行緒即可),Simple 用 64 個執行緒(高頻寬,需要更多執行緒並行搬運)。

協議緩衝區

📎 src/include/comm.h:691-691定義了buffSizes[NCCL_NUM_PROTOCOLS]——每種協議有獨立的緩衝區大小。

協議相關的 FIFO 結構

📎 src/include/comm.h:59-83定義了ncclSendMem和ncclRecvMem:

c
struct ncclSendMem {
  union {
    struct {
      uint64_t head;
      char pad1[CACHE_LINE_SIZE - sizeof(uint64_t)];
      void* ptrExchange;
      uint64_t redOpArgExchange[2];
      char pad2[CACHE_LINE_SIZE - sizeof(void*) - 2 * sizeof(uint64_t)];
      int offsFifo[NCCL_STEPS];
    };
    char pad3[MEM_ALIGN];
  };
};

struct ncclRecvMem {
  union {
    struct {
      uint64_t tail;
      char pad1[CACHE_LINE_SIZE - sizeof(uint64_t)];
      struct ncclConnFifo connFifo[NCCL_STEPS];
      int flush; // For GDRCopy-based flush
    };
    char pad4[MEM_ALIGN];
  };
};
〔設計推斷與架構權衡〕

ncclSendMem和ncclRecvMem是發送和接收的共享記憶體結構。head和tail是環形緩衝區的讀寫指標,pad1確保它們在不同快取行。connFifo陣列儲存每個步驟的連接資訊(模式、偏移、大小、指標),定義在📎 src/include/collectives.h:72-77:

c
struct ncclConnFifo {
  int mode;
  ssize_t offset;
  ssize_t size;
  void* ptr;
};

協議選擇邏輯

〔設計推斷與架構權衡〕

協議選擇由 tuning 模組完成,考慮因素包括:

  • 訊息大小:小訊息用 LL,中等用 LL128,大訊息用 Simple。
  • 拓撲結構:NVLink 連接適合 LL128,網路連接適合 Simple。
  • 硬體能力:某些 GPU 架構對特定協議有優化。

場景驅動 Walkthrough:LL 協議的資料搬運

假設使用 LL 協議傳輸 1KB 資料。

第一步:資料寫入發送緩衝區

主機側將資料寫入sendbuff,然後更新ncclSendMem.head指標,通知 GPU kernel 有新資料。

第二步:GPU kernel 讀取資料

GPU kernel 輪詢head指標,發現新資料後,從sendbuff讀取資料。

第三步:資料傳輸

GPU kernel 透過 NVLink 或網路將資料發送到目標 rank。

第四步:目標 rank 接收資料

目標 rank 的 GPU kernel 將資料寫入recvbuff,然後更新ncclRecvMem.tail指標。

第五步:主機側讀取資料

主機側輪詢tail指標,發現新資料後,從recvbuff讀取資料。

並發控制與硬體互動

LL 協議的低延遲機制

〔設計推斷與架構權衡〕

LL 協議使用輪詢(Polling)而非中斷來檢測資料到達。GPU kernel 不斷讀取head指標,一旦發現變化立即處理。這比中斷方式延遲更低,但會佔用 GPU 計算資源。

LL128 協議的 128 位元組對齊

〔設計推斷與架構權衡〕

LL128 協議要求資料按 128 位元組對齊,這樣每次傳輸正好填滿一個快取行。對齊的好處是:

  • 減少部分快取行寫入(Partial Cache Line Write)
  • 提高記憶體頻寬利用率
  • 簡化硬體處理邏輯

Simple 協議的批量傳輸

〔設計推斷與架構權衡〕

Simple 協議使用批量傳輸模式:積累一定量的資料後一次性發送,減少同步次數。這適合大訊息場景,因為同步開銷被分攤到大量資料上。

生產避坑指南

踩坑場景一:協議與訊息大小不匹配

〔設計推斷與架構權衡〕

如果強制使用 LL 協議傳輸大訊息,效能會急劇下降。 因為 LL 協議的設計目標是低延遲,不是高頻寬。大訊息應該用 Simple 協議。

踩坑場景二:LL128 對齊問題

〔設計推斷與架構權衡〕

如果資料沒有按 128 位元組對齊,LL128 協議會回退到 LL 或 Simple,導致效能不穩定。 建議確保發送緩衝區和接收緩衝區都按 128 位元組對齊。

踩坑場景三:協議切換開銷

〔設計推斷與架構權衡〕

在執行時動態切換協議會帶來額外開銷。 NCCL 在初始化時確定協議,執行時不再切換。如果需要切換,必須重新初始化通訊域。

2.5 傳輸層 transport:P2P/SHM/NET/CollNet 底層搬運通道

直覺模型

從 A 點到 B 點可以走路、騎車、坐地鐵或打車,NCCL 的傳輸層就是這些不同的「出行方式」。上層不關心具體怎麼走,只關心能不能送到。P2P 是「走路」(同機 GPU 直連),SHM 是「騎車」(共享記憶體),NET 是「坐地鐵」(網路),CollNet 是「打車」(網卡卸載)。

如果沒有傳輸層抽象,上層演算法需要針對每種物理鏈路寫不同的程式碼,無法復用。

資料結構與記憶體佈局

傳輸層列舉

📎 src/include/transport.h:18-23定義了傳輸層類型:

c
#define NTRANSPORTS 4
#define TRANSPORT_UNDEFINED -1
#define TRANSPORT_P2P 0
#define TRANSPORT_SHM 1
#define TRANSPORT_NET 2
#define TRANSPORT_COLLNET 3

傳輸層介面

📎 src/include/transport.h:129-146定義了ncclTransportComm——傳輸層的通訊介面:

c
struct ncclTransportComm {
  ncclResult_t (*setup)(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo*, struct ncclPeerInfo*,
                        struct ncclConnect*, struct ncclConnector*, int channelId, int connIndex);
  ncclResult_t (*connect)(struct ncclComm* comm, struct ncclConnect*, int nranks, int rank, struct ncclConnector*);
  ncclResult_t (*free)(struct ncclComm* comm, struct ncclConnector*);
  ncclResult_t (*proxySharedInit)(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState,
                                  int nChannels);
  ncclResult_t (*proxySetup)(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState, void* reqBuff,
                             int reqSize, void* respBuff, int respSize, int* done);
  ncclResult_t (*proxyConnect)(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState, void* reqBuff,
                               int reqSize, void* respBuff, int respSize, int* done);
  ncclResult_t (*proxyFree)(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState);
  ncclResult_t (*proxyProgress)(struct ncclProxyState* proxyState, struct ncclProxyArgs*);
  ncclResult_t (*proxyRegister)(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState,
                                void* reqBuff, int reqSize, void* respBuff, int respSize, int* done);
  ncclResult_t (*proxyDeregister)(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState,
                                  void* reqBuff, int reqSize, int* done);
};

關鍵回呼解析

  • setup:建立連接前的準備工作,交換連接參數。
  • connect:實際建立連接。
  • free:釋放連接資源。
  • proxySharedInit:初始化 proxy 執行緒共享資源。
  • proxySetup / proxyConnect:proxy 執行緒側的連線建立。
  • proxyProgress:proxy 執行緒推進資料傳輸。
  • proxyRegister / proxyDeregister:記憶體註冊和註銷。

傳輸層結構體

📎 src/include/transport.h:148-154定義了ncclTransport:

c
struct ncclTransport {
  const char name[8];
  ncclResult_t (*canConnect)(int*, struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo*,
                             struct ncclPeerInfo*);
  struct ncclTransportComm send;
  struct ncclTransportComm recv;
};
〔設計推斷與架構權衡〕

name是傳輸層名稱(如 "P2P"、"SHM"、"NET"),canConnect判斷兩個 rank 之間是否可以使用該傳輸層,send和recv分別是發送和接收方向的通訊介面。

傳輸層實例

📎 src/include/transport.h:36-36宣告了四個傳輸層實例:

c
extern struct ncclTransport p2pTransport;
extern struct ncclTransport shmTransport;
extern struct ncclTransport netTransport;
extern struct ncclTransport collNetTransport;

📎 src/include/transport.h:36-36定義了傳輸層陣列:

c
extern struct ncclTransport* ncclTransports[];

對等節點資訊

📎 src/include/transport.h:46-74定義了ncclPeerInfo——rank 之間交換的元資料:

c
struct ncclPeerInfo {
  int rank;
  int cudaDev;
  int nvmlDev;
  int gdrSupport;
  uint64_t hostHash;
  uint64_t pidHash;
  dev_t shmDev;
  int64_t busId;
  cudaUUID_t gpuUuid;
  struct ncclComm* comm;
  int cudaCompCap;
  int gpuCftSupport;
  size_t totalGlobalMem;
  // MNNVL support
  nvmlGpuFabricInfoV_t fabricInfo;
  int fabricHandleSupport;
  int cuMemSupport;
  int version;
  uint64_t supportedGinTypeBitMask;
  bool crossNicSupport;
  bool rmaPluginAvailable;
  bool cuMemGdrSupport;
  int mloPart; // MLOPart partition index, or -1 if not an MLOPart GPU
  int cudaDriverVersion;
  bool gpuCftMulticastSupport;
  bool gpuCftCountedSupport;
  uint32_t gitVersionHash;
};
〔設計推斷與架構權衡〕

這些欄位用於判斷兩個 rank 之間可以使用哪種傳輸層:

  • hostHash相同 → 同一主機 → 可用 P2P 或 SHM
  • hostHash不同 → 不同主機 → 必須用 NET
  • gdrSupport→ 是否支援 GPUDirect RDMA
  • cudaCompCap→ GPU 運算能力,影響協定選擇

場景驅動 Walkthrough:建立 P2P 連線

假設兩個 rank 在同一主機內,NCCL 選擇 P2P 傳輸層。

第一步:交換 PeerInfo

兩個 rank 透過 bootstrap 通道交換ncclPeerInfo,確認彼此在同一主機、GPU 支援 P2P。

第二步:呼叫 canConnect

📎 src/include/transport.h:148-154的canConnect回呼被呼叫,檢查拓撲圖確認兩個 GPU 之間有 NVLink 或 PCIe 連線。

第三步:呼叫 setup

p2pTransport.send.setup和p2pTransport.recv.setup被呼叫,準備連線參數(如 IPC 句柄)。

第四步:呼叫 connect

p2pTransport.send.connect和p2pTransport.recv.connect被呼叫,實際建立連線。

第五步:註冊記憶體

如果需要 RDMA,呼叫proxyRegister註冊發送和接收緩衝區。

並發控制與硬體互動

P2P 傳輸層

〔設計推斷與架構權衡〕

P2P 使用 CUDA IPC(Inter-Process Communication)機制,允許一個 GPU 直接存取另一個 GPU 的顯存。這需要:

  • 兩個 GPU 在同一 PCIe 域或 NVLink 域
  • 作業系統支援 CUDA IPC
  • 足夠的權限

SHM 傳輸層

〔設計推斷與架構權衡〕

SHM 使用主機共享記憶體作為中轉。當兩個 GPU 之間沒有直接連線時,資料先拷貝到主機記憶體,再拷貝到目標 GPU。這比 P2P 慢,但相容性更好。

NET 傳輸層

〔設計推斷與架構權衡〕

NET 使用網路裝置(InfiniBand 或 RoCE)傳輸資料。這需要:

  • 網路裝置支援 GPUDirect RDMA(可選,但推薦)
  • 正確的網路配置(IP 位址、子網路遮罩等)
  • 足夠的網路頻寬

CollNet 傳輸層

〔設計推斷與架構權衡〕

CollNet 利用網卡的集合通訊卸載能力(如 NVIDIA SHARP)。網卡直接執行歸約操作,減少 GPU 的運算負擔。這需要:

  • 支援 SHARP 的網卡
  • 正確的 SHARP 配置

生產避坑指南

踩坑場景一:P2P 不可用

〔設計推斷與架構權衡〕

如果兩個 GPU 之間沒有 NVLink 且 PCIe 拓撲不支援 P2P,NCCL 會回退到 SHM。 這會導致效能下降。可以透過NCCL_P2P_DISABLE=1強制停用 P2P,觀察效能變化。

踩坑場景二:網路配置錯誤

〔設計推斷與架構權衡〕

如果網路裝置的 IP 位址配置錯誤,NET 傳輸層無法建立連線。 常見錯誤包括:子網路遮罩錯誤、路由表缺失、防火牆阻擋。建議用ibstat和ibping檢查 InfiniBand 連線。

踩坑場景三:GPUDirect RDMA 未啟用

〔設計推斷與架構權衡〕

如果gdrSupport為 0,NET 傳輸層會回退到「先拷貝到主機記憶體再發送」模式,延遲顯著增加。 檢查nvidia-peermem模組是否載入,以及網卡驅動是否支援 GPUDirect。

2.6 五件套如何組合:一次通訊的完整生命週期

組合關係圖

mermaid
flowchart TD
    api["ncclAllReduce(sendbuff, recvbuff, count, ...)"] --> comm_lookup["查找 ncclComm"]
    comm_lookup --> task_create["创建 ncclTaskColl"]
    task_create --> tuning{"tuning 模块选择算法和协议"}
    tuning -->|"小消息"| tree_ll["Tree + LL"]
    tuning -->|"中等消息"| ring_ll128["Ring + LL128"]
    tuning -->|"大消息"| ring_simple["Ring + Simple"]
    tuning -->|"NVLS 可用"| nvls["NVLS + Simple"]
    tree_ll --> channel_assign["分配通道"]
    ring_ll128 --> channel_assign
    ring_simple --> channel_assign
    nvls --> channel_assign
    channel_assign --> transport_select{"选择传输层"}
    transport_select -->|"同机 GPU 直连"| p2p["P2P"]
    transport_select -->|"同机无直连"| shm["SHM"]
    transport_select -->|"跨机"| net["NET"]
    transport_select -->|"CollNet 可用"| collnet["CollNet"]
    p2p --> kernel_launch["启动 GPU kernel"]
    shm --> kernel_launch
    net --> kernel_launch
    collnet --> kernel_launch
    kernel_launch --> execute["执行通信"]
    execute --> complete["完成,更新 asyncResult"]

完整生命週期

階段一:API 呼叫

使用者呼叫ncclAllReduce,傳入發送緩衝區、接收緩衝區、元素數、資料類型、歸約操作、通訊域、CUDA stream。

階段二:任務建立

NCCL 建立ncclTaskColl結構體(📎 src/include/comm.h:212-273),填充func(AllReduce)、sendbuff、recvbuff、count、datatype、opHost等欄位。

階段三:演算法和協定選擇

Tuning 模組根據訊息大小、拓撲結構、硬體能力選擇演算法(Ring/Tree/NVLS)和協定(LL/LL128/Simple)。選擇結果寫入ncclTaskColl的algorithm和protocol欄位(📎 src/include/comm.h:227-227)。

階段四:通道分配

根據演算法和協定,確定使用的通道數和通道範圍。nChannels、channelLo、channelHi欄位被設定(📎 src/include/comm.h:254-257)。

階段五:傳輸層選擇

根據拓撲圖,為每對 rank 選擇傳輸層(P2P/SHM/NET/CollNet)。連線資訊儲存在channels[i].peers[j]中。

階段六:Kernel 啟動

NCCL 建構ncclKernelPlan(📎 src/include/comm.h:357-410),包含工作佇列、清理佇列、任務佇列等。然後啟動 GPU kernel。

階段七:執行通訊

GPU kernel 讀取工作 FIFO,執行資料傳輸和歸約操作。Proxy 執行緒非同步推進網路 I/O。

階段八:完成

所有通道完成後,asyncResult被設定為ncclSuccess。使用者可以透過ncclCommGetAsyncError查詢狀態。

設計思考

為什麼需要五件套?

〔設計推斷與架構權衡〕

這五個抽象分別解決了不同維度的問題:

  • ncclComm:解決「誰和誰通訊」的問題。
  • channel:解決「如何並行」的問題。
  • algorithm:解決「用什麼拓撲」的問題。
  • protocol:解決「用什麼策略」的問題。
  • transport:解決「走什麼物理鏈路」的問題。

它們正交組合,讓 NCCL 能夠適應各種硬體配置和訊息大小,而不需要為每種組合寫專門的程式碼。

組合的靈活性

〔設計推斷與架構權衡〕

五件套的組合數量是:

  • 演算法:5 種(Tree/Ring/CollNet/NVLS/PAT)
  • 協定:3 種(LL/LL128/Simple)
  • 傳輸層:4 種(P2P/SHM/NET/CollNet)

本章思考與自測

Q1: 如果將📎 src/include/comm.h:731-731中的intraPad1[64 - sizeof(uint64_t)]改為intraPad1[0](即去掉快取行填充),在多行程場景下會出現什麼效能問題?為什麼?

參考解析:

去掉填充後,intraBarrierPhase、intraBarrierCounter、intraBarrierGate三個欄位會緊密排列在記憶體中,很可能共享同一個快取行(通常 64 位元組)。

在多行程場景下,每個行程有自己的ncclComm副本,但intraComm0指向的 leader 通訊域的intraBarrierCounter和intraBarrierGate會被所有行程讀寫。當行程 A 呼叫ncclCommIntraBarrierIn更新intraBarrierCounter(📎 src/include/comm.h:943-959)時,會導致行程 B 的intraBarrierGate快取行失效。行程 B 在ncclCommIntraBarrierOut中輪詢intraBarrierGate(📎 src/include/comm.h:962-977),每次快取失效都要重新從記憶體載入,延遲從奈秒級上升到微秒級。

這就是偽共享(False Sharing)問題。填充 56 位元組確保每個欄位獨占一個快取行,消除偽共享。

Q2: 如果將📎 src/include/collectives.h:106-108的incRefCount從memory_order_relaxed改為memory_order_seq_cst,會有什麼影響?為什麼作者選擇relaxed?

參考解析:

memory_order_seq_cst會強制全域順序一致性,每次增加引用計數都要插入記憶體屏障,導致效能下降。

incRefCount只需要保證原子性,不需要同步其他記憶體操作。因為增加引用計數不會觸發物件銷毀,也不會依賴其他執行緒的寫操作。memory_order_relaxed正好滿足這個需求——只保證原子性,不插入屏障。

相比之下,decRefCount(📎 src/include/collectives.h:109-111)使用memory_order_release,因為減少引用計數可能觸發物件銷毀,需要確保之前的寫操作對其他執行緒可見。

這是 C++ 記憶體模型的經典應用:根據操作語義選擇最弱的記憶體序,在保證正確性的前提下最大化效能。

Q3: 如果將📎 src/include/channel.h:32-32的reverseBits(base, log2Up(comm->p2pnChannels))改為直接返回base % comm->p2pnChannels,在什麼場景下會導致效能下降?為什麼?

參考解析:

reverseBits是位反轉操作,用於打散通道分配。直接取模會導致通道分配呈現規律性:round 0 用通道 0,round 1 用通道 1,...,round N 用通道 N%p2pnChannels。

在多節點場景下,如果多個 rank 的 P2P 通訊同時進行,規律性的通道分配會導致熱點集中——某些通道被多個 rank 同時使用,而其他通道閒置。這會造成鏈路壅塞,降低整體頻寬利用率。

reverseBits打散了通道分配,讓不同 round 使用看似隨機的通道,均勻分布負載。這是負載均衡的經典手法。

另外,reverseBits是純位操作,比取模運算更快(取模需要除法指令,位操作只需幾條指令)。

---

下一章我們將深入ncclCommInitRank的內部實現,看看 NCCL 如何從一個空的ncclComm結構體開始,逐步建立拓撲圖、初始化通道、建立傳輸連接,最終構建出一個可用的通訊域。本章建立的五件套心智模型,將在下一章中逐一落地。

這五個抽象並非孤立存在:通訊域是容器,通道是並行執行的單位,演算法決定資料如何歸約,協定規定資料如何編碼,傳輸層負責資料如何移動。它們的組合——5 個維度、每個維度 3 到 4 種選擇——構成了 NCCL 效能調優的搜尋空間。那麼,這個通訊域物件究竟是如何從零開始被構建出來的?下一章我們將深入 ncclCommInitRank 的呼叫鏈,看 NCCL 如何在初始化階段完成裝置探測、拓撲發現與通道分配,並揭示 comm->rank、comm->nRanks、comm->channels 等關鍵欄位的賦值時機。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 03

第 3 章:初始化入局:ncclCommInitRank 如何把一群孤立行程建立成通訊域

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 3 章 / 共 25 章

上一章我們建立了貫穿全書的五個核心抽象:ncclComm、channel、algorithm、protocol 和 transport,它們共同構成了「一次通訊 = 若干 channel × 一個 algorithm × 一個 protocol × 若干 transport」的公共詞彙表。現在,我們要回答一個更根本的問題:這個 ncclComm 物件究竟是如何從無到有構建出來的?當你呼叫 ncclCommInitRank 時,NCCL 需要在幾百毫秒內完成一系列複雜操作:確認所有 rank 到齊、交換裝置資訊、探測機器拓撲、計算資料路徑、分配 GPU 記憶體與主機記憶體,最終將這一切打包成一個 ncclComm 物件。本章將沿著這條呼叫鏈,從 API 入口一路下鑽到 initTransportsRank 的最後一根微血管。

3.1 API 入口:ncclCommInitRank 的同步外殼與非同步核心

直覺模型

ncclCommInitRank表面上是「建一個通訊域」,實際上它做的是「發起一個背景任務,然後(預設情況下)等它完成」。這就像你去餐廳點餐:點餐這個動作(API 呼叫)瞬間返回,但廚房做菜(真正的初始化)是在背景進行的。預設的「阻塞模式」只是讓你在櫃檯前等到菜做好,而「非阻塞模式」則給你一個取餐號,你可以先去幹別的。

如果沒有這層非同步設計,NCCL 在初始化期間就無法與 CUDA Graph 捕獲、多通訊域並行初始化等場景配合——所有初始化都會變成串列的、無法與使用者程式碼重疊的阻塞操作。

資料結構與記憶體佈局

先看 API 入口本身。ncclCommInitRank是一個極薄的同步外殼:

📎 src/init.cc:2946-2970

它做了四件事:呼叫ncclInitEnv()載入環境變數外掛、開啟 NVTX 效能標記、讀取當前 CUDA 裝置號、然後呼叫ncclGroupStartInternal()進入 group 語義,最後把實際工作委託給ncclCommInitRankDev。

注意ncclGroupStartInternal() / ncclGroupEndInternal()這一對呼叫——即使你只初始化一個通訊域,NCCL 也把它包在 group 語義裡。這是為了統一處理「使用者在一個 group 裡初始化多個通訊域」的場景,避免為單通訊域和多通訊域寫兩套程式碼路徑。

真正的參數校驗和物件分配在ncclCommInitRankDev裡:

📎 src/init.cc:2851-2943

這個函式是整條鏈路的「總排程台」。它先做參數校驗(nId範圍、nranks/myrank合法性),然後分配ncclComm結構體本身,以及三個與中止機制相關的欄位:abortFlag(主機側原子標誌)、abortFlagDev(裝置側可見的固定記憶體副本)、abortFlagRefCount(引用計數,因為 split 出來的子通訊域可能共享父通訊域的 abortFlag)。

這裡有一個值得注意的細節——comm->startMagic = comm->endMagic = NCCL_MAGIC:

📎 src/init.cc:2886-2886

這對 magic 值像「封條」一樣夾在ncclComm結構體的首尾。任何越界寫入或結構體損壞都會破壞這對 magic,後續操作可以透過校驗它們來檢測記憶體踩踏。這是一種廉價但有效的記憶體完整性防護。

Step-by-Step Walkthrough

當ncclCommInitRankDev走到最後,它構造一個ncclCommInitRankAsyncJob並啟動非同步任務:

📎 src/init.cc:2896-2929

job結構體承載了所有初始化所需的參數。注意job->commId是拷貝出來的,而不是直接引用使用者傳入的commId:

📎 src/init.cc:2903-2910

為什麼要拷貝?原始碼註解給出了答案:ncclUniqueId和ncclBootstrapHandle的對齊要求不同,使用者傳入的陣列可能沒有正確對齊到ncclBootstrapHandle所需的邊界。拷貝到新分配的記憶體可以保證對齊。這是一個典型的「ABI 相容性陷阱」——使用者看到的是ncclUniqueId,內部要當ncclBootstrapHandle用,兩者大小相同但對齊不同。

最後,根據ncclParamEnqueueRearchEnable()的值,任務要麼進入管理佇列,要麼直接透過ncclAsyncLaunch啟動:

📎 src/init.cc:2922-2929

ncclAsyncLaunch會建立一個新執行緒執行ncclCommInitRankFunc。如果是阻塞模式(預設),呼叫方會在ncclGroupEndInternal()裡等待這個執行緒完成;如果是非阻塞模式,呼叫方立即返回,使用者後續透過ncclCommGetAsyncError輪詢狀態。

設計思考

這裡的設計核心是「同步 API + 非同步實作」。為什麼不讓ncclCommInitRank直接同步執行所有初始化?因為 NCCL 需要支援ncclCommInitRankConfig的非阻塞模式,而非阻塞模式要求初始化在背景執行緒執行。如果同步路徑和非同步路徑是兩套程式碼,維護成本會翻倍。統一走非同步、同步路徑只是「啟動後立即等待」,程式碼只有一份。

mermaid
flowchart TD
    api["ncclCommInitRank(newcomm, nranks, commId, myrank)"]
    env["ncclInitEnv() 加载环境变量插件"]
    group["ncclGroupStartInternal()"]
    dev["ncclCommInitRankDev(...)"]
    check{"nId/nranks/myrank 合法?"}
    alloc["ncclCalloc 分配 comm + abortFlag"]
    parse["parseCommConfig() 解析配置"]
    job["构造 ncclCommInitRankAsyncJob"]
    copyid["拷贝 commId 保证对齐"]
    enq{"ncclParamEnqueueRearchEnable()?"}
    mgmt["ncclMgmtTaskEnqueue()"]
    async["ncclAsyncLaunch() 启动后台线程"]
    func["ncclCommInitRankFunc() 执行初始化"]
    fail["返回 ncclInvalidArgument"]

    api --> env --> group --> dev --> check
    check -->|否| fail
    check -->|是| alloc --> parse --> job --> copyid --> enq
    enq -->|是| mgmt --> func
    enq -->|否| async --> func

3.2 Bootstrap:rank 之間的第一條控制通道

直覺模型

Bootstrap 是 NCCL 的「會前微信群」。在正式通訊開始之前,所有 rank 需要先建立一條控制通道,用來交換「我是誰、我在哪台機器、我的 GPU 是什麼型號、我的網卡位址是什麼」這些元資料。沒有 bootstrap,rank 之間就是一群互不相識的陌生人,無法協調任何通訊。

如果 bootstrap 失敗或逾時,整個通訊域初始化就會卡死——這是生產環境中最常見的 NCCL 掛起原因之一。

資料結構與記憶體佈局

Bootstrap 的核心狀態保存在bootstrapState結構體中:

📎 src/bootstrap.cc:527-546

這個結構體有幾個關鍵欄位值得展開:

  • ring:一個聯合體,要麼是網路設備句柄(net.sendComm/net.recvComm),要麼是一對 socket(socket.send/socket.recv)。這對應兩種 bootstrap 模式:基於 socket 的預設模式和基於網路設備的NCCL_OOB_NET_ENABLE模式。
  • listen:監聽端資訊,同樣有網路和 socket 兩種形態。
  • peerP2pAddresses / peerProxyAddresses:所有 rank 的 P2P 位址和 proxy 位址陣列,透過 ring allgather 填充。
  • unexpectedConnections:一個鏈結串列,快取「收到了但還沒被匹配」的連線。這是 bootstrap 協定的一個關鍵設計——因為接收方無法預知誰會先連過來,所以必須先把不匹配的連線存起來。
  • asyncSendQueue + asyncSendLock + asyncSendCond:非同步發送佇列及其同步原語,用於 TLS 加密模式下的並行發送。

bootstrapState的分配發生在bootstrapInit開頭:

📎 src/bootstrap.cc:769-776

注意comm->bootstrap = state這一行——bootstrap 狀態被掛到通訊域上,後續所有 bootstrap 操作都透過comm->bootstrap存取。

Step-by-Step Walkthrough

bootstrapInit是 bootstrap 的主幹函式。讓我們按執行順序拆解:

第一步:確定 magic 值。magic 是 bootstrap 通訊的「暗號」,只有持有相同 magic 的 rank 才能互相連線。

📎 src/bootstrap.cc:778-788

如果是正常初始化(handles != NULL),magic 來自第一個 handle;如果是 split/grow(parent != NULL),magic 透過hashCombine(parent->magic, parent->childCount)派生。這保證了每個子通訊域有唯一的 magic。

第二步:建立監聽 socket。每個 rank 需要兩個監聽端點:一個用於 ring 鄰居連線(STATE_LISTEN(state, socket)),一個用於 root 連線(listenSockRoot):

📎 src/bootstrap.cc:797-831

這裡有一個關鍵的分工:ring 監聽 socket 使用comm->magic,而 root 監聽 socket 使用BOOTSTRAP_HANDLE(handles, curr_root)->magic。為什麼?因為 root 是全域協調者,所有 rank 都要連它,所以它用統一的 magic;而 ring 鄰居是點對點的,用通訊域自己的 magic 就夠了。

第三步:錯峰連線。當 rank 數量很大時,所有 rank 同時連 root 會造成連線風暴。NCCL 用NCCL_UID_STAGGER_RATE和NCCL_UID_STAGGER_THRESHOLD來控制錯峰:

📎 src/bootstrap.cc:833-843

當某個 root 負責的 rank 數超過閾值(預設 256)時,每個 rank 根據自己在 root 下的局部 ID 計算延遲微秒數,然後 sleep。這是一個簡單但有效的「令牌桶」式限流。

第四步:向 root 發送自己的連線資訊。每個 rank 把自己的監聽位址發給 root:

📎 src/bootstrap.cc:845-867

root 收到所有 rank 的資訊後,會做一次「環形配對」——把 rank i 的位址發給 rank i-1,把 rank i+1 的位址發給 rank i。這樣每個 rank 就知道了自己 ring 上的前後鄰居。

第五步:建立 ring 連線。每個 rank 連線自己的「下一個」鄰居,同時接受「上一個」鄰居的連線:

📎 src/bootstrap.cc:885-894

這裡socketRingConnect內部使用了bootstrapConcurrent——在 TLS 加密模式下,connect 和 accept 必須並行執行,否則會死鎖(因為 TLS 握手需要雙方同時參與)。非加密模式下則串行執行 connect 再 accept。

第六步:AllGather 所有位址。ring 建立後,透過ringAllInfo把所有 rank 的 P2P 位址、proxy 位址、UDS 位址做一次 allgather:

📎 src/bootstrap.cc:934-938

ringAllInfo內部呼叫bootstrapAllGather,後者在 socket 模式下使用socketRingAllGather——一個雙向 ring allgather 演算法,N 個 rank 只需要 N/2 步:

📎 src/bootstrap.cc:1363-1412

這個雙向演算法是 bootstrap 效能的關鍵最佳化。傳統的單向 ring allgather 需要 N-1 步,雙向版本把步數減半。每一步同時向兩個方向發送和接收資料,用socketDoubleSendRecv把 4 個操作(2 發 2 收)打包成一次系統呼叫。

並行控制與底層互動

Bootstrap 的並行控制有幾個層次:

第一層:abort 檢查。所有阻塞迴圈都定期檢查 abortFlag:

📎 src/bootstrap.cc:150-159

BOOTSTRAP_N_CHECK_ABORT設為 10000,意味著每 10000 次迴圈檢查一次 abort 標誌。這個數字是效能與回應性的折中——檢查太頻繁會影響效能,檢查太少會導致 abort 回應延遲。

第二層:非同步發送佇列。在 TLS 加密模式下,bootstrapSend不能同步執行(因為 TLS 握手需要接收方也參與),所以 NCCL 把發送操作放到獨立執行緒:

📎 src/bootstrap.cc:1161-1217

這裡有一個精妙的順序保證機制。bootstrapAsyncSendMain在發送前會檢查佇列中是否有「更早的、發往同一 (peer, tag) 的發送」:

📎 src/bootstrap.cc:1124-1152

為什麼要保證同一 (peer, tag) 的發送順序?原始碼註解解釋得很清楚:接收方按 (peer, tag) 匹配連接,如果兩個發往同一 (peer, tag) 的訊息到達順序顛倒,接收方會把它們匹配錯。NVLS 初始化期間會多次向同一 peer 用同一 tag 廣播,所以這個順序保證是必須的。

第三層:意外連接佇列。接收方無法預知誰會先連過來,所以socketAccept會把不匹配的連接存入unexpectedConnections鏈結串列:

📎 src/bootstrap.cc:1276-1300

這個設計解決了一個經典的分散式問題:多個 rank 可能同時向你發起連接,但你的bootstrapRecv呼叫順序是固定的。如果不匹配的連接被直接丟棄,發送方會逾時;如果阻塞等待,又可能死鎖。存入佇列是最安全的做法。

生產避坑指南

坑一:bootstrap 逾時導致初始化掛起。如果某個 rank 因為網路問題無法連接到 root,其他所有 rank 都會在ncclSocketAccept或ncclSocketRecv上無限等待。NCCL 沒有內建的 bootstrap 逾時機制,唯一的逃生通道是 abortFlag。生產環境中建議設定NCCL_UID_STAGGER_RATE來緩解大規模叢集的連接風暴。

坑二:NCCL_COMM_ID與多 handle 衝突。當使用者設定NCCL_COMM_ID環境變數時,NCCL 會強制把nId降為 1:

📎 src/init.cc:2912-2921

這意味著ncclCommInitRankScalable的多 handle 特性會被靜默停用。如果你在用 scalable 初始化又設了NCCL_COMM_ID,行為會和你預期的不一樣。

坑三:TLS 模式下的死鎖。在 TLS 加密模式下,如果 connect 和 accept 不並行執行,雙方都會卡在 TLS 握手。bootstrapConcurrent就是為了解決這個問題:

📎 src/bootstrap.cc:648-669

非加密模式下串行執行(先 send 後 recv),加密模式下啟動一個執行緒處理 send,主執行緒處理 recv。

mermaid
sequenceDiagram
    participant R0 as Rank 0
    participant Root as Bootstrap Root
    participant R1 as Rank 1
    participant R2 as Rank 2

    R0->>Root: sendToRoot(extInfo{rank=0, listenAddr})
    R1->>Root: sendToRoot(extInfo{rank=1, listenAddr})
    R2->>Root: sendToRoot(extInfo{rank=2, listenAddr})
    Note over Root: 收集所有 rank 的监听地址
    Root-->>R0: rootSend(rank2.addr) 下一个邻居
    Root-->>R1: rootSend(rank0.addr) 下一个邻居
    Root-->>R2: rootSend(rank1.addr) 下一个邻居
    R0->>R1: socketRingConnect(connect to next)
    R1->>R2: socketRingConnect(connect to next)
    R2->>R0: socketRingConnect(connect to next)
    Note over R0,R2: Ring 建立完成
    R0->>R1: socketRingAllGather 双向交换
    R1->>R2: socketRingAllGather 双向交换
    R2->>R0: socketRingAllGather 双向交换
    Note over R0,R2: 所有地址交换完成

3.3 commAlloc:通訊域物件的記憶體骨架

直覺模型

commAlloc是通訊域的「毛胚屋交付」——它分配結構體記憶體、初始化所有欄位到安全預設值、建立必要的 CUDA 物件和同步原語,但還沒有填充拓撲資訊、通道配置、傳輸連接這些「精裝修」內容。如果把ncclComm比作一棟大樓,commAlloc就是打地基和澆築框架,initTransportsRank才是內部裝修。

如果沒有commAlloc的初始化,後續程式碼存取未初始化的欄位會導致不可預測的行為——比如comm->channels[c].id如果是隨機值,通道初始化邏輯就會誤判通道狀態。

資料結構與記憶體佈局

commAlloc的簽名和開頭校驗:

📎 src/init.cc:512-526

它首先校驗ndev和rank的合法性,然後建構兩個記憶體堆疊(memPermanent和memScoped),設定rank和nRanks。這兩個記憶體堆疊是 NCCL 的記憶體管理基礎設施——memPermanent用於生命週期與通訊域相同的分配,memScoped用於臨時分配。

接下來是 CUDA 裝置探測:

📎 src/init.cc:528-531

cudaGetDevice取得目前裝置號,ncclCudaCompCap取得計算能力。原始碼註解說得很直白:"Try to create a CUDA object right away. If there is something wrong with the device we're on, better know it early."——儘早暴露裝置問題,避免在初始化後期才發現。

然後是共享資源的分配或繼承:

📎 src/init.cc:533-555

這裡有一個重要的分支:如果parent == NULL || !parent->shareResources,就建立新的ncclSharedResources;否則繼承父通訊域的共享資源並增加引用計數。ncclSharedResources包含裝置流、主機流、啟動事件、scratch 事件等——這些資源在 split 場景下可以被子通訊域複用,避免重複建立。

注意sharedRes->refCount = 1這一行——初始引用計數為 1,每次 split 共享時遞增,最後一個引用釋放時才真正銷毀。

接下來是網路、RMA、GIN 的初始化:

📎 src/init.cc:547-549

這三個子系統分別負責網路傳輸、遠端記憶體存取、GPU 發起的網路通訊。它們的初始化順序有講究——ncclNetInit必須先於ncclRmaInit,因為 RMA 依賴網路外掛。

記憶體管理器的初始化:

📎 src/init.cc:567-576

同樣有共享/新建兩種路徑。ncclMemManager負責管理 CUDA 記憶體池和註冊快取。

通道初始化標記:

📎 src/init.cc:607-608

這一行把所有通道的id設為 -1,表示「未初始化」。後續setupChannel會檢查這個值來決定是否需要初始化。

中斷佇列的建構:

📎 src/init.cc:619-632

NCCL 使用侵入式佇列(intrusive queue)來管理各種任務。這些佇列在commAlloc階段全部建構為空,後續任務入佇列時直接使用。

CUDA 記憶體池的建立:

📎 src/init.cc:636-652

如果裝置支援記憶體池(cudaDevAttrMemoryPoolsSupported),就建立一個 pinned 類型的記憶體池,並把釋放閾值設為最大值(~uint64_t(0)),意思是「永遠不自動釋放」。這是為了避免 CUDA 執行階段在 NCCL 不知情的情況下回收記憶體。

Step-by-Step Walkthrough

讓我們追蹤一個具體的初始化場景:單機 8 卡,每個行程一個 rank,正常初始化。

1. commAlloc(comm, NULL, 8, rank)被呼叫,parent == NULL。

2. 校驗通過,comm->rank = rank,comm->nRanks = 8。

3. cudaGetDevice傳回目前裝置號,comm->compCap被設定。

4. 建立新的ncclSharedResources,引用計數為 1。

5. ncclNetInit初始化網路外掛(可能是 Socket 或 IB)。

6. ncclMemManagerInit建立記憶體管理器。

7. getBusId取得 PCI 匯流排 ID,ncclNvmlDeviceGetHandleByPciBusId取得 NVML 控制代碼。

8. dmaBufSupported偵測 DMA-BUF 支援。

9. 分配connectSend / connectRecv位圖陣列。

10. 所有通道id設為 -1。

11. 建構所有中斷佇列。

12. 建立 CUDA 記憶體池。

設計思考

commAlloc中最值得玩味的設計是「儘早失敗」原則。它在函式開頭就呼叫cudaGetDevice,而不是等到後面需要裝置資訊時再呼叫。這樣做的好處是:如果裝置有問題(例如被其他行程獨占),錯誤會在初始化早期就暴露,而不是在分配了大量記憶體之後才發現。

另一個設計是preconnectNext的初始化:

📎 src/init.cc:598-598

reinterpret_cast<struct ncclComm*>(0x1)是一個哨兵值,用於標記「下一個預連接」的狀態。這種用非法指標值作為狀態標記的手法在系統程式設計中很常見——它比額外的布林欄位更省記憶體,但需要小心不要解引用。

3.4 initTransportsRank:拓撲探索與通道分配

直覺模型

initTransportsRank是初始化的「心臟」。它做三件大事:透過兩次 AllGather 交換所有 rank 的裝置資訊和拓撲資訊;根據這些資訊計算 ring/tree/collnet/nvls 等演算法的圖結構;最後建立所有傳輸連線。如果把通訊域比作一個城市的交通系統,initTransportsRank就是規劃所有道路、立交橋和公車路線的過程。

如果沒有這一步,NCCL 就不知道資料該走哪條路——它可能讓資料繞遠路,或者根本找不到可達的路徑。

資料結構與記憶體佈局

initTransportsRank的區域變數非常多,我們挑關鍵的看:

📎 src/init.cc:1163-1179

這裡把comm->graphs陣列中的各個圖結構取出來,建立別名。graphs陣列按演算法索引,注意nvlsGraph被用了兩次(NVLS 和 NVLSTree 共享同一個圖結構)。

兩個關鍵的臨時結構體:

📎 src/init.cc:1181-1206

graphInfo儲存單個 rank 對某個演算法的圖資訊(通道數、頻寬、類型等),allGatherInfo是 AllGather 的資料單元,包含所有演算法的圖資訊加上拓撲 rank 資訊。

Step-by-Step Walkthrough

階段一:AllGather1——交換裝置資訊。

📎 src/init.cc:1234-1239

每個 rank 呼叫fillInfo填充自己的ncclPeerInfo,然後透過bootstrapAllGather交換。fillInfo填充的資訊包括:rank 號、CUDA 裝置號、NVML 裝置號、NCCL 版本、git hash、主機 hash、行程 hash、GPU UUID、匯流排 ID、顯示記憶體大小、驅動版本等。

📎 src/init.cc:888-982

注意info->hostHash = getHostHash() + commHash和info->pidHash = getPidHash() + commHash——host hash 和 pid hash 都加上了 commHash。這是為了區分同一台機器上的不同通訊域。

AllGather 完成後,每個 rank 走訪所有 peer 的資訊,計算全域屬性:

📎 src/init.cc:1250-1303

這個迴圈做了很多事:偵測版本不符、統計節點數、計算cuMemSupport的交集、偵測是否有多個 rank 使用同一個 GPU、計算 GIN 類型遮罩的交集等。注意nNodes的統計方式——每當遇到不同 hostHash 就遞增,這假設 rank 是按節點連續排列的。

階段二:拓撲探索。

📎 src/init.cc:1390-1403

這六步是拓撲探索的核心流程:ncclTopoGetSystem列舉系統裝置建構拓撲圖,ncclTopoComputePaths計算 GPU 到 NIC 的路徑,ncclTopoTrimSystem移除不可達裝置,再次計算路徑,ncclTopoSearchInit初始化搜尋狀態,最後列印拓撲。

階段三:圖計算。

📎 src/init.cc:1421-1468

依次計算 ring、tree、collnet chain、collnet direct、nvls 五種圖。每種圖有不同的 pattern 和通道數約束。注意treeGraph->minChannels = ringGraph->nChannels——tree 的通道數被約束為與 ring 相同,這是為了保證不同演算法之間的通道對齊。

階段四:AllGather3——交換圖資訊。

📎 src/init.cc:1490-1533

每個 rank 把自己的圖資訊填入allGather3Data[rank],然後再次bootstrapAllGather。這次交換的資訊包括:每種演算法的 pattern/nChannels/bwIntra/bwInter/typeIntra/typeInter/crossNic、CPU 架構、P2P 通道數、網路裝置數、CollNet 裝置數等。

AllGather3 完成後,每個 rank 走訪所有 peer 的圖資訊,取最小值/最大值來對齊:

📎 src/init.cc:1687-1703

注意這裡的對齊策略:nChannels、sameChannels、bwIntra、bwInter取最小值,typeIntra、typeInter、crossNic取最大值。為什麼?因為通道數和頻寬受限於最弱的鏈路,而類型和 crossNic 需要取聯集以確保相容性。

階段五:建立傳輸連線。

📎 src/init.cc:1811-1892

這裡有兩個分支:runtimeConn為真時只做通道 setup 不做連線(延遲到執行時連線),否則立即建立所有連線。連線順序是:ring → tree → NVLS → PAT → NVLS tree → CollNet。

並行控制與硬體互動

initTransportsRank中有幾個值得注意的並行/硬體互動點:

CPU 親和性設定:

📎 src/init.cc:1406-1412

NCCL 把當前執行緒綁定到 GPU 附近的 CPU 核心,確保主機記憶體分配是本地 NUMA 節點的。這減少了跨 NUMA 存取的延遲。

NVLS 初始化:

📎 src/init.cc:1419-1419

ncclNvlsInit檢測 NVLink SHARP 支援。NVLS 允許交換器直接執行 reduce 操作,大幅降低 AllReduce 延遲。

Proxy 執行緒建立:

📎 src/init.cc:1780-1786

Proxy 執行緒負責非同步推進網路 I/O。它在initTransportsRank中被建立,之後所有網路操作都透過 proxy 進行。

生產避坑指南

坑一:網路裝置數不匹配。如果不同 rank 的本地網卡數量不同,NCCL 會報錯:

📎 src/init.cc:1576-1596

除非設定NCCL_IGNORE_NET_MISMATCH=1。這在異構叢集中很常見——有些節點有 8 張網卡,有些只有 4 張。忽略不匹配可能導致效能下降,因為通道數會被最弱的節點限制。

坑二:多 rank 共用同一 GPU。如果兩個 rank 的 GPU UUID 相同,NCCL 會拒絕初始化:

📎 src/init.cc:1291-1296

除非設定NCCL_MULTI_RANK_GPU_ENABLE=1。這個檢查防止了使用者誤配置導致的效能問題。

坑三:CollNet 節點數不足。CollNet 需要至少NCCL_COLLNET_NODE_THRESHOLD個節點才能啟用:

📎 src/init.cc:1720-1728

預設閾值是 2。單節點環境下 CollNet 會被自動停用。

mermaid
flowchart TD
    start["initTransportsRank(comm, parent, timers)"]
    ag1["AllGather1: fillInfo + bootstrapAllGather"]
    check_ver{"版本匹配?"}
    fail_ver["返回 ncclInvalidUsage"]
    topo["ncclTopoGetSystem + ComputePaths + TrimSystem"]
    graphs["计算 ring/tree/collnet/nvls 图"]
    ag3["AllGather3: 交换图信息"]
    align["对齐 nChannels/bwIntra/bwInter"]
    setup["setupChannel 初始化所有通道"]
    conn_ring["ncclTransportRingConnect"]
    conn_tree["ncclTransportTreeConnect"]
    conn_nvls["ncclNvlsSetup + ncclNvlsBufferSetup"]
    conn_collnet{"collnetEnable?"}
    conn_collnet_yes["ncclCollNetSetup + BufferSetup"]
    devcomm["devCommSetup 映射到设备"]
    barrier["bootstrapIntraNodeBarrier"]
    done["初始化完成"]

    start --> ag1 --> check_ver
    check_ver -->|否| fail_ver
    check_ver -->|是| topo --> graphs --> ag3 --> align --> setup
    setup --> conn_ring --> conn_tree --> conn_nvls --> conn_collnet
    conn_collnet -->|是| conn_collnet_yes --> devcomm
    conn_collnet -->|否| devcomm
    devcomm --> barrier --> done

3.5 NCCL_PARAM:環境變數體系的編譯期魔法

直覺模型

NCCL_PARAM是 NCCL 的「配置開關工廠」。它用巨集在編譯期生成一個函式,執行時第一次呼叫時讀取環境變數並快取結果。這就像家裡的電燈開關——你撥一下(呼叫函式),燈就亮了(回傳配置值),之後開關狀態被記住,不需要每次都重新撥。

如果沒有這套機制,NCCL 就需要在每個使用配置的地方手動呼叫getenv並解析字串,程式碼會變得極其冗長且容易出錯。

資料結構與記憶體佈局

NCCL_PARAM巨集的定義:

📎 src/include/param.h:22-31

這個巨集展開後生成一個函式ncclParam##name(),內部有三個靜態變數:

  • uninitialized = INT64_MIN:哨兵值,表示「尚未初始化」。
  • noCache:三態標誌,-1 表示未初始化,0 表示快取,1 表示不快取。
  • cache:快取的值,初始為uninitialized。

函式邏輯是:如果cache還是uninitialized,呼叫ncclLoadParam載入;否則直接回傳cache。COMPILER_EXPECT(..., false)告訴編譯器這個分支很少走,優化熱路徑。

ncclLoadParam的實作:

📎 src/misc/param.cc:78-108

它用互斥鎖保護整個載入過程,先檢查noCache策略,再檢查快取是否有效,然後讀取環境變數並解析。解析失敗時使用預設值並列印警告。

Step-by-Step Walkthrough

以NCCL_PARAM(BuffSize, "BUFFSIZE", -2)為例:

📎 src/init.cc:1007-1007

巨集展開後生成:

cpp
int64_t ncclParamBuffSize() {
  constexpr int64_t uninitialized = INT64_MIN;
  static int8_t noCache = -1;
  static_assert(-2 != uninitialized, "...");
  static int64_t cache = uninitialized;
  if (COMPILER_EXPECT(COMPILER_ATOMIC_LOAD(&cache, std::memory_order_relaxed) == uninitialized, false)) {
    return ncclLoadParam("NCCL_BUFFSIZE", -2, uninitialized, &cache, &noCache);
  }
  return cache;
}

第一次呼叫時,cache == uninitialized,進入ncclLoadParam。它讀取NCCL_BUFFSIZE環境變數,如果沒設定就回傳預設值 -2。然後根據noCache策略決定是否快取。

noCache策略由ncclParamIsCacheDisabled決定:

📎 src/misc/param.cc:74-76

如果環境變數名稱匹配某個模式(比如以_結尾),就不快取,每次都重新讀取。這允許使用者在執行時動態修改某些配置。

設計思考

這套設計的精妙之處在於「零成本抽象」:熱路徑上只有一次原子載入和比較,沒有鎖、沒有字串解析。冷路徑(首次載入)才付出完整代價。COMPILER_EXPECT提示編譯器把熱路徑放在指令快取的前面,進一步提高效能。

另一個設計是noCache的三態設計。-1 表示「還沒決定」,0 表示「快取」,1 表示「不快取」。這個決定只在首次載入時做一次,之後不再改變。

生產避坑指南

坑一:環境變數拼寫錯誤。如果使用者寫了NCCL_BUFSIZE而不是NCCL_BUFFSIZE,NCCL 不會報錯,只會使用預設值。建議用NCCL_DEBUG=ENV查看所有被識別的環境變數。

坑二:NCCL_CONF_FILE的載入順序。NCCL 會依次載入$NCCL_CONF_FILE(或~/.nccl.conf)和/etc/nccl.conf:

📎 src/misc/param.cc:52-67

後載入的檔案會覆蓋先載入的。如果兩個檔案都設定了同一個變數,/etc/nccl.conf的值會生效。

坑三:noCache變數的執行緒安全。原始碼註解說 "noCache is only load/stored within the mutex, no need for atomic":

📎 src/misc/param.cc:74-76

這意味著noCache的讀寫都在互斥鎖保護下,不需要原子操作。但cache的讀取是無鎖的(熱路徑),所以用原子載入。

3.6 devCommSetup:把通訊域映射到裝置

直覺模型

devCommSetup是通訊域的「裝置側投影」。GPU kernel 執行在裝置上,無法直接存取主機記憶體中的ncclComm結構體。所以 NCCL 需要把通訊域的關鍵欄位複製到裝置可存取的記憶體中,形成ncclDevComm。這就像把公司的通訊錄複印一份放到每個員工的工位上——員工不用每次都跑去找櫃檯問同事電話。

如果沒有devCommSetup,GPU kernel 就無法知道自己的 rank、通道配置、緩衝區大小等資訊,集合通訊 kernel 根本無法啟動。

資料結構與記憶體佈局

devCommSetup使用一個臨時結構體ncclKernelCommAndChannels來打包要複製到裝置的資料:

📎 src/init.cc:712-746

這個結構體包含ncclDevComm(裝置側通訊域)和通道陣列。函式先把主機側的資料填入臨時結構體,然後一次性cudaMemcpyAsync到裝置。

關鍵欄位的填充:

📎 src/init.cc:734-746

注意comm->devComm = &devCommAndChans->comm——主機側的comm->devComm指向裝置記憶體中的ncclDevComm。後續 kernel 啟動時會把comm->devComm作為參數傳入。

通道資訊的填充:

📎 src/init.cc:829-843

每個通道的 peers、ring、tree、collnetChain、collnetDirect、nvls 指標都被拷貝到裝置側。注意ring.userRanks需要額外的一次cudaMemcpyAsync,因為它是一個陣列。

Step-by-Step Walkthrough

1. 取得裝置串流:ncclStrongStreamAcquire取得一個強串流(strong stream),確保後續的非同步拷貝有序執行。

2. 分配裝置記憶體:ncclCudaCallocAsync分配devCommAndChans。

3. 填充主機側暫存結構體:設定 rank、nRanks、node、nNodes、abortFlag、buffSizes 等。

4. 分配並拷貝rankToLocalRank陣列。

5. 計算workFifoBytes:根據 CC(Confidential Computing)狀態決定。

6. 分配 workFifo 緩衝區:GDR 模式用ncclGdrCudaCalloc,否則用ncclCudaHostCalloc。

7. 分配 profiler 計數器。

8. 分配進度計數器(如果啟用)。

9. 填充通道資訊。

10. 一次性拷貝到裝置:ncclCudaMemcpyAsync(devCommAndChans, &tmpCommAndChans, 1, deviceStream)。

11. 釋放強串流並同步。

設計思考

devCommSetup中最值得注意的設計是「批次拷貝」。NCCL 沒有為每個欄位單獨呼叫cudaMemcpy,而是把所有欄位打包到一個暫存結構體,用一次cudaMemcpyAsync完成。這大幅減少了 CUDA API 呼叫次數和同步開銷。

另一個設計是workFifoBytes的 CC 處理:

📎 src/init.cc:750-763

在 CC(Confidential Computing)模式下,workFifoBytes被設為 0,因為 GDR 拷貝在 CC 模式下不可用。這是一個硬體限制的優雅降級。

生產避坑指南

坑一:devCommSetup必須在 barrier 之前呼叫。原始碼註解解釋了原因:

📎 src/init.cc:1950-1952

如果在 barrier 之後呼叫,可能有執行緒已經開始啟動 NCCL kernel,而此時裝置記憶體還沒分配完,會導致死鎖。

坑二:workFifoBytes必須是 2 的冪。如果不是,NCCL 會警告並使用預設值:

📎 src/init.cc:757-762

本章思考與自測

Q1: 如果將📎 src/init.cc:1291-1296中偵測「多個 rank 使用同一 GPU」的邏輯去掉,在什麼場景下會導致問題?為什麼 NCCL 預設拒絕這種配置?

參考解析:

這段程式碼偵測同一主機上兩個 rank 的 GPU UUID 是否相同。如果相同且NCCL_MULTI_RANK_GPU_ENABLE=0(預設),就返回ncclInvalidUsage。

去掉這個檢查後,多個 rank 會共享同一個 GPU。這會導致:

1. P2P 傳輸衝突:NCCL 的 P2P 傳輸假設每個 rank 獨佔一個 GPU。如果兩個 rank 共享 GPU,它們會同時向同一個 GPU 的同一塊緩衝區寫入資料,導致資料競爭和結果錯誤。

2. 通道分配衝突:comm->channels中的通道資源(緩衝區、FIFO)是按 rank 分配的。共享 GPU 的 rank 會爭搶同一份資源。

3. 效能災難:即使沒有正確性問題,兩個 rank 共享一個 GPU 的算力和顯示記憶體頻寬,效能會急劇下降。

NCCL 預設拒絕這種配置是為了「快速失敗」——與其讓使用者在一個錯誤配置上浪費數小時除錯,不如在初始化時就明確報錯。NCCL_MULTI_RANK_GPU_ENABLE=1是給那些明確知道自己在做什麼的使用者(比如 MPS 場景)準備的逃生通道。

Q2: 如果將📎 src/bootstrap.cc:1129-1134中等待「同一 (peer, tag) 的更早發送」的邏輯去掉,在什麼場景下會導致接收方匹配錯誤?

參考解析:

這段程式碼在非同步發送執行緒中等待,直到佇列中沒有更早的、發往同一 (peer, tag) 的發送。

去掉這個等待後,兩個發往同一 (peer, tag) 的發送可能並行執行,到達接收方的順序不確定。接收方的socketAccept按 (peer, tag) 匹配連線:

📎 src/bootstrap.cc:1291-1292

如果發送方 A 先呼叫bootstrapSend但後到達,發送方 B 後呼叫但先到達,接收方會把 B 的訊息當作 A 的回應。這會導致資料錯位——接收方以為收到的是第一個請求的回應,實際上是第二個請求的。

原始碼註解明確指出了這個場景:「NVLS setup broadcasts to the same peers with the same tag several times during init」。NVLS 初始化期間會多次向同一 peer 用同一 tag 廣播,如果順序顛倒,NVLS 配置會完全錯亂。

這個順序保證的代價是:同一 (peer, tag) 的發送被串行化。但不同 (peer, tag) 的發送仍然並行,所以整體吞吐量不受影響。

Q3: 如果將📎 src/init.cc:1691-1697中對齊策略從「nChannels 取 min、typeIntra 取 max」改為「全部取 min」或「全部取 max」,會分別導致什麼問題?

參考解析:

當前策略是:nChannels、sameChannels、bwIntra、bwInter取 min,typeIntra、typeInter、crossNic取 max。

如果全部取 min:typeIntra和typeInter取 min 會導致某些 rank 的傳輸類型被降級。比如 rank A 支援 P2P(typeIntra=P2P),rank B 只支援 SHM(typeIntra=SHM),取 min 後所有 rank 都用 SHM。但 SHM 的列舉值可能比 P2P 小,取 min 會選到錯誤的類型。實際上typeIntra是一個位元遮罩或列舉,取 max 是為了選擇「能力最強」的類型。

如果全部取 max:nChannels取 max 會導致某些 rank 被分配超過其能力的通道數。比如 rank A 只能支援 4 個通道,rank B 支援 8 個,取 max 後所有 rank 都嘗試用 8 個通道,rank A 會失敗或效能下降。bwIntra取 max 會導致頻寬估計過於樂觀,tuning 模組可能選擇不適合的演算法。

這個對齊策略的本質是:資源約束取交集(min),能力列舉取聯集(max)。通道數和頻寬是「上限」約束,必須取最保守的值;傳輸類型是「能力」列舉,取最大值確保所有 rank 都能找到相容的傳輸方式。

下一章我們將深入拓撲發現與圖搜尋,看 NCCL 如何列舉機器裡的 GPU、網卡、PCI 交換機,構建出一張完整的拓撲圖,並在這張圖上搜尋最優的 ring 和 tree 結構。本章建立的 bootstrap 通訊、commAlloc 記憶體骨架、initTransportsRank 主幹流程,將在下一章中逐一展開其拓撲細節。

至此,我們已經完整走過了 ncclCommInitRank 的呼叫鏈,看清了 ncclComm 物件從零構建的全過程。但初始化過程中有一個關鍵環節我們只是匆匆掠過:NCCL 是如何探測機器內部的 GPU 和網卡,並據此決定資料該走哪條路的?這正是下一章要深入的主題——拓撲發現與圖搜尋。我們將拆解 src/graph/topo.cc 如何列舉 PCI/NVLink/網卡裝置並構建拓撲圖,src/graph/search.cc 如何在該圖上搜尋最優路徑,以及 src/graph/rings.cc 與 trees.cc 如何將搜尋結果具體化為 Ring 與 Tree 演算法拓撲。理解了這套機制,你就能明白為什麼 NCCL 能在不同機器上自動選到合適的演算法。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 04

第 4 章:拓撲發現與圖搜尋:NCCL 如何「看清」多 GPU 系統的物理互聯

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 4 章 / 共 25 章

上一章我們沿 ncclCommInitRank 的呼叫鏈逐層下鑽,看到了 comm->topo 欄位被填充的時機,但並未展開它內部的結構。那麼,NCCL 究竟是如何「看見」機器裡的 GPU 和網卡,並將它們組織成可用的拓撲資訊的?本章將拆解這一過程的三個關鍵環節:topo.cc 負責將物理裝置列舉成一張圖,search.cc 在這張圖上搜尋最優路徑,rings.cc 和 trees.cc 則把搜尋結果具體化為 Ring 與 Tree 兩種演算法拓撲。理解這三者的配合,才能明白 NCCL 為何能在不同機器上自動選到合適的演算法。

拓撲圖:把機器畫成一張「地鐵線路圖」

直覺模型

想像你是一個剛到陌生城市的快遞員。你要把包裹從 A 點送到 B 點,但你不知道哪條路最快。你需要一張地圖——上面標著所有站點(GPU、網卡、CPU、PCI 交換機)以及站點之間的連接(NVLink、PCIe、網路)。NCCL 的拓撲圖就是這張地圖。

如果沒有這張圖,NCCL 只能盲目地假設「所有 GPU 之間頻寬相同」,在 8 卡 NVLink 全互聯的機器上或許還能湊合,但一旦遇到跨 NUMA、跨 PCI 交換機、混合 NVLink + PCIe 的複雜拓撲,就會選錯路徑,把本該走 NVLink 的資料塞進慢速 PCIe,效能直接腰斬。

資料結構與記憶體佈局

拓撲圖的核心是ncclTopoSystem,它按節點類型分組儲存所有裝置。節點類型定義在topoNodeTypeStr陣列裡:

📎 src/graph/topo.cc:33-35

c
const char* topoNodeTypeStr[] = {"GPU", "PCI", "NVS", "CPU", "NIC", "NET", "GIN", "RMA", "DEV", "CXB"};
const char* topoLinkTypeStr[] = {"LOC", "NVL", "", "C2C", "PCI", "", "", "", "", "SYS", "NET"};
const char* topoPathTypeStr[] = {"LOC", "NVL", "NVB", "C2C", "PIX", "PXB", "P2C", "PXN", "PHB", "SYS", "NET", "DIS"};

這三個陣列分別定義了節點類型、鏈路類型和路徑類型的字串表示。注意topoPathTypeStr的順序——它同時充當了路徑品質的排序:索引越小,路徑越快。LOC(本地)最快,DIS(斷開)最慢。這個順序在後續搜尋中會被反覆用來比較路徑優劣。

每個節點由ncclTopoNode表示,建立時根據類型初始化不同的欄位。以 GPU 節點為例:

📎 src/graph/topo.cc:105-141

c
ncclResult_t ncclTopoCreateNode(struct ncclTopoSystem* system, struct ncclTopoNode** node, int type, uint64_t id) {
  if (system->nodes[type].count == NCCL_TOPO_MAX_NODES) {
    WARN("Error : tried to create too many nodes of type %d", type);
    return ncclInternalError;
  }
  struct ncclTopoNode* n = system->nodes[type].nodes + system->nodes[type].count;
  system->nodes[type].count++;
  n->type = type;
  n->id = id;
  if (type == GPU) {
    n->gpu.dev = NCCL_TOPO_UNDEF;
    n->gpu.rank = NCCL_TOPO_UNDEF;
    n->gpu.cudaCompCap = NCCL_TOPO_UNDEF;
    n->gpu.mloPart = NCCL_TOPO_UNDEF;
  } else if (type == CPU) {
    ...

這裡有幾個關鍵設計點。第一,節點儲存在一個預先分配的陣列中(system->nodes[type].nodes),而不是鏈結串列。這意味著節點在記憶體中是連續排列的,遍歷時快取友好。第二,NCCL_TOPO_MAX_NODES是一個硬上限,超過就報錯——這是為了防止拓撲異常時無限增長。第三,每個節點有一個id欄位,它是一個 64 位元整數,高 32 位元是 systemId(標識哪台主機),低 32 位元是 localId(主機內的裝置編號)。

節點之間的連接由ncclTopoLink表示。ncclTopoConnectNodes負責建立雙向連接:

📎 src/graph/topo.cc:179-204

c
ncclResult_t ncclTopoConnectNodes(struct ncclTopoNode* node, struct ncclTopoNode* remNode, int type, float bw) {
  // Aggregate links into higher bw for NVLink
  struct ncclTopoLink* link;
  for (link = node->links; link - node->links != NCCL_TOPO_MAX_LINKS && link->remNode; link++) {
    if (link->remNode == remNode && link->type == type) break;
  }
  if (link - node->links == NCCL_TOPO_MAX_LINKS) {
    WARN("Error : too many Topo links (max %d)", NCCL_TOPO_MAX_LINKS);
    return ncclInternalError;
  }
  if (link->remNode == NULL) node->nlinks++;
  link->type = type;
  link->remNode = remNode;
  link->bw += bw;

  // Sort links in BW descending order
  struct ncclTopoLink linkSave;
  memcpy(&linkSave, link, sizeof(struct ncclTopoLink));
  while (link != node->links) {
    if ((link - 1)->bw >= linkSave.bw) break;
    memcpy(link, link - 1, sizeof(struct ncclTopoLink));
    link--;
  }
  memcpy(link, &linkSave, sizeof(struct ncclTopoLink));
  return ncclSuccess;
}

這個函式做了三件事。第一,查找是否已存在到同一目標、同一類型的鏈路——如果存在,就把頻寬累加(link->bw += bw)。這處理的是多條 NVLink 連到同一個 GPU 的情況:4 條 NVLink 各 25 GB/s,聚合後就是 100 GB/s。第二,如果沒找到,就新增一條鏈路。第三,插入後按頻寬降序排列,這樣後續遍歷時優先看到高頻寬鏈路。

〔設計推斷與架構權衡〕

頻寬降序排列的設計動機是讓搜尋演算法儘早發現高頻寬路徑,從而更快收斂到較優解。搜尋有逾時限制(後面會看到NCCL_SEARCH_TIMEOUT),排序能讓有限的時間預算花在更有希望的路徑上。

場景驅動的 Step-by-Step Walkthrough

現在代入一個具體場景:一台 8 卡 A100 伺服器,每張卡透過 NVLink 全互聯,另有 4 張 Mellanox ConnectX-6 網路卡插在 PCIe 插槽上。NCCL 初始化時,ncclTopoGetSystem被呼叫,它從 XML 檔案(由nvidia-topologyd或 NCCL 自己生成)讀取裝置資訊,然後構建拓撲圖。

第一步,解析 CPU 節點。ncclTopoAddCpu從 XML 中讀取 CPU 的架構、廠商、型號,並建立 CPU 節點:

📎 src/graph/topo.cc:806-875

c
ncclResult_t ncclTopoAddCpu(struct ncclXmlNode* xmlCpu, struct ncclTopoSystem* system) {
  int numaId;
  NCCLCHECK(xmlGetAttrInt(xmlCpu, "numaid", &numaId));
  int systemId;
  NCCLCHECK(ncclGetSystemId(system, xmlCpu, &systemId));
  struct ncclTopoNode* cpu;
  NCCLCHECK(ncclTopoCreateNode(system, &cpu, CPU, NCCL_TOPO_ID(systemId, numaId)));
  ...
  for (int s = 0; s < xmlCpu->nSubs; s++) {
    struct ncclXmlNode* node = xmlCpu->subs[s];
    if (strcmp(node->name, "pci") == 0) NCCLCHECK(ncclTopoAddPci(node, system, cpu, systemId, numaId));
    if (strcmp(node->name, "nic") == 0) {
      ...
    }
  }
  return ncclSuccess;
}

CPU 節點是拓撲樹的根。每個 CPU 下面掛著 PCI 子樹和 NIC 節點。ncclTopoAddPci遞迴處理 PCI 樹,遇到 GPU 就建立 GPU 節點,遇到 NIC 就建立 NIC 節點。

第二步,新增 NVLink 連接。注意ncclTopoAddGpu只讀取 GPU 的基本屬性,註解明確說 "Do not go any further, nvlinks will be added in a second pass":

📎 src/graph/topo.cc:590-598

c
ncclResult_t ncclTopoAddGpu(struct ncclXmlNode* xmlGpu, struct ncclTopoSystem* system, struct ncclTopoNode* gpu) {
  NCCLCHECK(xmlGetAttrInt(xmlGpu, "rank", &gpu->gpu.rank));
  NCCLCHECK(xmlGetAttrInt(xmlGpu, "sm", &gpu->gpu.cudaCompCap));
  NCCLCHECK(xmlGetAttrInt(xmlGpu, "dev", &gpu->gpu.dev));
  NCCLCHECK(xmlGetAttrInt(xmlGpu, "gdr", &gpu->gpu.gdrSupport));
  NCCLCHECK(xmlGetAttrIntDefault(xmlGpu, "mlopart", &gpu->gpu.mlopart, NCCL_TOPO_UNDEF));
  // Do not go any further, nvlinks will be added in a second pass
  return ncclSuccess;
}

為什麼要分兩遍?因為 NVLink 是 GPU 之間的連接,需要兩端 GPU 節點都已存在才能建立鏈路。第一遍建立所有節點,第二遍ncclTopoAddNvLinks再連接它們。

第三步,處理網路裝置。ncclTopoAddNic遍歷 NIC 下的 net/gin/rma 子節點,分別呼叫對應的新增函式。以ncclTopoAddNet為例:

📎 src/graph/topo.cc:461-503

c
static ncclResult_t ncclTopoAddNet(struct ncclXmlNode* xmlNet, struct ncclXmlNode* parent,
                                   struct ncclTopoSystem* system, struct ncclTopoNode* nic, int systemId) {
  int dev;
  NCCLCHECK(xmlGetAttrInt(xmlNet, "dev", &dev));
  int64_t netId = NCCL_TOPO_ID(systemId, dev);
  struct ncclTopoNode* net;
  NCCLCHECK(ncclTopoCreateNode(system, &net, NET, netId));
  net->net.dev = dev;
  int mbps;
  NCCLCHECKNOWARN(xmlGetAttrIntDefault(xmlNet, "speed", &mbps, 0), NCCL_GRAPH);
  if (mbps <= 0) mbps = 10000; // Some NICs define speed = -1
  net->net.bw = mbps / 8000.0;
  ...
  NCCLCHECK(ncclTopoConnectNodes(nic, net, LINK_NET, net->net.bw));
  NCCLCHECK(ncclTopoConnectNodes(net, nic, LINK_NET, net->net.bw));
  return ncclSuccess;
}

注意mbps / 8000.0這個轉換:mbps 是兆位元每秒,除以 8000 得到 GB/s(因為 1 GB/s = 8000 Mbps)。如果網路卡報告 speed = -1(某些虛擬網路卡會這樣),就預設 10000 Mbps = 1.25 GB/s。

第四步,收尾處理。ncclTopoGetSystemFromXml在完成所有節點和鏈路新增後,還會做幾件清理工作:

📎 src/graph/topo.cc:1080-1088

c
  NCCLCHECK(ncclTopoAddNvLinks(topNode, *topoSystem, NULL, 0));
  NCCLCHECK(ncclTopoAddC2c(topNode, *topoSystem, NULL, 0));
  NCCLCHECK(ncclTopoAddPciLinks(topNode, *topoSystem, NULL, 0));

  NCCLCHECK(ncclTopoFlattenBcmSwitches(*topoSystem));
  NCCLCHECK(ncclTopoConnectCpus(*topoSystem));
  NCCLCHECK(ncclTopoSortSystem(*topoSystem));

ncclTopoFlattenBcmSwitches處理 Broadcom Gen4 PCIe 交換機的特殊情況——它們把自己呈現為兩層交換機,但實際是全頻寬的,需要「壓平」以避免搜尋演算法被誤導。ncclTopoConnectCpus把所有 CPU 節點互相連接(跨 NUMA 存取走 SYS 鏈路)。ncclTopoSortSystem對鏈路排序,讓 PCI 下行鏈路排在前面,方便遍歷。

設計思考與生產踩坑

〔設計推斷與架構權衡〕

為什麼用 XML 作為中間格式?因為拓撲發現需要跨行程共享——每個 rank 只探測自己管理的 GPU,然後透過 bootstrap 交換 XML,最後融合成完整拓撲。XML 是自描述的文字格式,便於除錯(可以 dump 出來看)和版本相容。

坑點一:ncclTopoGetNode找不到節點時不報錯。看這個函式:

📎 src/graph/topo.cc:95-103

c
ncclResult_t ncclTopoGetNode(struct ncclTopoSystem* system, struct ncclTopoNode** node, int type, uint64_t id) {
  for (int i = 0; i < system->nodes[type].count; i++) {
    if (system->nodes[type].nodes[i].id == id) {
      *node = system->nodes[type].nodes + i;
      return ncclSuccess;
    }
  }
  return ncclSuccess;
}

如果沒找到,它回傳ncclSuccess但*node保持不變(呼叫者通常初始化為 NULL)。呼叫者必須自己檢查*node == NULL。這種設計容易漏檢——如果呼叫者忘了檢查,後續解引用就會崩潰。

坑點二:ncclTopoConnectNodes的頻寬累加可能導致溢位。如果同一對節點之間有大量鏈路(比如 NVSwitch 場景),link->bw += bw可能累加到很大。雖然 float 的精度足夠,但如果鏈路數量異常多,排序邏輯可能出問題。

坑點三:ncclTopoRemoveNode的指標修正。刪除節點時,所有指向被刪節點的鏈路都要移除,且指向被刪節點之後節點的指標要前移:

📎 src/graph/topo.cc:143-177

c
ncclResult_t ncclTopoRemoveNode(struct ncclTopoSystem* system, int type, int index) {
  struct ncclTopoNode* delNode = system->nodes[type].nodes + index;
  for (int t = 0; t < NCCL_TOPO_NODE_TYPES; t++) {
    if (delNode->paths[t] != nullptr) {
      WARN("Cannot remove topology node %d/%lx while paths are computed", type, delNode->id);
      return ncclInternalError;
    }
    for (int n = 0; n < system->nodes[t].count; n++) {
      struct ncclTopoNode* node = system->nodes[t].nodes + n;
      if (node == delNode) continue;
      for (int l = 0; l < node->nlinks; l++) {
        while (l < node->nlinks && node->links[l].remNode == delNode) {
          memmove(node->links + l, node->links + l + 1, (node->nlinks - l - 1) * sizeof(struct ncclTopoLink));
          node->nlinks--;
        }
        if (l < node->nlinks && node->links[l].remNode->type == type && node->links[l].remNode >= delNode) {
          node->links[l].remNode--;
        }
      }
    }
  }
  ...

這裡有個微妙之處:node->links[l].remNode--是在修正指標。因為節點儲存在連續陣列中,刪除一個節點後,後面的節點位址都會前移一個sizeof(struct ncclTopoNode)。所以所有指向被刪節點之後節點的指標都要減一。這個操作在memmove之前執行,順序很關鍵。

路徑搜尋:在圖上找「最優路線」

直覺模型

有了地圖還不夠,你還需要一個導航演算法。NCCL 的路徑搜尋分兩層:第一層是預處理,計算所有節點對之間的最短路徑(BFS);第二層是圖搜尋,在預處理結果上嘗試不同的 Ring/Tree 結構,找到頻寬最高的那個。

如果沒有路徑搜尋,NCCL 只能硬編碼「GPU 0 連 GPU 1 連 GPU 2...」這種固定順序,在非均勻拓撲上會選到慢速路徑。

資料結構與記憶體佈局

路徑搜尋的核心資料結構是ncclTopoLinkList,它儲存從某個源節點到某個目標節點的完整路徑:

c
struct ncclTopoLinkList {
  struct ncclTopoLink* list[NCCL_TOPO_MAX_HOPS];  // 路径上的链路
  int count;      // 跳数
  float bw;       // 瓶颈带宽
  int type;       // 路径类型(PATH_LOC, PATH_NVL, ...)
  int capacity;   // list 数组的容量
};

每個節點有一個paths[type]陣列,儲存到所有該類型節點的路徑。比如 GPU 節點的paths[NET]儲存到所有網卡的路徑。

路徑計算由ncclTopoSetPaths完成,它是一個 BFS:

📎 src/graph/paths.cc:52-147

c
static ncclResult_t ncclTopoSetPaths(struct ncclTopoNode* baseNode, struct ncclTopoSystem* system) {
  if (baseNode->paths[baseNode->type] == NULL) {
    NCCLCHECK(ncclCalloc(baseNode->paths + baseNode->type, system->nodes[baseNode->type].count));
    for (int i = 0; i < system->nodes[baseNode->type].count; i++) baseNode->paths[baseNode->type][i].type = PATH_DIS;
  }

  // breadth-first search to set all paths to that node in the system
  struct ncclTopoNodeList nodeList;
  struct ncclTopoNodeList nextNodeList = {{0}, 0};
  nodeList.count = 1;
  nodeList.list[0] = baseNode;
  ...
  while (nodeList.count) {
    nextNodeList.count = 0;
    for (int n = 0; n < nodeList.count; n++) {
      struct ncclTopoNode* node = nodeList.list[n];
      struct ncclTopoLinkList* path;
      NCCLCHECK(getPath(system, node, baseNode->type, baseNode->id, &path));
      for (int l = 0; l < node->nlinks; l++) {
        struct ncclTopoLink* link = node->links + l;
        struct ncclTopoNode* remNode = link->remNode;
        ...
        float bw = std::min(path->bw, link->bw);
        ...
        // Update if better path type, OR same type with higher bw, OR same type/bw with strickly fewer hops.
        if (newType < remPath->type || (newType == remPath->type && remPath->bw < bw) ||
            (newType == remPath->type && remPath->bw == bw && remPath->count > (path->count + 1))) {
          ...
          remPath->bw = bw;
          remPath->type = newType;
          ...
        }
      }
    }
    memcpy(&nodeList, &nextNodeList, sizeof(nodeList));
  }
  return ncclSuccess;
}

BFS 從baseNode出發,逐層擴展。每到達一個新節點,就計算路徑的瓶頸頻寬(std::min(path->bw, link->bw))和路徑類型。路徑類型的計算有幾個特殊規則:

  • 如果經過兩個 PCI 交換器,類型升級為PATH_PXB
  • 如果經過 CPU,類型升級為PATH_PHB
  • 如果經過 DEV 節點且是 NVLink,類型升級為PATH_NVB

更新條件是「更優路徑」:類型更好,或類型相同但頻寬更高,或類型頻寬相同但跳數更少。

場景驅動的 Step-by-Step Walkthrough

現在看第二層搜尋。ncclTopoCompute是入口,它嘗試不同的參數組合,呼叫ncclTopoSearchRec進行搜尋。

搜尋的核心是遞迴函式ncclTopoSearchRecGpu。它從某個 GPU 出發,嘗試走到下一個 GPU,直到走完所有 GPU 形成一條路徑:

📎 src/graph/search.cc:639-756

c
ncclResult_t ncclTopoSearchRecGpu(struct ncclTopoSystem* system, struct ncclTopoGraph* graph,
                                  struct ncclTopoGraph* saveGraph, struct ncclTopoNode* gpu, int step, int backToNet,
                                  int backToFirstRank, int forcedOrder, int* time) {
  if ((*time) <= 0) return ncclSuccess;
  (*time)--;
  ...
  if (step == ngpus) {
    // Determine whether we found a better solution or not
    int copy = 0;
    graph->nChannels++;
    NCCLCHECKGOTO(ncclTopoCompareGraphs(system, graph, saveGraph, &copy), ret, exit);
    if (copy) {
      memcpy(saveGraph, graph, sizeof(struct ncclTopoGraph));
      if (graph->nChannels == graph->maxChannels) *time = -1;
    }
    if (graph->nChannels < graph->maxChannels) {
      NCCLCHECKGOTO(ncclTopoSearchRec(system, graph, saveGraph, time), ret, exit);
    }
    graph->nChannels--;
    ret = ncclSuccess;
    goto exit;
  }
  graph->intra[graph->nChannels * ngpus + step] = gpu->gpu.rank;
  g = gpu - system->nodes[GPU].nodes;
  if (step == backToNet) {
    // first get back to NIC
    ...
  } else if (graph->pattern == NCCL_TOPO_PATTERN_NVLS) {
    ...
  } else if (step < system->nodes[GPU].count - 1) {
    // Go to next GPU
    ...
  } else if (step == backToFirstRank) {
    // Find first GPU and loop back to it
    ...
  } else {
    // Next path
    NCCLCHECKGOTO(ncclTopoSearchRecGpu(system, graph, saveGraph, gpu, ngpus, -1, -1, forcedOrder, time), ret, exit);
  }
  ...
}

這個函式有幾個關鍵分支:

1. step == ngpus:已經走完所有 GPU,形成了一條完整路徑。此時遞增nChannels,比較當前圖和儲存的最優圖,如果更好就儲存。然後遞迴呼叫ncclTopoSearchRec嘗試搜尋下一個 channel。

2. step == backToNet:需要回到網卡。這發生在 Ring 模式(最後一個 GPU 要連回起始網卡)或 Tree 模式(第一個 GPU 要連到網卡)。

3. step < ngpus - 1:繼續走下一個 GPU。這裡會呼叫ncclTopoSearchNextGpuSort對候選 GPU 排序。

4. step == backToFirstRank:Ring 模式下,最後一個 GPU 要連回第一個 GPU。

5. else:路徑結束,進入下一輪。

ncclTopoSearchNextGpuSort決定嘗試下一個 GPU 的順序:

📎 src/graph/search.cc:254-327

c
ncclResult_t ncclTopoSearchNextGpuSort(struct ncclTopoSystem* system, struct ncclTopoGraph* graph,
                                       struct ncclTopoNode* gpu, int* next, int* countPtr, int sortNet) {
  const uint64_t flag = 1ULL << (graph->nChannels);
  int ngpus = system->nodes[GPU].count;
  struct ncclTopoLinkList* paths = gpu->paths[GPU];
  ...
  for (int i = 1; i < ngpus; i++) {
    int g = (start + i) % ngpus;
    if (paths[g].count == 0) continue; // There is no path to that GPU
    if (system->nodes[GPU].nodes[g].used & flag) continue;
    scores[count].g = g;
    scores[count].startIndex = i;
    scores[count].intraNhops = paths[g].count;
    scores[count].intraBw = paths[g].bw;
    if (netPaths) {
      scores[count].interNhops = netPaths[g].count;
      scores[count].interPciBw = gpuPciBw(system->nodes[GPU].nodes + g);
      scores[count].interBw = netPaths[g].bw;
    }
    count++;
  }

  // Sort GPUs
  qsort(scores, count, sizeof(struct ncclGpuScore), cmpScore);
  ...
}

它給每個候選 GPU 打分,排序規則是:先比 interBw(到網卡的頻寬),再比 interPciBw,再比 interNhops,再比 intraBw,最後比 intraNhops。這個優先級反映了 NCCL 的優化目標:跨機通訊是瓶頸,所以優先選到網卡頻寬高的 GPU。

設計思考與生產踩坑

為什麼搜尋有超時?看這些常數:

📎 src/graph/search.cc:329-330

c
#define NCCL_SEARCH_GLOBAL_TIMEOUT (1ULL << 19)
#define NCCL_SEARCH_TIMEOUT (1 << 14)
#define NCCL_SEARCH_TIMEOUT_TREE (1 << 14)
#define NCCL_SEARCH_TIMEOUT_SAMECHANNELS (1 << 8)

搜尋空間是指數級的——每個 channel 有 O(ngpus!) 種排列。8 卡機器就是 40320 種,16 卡就是 2 兆種。必須限制搜尋時間。NCCL_SEARCH_TIMEOUT是 16384 次迭代,NCCL_SEARCH_GLOBAL_TIMEOUT是 524288 次。超時後返回當前最優解。

坑點一:ncclTopoFollowPath的頻寬扣減是全域副作用。看這個函式:

📎 src/graph/search.cc:127-173

c
static ncclResult_t ncclTopoFollowPath(struct ncclTopoSystem* system, struct ncclTopoGraph* graph, int type1,
                                       int index1, int type2, int index2, float mult, struct ncclTopoNode** node) {
  ...
  bw *= mult;
  // Check there is enough bandwidth on paths.
  int step = 0;
  NCCLCHECK(followPath(path, node1, path->count, bw, &step));
  if (step < path->count) goto rewind;
  // Enough bandwidth : return destination node.
  graph->nHops += mult * path->count;
  *node = system->nodes[type2].nodes + index2;
  return ncclSuccess;
rewind:
  // Not enough bandwidth : rewind and exit.
  NCCLCHECK(followPath(path, node1, step, -bw, &step));
  return ncclSuccess;
}

followPath會修改路徑上每條鏈路的bw(扣減已用頻寬)。如果搜尋失敗,必須呼叫followPath用-bw恢復。這個「扣減-恢復」模式在遞迴搜尋中很容易出錯——如果某個分支忘記恢復,後續搜尋就會看到錯誤的頻寬。

坑點二:ncclTopoCompareGraphs的比較邏輯很微妙。它優先比較nChannels * bwIntra,但還有一堆特殊情況:

📎 src/graph/search.cc:446-477

c
ncclResult_t ncclTopoCompareGraphs(struct ncclTopoSystem* system, struct ncclTopoGraph* graph,
                                   struct ncclTopoGraph* refGraph, int* copy) {
  // 1. Try to get the same nChannels between Rings and Trees
  if (graph->nChannels < graph->minChannels) return ncclSuccess;
  const bool evenReference = refGraph->nChannels > 0 && !(refGraph->nChannels & 1);
  const bool evenReferenceIsBetter = refGraph->nChannels * refGraph->bwIntra >= graph->nChannels * graph->bwIntra;
  // Favor an even number of channels when aggregate bandwidth is equal or better.
  if (graph->pattern != NCCL_TOPO_PATTERN_NVLS && evenReference && (graph->nChannels & 1) &&
      graph->nChannels < system->nodes[NET].count && evenReferenceIsBetter)
    return ncclSuccess;
  ...
〔設計推斷與架構權衡〕

為什麼要偏好偶數 channel? 因為 Ring 演算法在偶數 channel 時能更好地配對——每個 channel 可以分成兩半,一半順時針一半逆時針,減少網路壅塞。

Ring 與 Tree:把搜尋結果變成演算法拓撲

直覺模型

搜尋演算法找到的是一組路徑,但演算法需要的是明確的「誰發給誰」的順序。Ring 把所有 rank 串成一個環,每個 rank 從上一個收、發給下一個。Tree 則是一棵樹,資料從根往下流或從葉子往上匯聚。

如果沒有這兩個模組,搜尋演算法就只是找到了一堆路徑,無法告訴 GPU kernel 具體怎麼發資料。

資料結構與記憶體佈局

Ring 的構建由ncclBuildRings完成:

📎 src/graph/rings.cc:29-74

c
ncclResult_t ncclBuildRings(int nrings, int* rings, int rank, int nranks, int* prev, int* next) {
  ncclResult_t ret = ncclSuccess;
  uint64_t* rankFound;
  int rankFoundSize = DIVUP(nranks, 64);
  NCCLCHECK(ncclCalloc(&rankFound, rankFoundSize));

  for (int r = 0; r < nrings; r++) {
    int current = rank;
    for (int i = 0; i < nranks; i++) {
      rankFound[current / 64] |= (1ULL << (current % 64));
      rings[r * nranks + i] = current;
      current = next[r * nranks + current];
    }
    ...
    if (current != rank) {
      WARN("Error : ring %d does not loop back to start (%d != %d)", r, current, rank);
      ret = ncclInternalError;
      goto end;
    }
    // Check that all ranks are there
    for (int i = 0; i < nranks; i++) {
      uint64_t bits = rankFound[i / 64], mask = 1ULL << (i % 64);
      // Fast check 64 ranks at a time
      if (mask == 1 && bits == 0xffffffffffffffff) {
        i += 63;
        continue;
      }
      if ((bits & mask) == 0) {
        WARN("Error : ring %d does not contain rank %d", r, i);
        ret = ncclInternalError;
        goto end;
      }
    }
    memset(rankFound, 0, rankFoundSize * sizeof(uint64_t));
  }
end:
  free(rankFound);
  return ret;
}

輸入是prev和next陣列(每個 rank 的前驅和後繼),輸出是rings陣列(每個 channel 的完整 rank 順序)。它從當前 rank 出發,沿著next指標走一圈,驗證是否回到起點,並檢查所有 rank 都被訪問到。

Tree 的構建由ncclGetBtree完成:

📎 src/graph/trees.cc:32-67

c
ncclResult_t ncclGetBtree(int nranks, int rank, int* u, int* d0, int* d1, int* parentChildType) {
  int up, down0, down1;
  int bit;
  for (bit = 1; bit < nranks; bit <<= 1) {
    if (bit & rank) break;
  }

  if (rank == 0) {
    *u = -1;
    *d0 = -1;
    // Child rank is > 0 so it has to be our child 1, not 0.
    *d1 = nranks > 1 ? bit >> 1 : -1;
    return ncclSuccess;
  }

  up = (rank ^ bit) | (bit << 1);
  // if smaller than the parent, we are his first child, otherwise we're his second
  if (up >= nranks) up = (rank ^ bit);
  *parentChildType = (rank < up) ? 0 : 1;
  *u = up;

  int lowbit = bit >> 1;
  // down0 is always within bounds
  down0 = lowbit == 0 ? -1 : rank - lowbit;

  down1 = lowbit == 0 ? -1 : rank + lowbit;
  // Make sure down1 is within bounds
  while (down1 >= nranks) {
    down1 = lowbit == 0 ? -1 : rank + lowbit;
    lowbit >>= 1;
  }
  *d0 = down0;
  *d1 = down1;

  return ncclSuccess;
}

這個函式用位運算構建二元樹。核心思想是:找到 rank 的最低非零位bit,父節點是(rank ^ bit) | (bit << 1),左子是rank - (bit >> 1),右子是rank + (bit >> 1)。註解裡的 ASCII 圖很清楚地展示了這個結構。

場景驅動的 Step-by-Step Walkthrough

以 8 卡 Ring 為例。假設搜索結果給出了每個 rank 的next指標:

code
rank 0 -> rank 1
rank 1 -> rank 2
...
rank 7 -> rank 0

ncclBuildRings從 rank 0 出發,依次訪問 1, 2, ..., 7,最後回到 0。生成的rings[0..7] = {0, 1, 2, 3, 4, 5, 6, 7}。

對於 Tree,ncclGetBtree為每個 rank 計算父節點和子節點。以 rank 1 為例:

  • bit= 1(最低非零位是第 0 位)
  • up = (1 ^ 1) | (1 << 1) = 0 | 2 = 2
  • up >= nranks? 2 < 8,所以up = 2
  • parentChildType = (1 < 2) ? 0 : 1 = 0(是父節點的第一個孩子)
  • lowbit = 0,所以down0 = -1
  • down1 = -1

所以 rank 1 的父節點是 rank 2,沒有子節點。這符合註釋裡的樹結構:rank 1 是葉子。

設計思考與生產踩坑

〔設計推斷與架構權衡〕

為什麼 Tree 用位運算而不是顯式建樹?因為每個 rank 只需要知道自己的父節點和子節點,不需要全局樹結構。位運算可以在 O(1) 時間內計算出這些信息,避免了存儲和同步整棵樹的開銷。

坑點一:ncclBuildRings的驗證可能被跳過。如果next數組有環(比如 rank 0 -> rank 1 -> rank 0),循環會在nranks次迭代後退出,但current != rank檢查會捕獲這個問題。但如果環的長度恰好是nranks的因子,且不包含所有 rank,rankFound檢查會捕獲。

坑點二:ncclGetDtree的奇數 rank 處理。對於奇數個 rank,第二棵樹是「移位」而不是「鏡像」:

📎 src/graph/trees.cc:90-112

c
ncclResult_t ncclGetDtree(int nranks, int rank, int* s0, int* d0_0, int* d0_1, int* parentChildType0, int* s1,
                          int* d1_0, int* d1_1, int* parentChildType1) {
  // First tree ... use a btree
  ncclGetBtree(nranks, rank, s0, d0_0, d0_1, parentChildType0);
  // Second tree ... mirror or shift
  if (nranks % 2 == 1) {
    // shift
    int shiftrank = (rank - 1 + nranks) % nranks;
    ...
  } else {
    // mirror
    int u, d0, d1;
    ncclGetBtree(nranks, nranks - 1 - rank, &u, &d0, &d1, parentChildType1);
    *s1 = u == -1 ? -1 : nranks - 1 - u;
    ...
  }
  return ncclSuccess;
}

雙二叉樹(Double Tree)是 NCCL 的 Tree 算法實現——兩棵樹同時工作,一棵負責前半段數據,一棵負責後半段,提高帶寬利用率。奇數 rank 時鏡像會導致 rank 映射不完整,所以改用移位。

三者的配合:從拓撲到算法

現在把三個模塊串起來。整個流程可以用一張圖表示:

mermaid
flowchart TD
    A["ncclTopoGetSystem()"] --> B["解析 XML,创建节点"]
    B --> C["ncclTopoConnectNodes() 建立链路"]
    C --> D["ncclTopoComputePaths() 计算所有路径"]
    D --> E{"ncclTopoCompute() 搜索"}
    E -->|"Ring 模式"| F["ncclTopoSearchRecNet()"]
    E -->|"Tree 模式"| G["ncclTopoSearchRecNet()"]
    F --> H["ncclTopoSearchRecGpu() 递归搜索"]
    G --> H
    H --> I{"找到更优解?"}
    I -->|"是"| J["memcpy 保存到 saveGraph"]
    I -->|"否"| K["继续尝试其他路径"]
    J --> L["ncclBuildRings() 或 ncclGetDtree()"]
    K --> H
    L --> M["生成最终算法拓扑"]

這張圖展示了從拓撲發現到算法生成的完整流程。注意ncclTopoSearchRecGpu是一個遞歸函數,它會不斷嘗試不同的 GPU 順序,直到超時或找到最優解。

再看一個更細粒度的時序圖,展示搜索過程中各模塊的交互:

mermaid
sequenceDiagram
    participant Init as ncclTopoCompute
    participant Search as ncclTopoSearchRec
    participant Net as ncclTopoSearchRecNet
    participant Gpu as ncclTopoSearchRecGpu
    participant Follow as ncclTopoFollowPath
    participant Compare as ncclTopoCompareGraphs

    Init->>Search: ncclTopoSearchRec(system, tmpGraph, graph, &time)
    Search->>Net: ncclTopoSearchRecNet(system, graph, saveGraph, backToNet, backToFirstRank, time)
    Net->>Net: ncclTopoSelectNets() 选择候选网卡
    Net->>Gpu: ncclTopoSearchTryGpu(..., NET, n, gpu)
    Gpu->>Follow: ncclTopoFollowPath(system, graph, NET, n, GPU, g, 1, &gpu)
    Follow-->>Gpu: 返回目标 GPU 节点
    Gpu->>Gpu: 递归 ncclTopoSearchRecGpu(step+1)
    Gpu->>Compare: ncclTopoCompareGraphs(system, graph, saveGraph, &copy)
    Compare-->>Gpu: copy=1 表示更优
    Gpu->>Gpu: memcpy(saveGraph, graph)
    Gpu->>Follow: ncclTopoFollowPath(..., -1, &gpu) 恢复带宽

這個時序圖展示了搜索的核心循環:選擇網卡 -> 嘗試 GPU -> 遞歸搜索 -> 比較結果 -> 恢復帶寬。

本章小結

本章拆解了 NCCL 拓撲感知的三個環節:

1. 拓撲發現(topo.cc):從 XML 讀取設備信息,創建 GPU/CPU/PCI/NIC 節點,建立 NVLink/PCIe/網絡鏈路,形成一張完整的拓撲圖。

2. 路徑搜索(search.cc + paths.cc):先用 BFS 預計算所有節點對之間的最短路徑,再用遞歸搜索嘗試不同的 Ring/Tree 結構,找到帶寬最高的方案。

3. 算法拓撲生成(rings.cc + trees.cc):把搜索結果轉換成具體的 rank 順序,Ring 用ncclBuildRings生成環,Tree 用ncclGetBtree生成二叉樹。

本章思考與自測

Q1: 如果把ncclTopoConnectNodes中的帶寬累加link->bw += bw改成link->bw = std::max(link->bw, bw),在什麼場景下會導致性能下降?為什麼?

參考解析:帶寬累加處理的是多條並行鏈路的情況。以 4 條 NVLink 各 25 GB/s 為例,累加後是 100 GB/s,取 max 後只有 25 GB/s。在ncclTopoSetPaths中,路徑帶寬是std::min(path->bw, link->bw),如果鏈路帶寬被低估,整條路徑的帶寬都會被低估。這會導致ncclTopoCompareGraphs選擇錯誤的圖——可能選了一個 channel 數更多但每個 channel 帶寬更低的方案,實際性能反而更差。具體場景:8 卡 A100 全 NVLink 互聯,每對 GPU 之間有 4 條 NVLink。累加得到 100 GB/s,取 max 得到 25 GB/s。搜索算法會認為 NVLink 和 PCIe Gen4 x16(約 25 GB/s)帶寬相同,可能選擇走 PCIe 的路徑。

Q2: ncclTopoSearchRecGpu中(*time)--在函數入口處執行。如果搜索超時(*time <= 0),函數直接返回。這個設計在什麼情況下會導致搜索陷入死循環?如何修復?

參考解析:(*time)--在入口處遞減,如果*time初始值為 0 或負數,函數直接返回,不會遞減。但如果*time是一個很大的正數,每次遞歸都會遞減,最終會到 0。問題在於:如果某個分支的遞歸深度很大,但每次遞減後*time仍然大於 0,搜索會繼續。真正的風險是ncclTopoSearchRec中的goto search循環——如果time在循環中沒有被正確重置,可能無限循環。看ncclTopoCompute中的globalTimeout邏輯:globalTimeout -= time在每次search標籤處執行,如果globalTimeout變成負數,會goto done。但如果time被重置為NCCL_SEARCH_TIMEOUT,globalTimeout可能永遠不會變成負數。修復方法是確保globalTimeout在每次搜索後都遞減,且有一個硬上限。

Q3: ncclTopoFollowPath在搜索失敗時會調用followPath(path, node1, step, -bw, &step)恢復帶寬。如果某個遞歸分支在恢復之前就返回了(比如NCCLCHECKGOTO跳轉到exit),會發生什麼?如何檢測這種問題?

參考解析:如果恢復被跳過,路徑上的鏈路帶寬會保持被扣減的狀態。後續搜索會看到錯誤的帶寬,可能錯過最優解。檢測方法:在ncclTopoCompute結束後,遍歷所有鏈路,檢查頻寬是否與初始值一致。如果發現不一致,說明有恢復遺漏。修復方法:使用 RAII 風格的守衛物件,在解構時自動恢復頻寬。或者,在每次搜尋前保存所有鏈路的頻寬快照,搜尋後恢復。NCCL 當前的做法是在每個ncclTopoFollowPath呼叫點手動配對正向和反向呼叫,這容易出錯。一個更健壯的設計是把頻寬扣減和恢復封裝成一個函式,確保成對出現。

下一章我們將深入 tuning 模組,看 NCCL 如何根據拓撲搜尋結果和訊息大小,在 Ring、Tree、CollNet 等演算法之間做出最終選擇。本章建立的拓撲圖、路徑搜尋結果和演算法模板,將成為 tuning 模組的輸入。

透過 topo.cc 的圖構建、search.cc 的路徑搜尋以及 rings.cc 和 trees.cc 的拓撲生成,NCCL 實現了用通用圖結構描述任意拓撲、用可配置搜尋演算法找到最優解、用簡單模板生成最終演算法的設計哲學。這套機制讓 NCCL 能在從 2 卡工作站到 10000 卡叢集的各種機器上自動選到合適的演算法。然而,拓撲圖只是提供了演算法的候選路徑,具體到一次通訊該走哪條路、用哪種協定,還需要更精細的決策。下一章我們將聚焦 src/tuning 目錄,看看 tuning 模組如何結合代價模型與演算法估計,在 Ring/Tree/NVLS/PAT 以及 LL/LL128/Simple 之間做出最終選擇。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 05

第 5 章:演算法與協定選型:tuning 模組如何決定通訊路徑

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 5 章 / 共 25 章

上一章我們拆解了 NCCL 的拓撲感知能力:從 src/graph/topo.cc 枚舉裝置構建拓撲圖,到 src/graph/search.cc 搜尋最優路徑,再到 rings.cc 與 trees.cc 將搜尋結果具體化為 Ring 與 Tree 演算法拓撲。但拓撲圖只回答了「資料能走哪條路」,它沒有回答「這次通訊應該走哪條路」。同一台機器上,一次 4KB 的 AllReduce 和一次 400MB 的 AllReduce,最優解可能完全不同:前者拼的是延遲,後者拼的是頻寬;前者可能選 Tree/LL,後者可能選 Ring/Simple 或者 NVLS。tuning 模組就是那個「拍板的人」。它的輸入是訊息大小、rank 數、拓撲圖(上一章的產物)和使用者環境變數;輸出是一個 ncclTuningResult_t,裡面寫著用哪個演算法(algo)、哪個協定(proto)、開多少 channel、用多少 warp。這一章我們按「總調度 → 代價模型 → 各演算法估計 → 收尾決策」的順序,把 src/tuning 目錄拆開。核心問題只有一個:NCCL 怎麼在幾十種 (演算法, 協定) 組合裡,用一套純 CPU 的數學模型,在微秒級時間內選出最快的那一個?

一、tuning.cc:總調度與決策主幹

直覺模型

把 tuning 模組想像成一家搬家公司。客戶(一次集合通訊)來了,說「我要搬 100MB 的貨,從 8 個倉庫搬到 8 個倉庫」。調度員(ncclTuningCompute)不會真的去搬一遍試試,而是拿出一張價目表(代價模型),對每種方案(Ring/LL、Tree/Simple、NVLS/Simple……)估算一個「預計耗時」,然後挑最短的那個報價給客戶。

如果沒有這個調度員,NCCL 就只能寫死「AllReduce 永遠用 Ring」,那在小訊息場景會被 Tree 吊打,在大規模 NVLink 場景會被 NVLS 吊打。代價就是效能將在特定場景下腰斬甚至更差。

資料結構與記憶體佈局

決策的載體是ncclTuningResult_t,候選集合是ncclTuningResultList_t(一個單向鏈結串列)。鏈結串列節點定義在tuning_int.h,但 push 邏輯在tuning.cc裡:

📎 src/tuning/tuning.cc:32-39

c
ncclResult_t ncclTuningResultListPushFront(struct ncclTuningResultList_t* list, struct ncclTuningResult_t result) {
  struct ncclTuningResultListNode* node = nullptr;
  NCCLCHECK(ncclCalloc(&node, 1));
  node->result = result;
  node->next = list->head;
  list->head = node;
  return ncclSuccess;
}
〔設計推斷與架構權衡〕

注意這裡是頭插法:每算出一個有效候選,就插到鏈結串列頭部。這意味著鏈結串列順序和 id 順序是反的。為什麼用鏈結串列而不是陣列? 因為候選數量在編譯期由NCCL_TUNING_COUNT決定,但實際有效的候選是動態的(受tuningMask、平台能力、使用者環境變數影響),鏈結串列允許「只把有效的掛上去」,避免遍歷時反覆判斷valid。代價是每次決策要ncclCalloc一次,但 tuning 發生在入隊路徑上、頻率不高,這點分配開銷可以接受。

ncclTuningResult_t裡最關鍵的兩個欄位是timeUs(預計耗時,微秒)和selectionTimeUs(用於選擇的耗時,可能被 tuner 外掛覆蓋)。選擇邏輯只看後者:

📎 src/tuning/tuning.cc:155-173

c
static ncclResult_t ncclTuningSelectBestTuning(struct ncclTuningResultList_t* tunings,
                                               struct ncclTuningResult_t* const bestTuning) {
  bestTuning->timeUs = FLT_MAX;
  float bestSelectionTimeUs = FLT_MAX;
  struct ncclTuningResultListNode* node = tunings->head;
  while (node != nullptr) {
    const struct ncclTuningResult_t& tuning = node->result;
    float selectionTimeUs = tuning.selectionTimeUs > 0.0f ? tuning.selectionTimeUs : tuning.timeUs;
    ...
    if (selectionTimeUs < bestSelectionTimeUs) {
      *bestTuning = tuning;
      bestSelectionTimeUs = selectionTimeUs;
    }
    node = node->next;
  }
  return ncclSuccess;
}

這裡有個細節:bestTuning->timeUs先被設成FLT_MAX,然後遍歷。如果鏈結串列為空(所有候選都無效),bestTuning會保持NCCL_TUNING_RESULT_INIT的初始值,algo/proto 都是UNDEF。這個「空結果」在呼叫方會被特殊處理——見後面的錯誤分支。

Step-by-Step Walkthrough:一次 AllReduce 的決策流

假設應用呼叫ncclAllReduce,訊息 1MB,8 個 rank 單機 NVLink。我們跟著ncclTuningCompute走一遍。

第 0 步:單 rank 短路。如果nRanks <= 1,根本不需要通訊,直接回傳 Ring/Simple,channel 數設 0:

📎 src/tuning/tuning.cc:191-200

c
  // Set tuning to Ring/Simple for single rank case
  if (input->comm->nRanks <= 1) {
    bestTuning.algo = NCCL_ALGO_RING;
    bestTuning.proto = NCCL_PROTO_SIMPLE;
    bestTuning.symKernelId = ncclSymkKernelId_Count;
    bestTuning.ceMethodId = ncclCeMethodId_Count;
    bestTuning.nChannels = 0;
    bestTuning.maxChannels = 0;
    bestTuning.nWarps = 0;
    bestTuning.forced = 0;
  } else {

這個短路很重要:單 rank 時任何演算法估計都會除以nRanks-1之類的量,容易出 NaN 或除零。先兜底,再算帳,是防禦式編程的典型。

第 1 步:列舉所有候選。進入ncclTuningComputeAllTunings,它遍歷NCCL_TUNING_COUNT個 id:

📎 src/tuning/tuning.cc:128-149

c
ncclResult_t ncclTuningComputeAllTunings(struct ncclTuningInput_t* const input,
                                         struct ncclTuningResultList_t* const tunings) {
  ncclResult_t ret = ncclSuccess;

  for (int i = 0; i < NCCL_TUNING_COUNT; i++) {
    struct ncclTuningResult_t tuning = NCCL_TUNING_RESULT_INIT;
    tuning.id = i;
    tuning.valid = 1;

    if (!(input->tuningMask & (1ULL << i))) {
      tuning.valid = 0;
      continue;
    }
    NCCLCHECK(ncclTuningExpandId(i, &tuning.algo, &tuning.proto, &tuning.symKernelId, &tuning.ceMethodId));
    NCCLCHECKGOTO(ncclTuningComputeTuning(i, input, &tuning), ret, fail);
    if (tuning.valid) NCCLCHECKGOTO(ncclTuningResultListPushFront(tunings, tuning), ret, fail);
  }
...
}

注意tuningMask是一個 64 位元遮罩,第 i 位表示「第 i 個 (algo, proto) 組合是否允許」。這個遮罩在更上層根據平台能力、使用者環境變數、函式類型算出來。遮罩是「粗篩」,代價模型是「精算」——先排除掉根本不可能的(比如 PCI 機器上不可能有 NVLS),再對剩下的算時間。

ncclTuningExpandId把一維 id 展開成 (algo, proto, symKernelId, ceMethodId)。這個映射關係必須和cost_model.cc裡的modelMap陣列嚴格一致,否則會算錯模型。

第 2 步:逐個算代價。 ncclTuningComputeTuning只有一行,轉交給代價模型:

📎 src/tuning/tuning.cc:339-343

c
ncclResult_t ncclTuningComputeTuning(int id, struct ncclTuningInput_t* const input,
                                     struct ncclTuningResult_t* const result) {
  NCCLCHECK(ncclTuningCostModelSimModel(id, input, result));
  return ncclSuccess;
}

第 3 步:tuner 外掛介入(可選)。如果使用者裝了 tuner 外掛(比如某些雲廠商的自研調優器),NCCL 會把所有候選的timeUs打包成一個二維表generalTable[algo][proto]交給外掛,讓外掛覆蓋:

📎 src/tuning/tuning.cc:203-230

c
    if (input->comm->tuner != NULL) {
      float generalTable[NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
      for (int i = 0; i < NCCL_NUM_ALGORITHMS; i++) {
        for (int j = 0; j < NCCL_NUM_PROTOCOLS; j++) {
          generalTable[i][j] = NCCL_TUNING_IGNORE;
        }
      }
      struct ncclTuningResultListNode* node = tunings.head;
      while (node != nullptr) {
        const struct ncclTuningResult_t& tuning = node->result;
        node = node->next;
        if (tuning.algo == NCCL_ALGO_UNDEF || tuning.proto == NCCL_PROTO_UNDEF) continue;
        generalTable[tuning.algo][tuning.proto] = tuning.timeUs;
      }
      node = tunings.head;
      int nMaxChannels = 0;
      NCCLCHECKGOTO(input->comm->tuner->getCollInfo(input->comm->tunerContext, input->func, input->nBytes,
                                                    input->numPipeOps, (float**)generalTable, NCCL_NUM_ALGORITHMS,
                                                    NCCL_NUM_PROTOCOLS, input->regBuff, &nMaxChannels),
                    ret, exit);
      while (node != nullptr) {
        struct ncclTuningResult_t& tuning = node->result;
        node = node->next;
        if (tuning.algo == NCCL_ALGO_UNDEF || tuning.proto == NCCL_PROTO_UNDEF) continue;
        tuning.maxChannels = nMaxChannels;
        tuning.timeUs = generalTable[tuning.algo][tuning.proto];
      }
    }

這裡NCCL_TUNING_IGNORE是一個哨兵值,表示「這個組合沒算過/不適用」。外掛可以只改它關心的格子,其他格子保持 IGNORE,NCCL 會跳過。

第 4 步:選最優。調ncclTuningSelectBestTuning,遍歷鏈結串列取selectionTimeUs最小的。

第 5 步:算 channel 數。選出演算法後,還要決定開多少 channel:

📎 src/tuning/tuning.cc:233-235

c
  if (bestTuning.algo != NCCL_ALGO_UNDEF && bestTuning.proto != NCCL_PROTO_UNDEF) {
    NCCLCHECKGOTO(ncclTuningGetChannels(input, &bestTuning), ret, exit);
  }

ncclTuningGetChannels在tuning_int.h裡,邏輯是根據訊息大小和演算法類型,在minChannels和maxChannels之間插值。channel 數直接影響頻寬:channel 越多,並行度越高,但每個 channel 的啟動開銷也越大。

第 6 步:CTA Policy 偏置(NVLS 優先)。如果使用者設了NCCL_CTA_POLICY_EFFICIENCY,且當前是 AllGather/ReduceScatter 且 buffer 已註冊,NCCL 會嘗試把結果改成 NVLS:

📎 src/tuning/tuning.cc:240-257

c
  if (input->comm->tuner == NULL && (input->CTAPolicy & NCCL_CTA_POLICY_EFFICIENCY) &&
      ncclGetEnv("NCCL_ALGO") == NULL && ncclGetEnv("NCCL_PROTO") == NULL && !input->comm->MNNVL &&
      (input->tuningMask & (1ull << (NCCL_ALGO_NVLS * NCCL_NUM_PROTOCOLS + NCCL_PROTO_SIMPLE)))) {
    if (input->regBuff && (input->func == ncclFuncAllGather || input->func == ncclFuncReduceScatter)) {
      if ((input->comm->nNodes > 1 && input->collNetSupport && input->nvlsSupport) ||
          (input->comm->nNodes == 1 && input->nvlsSupport)) {
        int recChannels;
        NCCLCHECKGOTO(ncclNvlsRegResourcesQuery(input->comm, input->func, &recChannels), ret, exit);
        if (recChannels <= bestTuning.nChannels) {
          bestTuning.algo = NCCL_ALGO_NVLS;
          ...

這段程式碼的註解很關鍵:EFFICIENCY 偏置必須在GetChannels之後跑,因為要用到bestTuning.nChannels;而且必須檢查tuningMask裡 NVLS 位是否被允許,否則會「復活」一個被上層排除的演算法。這是典型的狀態依賴順序陷阱。

第 7 步:對稱 kernel 回退。如果選中的是對稱 kernel(symKernelId),但 buffer 沒註冊、或者平台不支援,需要回退到普通 kernel。這段邏輯在tuning.cc:258-298,是整章最繞的地方,我們放到第五節專門講。

第 8 步:無解報錯。如果所有候選都無效,algo/proto 都是 UNDEF,NCCL 會打一條 WARN,並根據使用者是否設了環境變數回傳不同錯誤碼:

📎 src/tuning/tuning.cc:308-329

c
  if ((bestTuning.algo == NCCL_ALGO_UNDEF || bestTuning.proto == NCCL_PROTO_UNDEF) &&
      bestTuning.symKernelId == ncclSymkKernelId_Count && bestTuning.ceMethodId == ncclCeMethodId_Count) {
    ...
    WARN("No algorithm/protocol nor symKernelId available for function %s with datatype %s.%s%s%s",
         ncclFuncToString(input->func), ncclDatatypeToString(input->datatype), ncclAlgoEnvStr, ncclProtoEnvStr,
         ncclSymKernelIdEnvStr);
    ret = (algoEnv || protoEnv || symKernelIdEnv) ? ncclInvalidUsage : ncclInternalError;
  }

為什麼區分錯誤碼?如果使用者設了NCCL_ALGO=ring但當前平台不支援 ring(比如某些特殊拓撲),那是使用者配置錯誤(ncclInvalidUsage);如果使用者沒設任何環境變數卻選不出演算法,那是NCCL 內部 bug(ncclInternalError)。這個區分對排障至關重要。

決策主幹流程圖

mermaid
flowchart TD
    start["ncclTuningCompute(input)"] --> check_rank{"comm->nRanks <= 1?"}
    check_rank -->|是| single["bestTuning = Ring/Simple<br/>nChannels = 0"]
    check_rank -->|否| enum["ncclTuningComputeAllTunings<br/>遍历 NCCL_TUNING_COUNT"]
    enum --> mask{"tuningMask & (1<<i)?"}
    mask -->|否| skip["tuning.valid = 0<br/>continue"]
    mask -->|是| expand["ncclTuningExpandId(i)"]
    expand --> sim["ncclTuningComputeTuning<br/>-> ncclTuningCostModelSimModel"]
    sim --> valid{"result.valid?"}
    valid -->|是| push["ncclTuningResultListPushFront"]
    valid -->|否| skip
    push --> tuner{"comm->tuner != NULL?"}
    tuner -->|是| plugin["tuner->getCollInfo<br/>覆盖 generalTable"]
    tuner -->|否| select
    plugin --> select["ncclTuningSelectBestTuning<br/>取 selectionTimeUs 最小"]
    select --> getch["ncclTuningGetChannels"]
    getch --> cta{"CTA_POLICY_EFFICIENCY<br/>且 NVLS 在 mask 内?"}
    cta -->|是| nvls["ncclNvlsRegResourcesQuery<br/>可能改写为 NVLS"]
    cta -->|否| symk
    nvls --> symk{"symKernelId 需要回退?"}
    symk -->|是| fallback["ncclTuningCompute(generalInput)<br/>回退普通 kernel"]
    symk -->|否| done
    fallback --> done["*result = bestTuning"]
    single --> done
    done --> undef{"algo/proto 仍 UNDEF?"}
    undef -->|是| warn["WARN + 返回<br/>InvalidUsage 或 InternalError"]
    undef -->|否| ret_ok["返回 ncclSuccess"]

---

二、cost_model.cc:模型註冊表與開關矩陣

直覺模型

cost_model.cc是 tuning 的總帳本。它維護一張modelMap表,每一行對應一個 (algo, proto) 組合,記錄「這個組合的初始化函式是誰、模擬函式是誰、對哪些函式啟用」。同時它負責解析使用者環境變數NCCL_ALGO/NCCL_PROTO/NCCL_SYM_KERNEL,把使用者的意圖翻譯成一張enabled[i][f]開關矩陣。

如果沒有這張表,每加一個新演算法就要改一遍 tuning 主流程,程式碼會爛成一鍋粥。表驅動讓「加演算法」變成「加一行」。

資料結構:modelMap 與開關矩陣

modelMap是一個靜態陣列,每個元素是ncclTuningModelEntry_t:

📎 src/tuning/cost_model.cc:230-277

c
static struct ncclTuningModelEntry_t modelMap[] = {
  {ncclTuningTreeModelInit, ncclTuningTreeModelSim, nullptr, {0, 0, 0, 0, 1}},       // Tree/LL
  {ncclTuningTreeModelInit, ncclTuningTreeModelSim, nullptr, {0, 0, 0, 0, 1}},       // Tree/LL128
  {ncclTuningTreeModelInit, ncclTuningTreeModelSim, nullptr, {0, 0, 0, 0, 1}},       // Tree/Simple
  {ncclTuningRingModelInit, ncclTuningRingModelSim, nullptr, {1, 1, 1, 1, 1}},       // Ring/LL
  ...
  {nullptr, nullptr, nullptr, {0}}, // CollNetDirect/LL, disabled as there is no implementation
  ...
};

每個 entry 有四個欄位:init(初始化,算好 latency/bandwidth 存到 comm 裡)、model(模擬,根據訊息大小算最終 timeUs)、finalize(清理)、enabled[5](對 Broadcast/Reduce/AllGather/ReduceScatter/AllReduce 五個函數是否啟用)。

注意enabled陣列的順序註解在 L234:Enable order: Broadcast, Reduce, AllGather, ReduceScatter, AllReduce。這個順序必須和ncclFunc_t列舉一致,否則會張冠李戴。

〔設計推斷與架構權衡〕

為什麼 init 和 sim 要分開?因為 init 裡算的東西(latency、bandwidth)只依賴 comm 的靜態屬性(拓撲、rank 數、compCap),和具體訊息大小無關。一次通訊裡可能連續調多次 tuning(比如 group 裡有多個 op),init 只跑一次,sim 每次跑。這是典型的「預計算 + 快速查詢」優化。

Step-by-Step:環境變數解析與開關矩陣構建

第 1 步:預設全開,LL128 特殊。 ncclTuningCostModelInit一開始把所有 proto 設成 1(啟用),但 LL128 設成 2:

📎 src/tuning/cost_model.cc:313-323

c
  for (int f = 0; f < NCCL_NUM_FUNCTIONS; f++) {
    for (int p = 0; p < NCCL_NUM_PROTOCOLS; p++) {
      protoEnable[f * NCCL_NUM_PROTOCOLS + p] = p == NCCL_PROTO_LL128 ? 2 : 1;
    }
    for (int a = 0; a < NCCL_NUM_ALGORITHMS; a++) {
      algoEnable[f * NCCL_NUM_ALGORITHMS + a] = 1;
    }
    for (int k = 0; k < ncclSymkKernelId_Count; k++) {
      symKernelIdEnable[f * ncclSymkKernelId_Count + k] = 1;
    }
  }

為什麼 LL128 是 2 而不是 1?因為 LL128 不是「預設啟用」,而是「有條件啟用」。2 是一個特殊標記,表示「使用者沒顯式要求,稍後由isLL128Enabled根據平台能力決定」。1 表示「無條件啟用」,0 表示「禁用」。這個三態設計在 L366 的判斷裡體現:

📎 src/tuning/cost_model.cc:364-370

c
      // Disable LL128 when 1) it is not supported on the platform, and 2) user did not explicitly request it.
      // protoEnable[..] == 2 indicates that user did not set NCCL_PROTO=LL128 explicitly.
      if (proto == NCCL_PROTO_LL128 && protoEnable[f * NCCL_NUM_PROTOCOLS + proto] == 2 &&
          !isLL128Enabled(comm->minCompCap, comm->maxCompCap, comm->graphs[algo].typeInter,
                          comm->graphs[algo].typeIntra, comm->nRanks, f, algo, comm->minDriverVersion)) {
        comm->tuningContext.enabled[i][f] = 0;
      }

第 2 步:解析使用者環境變數。如果使用者設了NCCL_ALGO或NCCL_SYM_KERNEL,先把 algo 和 symKernel 全清零(因為使用者指定了白名單):

📎 src/tuning/cost_model.cc:327-345

c
  if ((algoStr && strlen(algoStr) > 0) || (symKernelIdStr && strlen(symKernelIdStr) > 0)) {
    std::fill_n(algoEnable, NCCL_NUM_FUNCTIONS * NCCL_NUM_ALGORITHMS, 0);
    std::fill_n(symKernelIdEnable, NCCL_NUM_FUNCTIONS * ncclSymkKernelId_Count, 0);
  }
  if (protoStr) {
    INFO(NCCL_ENV, "NCCL_PROTO set by environment to %s", protoStr);
    NCCLCHECK(parseList(protoStr, ncclFuncStr, NCCL_NUM_FUNCTIONS, ncclProtoStr, NCCL_NUM_PROTOCOLS, protoEnable,
                        comm->tuningContext.forced));
  }

注意 proto 沒有清零——因為 proto 的預設值是 1/2,使用者設NCCL_PROTO=LL時,parseList會把 LL 設成 1、其他設成 0(因為unset邏輯)。這個不對稱是刻意的:algo 預設全開但使用者指定後要收窄,proto 的收窄由parseList內部處理。

第 3 步:parseList 的語法。這個函數支援相當複雜的語法,註解裡給了例子:

📎 src/tuning/cost_model.cc:14-32

c
// Parse a map of prefixes to a list of elements. The first prefix is
// optional and, if not present, the list of elements will be applied
// to all prefixes. Only the first list of elements can lack a
// prefix. Prefixes (if present) are followed by a colon. Lists of
// elements are comma delimited. Mappings of prefix to the lists of
// elements are semi-colon delimited.
//
// For example:
//
//     NCCL_ALGO="ring,collnetdirect;allreduce:tree,collnetdirect;broadcast:ring"
// Enable ring and collnetdirect for all functions, then select tree
// and collnetdirect for allreduce and ring for broadcast.

^前綴表示「取反」:

📎 src/tuning/cost_model.cc:59-67

c
    int unset, set;
    if (elemList[0] == '^') {
      unset = 1;
      set = 0;
      elemList++;
    } else {
      unset = 0;
      set = 1;
    }

所以NCCL_PROTO="^LL128;allreduce:LL128"的意思是:全域禁用 LL128,但 AllReduce 例外啟用 LL128。

第 4 步:合併 enabled 矩陣。最後遍歷所有 model,把model->enabled[f]和使用者開關做與運算:

📎 src/tuning/cost_model.cc:371-383

c
      //  Check the user env vars only for functions that have a forced configuration and not already disabled.
      if (comm->tuningContext.forced[f] == 0 || comm->tuningContext.enabled[i][f] == 0) continue;
      comm->tuningContext.enabled[i][f] = 0;
      ...
      if (((algo != NCCL_ALGO_UNDEF && algoEnable[f * NCCL_NUM_ALGORITHMS + algo] != 0) &&
           (proto != NCCL_PROTO_UNDEF && protoEnable[f * NCCL_NUM_PROTOCOLS + proto] != 0)) ||
          (symKernelId != ncclSymkKernelId_Count && symKernelIdEnable[f * ncclSymkKernelId_Count + symKernelId] != 0)) {
        comm->tuningContext.enabled[i][f] = 1;
      }

邏輯是:只有當使用者對某個函數設了 forced 配置時,才用使用者配置覆蓋模型預設值。如果使用者沒設,forced[f] == 0,直接continue,保留模型自己的enabled。這是「使用者顯式指定 > 模型預設」的優先級。

模型仿真的統一入口

所有模型最終都通過ncclTuningCostModelSimModel調用:

📎 src/tuning/cost_model.cc:470-497

c
ncclResult_t ncclTuningCostModelSimModel(int id, struct ncclTuningInput_t* const input,
                                         struct ncclTuningResult_t* const result) {
  struct ncclTuningModelEntry_t* model = nullptr;
  ncclResult_t ret = ncclSuccess;
  result->forced = input->comm->tuningContext.forced[input->func];
  NCCLCHECKGOTO(getModelEntry(id, &model), ret, not_valid);
  if (model == nullptr) {
    ret = ncclInternalError;
    goto not_valid;
  }
  if (input->comm->tuningContext.enabled[id][input->func] == 0) {
    goto not_valid;
  }
  if (model->model != nullptr) {
    NCCLCHECKGOTO(model->model(input, result), ret, not_valid);
    if (result->timeUs <= 0.0) {
      goto not_valid;
    }
  } else {
    goto not_valid;
  }
exit:
  return ret;
not_valid:
  result->timeUs = NCCL_TUNING_IGNORE;
  result->valid = 0;
  goto exit;
}

三層過濾:id 越界 → 模型禁用 → 模型返回非正時間,任何一層不過都走not_valid,把timeUs設成NCCL_TUNING_IGNORE(一個負數哨兵)、valid = 0。調用方看到valid == 0就不會把它掛進候選鏈表。

設計思考

modelMap的註解裡有一句關鍵警告:

📎 src/tuning/cost_model.cc:229

c
// IMPORTANT: this table need must be consistent with the algRegistry in src/config/algorithm_registry.cc
〔設計推斷與架構權衡〕

這意味著modelMap的下標順序必須和algorithm_registry.cc裡的算法註冊順序嚴格一致。如果有人在 registry 裡插了一個新算法但忘了改modelMap,所有 id 都會錯位,tuning 會選出一個完全錯誤的算法。這是表驅動設計的經典陷阱:隱式契約。更健壯的做法是用列舉名做 key 而不是下標,但那樣會犧牲一點編譯期優化。

---

三、ring.cc:Ring 算法的代價估計

直覺模型

Ring 算法把 N 個 rank 排成一個環,數據沿著環一圈一圈傳。它的代價模型要回答兩個問題:每步傳多少數據(帶寬)、一共要多少步(延遲)。

Ring 的直覺是「流水線」:想像 N 個人站成一圈傳水桶,每個人接到桶後倒一點水再傳給下一個人。桶轉一圈,所有人的水都混勻了。桶轉得越快(帶寬高)、圈越小(步數少),整體越快。

數據結構:latency/bandwidth 表

Ring 模型不引入新結構,它把估計結果寫進comm->tuningContext.generalLatencies[c][algo][proto]和generalBandwidths[c][algo][proto]。這兩個是三維陣列:函數 × 算法 × 協議。

初始化時先全部設成 -1.0(哨兵,表示「沒算過」):

📎 src/tuning/ring.cc:31-33

c
  for (int c = 0; c < NCCL_NUM_FUNCTIONS; c++) {
    comm->tuningContext.generalLatencies[c][algo][proto] = -1.0;
    comm->tuningContext.generalBandwidths[c][algo][proto] = -1.0;

-1.0 這個哨兵在 sim 階段被檢查:

📎 src/tuning/ring.cc:94-97

c
  if (inputs->comm->tuningContext.generalBandwidths[inputs->func][tuning->algo][tuning->proto] == -1.0f) {
    tuning->valid = 0;
    return ncclSuccess;
  }

為什麼用 -1.0 而不是 0?因為 0 是一個合法的帶寬值(雖然物理上不可能),而 -1.0 明確表示「未初始化」。浮點比較用==在這裡是安全的,因為 -1.0 是精確可表示的。

Step-by-Step:Ring 帶寬估計

第 1 步:確定用 intra 還是 inter 帶寬。單機(nNodes==1)用 intra,多機用 inter:

📎 src/tuning/ring.cc:34-37

c
    int nSteps = ncclTuningGetNsteps(c, comm->nRanks);
    float bw = (comm->nNodes == 1 || (comm->nNodes <= 2 && comm->minCompCap < 100)) ? comm->graphs[algo].bwIntra :
                                                                                      comm->graphs[algo].bwInter;
    float busBw = bw * comm->graphs[algo].nChannels;

nSteps是算法需要的步數,對 Ring 來說 AllReduce 是2*(nRanks-1),其他是nRanks-1。busBw是「總線帶寬」= 單鏈路帶寬 × channel 數。

第 2 步:按協議打折。LL 協定只用了頻寬的一半(因為 LL 的 flag 開銷),LL128 用 92%(120/128):

📎 src/tuning/ring.cc:38-42

c
    if (proto == NCCL_PROTO_LL) {
      busBw = std::min(llMaxBw, busBw * .5);
    }
    if (proto == NCCL_PROTO_LL128)
      busBw = std::min(busBw * (0.92 /*120.0/128.0*/), comm->graphs[algo].nChannels * perChMaxRingLL128Bw);

0.92 = 120/128是因為 LL128 每 128 位元組裡有 8 位元組是 flag,有效載荷只有 120 位元組。這個數字直接來自協定設計。

第 3 步:算有效頻寬。注意這裡乘了nRanks / nSteps:

📎 src/tuning/ring.cc:44-46

c
    comm->tuningContext.generalLatencies[c][algo][proto] =
      comm->tuningContext.tuningConstants.baseLatencies[algo][proto];
    comm->tuningContext.generalBandwidths[c][algo][proto] = busBw * comm->nRanks / nSteps;

為什麼乘nRanks / nSteps?這是 Ring 演算法的核心特性:每個 rank 實際搬運的資料量是nBytes * nSteps / nRanks(因為資料要繞環多圈)。所以「有效頻寬」= 總線頻寬 × nRanks / nSteps。對 AllReduce,nSteps = 2(nRanks-1),所以有效頻寬 ≈ busBw/2。

第 4 步:算延遲。延遲分 intra 和 inter 兩部分:

📎 src/tuning/ring.cc:48-63

c
    int intraHw, interHw;
    ncclTuningGetHwIndexes(comm, algo, &intraHw, &interHw);
    int hwLevel = comm->nNodes == 1 ? intraHw : interHw;

    float intraLat = comm->tuningContext.tuningConstants.hwLatencies[intraHw][algo][proto];
    // Preserve the pre-refactor model: with one rank per node, Ring inter-node steps use the exposed Tree NET latency.
    float interLat;
    if (comm->nNodes == 1) {
      interLat = intraLat;
    } else if (comm->maxLocalRanks == 1) {
      interLat = comm->tuningContext.tuningConstants.hwLatencies[NCCL_HW_NET][NCCL_ALGO_TREE][proto];
    } else {
      interLat = comm->tuningContext.tuningConstants.hwLatencies[interHw][algo][proto];
    }
    interLat += comm->graphs[algo].latencyInter;
    if (proto == NCCL_PROTO_SIMPLE) interLat += comm->graphs[algo].latencyInter;

注意 L57-58 的特殊處理:當maxLocalRanks == 1(每個節點只有 1 個 rank)時,Ring 的 inter-node 延遲用Tree 的 NET 延遲。註解說這是「preserve the pre-refactor model」——即為了保持和重構前行為一致,刻意保留的一個「怪癖」。這種歷史包袱在成熟系統裡很常見,讀原始碼時看到「preserve」字樣要格外小心,它往往意味著這裡有個不能動的相容性約束。

第 5 步:按函式類型累加。Reduce/Broadcast 和 AllReduce/AllGather/ReduceScatter 的延遲模型不同:

📎 src/tuning/ring.cc:65-87

c
    if ((c == ncclFuncReduce || c == ncclFuncBroadcast)) {
      float lat = comm->tuningContext.tuningConstants.hwLatencies[hwLevel][algo][proto];
      if (comm->graphs[algo].sameChannels) {
        comm->tuningContext.generalLatencies[c][algo][proto] += lat;
      } else {
        if (proto == NCCL_PROTO_SIMPLE)
          lat =
            comm->tuningContext.tuningConstants
              .hwLatencies[hwLevel][NCCL_ALGO_TREE][proto]; // Add some chunk latency, waiting for proper chunk modeling
        comm->tuningContext.generalLatencies[c][algo][proto] += nSteps * lat;
      }
    } else {
      // Inter-node rings still have to launch nsteps * net overhead.
      float netOverhead = 0.0;
      if (comm->nNodes > 1) {
        netOverhead = getNetOverhead(comm);
        if (proto == NCCL_PROTO_SIMPLE) netOverhead *= 3;
      }
      intraLat = std::max(intraLat, netOverhead);
      int nInterSteps = comm->nNodes == 1 ? 0 : c == ncclFuncAllReduce ? 2 * (comm->nNodes - 1) : comm->nNodes - 1;
      comm->tuningContext.generalLatencies[c][algo][proto] +=
        (nSteps - nInterSteps) * intraLat + nInterSteps * interLat;
    }

sameChannels是一個拓撲屬性,表示「環上的 intra 和 inter 步是否用同一組 channel」。如果不同,延遲要乘nSteps(每步都要等)。netOverhead是網路 post 開銷,Simple 協定要乘 3(因為 Simple 有三次網路往返:send、recv、ack)。

生產避坑:Ring/Simple 的 plateau 效應

ncclTuningRingModelSim裡有一段專門處理「plateau」的程式碼:

📎 src/tuning/ring.cc:105-137

c
  // Update Ring/Simple latency for multi-node AllReduce and
  // single NVL Domain AllReduce/AllGather/ReduceScatter for Blackwell
  bool isBlackwellNvLink =
    inputs->comm->minCompCap >= 100 && inputs->comm->graphs[NCCL_ALGO_RING].typeIntra == PATH_NVL;
  bool ringSimplePlateau =
    (inputs->comm->nNodes > 1 && inputs->func == ncclFuncAllReduce) ||
    (inputs->comm->nNodes == 1 && isBlackwellNvLink &&
     (inputs->func == ncclFuncAllReduce || inputs->func == ncclFuncAllGather || inputs->func == ncclFuncReduceScatter));
  size_t bytesPerRankPerChannel = inputs->nBytes / (inputs->comm->nChannels * inputs->comm->nRanks);

  if (tuning->algo == NCCL_ALGO_RING && tuning->proto == NCCL_PROTO_SIMPLE && ringSimplePlateau &&
      bytesPerRankPerChannel >= 64) {
    float plateauFactor = inputs->comm->minCompCap < 80 ? 1.9 : 1.4;
    ...
    lat *= plateauFactor; // Plateau effect of ring
  }
〔設計推斷與架構權衡〕

什麼是 plateau?在 Ring/Simple 裡,當訊息大到一定程度,延遲不再隨訊息線性增長,而是「卡」在一個平台上——因為此時瓶頸從「啟動開銷」變成了「頻寬」,而頻寬已經飽和。這個現象在 Blackwell NVLink 上尤其明顯(因為 NVLink 頻寬太高,延遲佔比更大)。程式碼用plateauFactor(1.4 或 1.9)乘到延遲上,模擬這個「延遲被放大」的效果。

bytesPerRankPerChannel >= 64是觸發條件:每個 rank 每個 channel 至少要傳 64 位元組,否則 plateau 不成立。這個 64 位元組來自 LL 協定的 flag 大小。

踩坑場景:如果你在 Blackwell 上跑一個 1MB 的 AllReduce,發現實際延遲比模型預測的高 40%,不要以為是 bug——這是 plateau 效應,模型已經把它算進去了。如果你手動改小plateauFactor,模型會低估延遲,導致選錯演算法。

---

四、tree.cc 與 nvls.cc:Tree 與 NVLS 的代價估計

直覺模型

Tree 演算法是「樹形廣播」:根節點把資料分給子節點,子節點再分給孫節點。它的優勢是步數少(log N 而不是 N),適合小訊息;劣勢是頻寬利用率低(每個非葉節點要轉發,實際有效頻寬只有一半)。

NVLS(NVLink SHARP)是「硬體多播」:交換器直接把資料複製給多個 GPU,不需要軟體轉發。它的優勢是頻寬高、延遲低,但需要特定硬體(Hopper 以上)和特定配置。

Tree 模型:只服務 AllReduce

Tree 模型有個硬性限制——只對 AllReduce 啟用:

📎 src/tuning/tree.cc:21-27

c
  for (int c = 0; c < NCCL_NUM_FUNCTIONS; c++) {
    if (c != ncclFuncAllReduce) {
      comm->tuningContext.generalLatencies[c][algo][proto] = -1.0;
      comm->tuningContext.generalBandwidths[c][algo][proto] = -1.0;
      enabled[c] = 0; // Hard disable
      continue;
    }
〔設計推斷與架構權衡〕

為什麼?因為 NCCL 的 Tree 實作只支援 AllReduce(其他集合操作沒有 Tree 版本)。這是一個實作約束,不是理論限制。enabled[c] = 0是「硬禁用」,比generalBandwidths = -1更徹底——前者直接讓ncclTuningCostModelSimModel在 L480 就返回not_valid,後者要到 sim 函式裡才檢查。

Tree 頻寬估計:

📎 src/tuning/tree.cc:28-43

c
    float bw = (comm->minCompCap < 100) ?
                 ((comm->nNodes <= 2) ? comm->graphs[algo].bwIntra : comm->graphs[algo].bwInter) :
                 std::min(comm->graphs[algo].bwInter, comm->graphs[algo].bwIntra);
    float busBw = bw * comm->graphs[algo].nChannels;
    if (c == ncclFuncAllReduce) busBw = std::min(busBw * .92, comm->graphs[algo].nChannels * perChMaxTreeBw);
    if (proto == NCCL_PROTO_LL) {
      busBw = std::min(busBw * 1.0 / 3.8, llMaxBw);
    }
    if (proto == NCCL_PROTO_LL128)
      busBw = std::min(busBw * (comm->nNodes == 1 ? 7.0 / 9.0 : 120.0 / 128.0),
                       comm->graphs[algo].nChannels * perChMaxTreeLL128Bw);
    if (comm->maxTreePattern == NCCL_TOPO_PATTERN_TREE) busBw *= .85;
〔設計推斷與架構權衡〕

注意 LL 協定的打折係數是1/3.8,比 Ring 的0.5更狠。為什麼 Tree 的 LL 效率更低?因為 Tree 的每個中間節點既要收又要發,LL 的 flag 開銷在雙向流量下被放大。1/3.8這個數字來自實測。

Tree 延遲估計:

📎 src/tuning/tree.cc:55-58

c
    if (c == ncclFuncAllReduce) {
      comm->tuningContext.generalLatencies[c][algo][proto] +=
        2 * ((comm->nRanks / comm->nNodes - 1) * intraLat + log2i(comm->nNodes) * interLat);
    }

2 *是因為 AllReduce = ReduceScatter + AllGather,兩趟。(nRanks/nNodes - 1)是節點內步數(每個節點內的 rank 數減一),log2i(nNodes)是節點間步數(樹的高度)。

Tree 的修正因子:Tree 模型在 sim 階段乘了一個treeCorrectionFactor:

📎 src/tuning/tree.cc:75-79

c
  int logSize = log2i(inputs->nBytes >> 6);
  float bw = inputs->comm->tuningContext.generalBandwidths[inputs->func][tuning->algo][tuning->proto];
  float lat = inputs->comm->tuningContext.generalLatencies[inputs->func][tuning->algo][tuning->proto];
  if (inputs->func == ncclFuncAllReduce && logSize >= 0 && logSize < 23)
    bw *= treeCorrectionFactor[tuning->proto][logSize];

treeCorrectionFactor是一個 3×24 的表:

📎 src/tuning/cost_model.cc:223-227

c
float treeCorrectionFactor[NCCL_NUM_PROTOCOLS][24] = {
  {1.0, 1.0, 1.0, 1.0, .9, .8, .7, .7, .7, .7, .6, .5, .4, .4, .5, .6, .7, .8, .9, 1.0, 1.0, 1.0, 1.0, 1.0},
  {1.0, 1.0, 1.0, 1.0, 1.0, .9, .8, .8, .8, .7, .6, .6, .6, .6, .6, .6, .8, .9, .9, .9, .9, 1.0, 1.0, 1.0},
  {.9, .9, .9, .9, .9, .9, .9, .8, .7, .6, .6, .5, .5, .5, .5, .6, .7, .8, .7, .7, .8, .9, .9, .9}
};

logSize = log2(nBytes >> 6),即訊息大小以 64 位元組為單位取 log2。表的下標 0-23 對應 64B 到 64B×2^23 ≈ 512MB。這個表是實測出來的「Tree 效率曲線」:小訊息時效率 1.0(延遲主導),中等訊息時效率掉到 0.4-0.5(頻寬沒打滿),大訊息時回到 1.0(頻寬打滿)。這個「中間凹陷」是 Tree 演算法的固有特性。

NVLS 模型:硬體多播的代價

NVLS 模型首先檢查硬體是否支援:

📎 src/tuning/nvls.cc:19-24

c
ncclResult_t ncclTuningNvlsModelInit(struct ncclComm* comm, int id, int enabled[NCCL_NUM_FUNCTIONS]) {
  ncclResult_t ret = ncclSuccess;
  if (!ncclNvlsTransportEnabled(comm)) {
    memset(enabled, 0, NCCL_NUM_FUNCTIONS * sizeof(int));
    return ncclSuccess;
  }

然後是一系列硬性約束:只支援 Simple 協定、單機不支援 NVLSTree、多機 NVLS 需要 CollNet:

📎 src/tuning/nvls.cc:28-41

c
  if ((algo == NCCL_ALGO_NVLS || algo == NCCL_ALGO_NVLS_TREE) && (proto != NCCL_PROTO_SIMPLE)) {
    memset(enabled, 0, NCCL_NUM_FUNCTIONS * sizeof(int));
    return ncclSuccess;
  }

  if (comm->nNodes == 1 && algo == NCCL_ALGO_NVLS_TREE) {
    memset(enabled, 0, NCCL_NUM_FUNCTIONS * sizeof(int));
    return ncclSuccess;
  }

  if (comm->config.collnetEnable == 0 && algo == NCCL_ALGO_NVLS && comm->nNodes > 1) {
    memset(enabled, 0, NCCL_NUM_FUNCTIONS * sizeof(int));
    return ncclSuccess;
  }

NVLS 頻寬估計用了一個效率因子:

📎 src/tuning/nvls.cc:12-17

c
static const float nvlsEfficiency[NCCL_NUM_COMPCAPS] = {
  0.0f, // Volta
  0.0f, // Ampere
  0.85f, // Hopper
  0.74f, // Blackwell
};
〔設計推斷與架構權衡〕

Hopper 是 0.85,Blackwell 反而降到 0.74。為什麼新一代硬體效率更低?因為 Blackwell 的 NVLink 頻寬更高,但 NVLS 的交換器處理能力沒有同比提升,導致相對效率下降。這個數字是實測的,不是理論值。

頻寬計算裡有個(nChannels - 1) / nChannels因子:

📎 src/tuning/nvls.cc:62-74

c
    int nSteps = ncclTuningGetNsteps(c, comm->nRanks);
    float intraBw = comm->graphs[algo].bwIntra * nvlsEfficiency[compCapIndex] * (comm->graphs[algo].nChannels - 1) /
                    comm->graphs[algo].nChannels;
    if (c == ncclFuncAllReduce) {
      intraBw *= 2.0f;
    } else {
      float ppn = comm->minLocalRanks;
      intraBw *= (ppn - 1) / ppn;
    }
    float interBw = comm->graphs[algo].bwInter * ((comm->nNodes <= 2 && algo == NCCL_ALGO_NVLS_TREE) ? 2 : 1);
    bw = std::min({intraBw, interBw,
                   algo == NCCL_ALGO_NVLS_TREE ? (float)perChMaxNVLSTreeBw : std::numeric_limits<float>::max()});
    bw = bw * comm->graphs[algo].nChannels;

(nChannels - 1) / nChannels是因為 NVLS 需要留一個 channel 做同步。(ppn - 1) / ppn是 AllGather/ReduceScatter 的額外開銷(每個 rank 要等前一個 rank 的資料)。

生產避坑:NVLS 的硬性約束

NVLS 模型在 sim 階段還有一層執行時檢查:

📎 src/tuning/nvls.cc:136-156

c
  int nvlsSupport = inputs->nvlsSupport;
  if (!nvlsSupport) {
    tuning->valid = 0;
    tuning->timeUs = -1.0;
    return ret;
  }
  if (inputs->func != ncclFuncAllReduce && inputs->comm->graphs[tuning->algo].nChannels > NCCL_MAX_NVLS_ARITY) {
    tuning->valid = 0;
    tuning->timeUs = -1.0;
    return ret;
  }
  if (inputs->func != ncclFuncAllReduce && inputs->comm->localRanks > NCCL_MAX_NVLS_ARITY) {
    tuning->valid = 0;
    tuning->timeUs = -1.0;
    return ret;
  }

NCCL_MAX_NVLS_ARITY是 NVLS 多播組能容納的最大 GPU 數。如果超過這個數,NVLS 不可用。踩坑場景:在一個 16 卡 NVLink 域裡跑 AllGather,如果NCCL_MAX_NVLS_ARITY是 8,NVLS 會被禁用,tuning 會回退到 Ring。如果你不知道這個限制,會以為「NVLS 明明硬體支援為什麼不用」。

---

五、對稱 kernel 回退與錯誤恢復鏈

直覺模型

對稱 kernel(symmetric kernel)是 NCCL 的新特性:當所有 rank 的 buffer 都註冊到對稱記憶體後,kernel 可以用更高效的指令存取對端記憶體。但如果 buffer 沒註冊,或者平台不支援,就必須回退到普通 kernel。這個回退邏輯是 tuning 裡最繞的部分。

Step-by-Step:回退決策

回退邏輯在tuning.cc:258-298。我們拆開看。

第 1 步:判斷是否需要回退。入口條件:

📎 src/tuning/tuning.cc:258-263

至此,tuning 模組的決策鏈條已經清晰:它接收拓撲圖與通訊參數,透過代價模型和演算法估計,在微秒級內輸出最優的 (演算法, 協定, channel, warp) 組合。但選型只是開始——這個決策結果如何被下游使用?下一章我們將進入 src/enqueue/enqueue.cc 的主幹,看一次 ncclAllReduce 呼叫如何經過參數校驗、演算法/協定確定、channel 切分,最終生成 ncclInfo 與 ncclTaskColl 結構。這是全書從「使用者視角」切換到「引擎視角」的關鍵一章,你將探明一次集合通訊呼叫在 host 側被翻譯成了什麼,以及它與後續 kernel 啟動之間的邊界。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 06

第 6 章:算子下發全景:ncclAllReduce 如何變成一個可執行的 kernel 任務

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 6 章 / 共 25 章

上一章我們走完了 tuning 模組,知道 NCCL 會在微秒級內為一次集合通訊選定 (演算法, 協定, channel, warp) 組合。但選型結果本身只是一堆數字——它需要被「翻譯」成 GPU kernel 能讀懂的任務描述物件,才能被真正執行。本章進入 src/enqueue/enqueue.cc 的主幹,回答一個核心問題:當使用者呼叫 ncclAllReduce 時,host 側到底發生了什麼?從 ncclAllReduce 到 ncclEnqueueCheck,經過參數校驗、演算法/協定確定、channel 切分,最終生成 ncclInfo 與 ncclTaskColl 結構。這是全書從「使用者視角」切換到「引擎視角」的關鍵一章。如果把 NCCL 比作一家餐廳,那麼 enqueue 模組就是「前台點單系統」:使用者(應用層)說「我要一份 AllReduce」,前台把它翻譯成廚房(GPU kernel)能執行的工單——幾號灶台、用什麼鍋、分幾批做。沒有這個翻譯層,廚房根本不知道要做什麼菜。

一、入口:ncclAllReduce 如何建構 ncclInfo

直覺模型

ncclAllReduce是使用者直接呼叫的 API 函式。它的職責極其單一:把使用者傳入的裸參數打包成一個ncclInfo結構體,然後交給ncclEnqueueCheck。這就像你去銀行櫃檯辦業務,櫃員先把你的需求填進一張標準表單,再轉交給後台系統。

如果沒有這一層,每個集合通訊 API 都要自己處理參數校驗、group 語意、profiler 埋點——程式碼會重複到無法維護。

資料結構:ncclInfo 的記憶體佈局

ncclInfo是貫穿整個 enqueue 流程的核心載體。它的定義在src/include/info.h:

📎 src/include/info.h:17-44

這個結構體有 20+ 個欄位,我們可以按功能分成四組:

欄位組欄位作用
集合通訊參數coll, sendbuff, recvbuff, count, datatype, op, root描述「做什麼」
通訊域與流comm, stream描述「在哪做」
演算法細節chunkSteps, sliceSteps描述「怎麼切分」
單邊操作peerWinOffset, peerWin, sigIdx, ctx, flags, nDesc, signalDescsRMA 專用
使用者配置collConfig從使用者 config 拷貝的私有副本

注意collConfig的註解:"A config copied from config passed by user so older user config can be safely accessed during synchronous host scheduling (never at launch/replay)" 📎 src/include/info.h:41-43。這是一個關鍵設計——使用者傳入的 config 指標可能在ncclGroupEnd之前就被銷毀,所以 NCCL 在ncclInfo裡做了一份拷貝。

Step-by-Step:ncclAllReduce 的呼叫鏈

我們以ncclAllReduce為例,追蹤從使用者呼叫到ncclInfo建構的完整路徑。

第 1 步:使用者呼叫 ncclAllReduce。入口在src/collectives.cc:

📎 src/collectives.cc:206-211

這裡做了三件事:

1. NVTX3_FUNC_WITH_PARAMS打 NVTX 標記(用於 Nsight 等工具視覺化)

2. 呼叫ncclAllReduceConfigImpl,傳入config = nullptr

3. 回傳結果

第 2 步:ncclAllReduceConfigImpl 建構 ncclInfo。這是關鍵的一步:

📎 src/collectives.cc:192-202

注意這裡用了 C 風格的聚合初始化:

c
struct ncclInfo info = {ncclFuncAllReduce, "AllReduce",
                        sendbuff, recvbuff, count, datatype, op, 0, comm, stream,
                        ALLREDUCE_CHUNKSTEPS, ALLREDUCE_SLICESTEPS};

欄位按ncclInfo的宣告順序一一對應。ALLREDUCE_CHUNKSTEPS和ALLREDUCE_SLICESTEPS定義在src/include/collectives.h:

📎 src/include/collectives.h:19-20

NCCL_STEPS是環形緩衝區裡的步數(通常為 8 或 16),所以 AllReduce 的 chunkSteps 是NCCL_STEPS/2,sliceSteps 是NCCL_STEPS/4。這意味著一個 chunk 包含 2 個 slice。

第 3 步:解析使用者 config。 ncclParseCollConfig把使用者傳入的ncclCollConfig_t*解析進info.collConfig。如果config == nullptr,這個欄位保持零初始化。

第 4 步:交給 ncclEnqueueCheck。這是 enqueue 模組的真正入口。

設計思考:為什麼用聚合初始化而不是逐欄位賦值?

〔設計推斷與架構權衡〕

聚合初始化有兩個好處:一是編譯器會檢查欄位數量是否匹配(少一個欄位會警告),二是程式碼更緊湊。但缺點是欄位順序必須與結構體宣告嚴格一致——如果有人在ncclInfo中間插入一個欄位,所有聚合初始化點都會靜默錯位。這是 NCCL 程式碼裡一個隱含的維護風險。

生產踩坑:config 生命週期

一個真實的踩坑場景:使用者這樣寫程式碼:

c
ncclCollConfig_t config = {...};
ncclAllReduceConfig(..., &config);
// config 在这里被销毁(比如是栈变量,函数返回了)

如果 NCCL 沒有在ncclInfo裡拷貝 config,那麼ncclGroupEnd時存取info.collConfig就會讀到已釋放的記憶體。src/include/info.h:41-43的註解正是為了說明這個設計——config 在 task append 階段就被解析並拷貝,之後不再依賴使用者指標。

---

二、ncclEnqueueCheck:參數校驗與 group 語意

直覺模型

ncclEnqueueCheck是 enqueue 模組的「總閘門」。所有集合通訊 API 最終都匯聚到這裡。它的職責是:校驗參數合法性、處理 group 語意、呼叫 taskAppend 生成任務。如果把它比作機場安檢,那麼每個 API 函式就是值機櫃檯——值機只是收行李,真正的安檢在ncclEnqueueCheck。

如果沒有這一層,每個 API 都要自己寫一遍參數校驗和 group 處理,程式碼會膨脹數倍,而且容易漏掉某個校驗。

Step-by-Step:ncclEnqueueCheck 的執行流程

📎 src/enqueue/enqueue.cc:3478-3527

我們逐步拆解:

第 1 步:CommCheck 校驗通信域。 CommCheck(info->comm, info->opName, "comm")檢查 comm 指標是否非空、是否已初始化。如果 comm 被 revoke(比如某個 rank 出錯),直接返回錯誤:

📎 src/enqueue/enqueue.cc:3480-3485

第 2 步:處理 profiler 深度。如果已經在 group 內部(profilerGroupDepth > 0),遞增深度計數。這是為了正確處理隱式的ncclGroupStartInternal/ncclGroupEndInternal調用。

第 3 步:進入內部 group。 ncclGroupStartInternal()是 NCCL 內部的 group 機制。關鍵點:即使用戶沒有顯式調用ncclGroupStart,NCCL 也會為每次 API 調用創建一個隱式 group。這保證了單次調用的原子性。

第 4 步:確保 comm 就緒。 ncclCommEnsureReady(info->comm)等待通信域初始化完成(比如 bootstrap 完成、連接建立)。

第 5 步:ArgsCheck 參數校驗。這是最複雜的校驗步驟:

📎 src/enqueue/enqueue.cc:3497-3503

注意checkMode的處理:如果是ncclCheckModeDebugGlobal,ArgsCheck會把 info 入隊,等ncclGroupEnd時做全局校驗(比如檢查所有 rank 的 count 是否一致)。

第 6 步:調用 taskAppend。這是核心轉換步驟:

📎 src/enqueue/enqueue.cc:3513

第 7 步:遞增 opCount。每次成功入隊後,comm->opCount++。這個計數器用於匹配 send/recv 操作,也是 profiler 的時間線依據。

第 8 步:退出 group。 ncclGroupEndInternal()如果 depth 降到 0,會觸發真正的 group 操作(調度、啟動 kernel)。

並發控制:group 語義與線程安全

〔設計推斷與架構權衡〕

ncclGroupStartInternal/ncclGroupEndInternal使用線程局部存儲(TLS)來維護 group 狀態。這意味著同一個線程內的多個 API 調用會被合併成一個 group,但不同線程的調用是獨立的。這是 NCCL 支持多線程調用的基礎。

一個容易踩的坑:如果用戶在ncclGroupStart和ncclGroupEnd之間調用了非 NCCL 的 CUDA API(比如cudaMemcpy),可能會導致 stream 順序問題。NCCL 的 group 機制假設 group 內的操作都在同一組 stream 上。

錯誤恢復鏈

ncclEnqueueCheck的錯誤處理有一個精巧的設計:

📎 src/enqueue/enqueue.cc:3524-3526

如果taskAppend失敗,且 comm 是非阻塞模式,會調用ncclCommSetAsyncError記錄錯誤。這樣後續的 API 調用會立即返回錯誤,而不是繼續嘗試。這是異步錯誤傳播機制。

---

三、taskAppend:任務分發的十字路口

直覺模型

taskAppend是 enqueue 模組的「交通樞紐」。它根據info->coll的值,把任務分發到不同的處理路徑:P2P、RMA、CE、或者普通集合通信。這就像一個郵局分揀中心——根據信封上的地址,把信件投到不同的郵筒。

如果沒有這個分發層,所有類型的操作都要擠在一個巨大的 if-else 裡,代碼會難以維護。

Step-by-Step:taskAppend 的分發邏輯

📎 src/enqueue/enqueue.cc:3337-3476

第 1 步:判斷是否啟用新架構。 ncclParamEnqueueRearchEnable()是一個環境變量開關(默認 0)。如果啟用,走rawTaskAppend路徑——這是 NCCL 正在開發的新任務模型。

第 2 步:P2P 分發。如果是 Send/Recv,調用p2pTaskAppend:

📎 src/enqueue/enqueue.cc:3343-3345

第 3 步:RMA 分發。如果是 PutSignal/Signal/WaitSignal,調用rmaTaskAppend:

📎 src/enqueue/enqueue.cc:3346-3347

第 4 步:空集合通信提前返回。 if (info->count == 0) return ncclSuccess;——count 為 0 的集合通信直接丟棄。

第 5 步:算法選擇校驗。 ncclCollConfigGetAlgMask校驗用戶傳入的算法選擇是否合法:

📎 src/enqueue/enqueue.cc:3357-3358

第 6 步:FP8 類型檢查。FP8 歸約需要 sm90+:

📎 src/enqueue/enqueue.cc:3360-3366

第 7 步:歸約操作轉換。 hostToDevRedOp把 host 側的ncclRedOp_t轉換成設備側的ncclDevRedOpFull:

📎 src/enqueue/enqueue.cc:3370-3371

第 8 步:單 rank 提前返回。如果comm->nRanks == 1,直接調用ncclLaunchOneRank執行本地歸約,不需要生成任務:

📎 src/enqueue/enqueue.cc:3373-3377

第 9 步:多 rank 路徑。這是最複雜的分支,包含 CE 路由、AllToAll/Gather/Scatter 降級、以及普通集合通信:

📎 src/enqueue/enqueue.cc:3378-3470

數據結構:ncclTaskColl 的字段

collTaskAppend是生成ncclTaskColl的地方。我們看它的核心邏輯:

📎 src/enqueue/enqueue.cc:2757-2851

關鍵字段賦值:

字段來源含義
funcinfo->coll集合通信類型
sendbuff/recvbuffinfo->sendbuff/recvbuff緩衝區指標
countinfo->count元素數量
datatypeinfo->datatype數據類型
trafficBytescount * elementSize * ncclFuncTrafficPerByte流量估算
opHost/opDevinfo->op/opDev歸約操作
chunkSteps/sliceStepsinfo->chunkSteps/sliceSteps切分步數
minCTAs/maxCTAs/nvlsCTAs配置解析資源上限
algMaskncclCollConfigGetAlgMask算法選擇掩碼

注意trafficBytes的計算:

📎 src/enqueue/enqueue.cc:2813

ncclFuncTrafficPerByte返回每種集合通信的流量倍數:

📎 src/enqueue/enqueue.cc:123-134

AllReduce 返回 2(因為要 reduce + broadcast),AllGather/ReduceScatter 返回 nRanks,其他返回 1。

設計思考:為什麼 AllGather/Broadcast 要轉成 int8?

📎 src/enqueue/enqueue.cc:2808-2812

AllGather 和 Broadcast 把 count 乘以 elementSize,然後把 datatype 改成ncclInt8。這是一個優化:這兩種操作不涉及歸約,所以不需要關心資料類型,統一按位元組處理可以簡化 kernel 邏輯。

生產踩坑:CTAPolicy 的解析順序

📎 src/enqueue/enqueue.cc:3390-3397

CTAPolicy 的解析有一個微妙的優先級:env > per-call > comm。而且NCCL_CTA_POLICY_ZERO優先於NCCL_CTA_POLICY_EFFICIENCY。如果使用者同時設定了這兩個標誌,ZERO 會生效。

一個真實的踩坑場景:使用者設定了NCCL_CTA_POLICY=EFFICIENCY,但發現 CE 路徑沒有被使用。原因是 CE 路由要求CTAPolicy & NCCL_CTA_POLICY_ZERO為真,而 EFFICIENCY 不滿足這個條件。

---

四、ncclPrepareTasks:從任務列表到調度佇列

直覺模型

ncclPrepareTasks是 enqueue 模組的「預處理器」。它把散亂的任務列表按 (func, op, datatype) 分桶,然後為每個桶計算演算法和協定。這就像一個圖書館管理員——先把還回來的書按類別分好,再決定每類書放在哪個書架。

如果沒有這一步,後續的scheduleCollTasksToPlan就要為每個任務單獨計算演算法,效率極低。

Step-by-Step:ncclPrepareTasks 的分桶邏輯

📎 src/enqueue/enqueue.cc:423-642

第 1 步:Broadcast 任務轉換。如果只有一個 broadcast peer,把 broadcast 任務轉成 coll 任務:

📎 src/enqueue/enqueue.cc:430-461

注意這裡把bcastTask的欄位拷貝到新的ncclTaskColl,並計算trafficBytes。然後從memPool_ncclTaskBcast釋放原任務。

第 2 步:按 (func, op, datatype) 分桶。任務從 sorter 出來是按 size 降序的,然後被分到tasksByFnOpTy陣列:

📎 src/enqueue/enqueue.cc:464-487

索引計算:((int)task->func * ncclNumDevRedOps + (int)task->opDev.op) * ncclNumTypes + (int)task->datatype。這是一個三維陣列的線性化。

第 3 步:聚合與演算法選擇。對每個桶,聚合大小相近的任務(4 倍以內),然後呼叫ncclGetAlgoInfo:

📎 src/enqueue/enqueue.cc:503-547

第 4 步:按 (collnet, nvls) 分桶。根據演算法類型,把任務分到collBins[2][2]:

📎 src/enqueue/enqueue.cc:517-544

第 5 步:拼接最終佇列。把四個桶拼接成planner->collTaskQueue:

📎 src/enqueue/enqueue.cc:553-557

資料結構:ncclTaskCollSorter

ncclTaskCollSorter是一個按trafficBytes排序的插入式排序器。ncclTaskCollSorterInsert把任務插入到正確位置,ncclTaskCollSorterDequeueAll按順序取出所有任務。

〔設計推斷與架構權衡〕

這個排序器的設計動機是:大任務優先調度。因為大任務的傳輸時間長,先啟動它們可以更好地重疊計算和通訊。

並發控制:runtimeConn 與連線建立

📎 src/enqueue/enqueue.cc:572-583

如果comm->runtimeConn為真(執行時連線模式),且某個演算法的 channel 還沒初始化,就標記algoNeedConnect。這會在後續觸發連線建立。

生產踩坑:聚合的邊界條件

📎 src/enqueue/enqueue.cc:507-508

聚合條件是aggEnd->trafficBytes < 4 * aggBeg->trafficBytes,且兩個任務都不設定aggIsolate。如果使用者設定了 per-call config(比如maxCTAs),aggIsolate會被設為 true,這個任務就不會被聚合。

一個真實的踩坑場景:使用者為某個 AllReduce 設定了maxCTAs=4,期望它只用 4 個 CTA。但由於聚合邏輯,這個任務可能和相鄰任務合併,導致實際使用的 CTA 數量不符合預期。解決方案是設定aggIsolate——NCCL 在collTaskAppend裡已經處理了這一點:

📎 src/enqueue/enqueue.cc:2821-2822

---

五、scheduleCollTasksToPlan:channel 切分與預算控制

直覺模型

scheduleCollTasksToPlan是 enqueue 模組的「調度器」。它把任務分配到具體的 channel,並計算每個 channel 的資料切分。這就像一個工廠的排產系統——決定每條生產線做什麼、做多少。

如果沒有這一步,GPU kernel 就不知道自己要處理哪部分資料。

Step-by-Step:channel 切分演算法

📎 src/enqueue/enqueue.cc:644-947

第 1 步:預算估算。先估算能放進這個 plan 的任務數量:

📎 src/enqueue/enqueue.cc:648-689

ncclTestBudget檢查工作位元組數是否超出預算:

📎 src/enqueue/enqueue.cc:343-349

第 2 步:計算每個 channel 的流量。根據 kind(collnet/nvls)計算trafficPerChannel:

📎 src/enqueue/enqueue.cc:701-707

第 3 步:Collnet 路徑。如果是 collnet 演算法,channel 分配比較簡單:

📎 src/enqueue/enqueue.cc:709-739

第 4 步:普通路徑的 cell 切分。這是最複雜的部分。NCCL 把資料切成 "cell",每個 cell 是一個最小傳輸單元:

📎 src/enqueue/enqueue.cc:740-845

關鍵變數:

  • cellSize:每個 cell 的位元組數,至少MinTrafficPerChannel(32KB)
  • cells:總 cell 數
  • cellsPerChannel:每個 channel 處理的 cell 數
  • cellsLo/cellsHi:首尾 channel 的 cell 數(可能不滿)

第 5 步:計算 chunkGrains。對每個 channel 段呼叫calcCollChunking:

📎 src/enqueue/enqueue.cc:811-825

第 6 步:生成 proxyOp。為每個 channel 生成 proxy 操作:

📎 src/enqueue/enqueue.cc:844-894

資料結構:ncclDevWorkColl

ncclDevWorkColl是裝置側的工作描述符。它的關鍵欄位:

欄位含義
sendbuff/recvbuff緩衝區指標
channelLo/channelHichannel 範圍
cbd.countLo/countMid/countHi各段元素數
cbd.chunkGrainsLo/Mid/Hi各段 chunk 粒度
direct直接標誌

並發控制:channelMask 的位元運算

📎 src/enqueue/enqueue.cc:897

這行程式碼用位元運算設定 channelMask:(2ull << channelHi) - (1ull << channelLo)。比如 channelLo=2, channelHi=5,結果是(2<<5) - (1<<2) = 64 - 4 = 60 = 0b111100,即 bit 2-5 被設定。

生產踩坑:預算溢出

📎 src/enqueue/enqueue.cc:792-794

如果預算不夠,直接返回ncclSuccess,讓外層迴圈建立新的 plan。這是一個優雅的降級策略——不報錯,只是分批處理。

一個真實的踩坑場景:如果NCCL_WORK_FIFO_BYTES設定得太小,會導致每個 plan 只能容納很少的任務,增加 kernel 啟動次數,降低效能。

---

六、finishPlan:從任務到 kernel 參數

直覺模型

finishPlan是 enqueue 模組的「打包器」。它把任務、batch、proxyOp 打包成 kernel 能直接讀取的參數結構。這就像快遞打包——把散件裝進箱子,貼上運單,等待發貨。

Step-by-Step:finishPlan 的打包邏輯

📎 src/enqueue/enqueue.cc:236-330

第 1 步:決定儲存類型。如果所有工作都能放進 kernel args,用ncclDevWorkStorageTypeArgs:

📎 src/enqueue/enqueue.cc:244-250

第 2 步:分配 kernelArgs。從記憶體堆疊分配:

📎 src/enqueue/enqueue.cc:251-255

第 3 步:Round-robin 放置 batch。每個 channel 的第一個 batch 必須放在batchZero[blockIdx.x]:

📎 src/enqueue/enqueue.cc:257-280

第 4 步:合併 proxyOp 佇列。按 opCount 合併排序:

📎 src/enqueue/enqueue.cc:282-329

資料結構:ncclDevKernelArgs

ncclDevKernelArgs是傳給 kernel 的參數結構。它包含:

  • comm:裝置側通訊器
  • channelMask:channel 位元遮罩
  • workStorageType:工作儲存類型
  • workBuf:工作緩衝區指標
  • workMask:工作緩衝區遮罩

生產踩坑:batch 順序

📎 src/enqueue/enqueue.cc:257-259

註解說得很清楚:"The first batch for each channel must be located at batchZero[blockIdx.x]"。如果這個順序錯了,kernel 會讀到錯誤的 batch,導致資料損壞。

---

本章小結

本章我們追蹤了從ncclAllReduce到ncclTaskColl的完整路徑:

1. ncclAllReduce建構ncclInfo,打包使用者參數

2. ncclEnqueueCheck校驗參數、處理 group 語義

3. taskAppend根據操作類型分發到不同路徑

4. collTaskAppend生成ncclTaskColl,解析配置

5. ncclPrepareTasks按 (func, op, datatype) 分桶,計算演算法

6. scheduleCollTasksToPlan切分 channel,生成ncclDevWorkColl

7. finishPlan打包成 kernel 參數

關鍵設計思想:

  • 分層解耦:每個函式只做一件事,透過ncclInfo和ncclTaskColl傳遞狀態
  • 預算控制:透過ncclTestBudget控制每個 plan 的大小
  • 聚合最佳化:大小相近的任務會被聚合,減少 kernel 啟動次數
  • 配置優先級:env > per-call > comm

下一章我們將進入task_sched,看 NCCL 如何編排多 channel 多 kernel 的執行順序。

本章思考與自測

Q1: 如果把collTaskAppend中的aggIsolate判斷去掉(即src/enqueue/enqueue.cc:2821-2822永遠返回 false),在什麼場景下會導致使用者設定的maxCTAs失效?為什麼?

參考解析:aggIsolate的作用是標記「這個任務不能被聚合」。如果去掉這個判斷,設定了 per-call config 的任務會和相鄰任務合併。在ncclPrepareTasks的聚合迴圈中(src/enqueue/enqueue.cc:507-508),聚合條件是aggEnd->trafficBytes < 4 * aggBeg->trafficBytes && !aggBeg->aggIsolate && !aggEnd->aggIsolate。如果aggIsolate永遠為 false,那麼即使任務設定了maxCTAs=4,它也可能和一個maxCTAs=32的任務合併。合併後的agg會取兩者的某種組合(具體取決於ncclGetAlgoInfo的實作),導致實際使用的 CTA 數量不符合使用者預期。

更嚴重的是,在scheduleCollTasksToPlan中(src/enqueue/enqueue.cc:665-666),taskAggIsolate用於確保配置了 per-call 資源的任務單獨佔一個 plan。如果這個判斷失效,多個任務會共享 plan 的 channel 預算,導致資源分配不符合預期。

Q2: 在ncclEnqueueCheck中,如果ncclGroupEndInternal()返回錯誤(比如某個 rank 的 ArgsCheck 失敗),但taskAppend已經成功執行了,會發生什麼?NCCL 如何保證狀態一致性?

參考解析:看src/enqueue/enqueue.cc:3513-3519的控制流:

c
NCCLCHECKGOTO(taskAppend(info->comm, info), ret, fail);
info->comm->opCount++;
exit:
  if (devOld != -1) CUDACHECK(cudaSetDevice(devOld));
  ncclGroupErrCheck(ret);
  NCCLCHECK(ncclGroupEndInternal());

如果taskAppend成功但ncclGroupEndInternal失敗,opCount已經遞增了。這會導致後續操作的 opCount 與對端不匹配,可能觸發 hang。

NCCL 的處理方式是:ncclGroupErrCheck(ret)會檢查是否有錯誤,如果有,會設定 comm 的錯誤狀態。後續的 API 呼叫會透過ncclCommGetAsyncError檢測到這個錯誤並立即返回。這是一種「快速失敗」策略——一旦出錯,整個 comm 進入錯誤狀態,不再嘗試恢復。

在生產環境中,這意味著一旦出現 group 錯誤,使用者需要銷毀並重建 communicator。

Q3: scheduleCollTasksToPlan中的 cell 切分演算法(src/enqueue/enqueue.cc:740-845)有一個邊界條件:當cellsLo == 0時,會跳過最少的 channel。如果這個跳過邏輯有 bug(比如channelId沒有正確遞增),會導致什麼後果?

參考解析:看src/enqueue/enqueue.cc:770-780:

c
if (cellsLo == 0) {
  // Least channel skipped. Make the next channel the new least.
  channelId += 1;
  if (nMidChannels == 0) {
    cellsLo = cellsHi;
    cellsHi = 0;
  } else {
    cellsLo = cellsPerChannel;
    nMidChannels -= 1;
  }
}

如果channelId沒有正確遞增,那麼下一個任務會從錯誤的 channel 開始分配。這會導致:

1. channel 重疊:兩個任務可能分配到同一個 channel 的同一段資料

2. 資料損壞:kernel 會重複處理或遺漏資料

3. 效能下降:channel 負載不均衡

更隱蔽的是,這種 bug 可能只在特定訊息大小下觸發(當cellsLo == 0時),難以復現。NCCL 透過plan->channelMask |= (2ull << devWork->channelHi) - (1ull << devWork->channelLo)來追蹤已使用的 channel,但這只是記錄,不能防止重疊。

至此,我們已經看清 ncclAllReduce 如何從使用者呼叫變成一串可執行的 kernel 任務:參數校驗、演算法/協定確定、channel 切分,最終生成 ncclInfo 與 ncclTaskColl。但任務被建立出來只是第一步——它們還需要被排程到多個 channel 上,生成 kernel 啟動參數,並在 group 語義下處理批次提交與依賴排序。下一章將深入 src/enqueue/task_sched 與 src/enqueue/task_prep,回答「為什麼一次 AllReduce 會啟動多個 kernel,它們之間的順序和依賴是怎麼保證的」,同時揭示 src/group.cc 中 ncclGroupStart/ncclGroupEnd 如何把多次 API 呼叫合併成一次提交。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 07

第 7 章:任務排程器:task_sched 如何編排多 channel 與 kernel 的執行順序

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 7 章 / 共 25 章

上一章我們把 ncclAllReduce 一路追到了 ncclTaskColl——任務描述物件已經躺在 comm->planner 裡了。但任務描述只是「工單」,還沒變成 GPU 上真正跑的 kernel。這一章要回答三個問題:多次 API 呼叫怎麼被攢起來一起提交?攢起來的任務怎麼被切到多個 channel 上?多個 kernel 之間的順序和依賴靠什麼保證?先給一個整體心智模型。把 NCCL 想像成一家餐廳:ncclGroupStart/ncclGroupEnd 是「購物車」,使用者把好幾道菜(多次集合通訊呼叫)丟進購物車;ncclGroupEnd 是「下單」,廚房才開始按訂單做菜。而 doLaunches 是「傳菜調度員」,它決定哪幾道菜先上、哪幾道菜可以並行做。沒有 group 語義,每道菜單獨下單,廚房每做一道就要重新點火(啟動 kernel),開銷巨大;沒有 doLaunches 的輪次調度,多 channel 的 kernel 會亂序啟動,導致資料依賴被破壞。

一、Group 語義的全域狀態:thread_local 變數與「購物車」模型

直覺模型

ncclGroupStart和ncclGroupEnd之間的所有通訊呼叫,不會立即啟動 kernel,而是被「攢」起來。攢在哪裡?攢在執行緒局部(thread_local)的全域變數裡。為什麼是 thread_local?因為 NCCL 假設同一個執行緒內的 group 呼叫是串行的,不同執行緒各自有獨立的購物車,互不干擾。如果這些狀態是全域變數而非 thread_local,兩個執行緒同時呼叫ncclGroupStart就會互相踩踏,導致一個執行緒的任務被另一個執行緒的ncclGroupEnd提交——這是災難性的。

資料結構與記憶體佈局

先看 group 的全域狀態定義。

📎 src/group.cc:34-34

cpp
thread_local int ncclGroupDepth = 0; // depth of ncclGroupStart nesting
thread_local ncclResult_t ncclGroupError = ncclSuccess;
thread_local struct ncclComm* ncclGroupCommHead[ncclGroupTaskTypeNum] = {nullptr};
thread_local struct ncclComm* ncclGroupCommPreconnectHead = nullptr;
thread_local struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> ncclAsyncJobs;
thread_local int ncclGroupBlocking = -1; /* default mode */

逐個欄位拆解:

  • ncclGroupDepth:嵌套深度。ncclGroupStart可以嵌套呼叫(雖然不常見),每次ncclGroupStart加一,ncclGroupEnd減一。只有減到 0 時才真正提交。這就像購物車可以嵌套——你在一個購物車裡又開了一個子購物車,只有最外層結算時才真正下單。
  • ncclGroupError:group 內任意一次呼叫出錯,錯誤被記錄在這裡,ncclGroupEnd時統一處理。這避免了「一次呼叫失敗後,後續呼叫還在往購物車裡加東西」的不一致狀態。
  • ncclGroupCommHead[ncclGroupTaskTypeNum]:按任務類型分組的通訊域鏈結串列頭。ncclGroupTaskTypeNum是任務類型數量(集合通訊、原始任務、管理任務、對稱註冊等)。每個類型一條鏈結串列,鏈結串列節點是ncclComm,透過comm->groupNext[type]串聯。為什麼按類型分?因為不同類型的任務提交時機和依賴關係不同——集合通訊任務需要先 preconnect,管理任務(如 destroy)需要最後執行。
  • ncclGroupCommPreconnectHead:需要預連接的通訊域鏈結串列。預連接是「提前把網路連接建好」,避免在 kernel 啟動時才建連接導致延遲。
  • ncclAsyncJobs:非同步任務佇列。有些任務(如ncclCommInitRank)是非同步的,它們被放進這個佇列,在ncclGroupEnd時統一啟動。
  • ncclGroupBlocking:阻塞模式標誌。-1表示還沒確定,0表示非阻塞,1表示阻塞。同一個 group 內不允許混用阻塞和非阻塞通訊域,否則報錯。

這裡有個關鍵設計:ncclGroupCommHead是陣列,每個元素是一條鏈結串列。鏈結串列節點透過comm->groupNext[type]串聯,而不是用獨立的鏈結串列節點結構。這意味著ncclComm結構體裡必須預留groupNext陣列欄位。這種「侵入式鏈結串列」的設計避免了額外的記憶體分配,但代價是ncclComm結構體變大。

場景驅動的 Step-by-Step Walkthrough

場景:使用者呼叫ncclGroupStart(),然後連續呼叫兩次ncclAllReduce(分別針對兩個不同的通訊域 commA 和 commB),最後呼叫ncclGroupEnd()。

第一步:ncclGroupStart做了什麼?

📎 src/include/group.h:63-66

cpp
inline ncclResult_t ncclGroupStartInternal() {
  ncclGroupDepth++;
  return ncclSuccess;
}

極其簡單:深度加一。沒有記憶體分配,沒有鎖,沒有系統呼叫。這就是為什麼ncclGroupStart幾乎零開銷。

第二步:ncclAllReduce在 group 內被呼叫時發生了什麼?

ncclAllReduce內部會呼叫ncclGroupCommJoin(comm, ncclGroupTaskTypeCollective),把通訊域加入 group 鏈結串列。

📎 src/include/group.h:80-116

cpp
inline void ncclGroupCommJoin(struct ncclComm* comm, int type) {
  if (comm->groupNext[type] == reinterpret_cast<struct ncclComm*>(NCCL_COMM_GROUP_INVALID)) {
    // Insert comm into ncclGroupCommHead adjacent to sibling comms. This preserves
    // the users program order yet insures siblings occur consecutively. This
    // is required by doLaunches() in "group.cc".
    struct ncclComm** pp = &ncclGroupCommHead[type];
    while (*pp != nullptr && comm->intraComm0 != (*pp)->intraComm0) pp = &(*pp)->groupNext[type];

    // didn't find its clique, we need to insert it with ascending order based on commHash
    if (*pp == nullptr) {
      pp = &ncclGroupCommHead[type];
      while (*pp != nullptr && (*pp)->commHash < comm->commHash) pp = &(*pp)->groupNext[type];
    }
    comm->groupNext[type] = *pp;
    *pp = comm;
    // Comms gets a new memory stack scope upon joining. Each task batched for
    // this comm is allocated there.
    if (type == ncclGroupTaskTypeCollective || type == ncclGroupTaskTypeRawTask) {
      // Initialize planner
      ncclMemoryStackPush(&comm->memScoped);
      ncclKernelPlanner::Peer* tmp = comm->planner.peers;
      ncclIntruQueue<ncclTaskRma, &ncclTaskRma::next>* tmpRmaQueues = comm->planner.rmaTaskQueues;
      int numRmaCtx = comm->config.numRmaCtx;
      memset(&comm->planner, 0, sizeof(comm->planner));
      comm->planner.peers = tmp;
      comm->planner.bcast_info.minBcastPeer = INT_MAX;
      comm->planner.bcast_info.maxBcastPeer = INT_MIN;
      comm->planner.rmaTaskQueues = tmpRmaQueues;
      if (comm->planner.rmaTaskQueues != NULL) {
        for (int i = 0; i < numRmaCtx; i++) {
          ncclIntruQueueConstruct(&comm->planner.rmaTaskQueues[i]);
        }
      }
    }
  }
  ncclGroupBlocking = comm->config.blocking;
}

這段程式碼有幾個精妙之處:

1. 冪等性檢查:if (comm->groupNext[type] == NCCL_COMM_GROUP_INVALID)確保同一個通訊域在同一個 group 內只被加入一次。如果使用者對同一個 comm 呼叫了兩次ncclAllReduce,第二次不會重複加入鏈結串列,但任務會被追加到comm->planner裡。

2. clique 排序:intraComm0是「全域實體」的標識。多個通訊域如果屬於同一個全域實體(比如透過ncclCommSplit分裂出來的),它們的intraComm0相同,被稱為一個 clique。程式碼先按intraComm0找到 clique,把 comm 插入到同 clique 的兄弟節點旁邊。如果沒找到 clique,就按commHash升序插入。這個排序是為了doLaunches能正確處理 clique 內的 barrier 同步。

3. 記憶體堆疊作用域:ncclMemoryStackPush(&comm->memScoped)為這個 comm 在 group 內分配一個新的記憶體堆疊作用域。所有為這個 comm 分配的任務(ncclTaskColl等)都從這個堆疊上分配。ncclGroupCommLeave時會ncclMemoryStackPop一次性釋放所有任務記憶體——這是「批量分配、批量釋放」的經典優化,避免了每個任務單獨malloc/free的開銷。

4. planner 重置:memset(&comm->planner, 0, sizeof(comm->planner))清空 planner,但保留了peers和rmaTaskQueues指標(先存到臨時變數,memset 後再恢復)。為什麼要保留?因為這兩個是預分配的陣列,不需要每次重新分配。bcast_info的 min/max 被重置為INT_MAX/INT_MIN,用於後續 broadcast 任務的合併優化。

第三步:ncclGroupEnd做了什麼?

📎 src/group.cc:1039-1164

ncclGroupEndInternal是核心。逐段解析:

📎 src/group.cc:1048-1061

cpp
if (ncclGroupDepth == 0) {
  WARN("ncclGroupEnd: not in a group call.");
  ret = ncclInvalidUsage;
  goto exit;
}
// ...
if ((--ncclGroupDepth) > 0) goto exit;

先檢查深度,然後減一。如果減一後還大於 0,說明還在嵌套的內層 group 裡,直接返回,不提交。只有減到 0 才繼續。

📎 src/group.cc:1063

cpp
if ((ret = ncclGroupError) != ncclSuccess) goto fail;

如果 group 內任何一次呼叫出過錯,直接跳到 fail 清理。

📎 src/group.cc:1084-1093

cpp
NEW_NOTHROW_GOTO(groupJob, ncclGroupJob, ret, fail);
ncclIntruQueueConstruct(&groupJob->asyncJobs);
groupJob->groupRefCount = 0;
groupJob->nonBlockingInit = false;
memcpy(groupJob->groupCommHead, ncclGroupCommHead, sizeof(ncclGroupCommHead));
groupJob->groupCommPreconnectHead = ncclGroupCommPreconnectHead;
groupJob->groupError = ncclSuccess;
groupJob->abortFlag = false;
groupJob->joined = false;
ncclIntruQueueTransfer(&groupJob->asyncJobs, &ncclAsyncJobs);

建立一個ncclGroupJob,把 thread_local 的 group 狀態「轉移」到 job 物件裡。ncclIntruQueueTransfer把ncclAsyncJobs佇列整體轉移到groupJob->asyncJobs。這一步很關鍵:thread_local 狀態是「臨時」的,job 物件是「持久」的,可以被非同步執行緒持有。

📎 src/group.cc:1095-1147

cpp
if (hasCommHead || !ncclIntruQueueEmpty(&groupJob->asyncJobs) || ncclGroupCommPreconnectHead != nullptr) {
  /* make sure ncclGroupBlocking has been set. */
  if (ncclGroupBlocking != 0 && ncclGroupBlocking != 1) {
    WARN("Invalid group blocking state %d", ncclGroupBlocking);
    ret = ncclInternalError;
    goto fail;
  }
  if (ncclGroupBlocking == 0) {
    /* nonblocking group */
    // ... 设置 async error 为 ncclInProgress,创建线程执行 groupLaunchNonBlocking
    groupJob->base.func = groupLaunchNonBlocking;
    STDTHREADCREATE_GOTO(groupJob->base.thread, ncclAsyncJobMain, ret, fail, &groupJob->base);
    groupJob->nonBlockingInit = true;
    ret = ncclInProgress;
  } else {
    /* blocking group */
    int savedDev;
    CUDACHECKGOTO(cudaGetDevice(&savedDev), ret, fail);
    NCCLCHECKGOTO(groupLaunch(&groupJob->base, internalSimInfoPtr), ret, fail);
    CUDACHECKGOTO(cudaSetDevice(savedDev), ret, fail);
    if (simInfo) memcpy((void*)simInfo, (void*)internalSimInfoPtr, realSize);
    delete groupJob;
  }
} else {
  // Free when not needed (single rank case)
  delete groupJob;
}

阻塞模式:直接在當前執行緒呼叫groupLaunch,同步完成。非阻塞模式:建立一個執行緒執行groupLaunchNonBlocking,立即返回ncclInProgress。使用者後續透過ncclCommGetAsyncError查詢進度。

注意cudaGetDevice/cudaSetDevice的保存和恢復:groupLaunch內部會切換 CUDA 裝置(因為不同 comm 可能在不同 GPU 上),執行完後恢復使用者原來的裝置。這是防止「NCCL 內部切換裝置後沒切回來」導致使用者後續 CUDA 呼叫跑錯裝置。

設計思考與生產踩坑

坑 1:阻塞和非阻塞通訊域混用。ncclAsyncLaunch裡有檢查:

📎 src/group.cc:55-64

cpp
/* check if there are blocking and nonblocking comms at the same time in group. */
if (comm->destroyFlag) {
  ncclGroupBlocking = 1;
} else if (ncclGroupBlocking == -1) {
  /* first met communicator */
  ncclGroupBlocking = comm->config.blocking;
} else if (ncclGroupBlocking != comm->config.blocking) {
  WARN("Blocking and nonblocking communicators are not allowed in the same group.");
  ret = ncclInvalidArgument;
}

為什麼不允許混用?因為阻塞 group 在當前執行緒同步執行,非阻塞 group 在獨立執行緒非同步執行。如果混用,無法確定ncclGroupEnd應該同步返回還是返回ncclInProgress。生產環境中,如果使用者不小心把阻塞和非阻塞 comm 放進同一個 group,會收到ncclInvalidArgument,但此時 group 狀態已經被污染,必須重新ncclGroupStart。

坑 2:ncclGroupError的傳播。如果 group 內某次呼叫失敗,ncclGroupError被設定,ncclGroupEnd會跳到 fail 分支執行groupCleanup。groupCleanup會遍歷所有 comm,釋放 planner 裡的 plan 記憶體、重置 planner、清理 rawTaskQueue。如果這一步沒做乾淨,下次ncclGroupStart時 planner 裡殘留舊資料,會導致任務重複提交或記憶體洩漏。

📎 src/group.cc:514-607

cpp
static void groupCleanup(struct ncclComm** groupCommHeadPtr,
                         struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next>* asyncJobsPtr,
                         ncclResult_t error) {
  struct ncclComm* comm;
  for (int type = 0; type < ncclGroupTaskTypeNum; ++type) {
    comm = groupCommHeadPtr[type];
    groupCommHeadPtr[type] = nullptr;
    while (comm != nullptr) {
      struct ncclComm* next = comm->groupNext[type];
      (void)ncclGroupCommLeave(comm, type);
      // We don't know if preconnect succeeded or happened at all, so clear
      // the flags that let `taskAppend()` skip over checking if preconnect
      // is needed.
      if (type == ncclGroupTaskTypeCollective || type == ncclGroupTaskTypeRawTask) {
        comm->preconnectNext = reinterpret_cast<struct ncclComm*>(0x1);
        for (int i = 0; i < comm->nRanks; i++) {
          comm->connectSend[i] = 0UL;
          comm->connectRecv[i] = 0UL;
        }
        // Reclaim abandoned kernel plan memory.
        while (!ncclIntruQueueEmpty(&comm->planner.planQueue)) {
          struct ncclKernelPlan* plan = ncclIntruQueueDequeue(&comm->planner.planQueue);
          if (!plan->persistent) {
            while (!ncclIntruQueueEmpty(&plan->proxyOpQueue)) {
              struct ncclProxyOp* pxop = ncclIntruQueueDequeue(&plan->proxyOpQueue);
              ncclMemoryPoolFree(&comm->memPool_ncclProxyOp, pxop);
            }
            ncclMemoryPoolFree(&comm->memPool_ncclKernelPlan, plan);
          }
        }
        // Reset comm->planner to empty.
        // ...
      }
      // ...
    }
  }
  // ...
}

注意comm->preconnectNext = reinterpret_cast<struct ncclComm*>(0x1)這一行。這是一個「哨兵值」,表示「這個 comm 需要重新 preconnect」。為什麼?因為 cleanup 時不知道 preconnect 是否成功,所以強制下次重新檢查。0x1這個值很巧妙——它不是一個合法的指標,但可以用來做「未初始化」標記。ncclGroupCommPreconnect裡檢查if (comm->preconnectNext == reinterpret_cast<struct ncclComm*>(0x1))來判斷是否需要加入 preconnect 鏈結串列。

---

二、任務準備:ncclPrepareTasks如何把任務描述變成可排程單元

直覺模型

ncclPrepareTasks是「備菜」環節。購物車裡的菜(任務描述)還是生的,需要先洗切配(確定演算法、協議、channel 切分),才能下鍋(啟動 kernel)。如果跳過這一步直接啟動 kernel,kernel 不知道資料怎麼切、走哪條路,會直接崩潰。

場景驅動的 Step-by-Step Walkthrough

ncclPrepareTasks在groupLaunchLegacy裡被呼叫:

📎 src/group.cc:705-746

cpp
static ncclResult_t ncclPrepareTasksAndCollPreconnect(
  struct ncclComm* comm, ncclSimInfo_t* simInfo,
  struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next>* asyncCollJobs) {
  if (ncclParamSingleProcMemRegEnable()) {
    // 单进程内存注册模式:把 prepare 和 preconnect 合并成一个异步 job
    struct ncclPrepareTasksAndCollPreconnectJob* job;
    NEW_NOTHROW(job, ncclPrepareTasksAndCollPreconnectJob);
    job->base.func = ncclPrepareTasksAndCollPreconnectFunc;
    // ...
    ncclIntruQueueEnqueue(asyncCollJobs, &job->base);
  } else {
    bool needConnect = false;
    bool algoNeedConnect[NCCL_NUM_ALGORITHMS];
    memset(algoNeedConnect, 0, sizeof(bool) * NCCL_NUM_ALGORITHMS);

    CUDACHECK(cudaSetDevice(comm->cudaDev));
    NCCLCHECK(ncclPrepareTasks(comm, algoNeedConnect, &needConnect, simInfo));

    if (comm->cuMemSupport && needConnect) {
      // 创建 preconnect job
      struct ncclPreconnectJob* job;
      NEW_NOTHROW(job, ncclPreconnectJob);
      job->base.func = ncclCollPreconnectFunc;
      // ...
      ncclIntruQueueEnqueue(asyncCollJobs, &job->base);
    }
  }
  return ncclSuccess;
}

ncclPrepareTasks的輸出是兩個東西:algoNeedConnect陣列(哪些演算法需要建立連線)和needConnect旗標(是否需要連線)。如果needConnect為真且支援 cuMem,就建立一個 preconnect job 非同步執行。

ncclPrepareTasks內部做了什麼?它遍歷comm->planner裡的任務,對每個任務確定演算法和協議,然後呼叫taskAppend把任務追加到 planner 的 plan 裡。這部分邏輯在上一章已經展開,這裡不再重複。

關鍵點:ncclPrepareTasks是按 comm 逐個呼叫的,但 preconnect 是按 clique 批次執行的。為什麼?看groupLaunchLegacy裡的註解:

📎 src/group.cc:818-834

cpp
do {
  // We need to preconnect connections for collectives clique by clique to avoid
  // race condition for split shared comms which can connect the same connections
  // at the same time.
  comm = cliqueHead;
  do {
    NCCLCHECKGOTO(ncclPrepareTasksAndCollPreconnect(comm, simInfo, &asyncCollJobs), ret, fail);
    comm = comm->groupNext[ncclGroupTaskTypeCollective];
  } while (comm != nullptr && comm->intraComm0 == cliqueHead->intraComm0);
  // connect
  NCCLCHECKGOTO(asyncJobLaunch(&asyncCollJobs, groupAbortFlag), ret, fail);
  // ...
  cliqueHead = comm;
} while (cliqueHead != nullptr);

註解說得很清楚:按 clique 逐個 preconnect,避免 split shared comms 同時連接同一組連線導致競態。如果兩個 comm 是從同一個父 comm split 出來的,它們可能共享一些連線。如果並行 preconnect,兩個執行緒可能同時嘗試建立同一個連線,導致重複連線或連線狀態不一致。按 clique 串行執行,保證同一時刻只有一個 clique 在建立連線。

並發控制與底層互動

asyncJobLaunch是非同步任務啟動的核心:

📎 src/group.cc:609-678

cpp
static ncclResult_t asyncJobLaunch(struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next>* asyncJobsMain,
                                   volatile bool* groupAbortFlag) {
  ncclResult_t ret = ncclSuccess;
  bool jobsDone = false;
  bool errorJobAbortFlag = false;

  if (!ncclIntruQueueEmpty(asyncJobsMain)) {
    struct ncclAsyncJob* job = ncclIntruQueueHead(asyncJobsMain);
    if (job->next == nullptr) {
      // 只有一个 job,直接在当前线程执行,避免线程创建开销
      job->isThreadMain = true;
      ncclAsyncJobMain(job);
      job->state = ncclGroupJobJoined;
      return job->result;
    }
    // 多个 job,每个创建一个线程
    do {
      STDTHREADCREATE(job->thread, ncclAsyncJobMain, job);
      job = job->next;
    } while (job != nullptr);

    do {
      jobsDone = true;
      job = ncclIntruQueueHead(asyncJobsMain);
      do {
        ncclGroupJobState_t state = COMPILER_ATOMIC_LOAD(&job->state, std::memory_order_acquire);
        if (state == ncclGroupJobRunning) {
          jobsDone = false;
        } else if (state == ncclGroupJobDone) {
          int err;
          if ((err = ncclThreadJoin(job->thread)) != ncclSuccess) {
            WARN("asyncJobLaunch: failed to join thread for job");
            ret = ncclSystemError;
          }
          job->state = ncclGroupJobJoined;
          if (job->result != ncclSuccess && ret == ncclSuccess) {
            ret = job->result;
            errorJobAbortFlag = true;
          }
        } else {
          // safety check
          if (state != ncclGroupJobJoined) {
            WARN("Async job state is %d, expected %d", state, ncclGroupJobJoined);
            if (ret == ncclSuccess) ret = ncclInternalError;
            errorJobAbortFlag = true;
          }
        }

        if (!job->destroyFlag &&
            (COMPILER_ATOMIC_LOAD(groupAbortFlag, std::memory_order_acquire) || errorJobAbortFlag == true)) {
          COMPILER_ATOMIC_STORE(job->abortFlag, uint32_t(1), std::memory_order_release);
          COMPILER_ATOMIC_STORE(job->abortFlagDev, uint32_t(1), std::memory_order_release);
          if (job->childAbortFlag) {
            COMPILER_ATOMIC_STORE(job->childAbortFlag, uint32_t(1), std::memory_order_release);
            COMPILER_ATOMIC_STORE(job->childAbortFlagDev, uint32_t(1), std::memory_order_release);
          }
        }

        job = job->next;
      } while (job != nullptr);
      // Let preconnect threads progress.
      if (jobsDone == false) std::this_thread::sleep_for(std::chrono::microseconds(1));
    } while (jobsDone == false);

    if (ret != ncclSuccess) goto fail;
  }

exit:
  return ret;
fail:
  goto exit;
}

這段程式碼有幾個關鍵設計:

1. 單 job 最佳化:如果佇列裡只有一個 job,不建立執行緒,直接在當前執行緒執行。這避免了執行緒建立和 join 的開銷。對於單 comm 的 group,這是常見情況。

2. 原子狀態機:job->state是一個原子變數,有三個狀態:ncclGroupJobRunning、ncclGroupJobDone、ncclGroupJobJoined。工作執行緒執行完後用COMPILER_ATOMIC_STORE(..., std::memory_order_release)設定為Done;主執行緒用COMPILER_ATOMIC_LOAD(..., std::memory_order_acquire)讀取。release/acquire 配對保證了工作執行緒的所有記憶體寫入對主執行緒可見。

3. 忙等待 + 微睡眠:主執行緒輪詢所有 job 的狀態,如果還有 job 在跑,sleep_for(1us)後繼續輪詢。為什麼用 1 微秒而不是條件變數?因為 preconnect 是短任務(通常幾十微秒到幾毫秒),條件變數的喚醒開銷可能比忙等待還大。1 微秒的睡眠避免了純自旋導致的 CPU 浪費。

4. 錯誤傳播與 abort:如果任何一個 job 失敗,errorJobAbortFlag被設定,後續所有 job 的abortFlag被原子設定為 1。工作執行緒在執行過程中會檢查abortFlag,如果發現被 abort,提前退出。這是「快速失敗」機制,避免一個 job 失敗後其他 job 還在傻跑。

Mermaid 圖:group 提交的控制流

mermaid
flowchart TD
    gs["ncclGroupStart()"] --> depth_inc["ncclGroupDepth++"]
    depth_inc --> api_calls["用户调用 ncclAllReduce 等"]
    api_calls --> join["ncclGroupCommJoin(comm, type)"]
    join --> check_dup{"comm->groupNext[type]<br/>== NCCL_COMM_GROUP_INVALID?"}
    check_dup -->|是| insert["插入 clique 链表<br/>ncclMemoryStackPush"]
    check_dup -->|否| skip["跳过(已加入)"]
    insert --> ge["ncclGroupEnd()"]
    skip --> ge
    ge --> depth_dec["--ncclGroupDepth"]
    depth_dec --> depth_zero{"depth == 0?"}
    depth_zero -->|否| ret_early["返回(嵌套内层)"]
    depth_zero -->|是| check_err{"ncclGroupError<br/>== ncclSuccess?"}
    check_err -->|否| fail_cleanup["groupCleanup()"]
    check_err -->|是| create_job["创建 ncclGroupJob<br/>转移 thread_local 状态"]
    create_job --> blocking{"ncclGroupBlocking?"}
    blocking -->|0 非阻塞| spawn_thread["STDTHREADCREATE<br/>groupLaunchNonBlocking"]
    blocking -->|1 阻塞| sync_launch["groupLaunch() 同步执行"]
    spawn_thread --> ret_progress["返回 ncclInProgress"]
    sync_launch --> ret_ok["返回 ncclSuccess"]
    fail_cleanup --> reset["groupLocalResetJobState()"]
    ret_progress --> reset
    ret_ok --> reset

---

三、doLaunches:多 channel 多 kernel 的輪次調度

直覺模型

doLaunches是「傳菜調度員」。廚房(GPU)有多個灶台(channel),每道菜(kernel plan)需要按順序上。但不同 comm 的菜可能可以並行上,同一個 comm 的菜必須按順序上。調度員要保證:同一個 clique 內的 comm 同步推進(用 barrier),不同 clique 之間可以獨立推進。

資料結構與記憶體佈局

doLaunches的核心資料結構是ncclKernelPlan和comm->planner.unlaunchedPlansHead。

📎 src/group.cc:427-503

cpp
ncclResult_t doLaunches(struct ncclComm* head, int taskType) {
  ncclResult_t result = ncclSuccess;
  struct ncclComm* cliqueHead = head;
  struct ncclComm* cliqueNextHead;
  bool useBarrier = ncclParamLaunchMode == ncclLaunchModeGroup;
  // This outer loop iterates over cliques of comms which are siblings of the
  // same global entity. We calculate a clique as all comms which have the same
  // `intraComm0` value.
  do {
    struct ncclComm* comm = cliqueHead;
    bool capturingYes = false, capturingNo = false;
    do {
      (ncclCudaGraphValid(comm->planner.capturingGraph) ? capturingYes : capturingNo) = true;
      CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), result, failure);
      NCCLCHECKGOTO(ncclLaunchPrepare(comm), result, failure);
      if (useBarrier) ncclCommIntraBarrierIn(comm, 1);
      comm = comm->groupNext[taskType];
    } while (comm != nullptr && comm != reinterpret_cast<struct ncclComm*>(NCCL_COMM_GROUP_INVALID) &&
             comm->intraComm0 == cliqueHead->intraComm0);
    cliqueNextHead = comm;

    if (capturingYes && capturingNo) {
      // We have entered barriers but are aborting without leaving them. Thus
      // these comms are permanently trashed. We need a good mechanism for
      // tracking and reporting that.
      WARN("Either none or all communicators in a ncclGroup() can be CUDA graph captured.");
      result = ncclInvalidUsage;
      goto failure;
    }

    while (true) {
      // Iterate rounds of launches for clique.
      bool moreRounds = false;
      comm = cliqueHead;
      do {
        // Iterate clique members.
        struct ncclComm* next = comm->groupNext[taskType];
        if (useBarrier) {
          // Barrier reduction result tells us if this was the final round.
          moreRounds = 0 != ncclCommIntraBarrierOut(comm);
        } else {
          moreRounds |= comm->planner.unlaunchedPlansHead != nullptr;
        }
        if (moreRounds) {
          // Pop next unlaunched kernel
          struct ncclKernelPlan* plan = comm->planner.unlaunchedPlansHead;
          if (plan != nullptr) {
            comm->planner.unlaunchedPlansHead = plan->next;
            CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), result, failure);
            NCCLCHECKGOTO(ncclLaunchKernelBefore_NoUncapturedCuda(comm, plan), result, failure);
            if (plan->isCeColl) {
              NCCLCHECKGOTO(ncclLaunchCeColl(comm, plan), result, failure);
            } else if (plan->isRma) {
              NCCLCHECKGOTO(ncclLaunchRma(comm, plan), result, failure);
            } else {
              NCCLCHECKGOTO(ncclLaunchKernel(comm, plan), result, failure);
            }
          }
          // Barrier reduction input indicates if we require further rounds.
          if (useBarrier) ncclCommIntraBarrierIn(comm, comm->planner.unlaunchedPlansHead != nullptr ? 1 : 0);
          if (plan != nullptr) {
            NCCLCHECKGOTO(ncclLaunchKernelAfter_NoCuda(comm, plan), result, failure);
          }
        } else {
          // Final round.
          CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), result, failure);
          NCCLCHECKGOTO(ncclLaunchFinish(comm), result, failure);
        }
        comm = next;
      } while (comm != reinterpret_cast<struct ncclComm*>(NCCL_COMM_GROUP_INVALID) && comm != cliqueNextHead);
      if (!moreRounds) break;
    }
    cliqueHead = cliqueNextHead;
  } while (cliqueHead != nullptr && cliqueHead != reinterpret_cast<struct ncclComm*>(NCCL_COMM_GROUP_INVALID));
failure:
  return result;
}

場景驅動的 Step-by-Step Walkthrough

場景:兩個 comm(commA 和 commB)屬於同一個 clique(intraComm0相同),每個 comm 有 3 個 kernel plan 待啟動。

第一層迴圈:遍歷 clique

外層do-while遍歷所有 clique。cliqueHead是當前 clique 的第一個 comm。內層do-while遍歷 clique 內的所有 comm(comm->intraComm0 == cliqueHead->intraComm0)。

對每個 comm:

  • cudaSetDevice(comm->cudaDev):切換到該 comm 對應的 GPU。
  • ncclLaunchPrepare(comm):準備啟動,包括設定 CUDA 流、檢查資源等。
  • ncclCommIntraBarrierIn(comm, 1):進入 barrier,初始值為 1。

第二層迴圈:輪次調度

while (true)迴圈執行「輪次」。每一輪,clique 內每個 comm 啟動一個 kernel plan。

關鍵在moreRounds的計算:

  • 有 barrier 模式(useBarrier == true):moreRounds = 0 != ncclCommIntraBarrierOut(comm)。ncclCommIntraBarrierOut是一個跨 comm 的 barrier 歸約操作。它等待 clique 內所有 comm 都呼叫了ncclCommIntraBarrierIn,然後返回所有輸入值的歸約結果(這裡是邏輯或)。如果任何一個 comm 還有未啟動的 plan,歸約結果為 1,moreRounds為 true,繼續下一輪。如果所有 comm 都沒有未啟動的 plan,歸約結果為 0,moreRounds為 false,進入 final round。
  • 無 barrier 模式:moreRounds |= comm->planner.unlaunchedPlansHead != nullptr。直接檢查每個 comm 是否還有未啟動的 plan。注意這裡用的是|=,只要有一個 comm 還有 plan,moreRounds就為 true。

為什麼需要 barrier?因為 clique 內的 comm 是「兄弟」,它們可能共享 GPU 資源或網路連接。如果一個 comm 啟動了 3 個 kernel,另一個只啟動了 1 個,先啟動完的 comm 會進入ncclLaunchFinish,釋放資源,而另一個 comm 還在用這些資源,導致 use-after-free。barrier 保證 clique 內所有 comm 同步推進:要麼都啟動第 N 輪,要麼都進入 final round。

kernel 啟動分支

📎 src/group.cc:477-483

cpp
if (plan->isCeColl) {
  NCCLCHECKGOTO(ncclLaunchCeColl(comm, plan), result, failure);
} else if (plan->isRma) {
  NCCLCHECKGOTO(ncclLaunchRma(comm, plan), result, failure);
} else {
  NCCLCHECKGOTO(ncclLaunchKernel(comm, plan), result, failure);
}

三種 plan 類型:

  • isCeColl:CollNet 集合通訊(用網卡卸載做集合通訊)。
  • isRma:RMA(Remote Memory Access)任務。
  • 預設:普通 GPU kernel。

每種類型的啟動函數不同,但都遵循「Before -> Launch -> After」的模式:

  • ncclLaunchKernelBefore_NoUncapturedCuda:啟動前準備(設定 kernel 參數、上傳到裝置等)。
  • ncclLaunchKernel:實際啟動 kernel(cudaLaunchKernel)。
  • ncclLaunchKernelAfter_NoCuda:啟動後清理(更新狀態、釋放臨時資源)。

Final round

當moreRounds為 false 時,執行ncclLaunchFinish(comm)。這一步做最終的清理:釋放 plan 記憶體、更新 comm 狀態、通知 proxy 執行緒等。

並發控制與硬體互動

ncclCommIntraBarrierIn/Out是 clique 內 comm 的同步原語。它的實作涉及原子操作和自旋等待。In把值寫入共享記憶體,Out等待所有 comm 都寫入後讀取歸約結果。這個 barrier 是跨行程的(如果 comm 在不同行程),底層可能用共享記憶體或網路。

為什麼用 barrier 而不是簡單的「檢查所有 comm 是否還有 plan」?因為「檢查」是非原子的:commA 檢查時 commB 還有 plan,commA 決定繼續;但 commB 在 commA 檢查後立即啟動完最後一個 plan,進入 final round。commA 還在啟動 kernel,commB 已經釋放了共享資源。barrier 把「檢查」和「決定」變成一個原子操作,消除了這個競態。

生產避坑指南

坑 1:CUDA graph capture 混用。

📎 src/group.cc:448-455

cpp
if (capturingYes && capturingNo) {
  // We have entered barriers but are aborting without leaving them. Thus
  // these comms are permanently trashed. We need a good mechanism for
  // tracking and reporting that.
  WARN("Either none or all communicators in a ncclGroup() can be CUDA graph captured.");
  result = ncclInvalidUsage;
  goto failure;
}

如果 clique 內一部分 comm 在 CUDA graph capture 模式下,另一部分不在,直接報錯。註解說「these comms are permanently trashed」——因為已經進入了 barrier 但沒有退出,這些 comm 的 barrier 狀態永遠不一致,後續無法再使用。這是一個不可恢復錯誤,使用者必須重建通訊域。生產環境中,如果使用者混用 graph capture 和非 capture 的 comm,會收到ncclInvalidUsage,但更嚴重的是 comm 已經損壞。

坑 2:useBarrier的配置依賴。useBarrier = ncclParamLaunchMode == ncclLaunchModeGroup。如果使用者設定了NCCL_LAUNCH_MODE=GROUP,走 barrier 路徑;否則走非 barrier 路徑。非 barrier 路徑下,moreRounds用|=累積,但每個 comm 獨立判斷。如果 commA 還有 plan 而 commB 沒有,commB 會進入 final round 執行ncclLaunchFinish,而 commA 還在啟動 kernel。這在某些場景下是安全的(comm 之間沒有共享資源),但如果共享了 proxy 執行緒或網路連接,可能導致問題。所以預設推薦用 barrier 模式。

---

四、groupLaunchLegacy的完整執行鏈

場景驅動的 Step-by-Step Walkthrough

groupLaunchLegacy是阻塞模式下的完整提交流程。按順序執行:

階段 1:P2P preconnect

📎 src/group.cc:756-774

cpp
if (!simInfo && groupCommPreconnectHeadMain != nullptr) {
  struct ncclComm* comm = groupCommPreconnectHeadMain;
  do {
    struct ncclPreconnectJob* job;
    NEW_NOTHROW_GOTO(job, ncclPreconnectJob, ret, fail);
    job->base.func = ncclP2PPreconnectFunc;
    // ...
    ncclIntruQueueEnqueue(asyncJobsMain, (struct ncclAsyncJob*)job);
    struct ncclComm* next = comm->preconnectNext;
    comm->preconnectNext = reinterpret_cast<struct ncclComm*>(0x1);
    comm = next;
  } while (comm != nullptr);
}
NCCLCHECKGOTO(asyncJobLaunch(asyncJobsMain, groupAbortFlag), ret, fail);

對每個需要 preconnect 的 comm 建立一個ncclP2PPreconnectFuncjob,然後批量啟動。ncclP2PPreconnectFunc內部呼叫ncclTransportP2pSetup建立 P2P 連接。

階段 2:對稱記憶體註冊

📎 src/group.cc:778-808

cpp
// only loop through sym alloc and register tasks
for (int type = ncclGroupTaskTypeSymRegister; type <= ncclGroupTaskTypeSymRegister; ++type) {
  if (groupCommHeadMain[type]) {
    // 按 clique 批量执行 ncclCommGroupRegisterSymmetric
  }
}

對稱記憶體註冊(ncclCommWindowRegister等)按 clique 批量執行。

階段 3:集合通訊 preconnect

📎 src/group.cc:810-870

cpp
if (groupCommHeadMain[ncclGroupTaskTypeCollective] != nullptr) {
  // 按 clique 逐个 prepare + preconnect
  // 然后 ncclTasksRegAndEnqueue
  // 然后 debug check
}

這是核心階段。按 clique 逐個呼叫ncclPrepareTasksAndCollPreconnect,然後asyncJobLaunch執行 preconnect。preconnect 完成後,呼叫ncclTasksRegAndEnqueue把任務註冊到 plan 並生成 kernel 啟動參數。

階段 4:doLaunches

📎 src/group.cc:872-874

cpp
if ((!simInfo) && (groupCommHeadMain[ncclGroupTaskTypeCollective] != nullptr)) {
  NCCLCHECKGOTO(doLaunches(groupCommHeadMain[ncclGroupTaskTypeCollective], ncclGroupTaskTypeCollective), ret, fail);
}

啟動所有 kernel plan。

階段 5:清理

📎 src/group.cc:876-903

cpp
while (!ncclIntruQueueEmpty(asyncJobsMain)) {
  struct ncclAsyncJob* job = ncclIntruQueueDequeue(asyncJobsMain);
  if (!job->destroyFlag && job->comm && !job->comm->config.blocking &&
      groupCommHeadMain[ncclGroupTaskTypeCollective] == nullptr) {
    (void)ncclCommSetAsyncError(job->comm, ret);
  }
  if (job->destructor) job->destructor((void*)job);
}

for (int type = 0; type < ncclGroupTaskTypeNum; ++type) {
  while (groupCommHeadMain[type] != nullptr) {
    struct ncclComm* comm = groupCommHeadMain[type];
    struct ncclComm* next = comm->groupNext[type];
    // Poll for callbacks sent to us from other threads.
    if (comm->reclaimSteps == GROUP_MAX_RECLAIM_STEPS) {
      NCCLCHECKGOTO(ncclCommPollCallbacks(comm, /*waitSome=*/false), ret, fail);
      comm->reclaimSteps = 0;
    } else {
      comm->reclaimSteps++;
    }
    (void)ncclGroupCommLeave(comm, type);
    if (!comm->config.blocking) {
      (void)ncclCommSetAsyncError(comm, ret);
    }
    groupCommHeadMain[type] = next;
  }
}

清理非同步 job,然後遍歷所有 comm 呼叫ncclGroupCommLeave。注意reclaimSteps的計數:每GROUP_MAX_RECLAIM_STEPS(10)次 group 呼叫,輪詢一次 callbacks。這是為了避免每次 group 都輪詢 callbacks 的開銷,同時保證 callbacks 不會無限堆積。

Mermaid 圖:groupLaunchLegacy的資料流

mermaid
flowchart LR
    subgraph input["输入"]
        preconnect["ncclGroupCommPreconnectHead"]
        coll["ncclGroupCommHead[Collective]"]
        sym["ncclGroupCommHead[SymRegister]"]
    end

    subgraph phase1["阶段1: P2P preconnect"]
        p2p_job["ncclPreconnectJob<br/>func=ncclP2PPreconnectFunc"]
        p2p_launch["asyncJobLaunch"]
    end

    subgraph phase2["阶段2: 对称内存注册"]
        sym_job["ncclGroupSymmetricJob<br/>func=ncclCommGroupRegisterSymmetric"]
    end

    subgraph phase3["阶段3: 集合通信 prepare+preconnect"]
        prep["ncclPrepareTasksAndCollPreconnect"]
        coll_job["ncclPreconnectJob<br/>func=ncclCollPreconnectFunc"]
        reg_enq["ncclTasksRegAndEnqueue"]
    end

    subgraph phase4["阶段4: kernel 启动"]
        do_launch["doLaunches<br/>轮次调度"]
        plan["ncclKernelPlan"]
        kernel["ncclLaunchKernel"]
    end

    preconnect --> p2p_job --> p2p_launch
    sym --> sym_job
    coll --> prep --> coll_job --> reg_enq
    reg_enq --> plan --> do_launch --> kernel

---

五、groupLaunchEnqueueRearch:新架構的排程器

直覺模型

groupLaunchEnqueueRearch是 NCCL 正在開發的新排程架構。它把任務準備、排程、啟動分成更細的階段,用非同步 job 佇列管理。目前排程器和啟動器模組「尚未實作」,回退到 legacy 的doLaunches。

📎 src/group.cc:991-996

cpp
// Schedule and launch tasks. Scheduler and launcher module of the enqueue framework
// is not yet implemented and falls back to the legacy launcher: a single phased
// doLaunches over the clique, run here on the user's thread.
if (!simInfo && groupCommHeadMain[ncclGroupTaskTypeRawTask] != nullptr) {
  NCCLCHECKGOTO(doLaunches(groupCommHeadMain[ncclGroupTaskTypeRawTask], ncclGroupTaskTypeRawTask), ret, fail);
}

新架構的執行流程:

1. 管理任務:ncclMgmtTaskJobFunc處理mgmtTaskQueue裡的任務(如 destroy)。

2. 任務準備:ncclTaskPrepareJobFunc呼叫ncclTaskPrepare。

3. 排程和啟動:回退到doLaunches。

新架構用ncclGroupJobLaunch替代asyncJobLaunch,增加了更嚴格的狀態檢查:

📎 src/group.cc:113-116

cpp
} else {
  /* safety check */
  assert(state == ncclGroupJobJoined);
}

legacy 版本用WARN而不是assert,新架構用assert。這說明新架構對狀態機的正確性要求更高。

設計思考

新架構的動機是解耦:legacy 的groupLaunchLegacy把所有階段揉在一個函式裡,難以維護和擴展。新架構把每個階段拆成獨立的 job 類型,透過佇列串聯。但目前排程器和啟動器還沒實作,所以只是「框架先行」。

ncclParamEnqueueRearchEnable()控制走新架構還是 legacy:

📎 src/group.cc:1031-1033

cpp
static ncclResult_t groupLaunch(struct ncclAsyncJob* job_, ncclSimInfo_t* simInfo = NULL) {
  return ncclParamEnqueueRearchEnable() ? groupLaunchEnqueueRearch(job_, simInfo) : groupLaunchLegacy(job_, simInfo);
}

使用者可以透過環境變數NCCL_ENQUEUE_REARCH_ENABLE切換。生產環境建議保持預設(legacy),因為新架構還在開發中。

---

六、非阻塞 group 與非同步錯誤處理

場景驅動的 Step-by-Step Walkthrough

非阻塞 group 的核心是ncclGroupJobComplete和ncclGroupJobAbort:

📎 src/group.cc:1166-1190

cpp
ncclResult_t ncclGroupJobComplete(struct ncclGroupJob* groupJob) {
  ncclResult_t ret = ncclSuccess;
  if (groupJob && groupJob->nonBlockingInit) {
    if (!COMPILER_ATOMIC_EXCHANGE(&groupJob->joined, true, std::memory_order_acq_rel)) {
      ret = ncclAsyncJobComplete(&groupJob->base);
    }
    if (ncclAtomicRefCountDecrement(&groupJob->groupRefCount) == 0) {
      delete groupJob;
    }
  }
  return ret;
}

ncclResult_t ncclGroupJobAbort(struct ncclGroupJob* groupJob) {
  if (groupJob && groupJob->nonBlockingInit) {
    if (!COMPILER_ATOMIC_EXCHANGE(&groupJob->joined, true, std::memory_order_acq_rel)) {
      COMPILER_ATOMIC_STORE(&groupJob->abortFlag, true, std::memory_order_relaxed);
      ncclAsyncJobComplete(&groupJob->base);
    }
    if (ncclAtomicRefCountDecrement(&groupJob->groupRefCount) == 0) {
      delete groupJob;
    }
  }
  return ncclSuccess;
}

關鍵設計:

1. joined原子標誌:用COMPILER_ATOMIC_EXCHANGE保證只有一個執行緒能執行 join 邏輯。如果兩個執行緒同時呼叫ncclGroupJobComplete,只有一個會真正 join,另一個直接跳過。這防止了 double-join。

2. 引用計數:groupRefCount記錄有多少個 comm 關聯到這個 group job。每個 comm 在ncclGroupEndInternal裡增加引用計數:

📎 src/group.cc:1108-1111

cpp
if (job->comm->groupJob == NULL) {
  job->comm->groupJob = groupJob;
  groupJob->groupRefCount++;
}

只有當所有 comm 都呼叫了ncclGroupJobComplete或ncclGroupJobAbort,引用計數減到 0,才刪除 group job。這保證了 group job 的生命週期覆蓋所有關聯的 comm。

3. abort 語意:ncclGroupJobAbort先設定abortFlag,然後 join。工作執行緒在執行過程中檢查abortFlag,如果發現被 abort,提前退出。這是「協作式取消」——不是強制殺死執行緒,而是讓執行緒自己檢查標誌後退出。

生產避坑指南

坑 3:非阻塞 group 的錯誤查詢。非阻塞 group 回傳ncclInProgress,使用者需要透過ncclCommGetAsyncError查詢進度。如果使用者忘記查詢,直接呼叫下一次通訊,可能遇到ncclInProgress錯誤。更嚴重的是,如果 group job 還在執行,使用者呼叫了ncclCommDestroy,會導致 use-after-free。NCCL 透過comm->groupJob指標和引用計數來防止這種情況:ncclCommDestroy會先檢查comm->groupJob,如果有未完成的 group job,會等待或報錯。

坑 4:ncclGroupJobComplete的回傳值。如果 group job 執行失敗,ncclAsyncJobComplete回傳錯誤碼。但ncclGroupJobComplete只在第一次呼叫時回傳這個錯誤碼,後續呼叫回傳ncclSuccess(因為joined已經是 true)。使用者必須在第一次呼叫時檢查回傳值,否則會遺失錯誤資訊。

---

本章小結

這一章我們拆解了 NCCL 從「任務描述」到「kernel 啟動」的完整排程鏈:

1. Group 語意:ncclGroupStart/ncclGroupEnd透過 thread_local 變數攢任務,ncclGroupEnd時統一提交。阻塞模式同步執行,非阻塞模式建立執行緒非同步執行。

2. 任務準備:ncclPrepareTasks確定演算法/協定,ncclPrepareTasksAndCollPreconnect按 clique 逐個 preconnect,避免 split comms 的競態。

3. 輪次排程:doLaunches按 clique 分組,用 barrier 同步 clique 內 comm,每輪啟動一個 kernel plan,直到所有 plan 啟動完畢。

4. 非同步任務:asyncJobLaunch用原子狀態機和忙等待管理非同步 job,支援快速失敗和 abort。

5. 新架構:groupLaunchEnqueueRearch是正在開發的新排程框架,目前回退到 legacy 的doLaunches。

下一章將進入 kernel 啟動的最後一哩路:ncclLaunchKernel如何把ncclKernelPlan變成 GPU 上真正執行的 kernel,以及裝置側如何讀取DevComm中繼資料。

本章思考與自測

Q1: 如果把ncclGroupCommJoin中的ncclMemoryStackPush(&comm->memScoped)去掉,會發生什麼?在什麼場景下會導致記憶體洩漏或資料損壞?

參考解析:ncclMemoryStackPush為 comm 在 group

至此,任務描述已經變成了可執行的啟動計畫:group 語意把多次 API 呼叫合併成一次提交,channel 切分把任務分配到多個執行流,doLaunches 的輪次調度則保證了 kernel 之間的順序與依賴。但計畫終究只是計畫,host 側的任務描述如何變成 GPU 上的一個 grid?下一章我們將深入 ncclLaunchKernel,看參數準備、kernel 變體選擇與 cudaLaunchKernel 呼叫,完成從 host 到 device 的最後一躍。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 08

第 8 章:Kernel 啟動與裝置端執行:從 host 側呼叫到 GPU 執行緒塊起跑

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 8 章 / 共 25 章

上一章我們拆解了任務如何被切分到多個 channel、如何生成 kernel 啟動參數,以及 group 語意下批次提交與依賴排序的機制。現在,啟動計畫已經就緒,但它還只是 host 側的資料結構。本章要回答的核心問題是:ncclKernelPlan如何變成 GPU 上一個真正在跑的 grid?我們將沿著ncclLaunchKernel的呼叫鏈,看參數如何被塞進 kernel args、kernel 變體如何被選中、cuLaunchKernelEx如何被呼叫,以及裝置側ncclKernelMain如何從共享記憶體裡把工作描述讀出來並分發到具體實現。

從 Plan 到 Grid:啟動路徑的全景

在深入細節之前,先建立一個整體心智模型。把ncclKernelPlan想像成一張「施工圖紙」:它記錄了這次要啟動幾個 channel(幾個 block)、每個 block 多少執行緒、要執行哪些 work、用哪個 kernel 函式。而ncclLaunchKernel就是「施工隊進場」的動作——它把圖紙上的資訊翻譯成 CUDA 驅動能理解的CUlaunchConfig,然後呼叫cuLaunchKernelEx把 grid 真正發射到 GPU 上。

如果沒有這一層,host 側的所有調度(上一章的 channel 切分、batch 組織、proxy op 排序)都只是紙上談兵,GPU 上不會有任何 kernel 運行,通訊永遠不會發生。這是端到端主幹的最後一環,也是 host 與 device 的分界線。

整個啟動路徑可以概括為三個階段:

1. 參數準備(finishPlan + uploadWork):把 work 結構體、batch 描述符、kernel args 組織到一塊連續記憶體裡,決定是放在 kernel 參數裡、FIFO 裡還是持久化緩衝區裡。

2. kernel 發射(ncclLaunchKernel):計算 grid/block 維度,組裝 launch attributes(CGA cluster、mem sync domain、launch completion event),呼叫cuLaunchKernelEx。

3. 裝置側入口(ncclKernelMain):每個 block 根據blockIdx.x確定自己的 channelId,從 args 或 FIFO 裡載入 work batch 到共享記憶體,然後透過ncclDevFuncTable分發到具體的演算法/協定實現。

下面這張圖展示了從 plan 到 grid 的完整控制流,包含關鍵的分支判斷:

mermaid
flowchart TD
    plan["ncclKernelPlan<br/>channelMask / workBytes / kernelFn"]
    finish["finishPlan()<br/>决定 workStorageType"]
    check_budget{"sizeof(args)+batchBytes<br/>+workBytes <= workArgsBytes?"}
    args_type["workStorageType = Args<br/>work 直接放 kernel 参数"]
    fifo_type["workStorageType = Fifo/Persistent<br/>work 放外部缓冲区"]
    upload["uploadWork()<br/>拷贝 work 到目标缓冲区"]
    launch["ncclLaunchKernel()<br/>组装 CUlaunchConfig"]
    check_cluster{"compCap >= 90<br/>且 clusterSize > 0?"}
    add_cluster["添加 CLUSTER_DIMENSION<br/>+ SPREAD 调度策略"]
    no_cluster["不添加 cluster 属性"]
    check_event{"userKernelEvent<br/>且 driver >= 12030?"}
    add_event["添加 LAUNCH_COMPLETION_EVENT"]
    no_event["无 completion event"]
    cu_launch["cuLaunchKernelEx()<br/>发射 grid 到 GPU"]

    plan --> finish --> check_budget
    check_budget -->|是| args_type
    check_budget -->|否| fifo_type
    args_type --> upload
    fifo_type --> upload
    upload --> launch --> check_cluster
    check_cluster -->|是| add_cluster
    check_cluster -->|否| no_cluster
    add_cluster --> check_event
    no_cluster --> check_event
    check_event -->|是| add_event
    check_event -->|否| no_event
    add_event --> cu_launch
    no_event --> cu_launch

這張圖錨定了本章的三個核心函式:finishPlan、uploadWork、ncclLaunchKernel。接下來我們逐個拆解。

參數準備:work 結構體如何找到自己的位置

直覺模型

finishPlan的角色類似於快遞分揀中心的「裝箱員」。它面對一堆零散的 work 結構體(每個 collective 或 p2p 操作對應一個),需要決定:這些 work 是塞進 kernel 參數這個「隨身背包」裡,還是放進 FIFO 這個「傳送帶」上,還是放進持久化緩衝區這個「倉庫」裡?

如果這個決策做錯了——比如 work 太大塞不進 kernel 參數卻硬塞——kernel 啟動會直接失敗。如果 work 放錯了位置,裝置側讀到的就是垃圾資料,通訊結果完全錯誤。

資料結構與記憶體佈局

先看ncclDevKernelArgs的結構,它是 host 和 device 之間的「信封」:

📎 src/include/device.h:514-522

c
struct alignas(16) ncclDevKernelArgs {
  struct ncclKernelComm* comm;      // 指向设备侧通信器元数据
  uint64_t channelMask;             // 哪些 channel 有工作
  enum ncclDevWorkStorageType workStorageType;  // work 存在哪里
  uint32_t workMask;                // FIFO 环形缓冲区的掩码
  void* workBuf;                    // work 缓冲区指针
  // struct ncclDevWorkBatch batches[];  // 紧随其后的是 batch 数组
};

這個結構體只有 5 個欄位,但每個欄位都承載著關鍵資訊。channelMask是一個 64 位元遮罩,每一位對應一個 channel,裝置側透過__popcll計算blockIdx.x對應的 channelId。workStorageType決定了裝置側從哪裡讀 work:Args表示 work 就在 kernel 參數裡,Fifo表示在環形緩衝區裡,Persistent表示在持久化緩衝區裡。

ncclDevWorkBatch是 batch 描述符,它告訴裝置側「這個 channel 的 work 在哪裡、有多少個」:

📎 src/include/device.h:400-421

c
struct alignas(16) ncclDevWorkBatch {
  union {
    struct {
      uint32_t nextJump:14, nextExtends:1;
      uint32_t workType:2, funcId : NCCL_DEV_WORK_BATCH_FUNC_ID_BITS, func : NCCL_DEV_WORK_BATCH_FUNC_BITS;
    };
    uint32_t flags;
  };
  uint32_t offsetBase;    // work 在 FIFO 中的起始偏移
  uint64_t offsetBitset;  // 哪些 work 属于这个 channel
};

offsetBitset是一個 64 位元遮罩,每一個位元對應一個 work 結構體。裝置側透過__popc和fns(find n-th set)指令來定位每個 work 的偏移。nextJump和nextExtends用於把多個 batch 串聯起來——當 work 太多裝不下一個 batch 時,會建立「擴展 batch」。

Step-by-Step Walkthrough

現在代入一個具體場景:一次 AllReduce 被切分到 4 個 channel,每個 channel 有 2 個 work 結構體,總共 8 個 work。

第一步:finishPlan決定儲存類型。

📎 src/enqueue/enqueue.cc:245-255

c
if (sizeof(ncclDevKernelArgs) + batchBytes + workBytes <= comm->workArgsBytes) {
  plan->workStorageType = ncclDevWorkStorageTypeArgs;
}
plan->kernelArgsSize = sizeof(struct ncclDevKernelArgs) + batchBytes;
plan->kernelArgsSize += (plan->workStorageType == ncclDevWorkStorageTypeArgs) ? workBytes : 0;
plan->kernelArgsSize = alignUp(plan->kernelArgsSize, 16);
plan->kernelArgs =
  (struct ncclDevKernelArgs*)ncclMemoryStackAlloc(&comm->memScoped, plan->kernelArgsSize, /*align=*/16);
plan->kernelArgs->comm = comm->devComm;
plan->kernelArgs->channelMask = plan->channelMask;
plan->kernelArgs->workStorageType = plan->workStorageType;

這裡的關鍵判斷是:如果sizeof(ncclDevKernelArgs) + batchBytes + workBytes能裝進comm->workArgsBytes(通常是 4KB),就把 work 直接放進 kernel 參數裡。否則,work 會被放到 FIFO 或持久化緩衝區,kernel 參數裡只放 batch 描述符。

〔設計推斷與架構權衡〕

為什麼優先放 kernel 參數? 因為 kernel 參數在 CUDA 驅動裡是透過常量記憶體(constant memory)傳遞的,裝置側讀取時走的是ld.param指令,比從全域記憶體讀取 FIFO 要快得多。對於小訊息(work 總量小),這能顯著降低延遲。

第二步:把 batch 按 channel 輪流放入 kernel args。

📎 src/enqueue/enqueue.cc:257-280

c
uint64_t hasBatchMask = plan->channelMask;
struct ncclDevWorkBatch* batchPrev[MAXCHANNELS] = {};
struct ncclDevWorkBatch* batchZero = (struct ncclDevWorkBatch*)(plan->kernelArgs + 1);
int batchIx = 0;
while (hasBatchMask != 0) {
  uint64_t tmpMask = hasBatchMask;
  do {
    int c = popFirstOneBit(&tmpMask);
    if (!ncclIntruQueueEmpty(&wipChannels[c].workBatchQueue)) {
      struct ncclWorkBatchList* batchNode = ncclIntruQueueDequeue(&wipChannels[c].workBatchQueue);
      if (batchPrev[c] != nullptr) {
        batchPrev[c]->nextJump = int(&batchZero[batchIx] - batchPrev[c]);
      }
      batchPrev[c] = &batchZero[batchIx];
      batchZero[batchIx++] = batchNode->batch;
    }
    if (ncclIntruQueueEmpty(&wipChannels[c].workBatchQueue)) {
      hasBatchMask ^= 1ull << c;
    }
  } while (tmpMask != 0);
}

這段程式碼的邏輯是「輪詢」:每一輪從每個還有 batch 的 channel 裡取一個 batch,按 channel 編號升序放入batchZero陣列。這樣做的目的是保證「每個 channel 的第一個 batch 位於batchZero[blockIdx.x]」——裝置側每個 block 透過blockIdx.x直接索引到自己的第一個 batch,不需要搜尋。

nextJump欄位記錄了同一個 channel 的下一個 batch 相對於當前 batch 的偏移。裝置側透過batchIx += batch.nextJump就能跳到下一個 batch,形成一個鏈結串列。

第三步:uploadWork把 work 拷貝到目標緩衝區。

📎 src/enqueue/enqueue.cc:1365-1430

c
static ncclResult_t uploadWork(struct ncclComm* comm, struct ncclKernelPlan* plan) {
  if (plan->isSymColl || plan->isCeColl || plan->isRma) return ncclSuccess;
  size_t workBytes = plan->workBytes;
  size_t batchBytes = plan->nWorkBatches * sizeof(struct ncclDevWorkBatch);
  void* fifoBufHost;
  uint32_t fifoCursor, fifoMask;
  switch (plan->workStorageType) {
  case ncclDevWorkStorageTypeArgs:
    plan->kernelArgs->workBuf = nullptr;
    fifoBufHost = (void*)plan->kernelArgs;
    fifoCursor = sizeof(ncclDevKernelArgs) + batchBytes;
    fifoMask = ~0u;
    break;
  case ncclDevWorkStorageTypeFifo:
    fifoBufHost = comm->workFifoBuf;
    fifoCursor = comm->workFifoProduced;
    fifoMask = comm->workFifoBytes - 1;
    NCCLCHECK(waitWorkFifoAvailable(comm, fifoCursor + workBytes));
    plan->kernelArgs->workBuf = comm->workFifoBufDev;
    break;
  // ...
  }
  plan->kernelArgs->workMask = fifoMask;
  // 修正 batch 的 offsetBase
  struct ncclDevWorkBatch* batchZero = (struct ncclDevWorkBatch*)(plan->kernelArgs + 1);
  for (int b = 0; b < plan->nWorkBatches; b++) {
    batchZero[b].offsetBase += fifoCursor;
  }
  // 拷贝 work 结构体
  struct ncclWorkList* workNode = ncclIntruQueueHead(&plan->workQueue);
  while (workNode != nullptr) {
    char* dst = (char*)fifoBufHost;
    char* src = (char*)(workNode + 1);
    for (int n = workNode->size; n != 0; n -= 16) {
      memcpy(COMPILER_ASSUME_ALIGNED(dst + (fifoCursor & fifoMask), 16), COMPILER_ASSUME_ALIGNED(src, 16), 16);
      fifoCursor += 16;
      src += 16;
    }
    workNode = workNode->next;
  }
  // ...
}

這裡有幾個關鍵點:

1. fifoCursor的語意:對於Args類型,它是相對於kernelArgs起始位址的偏移;對於Fifo類型,它是相對於 FIFO 基底位址的偏移;對於Persistent類型,它從 0 開始。

2. offsetBase的修正:finishPlan裡 batch 的offsetBase是相對於 plan 的 work 起始位置的(從 0 開始)。uploadWork需要把它轉換成相對於實際儲存位置的偏移。對於Args類型,加上sizeof(ncclDevKernelArgs) + batchBytes;對於Fifo類型,加上comm->workFifoProduced。

3. 16 位元組對齊拷貝:work 結構體都是 16 位元組對齊的(alignas(16)),所以拷貝時按 16 位元組為單位。COMPILER_ASSUME_ALIGNED告訴編譯器這個位址是 16 位元組對齊的,讓編譯器生成更高效的向量化指令。

4. FIFO 等待:對於Fifo類型,waitWorkFifoAvailable會自旋等待 FIFO 有足夠空間。這個等待會檢查comm->abortFlag,避免在 abort 時死鎖。

設計思考與生產踩坑

〔設計推斷與架構權衡〕

為什麼要有三種儲存類型?這是空間和延遲的權衡:

  • Args:最快(常量記憶體),但容量有限(4KB)。適合小訊息、少量 work。
  • Fifo:容量大(環形緩衝區),但裝置側讀取要走全域記憶體。適合中等訊息。
  • Persistent:用於 CUDA Graph 捕獲場景。因為 graph 捕獲時不能做cudaMemcpy,所以需要預先分配持久化緩衝區,把 work 拷貝進去,然後讓 kernel 從那裡讀。

踩坑點 1:FIFO 溢出導致死鎖。如果waitWorkFifoAvailable沒有檢查abortFlag,當 FIFO 滿且消費者(GPU kernel)因為某種原因停止消費時,host 會永遠自旋。原始碼裡📎 src/enqueue/enqueue.cc:1333-1349明確檢查了 abort flag:

c
if (COMPILER_ATOMIC_LOAD(comm->abortFlag, std::memory_order_acquire)) {
  return ncclInternalError;
}

踩坑點 2:offsetBitset溢出。 offsetBitset是 64 位元的,最多支援 64 個 work 在一個 batch 裡。如果超過 64 個,1ull << (offset / workSize)會溢出。原始碼裡透過NCCL_MAX_DEV_WORK_BATCH_BYTES限制了 batch 的大小(1024 位元組),而最小的 work 結構體是ncclDevWorkColl(約 80 位元組),所以最多 12 個 work,不會溢出。

踩坑點 3:Persistent 模式下的記憶體洩漏。在uploadWork的Persistent分支裡,fifoBufHost是透過ncclOsAlignedAlloc分配的,需要在uploadWork_cleanup_fn裡釋放。如果cudaMemcpyAsync失敗,fail標籤會檢查cleanup是否為 null,如果為 null 就直接釋放fifoBufHost。這個錯誤恢復鏈在📎 src/enqueue/enqueue.cc:1483-1485可以看到。

Kernel 發射:從 CUlaunchConfig 到 cuLaunchKernelEx

直覺模型

ncclLaunchKernel的角色類似於「火箭發射控制台」。它接收一個已經裝好燃料(work 資料)的 plan,計算出火箭的飛行參數(grid/block 維度),設定好各種發射選項(cluster、mem sync domain、completion event),然後按下發射按鈕(cuLaunchKernelEx)。

如果這個環節出錯——比如 grid 維度算錯了——GPU 上會啟動錯誤數量的 block,導致部分 channel 的工作永遠不會被執行,通訊掛起。

資料結構與記憶體佈局

CUlaunchConfig是 CUDA 驅動 API 的啟動配置結構體,NCCL 在堆疊上建構它:

📎 src/enqueue/enqueue.cc:1916-1917

c
CUlaunchConfig launchConfig = {0};
CUlaunchAttribute launchAttrs[6] = {};
int attrs = 0;

launchAttrs是一個最多 6 個元素的陣列,每個元素是一個CUlaunchAttribute。NCCL 根據硬體能力和驅動版本,有條件地加入不同的屬性:

  • CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION:CGA cluster 維度(sm90+)
  • CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE:cluster 排程策略
  • CU_LAUNCH_ATTRIBUTE_MEM_SYNC_DOMAIN:記憶體同步域(CUDA 12.0+)
  • CU_LAUNCH_ATTRIBUTE_LAUNCH_COMPLETION_EVENT:啟動完成事件(CUDA 12.3+)
  • CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION:程式化流序列化(sym kernel)
  • CU_LAUNCH_ATTRIBUTE_NVLINK_UTIL_CENTRIC_SCHEDULING:NVLink 利用率中心排程(CUDA 13.0+)

Step-by-Step Walkthrough

第一步:計算 grid 和 block 維度。

📎 src/enqueue/enqueue.cc:1889-1893

c
int nChannels = countOneBits(plan->channelMask);
void* sym = plan->kernelFn;
dim3 grid = {(unsigned)nChannels, 1, 1};
dim3 block = {(unsigned)plan->threadPerBlock, 1, 1};
int smem = plan->isSymColl ? plan->kernelDynSmem : ncclShmemDynamicSize(comm->cudaArch);

nChannels是channelMask中置位的個數,也就是這個 plan 要啟動多少個 block。每個 block 負責一個 channel。threadPerBlock是在scheduleCollTasksToPlan裡透過plan->threadPerBlock = std::max(plan->threadPerBlock, task->nWarps * WARP_SIZE)計算出來的,取所有 task 中最大的nWarps * 32。

smem是動態共享記憶體大小。對於普通 kernel,它是ncclShmemDynamicSize(comm->cudaArch),這是一個編譯期常數,取決於架構(sm70+ 是ncclShmemScratchWarpSize * (NCCL_MAX_NTHREADS / WARP_SIZE))。對於 sym kernel,它是plan->kernelDynSmem,因為 sym kernel 的共享記憶體需求可能不同。

第二步:組裝 kernel 參數。

📎 src/enqueue/enqueue.cc:1902-1903

c
void* extra[] = {CU_LAUNCH_PARAM_BUFFER_POINTER, plan->kernelArgs, CU_LAUNCH_PARAM_BUFFER_SIZE, &plan->kernelArgsSize,
                 CU_LAUNCH_PARAM_END};

這是 CUDA 驅動 API 的一種參數傳遞方式:CU_LAUNCH_PARAM_BUFFER_POINTER告訴驅動「參數不是一個個傳的,而是一個連續的記憶體塊」,CU_LAUNCH_PARAM_BUFFER_SIZE告訴驅動這個塊的大小。這樣做的好處是 NCCL 可以把ncclDevKernelArgs和後面的 batch 陣列一次性傳進去,不需要逐個參數打包。

第三步:加入 launch attributes。

📎 src/enqueue/enqueue.cc:1929-1936

c
if (clusterSize) {
  if (grid.x % clusterSize) clusterSize = 1;
  launchAttrs[attrs].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
  launchAttrs[attrs++].value.clusterDim = {clusterSize, 1, 1};
  launchAttrs[attrs].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE;
  launchAttrs[attrs++].value.clusterSchedulingPolicyPreference = CU_CLUSTER_SCHEDULING_POLICY_SPREAD;
}

CGA(Cooperative Group Array)是 sm90 引入的硬體特性,允許把多個 block 組成一個 cluster,cluster 內的 block 可以保證同時排程到一組 SM 上,並且可以互相存取共享記憶體。NCCL 用這個特性來實現 NVLS 等需要跨 block 同步的演算法。

注意if (grid.x % clusterSize) clusterSize = 1;這個保護:cluster 維度必須能整除 grid 維度,否則驅動會報錯。如果grid.x不能被clusterSize整除,就退化為不使用 cluster。

第四步:加入 launch completion event。

📎 src/enqueue/enqueue.cc:1944-1964

c
#if CUDART_VERSION >= 12030
enum ncclImplicitOrder implicitOrder;
NCCLCHECKGOTO(getImplicitOrder(&implicitOrder, comm, plan->persistent, driverVersion), ret, do_return);
if (implicitOrder == ncclImplicitOrderLaunch) {
  launchAttrs[attrs].id = CU_LAUNCH_ATTRIBUTE_LAUNCH_COMPLETION_EVENT;
  launchAttrs[attrs].value.launchCompletionEvent.event = comm->sharedRes->launchEvent;
  launchAttrs[attrs].value.launchCompletionEvent.flags = 0;
  attrs++;
  if (userKernelEvent) {
    NCCLCHECKGOTO(ncclUncapturedStreamPoolAcquire(&comm->sharedRes->uncapturedStreamPool, &relayStream), ret, do_return);
    relayUserLaunchCompletionEvent = true;
    userKernelEventArmed = true;
  }
} else if (userKernelEvent && driverVersion >= 12030) {
  launchAttrs[attrs].id = CU_LAUNCH_ATTRIBUTE_LAUNCH_COMPLETION_EVENT;
  launchAttrs[attrs].value.launchCompletionEvent.event = plan->launchCompletionEvent;
  launchAttrs[attrs].value.launchCompletionEvent.flags = 0;
  attrs++;
  userKernelEventArmed = true;
}
#endif

CU_LAUNCH_ATTRIBUTE_LAUNCH_COMPLETION_EVENT是 CUDA 12.3 引入的特性:驅動會在 kernel 真正開始執行時(而不是在 host 側呼叫返回時)記錄一個事件。這對於實現「隱式順序」(implicit order)至關重要——NCCL 需要保證多個 kernel 按順序執行,但又不希望 host 側阻塞等待。

getImplicitOrder的邏輯是:如果使用者設定了launchOrderImplicit,並且驅動版本足夠新,就使用ncclImplicitOrderLaunch(用 launch event 排序);否則使用ncclImplicitOrderSerial(用 completion event 排序,即串行執行)。

第五步:呼叫cuLaunchKernelEx。

📎 src/enqueue/enqueue.cc:1978-1996

c
launchConfig.gridDimX = grid.x;
launchConfig.gridDimY = grid.y;
launchConfig.gridDimZ = grid.z;
launchConfig.blockDimX = block.x;
launchConfig.blockDimY = block.y;
launchConfig.blockDimZ = block.z;
launchConfig.sharedMemBytes = smem;
launchConfig.attrs = launchAttrs;
launchConfig.numAttrs = attrs;
launchConfig.hStream = launchStream;
if (userKernelEvent && !userKernelEventArmed) {
  WARN("CUDA launch-completion events require CUDA 12.3 or newer; recording the user event before launch");
  CUDACHECKGOTO(cudaEventRecord(plan->launchCompletionEvent, launchStream), ret, do_return);
}
CUCHECKGOTO(cuLaunchKernelEx(&launchConfig, fn, nullptr, extra), ret, do_return);
if (relayUserLaunchCompletionEvent) {
  CUDACHECKGOTO(cudaStreamWaitEvent(relayStream, comm->sharedRes->launchEvent, 0), ret, do_return);
  CUDACHECKGOTO(cudaEventRecord(plan->launchCompletionEvent, relayStream), ret, do_return);
}

cuLaunchKernelEx是 CUDA 12.0 引入的新 API,支援 launch attributes。對於舊驅動(< 11.8),NCCL 會退回到cuLaunchKernel:

📎 src/enqueue/enqueue.cc:1998-2007

c
} else {
  // Standard kernel launch
  if (userKernelEvent) {
    WARN("CUDA launch-completion events require CUDA 12.3 or newer; recording the user event before launch");
    CUDACHECKGOTO(cudaEventRecord(plan->launchCompletionEvent, launchStream), ret, do_return);
  }
  CUCHECKGOTO(cuLaunchKernel(fn, grid.x, grid.y, grid.z, block.x, block.y, block.z, smem, launchStream, nullptr,
                             extra),
              ret, do_return);
}

並發控制與硬體互動

Launch completion event 的 relay 機制。當使用ncclImplicitOrderLaunch且使用者提供了launchCompletionEvent時,NCCL 不能直接把使用者的 event 傳給驅動,因為驅動只支援一個 launch completion event。NCCL 的做法是:

1. 把comm->sharedRes->launchEvent傳給驅動。

2. 在relayStream上等待launchEvent。

3. 在relayStream上記錄使用者的 event。

這樣使用者的 event 會在 kernel 真正開始執行後觸發,而不是在 host 側呼叫返回時觸發。

Mem Sync Domain。 📎 src/enqueue/enqueue.cc:1938-1942在 sm90+ 上,NCCL 設定CU_LAUNCH_ATTRIBUTE_MEM_SYNC_DOMAIN為cudaLaunchMemSyncDomainRemote。這是 Hopper 架構引入的記憶體同步域機制,用於隔離不同 kernel 的記憶體屏障,減少不必要的同步開銷。

生產踩坑指南

踩坑點 1:cluster 維度不整除導致啟動失敗。如果grid.x不能被clusterSize整除,驅動會返回CUDA_ERROR_INVALID_VALUE。原始碼裡透過if (grid.x % clusterSize) clusterSize = 1;做了保護,但這也意味著 cluster 特性被靜默停用了。如果使用者期望 cluster 帶來的效能提升,需要檢查cgaClusterSize和nChannels的關係。

踩坑點 2:驅動版本不滿足導致 kernel 不可用。 ncclInitKernelsForDevice會在初始化時檢查每個 kernel 的驅動要求:

📎 src/enqueue/enqueue.cc:71-76

c
for (int k = 0; k < kcount; k++) {
  if (kptrs[k] != nullptr && driverVersion < krequires[k]) {
    INFO(NCCL_INIT, "Skipping %skernel %d which requires driver %d", sym ? "symmetric " : "", k, krequires[k]);
    kptrs[k] = nullptr;
    if (kptrsProfile != nullptr) kptrsProfile[k] = nullptr;
  }

如果驅動版本不夠,kernel 指標會被置為 null。後續如果排程器選中了這個 kernel,cuLaunchKernelEx會失敗。NCCL 的 tuner 應該會避免選擇不可用的 kernel,但如果使用者強制指定了演算法(NCCL_ALGO),可能會觸發這個問題。

踩坑點 3:launchCompletionEvent在舊驅動上的行為。如果驅動版本 < 12.3,NCCL 會在 kernel 啟動前記錄 event,這意味著 event 會在 kernel 開始執行前就觸發,而不是在 kernel 真正開始執行時。這可能導致使用者程式的時序假設失效。

裝置側入口:從 blockIdx 到具體實現

直覺模型

ncclKernelMain是 GPU 上每個 block 的「入口大廳」。當一個 block 被調度到 SM 上開始執行時,它首先進入這個大廳,完成三件事:確定自己的身份(我是哪個 channel)、領取自己的任務(載入 work batch)、然後去對應的窗口辦事(呼叫具體的演算法實現)。

如果沒有這個入口,每個 kernel 變體都需要自己處理「我是誰、我要幹什麼」的問題,程式碼會大量重複。ncclKernelMain透過模板參數SpecializedFnId和SpecializedRunWorkBatch實現了「通用入口 + 特化執行」的模式。

資料結構與記憶體佈局

裝置側的共享記憶體佈局是理解ncclKernelMain的關鍵。ncclShmemData是所有 block 共享的「工作台」:

📎 src/device/common.h:48-72

c
struct ncclShmemData {
  struct ncclDevKernelArgs args;
  int channelId;
  int aborted;
  alignas(16) struct ncclKernelComm comm;
  alignas(16) struct ncclDevChannel channel;

  int batchIx, nextBatchIx;
  enum ncclDevWorkType workType;
  uint8_t directMode;
  uint16_t funcId;
  int nWorks;
  int workSize;
  uint64_t workCounter;
  bool profilerEnabled;
  uint8_t func;
  struct ncclShmemGroup groups[NCCL_MAX_GROUPS];

  alignas(16) char workStorage[ncclMaxDevWorkBatchBytes()];

  alignas(16) union {
    unpackShmem unpack;
  } devicePlugin;
};

這個結構體的佈局經過精心設計:

  • args放在最前面,因為它是從 kernel 參數拷貝過來的,需要 16 位元組對齊。
  • comm和channel也是 16 位元組對齊的,因為它們是透過copyToShmem16用向量化指令拷貝的。
  • workStorage是 work 結構體的臨時存放區,大小是ncclMaxDevWorkBatchBytes()(sm90+ 是 16KB)。
  • groups陣列用於儲存每個 group 的連接資訊,NCCL_MAX_GROUPS是 16。

Step-by-Step Walkthrough

第一步:拷貝 kernel args 到共享記憶體。

📎 src/device/common.h:426-428

c
if (tid < sizeof(ncclDevKernelArgs) / sizeof(uint32_t)) {
  ((uint32_t*)&ncclShmem.args)[tid] = ((uint32_t*)args)[tid];
}

這裡用前sizeof(ncclDevKernelArgs) / 4個執行緒,每個執行緒拷貝一個 32 位元字。為什麼要拷貝到共享記憶體?因為 kernel 參數在常量記憶體裡,存取速度雖然快,但每個執行緒都要存取時會有廣播開銷。拷貝到共享記憶體後,所有執行緒存取的是同一塊共享記憶體,效率更高。

第二步:確定 channelId。

📎 src/device/common.h:430-437

c
if (tid < MAXCHANNELS && (args->channelMask & (1ull << tid))) {
  int n = __popcll(args->channelMask & ((1ull << tid) - 1));
  if (blockIdx.x == n) ncclShmem.channelId = tid;
}
__syncthreads();

這段程式碼的邏輯是:對於每個置位的 channel(args->channelMask & (1ull << tid)),計算它前面有多少個置位的 channel(__popcll),如果這個數量等於blockIdx.x,那麼當前 block 就負責這個 channel。

舉個例子:channelMask = 0b1011(channel 0、1、3 有工作)。blockIdx.x = 0的 block 負責 channel 0(前面有 0 個置位),blockIdx.x = 1的 block 負責 channel 1(前面有 1 個置位),blockIdx.x = 2的 block 負責 channel 3(前面有 2 個置位)。

第三步:載入 comm 和 channel 到共享記憶體。

📎 src/device/common.h:446-478

c
switch (tid / WARP_SIZE) {
case 0:
  {
    void* dst = &ncclShmem.comm;
    void* src = ncclShmem.args.comm;
    int bytes = sizeof(ncclKernelComm);
    static_assert(sizeof(ncclKernelComm) <= 16 * WARP_SIZE,
                  "ncclKernelComm cannot be loaded by a single warp in one insn.");
    copyToShmem16(tid, dst, src, bytes);
  }
  break;
case 1:
  {
    void* dst = &ncclShmem.channel;
    void* src = &((ncclKernelCommAndChannels*)ncclShmem.args.comm)->channels[ncclShmem.channelId];
    int bytes = sizeof(ncclDevChannel);
    static_assert(sizeof(ncclDevChannel) <= 16 * WARP_SIZE,
                  "ncclDevChannel cannot be loaded by a single warp in one insn.");
    copyToShmem16(tid - WARP_SIZE, dst, src, bytes);
  }
  break;
default:
  {
    int subtid = tid - 2 * WARP_SIZE;
    int subtn = tn - 2 * WARP_SIZE;
    loadWorkBatchToShmem(subtid, subtn, args, /*batchIx=*/blockIdx.x);
  }
  break;
}
__syncthreads();

這裡把執行緒分成三組:

  • 第 0 個 warp:載入ncclKernelComm(通訊器元資料)到共享記憶體。
  • 第 1 個 warp:載入當前 channel 的ncclDevChannel(channel 元資料)到共享記憶體。
  • 其餘 warp:載入 work batch 到共享記憶體。

copyToShmem16是一個用內聯 PTX 實現的 16 位元組拷貝函式:

📎 src/device/common.h:131-139

c
inline __device__ void copyToShmem16(int tid, void* dst, void const* src, int bytes) {
  int offset = 16 * tid;
  if (offset < bytes) {
    uint64_t a = 0, b = 0;
    asm volatile("ld.v2.u64 {%0,%1},[%2];" : "=l"(a), "=l"(b) : "l"((char const*)src + offset) : "memory");
    uint32_t udst = (uint32_t)__cvta_generic_to_shared(dst);
    asm volatile("st.shared.v2.u64 [%0],{%1,%2};" ::"r"(udst + offset), "l"(a), "l"(b) : "memory");
  }
}

它用ld.v2.u64從全域記憶體載入 16 位元組,用st.shared.v2.u64儲存到共享記憶體。__cvta_generic_to_shared把通用位址轉換成共享記憶體位址(共享記憶體位址空間是 32 位的)。

第四步:載入 work batch。

loadWorkBatchToShmem是最複雜的部分。它的任務是把 batch 描述符指向的 work 結構體從全域記憶體(或 kernel 參數)拷貝到共享記憶體的workStorage裡。

📎 src/device/common.h:142-260

c
__device__ __forceinline__ void loadWorkBatchToShmem(int tid, int tn, struct ncclDevKernelArgs const* args,
                                                     int batchIx) {
  int lane = tid % WARP_SIZE;
  int workCursor = 0;
  while (true) {
    struct ncclDevWorkBatch batch = ((struct ncclDevWorkBatch*)(args + 1))[batchIx];

    uint8_t* fnsOfBitset = (uint8_t*)ncclScratchForWarp(threadIdx.x / WARP_SIZE);
    __syncwarp();
    if (uint32_t(batch.offsetBitset) & (1u << lane)) {
      int nWorksBelow = __popc(uint32_t(batch.offsetBitset) & ((1u << lane) - 1));
      fnsOfBitset[nWorksBelow] = lane;
    }
    int nWorksLow32 = __popc(uint32_t(batch.offsetBitset));
    if (uint32_t(batch.offsetBitset >> 32) & (1u << lane)) {
      int nWorksBelow = nWorksLow32;
      nWorksBelow += __popc(uint32_t(batch.offsetBitset >> 32) & ((1u << lane) - 1));
      fnsOfBitset[nWorksBelow] = 32 + lane;
    }
    int nWorks = nWorksLow32 + __popc(uint32_t(batch.offsetBitset >> 32));
    __syncwarp();
    // ...
  }
}

這段程式碼的核心是計算fnsOfBitset:對於offsetBitset中第 n 個置位,它的位索引是多少。PTX 有fns指令可以做這個,但它展開成很多 SASS 指令。NCCL 的做法是用共享記憶體:每個 lane 檢查自己的位是否置位,如果是,計算它前面有多少個置位,然後把自己的 lane 編號寫到fnsOfBitset[nWorksBelow]。

接下來是實際的拷貝:

📎 src/device/common.h:209-241

c
if (tid < nPacks) {
  int srcWork = fnsOfBitset[dstWork];
  ulonglong2 tmp;
  if (ncclShmem.args.workStorageType == ncclDevWorkStorageTypeArgs) {
    char* src = (char*)args + (batch.offsetBase + srcWork * workSize + packInWork * 16);
    tmp = *(ulonglong2*)src; // becomes ld.param.v2.u64
  } else {
    char* src = (char*)ncclShmem.args.workBuf +
                ((batch.offsetBase + srcWork * workSize + packInWork * 16) & ncclShmem.args.workMask);
    tmp = *(ulonglong2*)src; // becomes ld.v2.u64
  }
  char* dst = ncclShmem.workStorage;
  dst += (workCursor + dstWork) * workSize + packInWork * 16;
  *(ulonglong2*)dst = tmp;
}

這裡有一個關鍵的優化:對於Args型別,原始碼直接寫(char*)args + offset,編譯器會識別出這是從 kernel 參數讀取,生成ld.param.v2.u64指令。對於Fifo型別,原始碼寫(char*)ncclShmem.args.workBuf + (offset & workMask),編譯器生成ld.v2.u64指令。

註解裡特別強調了不能把這兩種情況合併:

📎 src/device/common.h:212-229

c
// The loads done in these two cases must be kept separate since we are
// relying on the compiler to use "ld.param" in the first one. The parameter
// space is not generically addressable, so any attempt to load through
// a pointer that *might* be parameter space backed will cause the
// compiler to spill the parameter struct (4K!) to each thread's local space
// before creating a pointer (to the spill) and decimate perf.

如果編譯器不能確定指標指向的是參數空間還是全域空間,它會把整個參數結構體(4KB)溢出到每個執行緒的本地記憶體,效能會急劇下降。

第五步:執行 work。

📎 src/device/common.h:481-497

c
while (ncclShmem.aborted == 0) {
  profiler(START);
  if (0 <= SpecializedFnId && ncclShmem.funcId == (unsigned)SpecializedFnId) {
    SpecializedRunWorkBatch().run();
  } else {
    ncclDevFuncTable[ncclShmem.funcId]();
  }

  if (ncclShmem.nextBatchIx == -1) break;
  int batchIx = ncclShmem.nextBatchIx;
  __syncthreads();
  profiler(STOP);
  if (ncclShmem.comm.progressCounters != nullptr) __syncthreads();
  loadWorkBatchToShmem(tid, tn, args, batchIx);
  __syncthreads();
}

這裡有一個重要的優化:如果SpecializedFnId匹配當前 batch 的funcId,直接呼叫SpecializedRunWorkBatch().run(),這是一個編譯期特化的函式,沒有函式指標呼叫的開銷。否則,透過ncclDevFuncTable[ncclShmem.funcId]()間接呼叫。

ncclDevFuncTable是一個裝置側的函式指標陣列,由generate.py生成:

📎 src/device/generate.py:261-270

python
out("__device__ ncclDevFuncPtr_t const ncclDevFuncTable[] = {\n")
index = 0
for fn in primary_funcs:
  sym = paste("_", "ncclDevFunc", *fn)
  cudart, arch = required_cuda(*fn)
  if (cudart, arch) != (0, 0):
    out("#if CUDART_VERSION >= %d && __CUDA_ARCH__ >= %d\n" % (cudart ,arch))
  out("/*%4d*/ %s,\n" % (index, sym))
  if (cudart, arch) != (0, 0):
    out("#else\n" "/*%4d*/ nullptr,\n" "#endif\n" % index)
  index += 1
out("nullptr};\n")

設計思考與生產踩坑

為什麼用__grid_constant__? 📎 src/device/common.h:19-24

c
#if __CUDA_ARCH__ >= 700
// __grid_constant__ appears to break cuda-gdb
#define NCCL_GRID_CONSTANT __grid_constant__
#else
#define NCCL_GRID_CONSTANT
#endif

__grid_constant__告訴編譯器這個參數是唯讀的,可以放在常量記憶體裡。這樣裝置側讀取時走ld.param指令,比從全域記憶體讀取快。註解裡提到它會破壞 cuda-gdb,所以只在 sm70+ 上啟用。

踩坑點 1:workStorage溢出。 workStorage的大小是ncclMaxDevWorkBatchBytes(),sm90+ 是 16KB。如果nWorks * workSize超過這個值,會寫越界。原始碼裡透過NCCL_MAX_DEV_WORK_BATCH_BYTES在 host 側限制了 batch 的大小,但裝置側沒有額外的檢查。如果 host 側的約束被繞過(比如透過修改環境變數),會導致共享記憶體越界。

踩坑點 2:__syncthreads()的缺失導致資料競爭。在loadWorkBatchToShmem之後,必須有一個__syncthreads()才能讓所有執行緒看到完整的workStorage。原始碼裡在📎 src/device/common.h:479有__syncthreads(); // publish ncclShmem。如果這個同步被去掉,某些執行緒可能會在workStorage還沒寫完時就開始讀取,導致讀到垃圾資料。

踩坑點 3:abort 檢查的時機。 while (ncclShmem.aborted == 0)只在每個 batch 開始時檢查 abort。如果某個 batch 執行時間很長,abort 信號可能要等很久才能生效。這是設計上的權衡:更頻繁的檢查會增加開銷,但回應更快。

Kernel 變體選擇:generate.py 如何生成 kernel 列表

直覺模型

generate.py的角色類似於「汽車工廠的生產線規劃師」。它面對一個巨大的組合空間(7 種集合操作 × 5 種歸約操作 × 12 種資料類型 × 7 種演算法 × 3 種協定),需要決定:哪些組合需要生成專門的 kernel?哪些可以共用一個通用 kernel?

如果每個組合都生成一個 kernel,編譯時間和二進位大小會爆炸。如果只生成一個通用 kernel,執行時會因為函式指標呼叫和分支判斷而變慢。generate.py的解決方案是「代表性 kernel」:為每個等價類生成一個 kernel,執行時透過函式指標表分發。

資料結構與記憶體佈局

generate.py生成三個關鍵檔案:

1. device_table.cu:裝置側的ncclDevFuncTable,把 funcId 映射到具體的裝置函式。

2. host_table.cc:host 側的ncclDevKernelList、ncclDevKernelForFunc、ncclDevFuncRowToId等表。

3. 各個<coll>_<op>_<ty>.cu:具體的 kernel 實作。

Step-by-Step Walkthrough

第一步:列舉所有函式行。

📎 src/device/generate.py:186-199

python
def enumerate_func_rows():
  yield ("SendRecv", None, None, None, None)
  for coll in ("AllGather", "Broadcast", "AllGatherV"):
    algos = algos_of_coll[coll]
    for algo in algos:
      for proto in all_protos:
        yield (coll, None, None, algo, proto)
  for coll in ("AllReduce", "Reduce", "ReduceScatter"):
    algos = algos_of_coll[coll]
    for redop in all_redops:
      for ty in all_tys:
        for algo in algos:
          for proto in all_protos:
            yield (coll, redop, ty, algo, proto)

這個列舉順序必須和ncclDevFuncId()的計算公式匹配:

📎 src/include/device.h:646-706

c
inline int ncclDevFuncId(int coll, int devRedOp, int type, int algo, int proto) {
  constexpr int NumTypes = ncclNumTypes;
  int row;
  do {
    row = 0; // ncclDevFuncIndex_P2p
    if (coll == ncclFuncSendRecv) break;
    row += 1;
    // ...
  } while (false);
  return ncclDevFuncRowToId[row];
}

ncclDevFuncId計算出的是「行號」,然後透過ncclDevFuncRowToId映射到「主函式 ID」。這個映射的原因是:很多行可能映射到同一個主函式(比如所有AllReduce Sum i32的行都映射到AllReduce Sum u32的主函式)。

第二步:計算主函式和 kernel 函式。

📎 src/device/generate.py:211-225

python
func_rows = [validate(*fn) for fn in enumerate_func_rows()]
primary_funcs = sorted(set(equivalent_primary(*fn) for fn in func_rows if fn is not None))
primary_to_index = {fn: i for (i,fn) in zip(range(len(primary_funcs)), primary_funcs)}
kernel_funcs = sorted(set(best_kernel(*fn) for fn in primary_funcs))

equivalent_primary把有符號整數映射到無符號整數(因為加法/乘法對兩者是一樣的):

📎 src/device/generate.py:158-166

python
def equivalent_primary(coll, redop, ty, algo, proto):
  if coll in ("AllReduce", "Reduce", "ReduceScatter"):
    if redop in ("Sum","Prod","PreMulSum","SumPostDiv") and ty[0]=="i":
      return (coll, redop, "u"+ty[1:], algo, proto)
    if redop=="MinMax" and ty[0]=="i" and ("NVLS" not in algo):
      return (coll, redop, "u"+ty[1:], algo, proto)
  return (coll, redop, ty, algo, proto)

best_kernel把多個主函式映射到同一個 kernel(比如所有AllGather的演算法都映射到AllGather RING LL):

📎 src/device/generate.py:171-183

python
def best_kernel(coll, redop, ty, algo, proto):
  def best(coll, redop, ty, algo, proto):
    if coll=="Nop": return ("Generic", None, None, None, None)
    if coll=="SendRecv": return ("SendRecv", None, None, None, None)
    if exact_kernel_names: return (coll, redop, ty, algo, proto)
    if coll in ("AllGather","Broadcast","AllGatherV"): return (coll, None, None, "RING", "LL")
    return (coll, "Sum", ty, ("TREE" if algo=="TREE" else "RING"), "LL")
  kfn = equivalent_primary(*best(coll, redop, ty, algo, proto))
  if not func_filter(*kfn): return ("Generic", None, None, None, None)
  return kfn

第三步:生成 kernel 定義。

📎 src/device/generate.py:458-480

python
(_, kfns) = name_to_kernels.get(name) or (None, [])
for kfn in kfns:
  (coll, redop, ty, algo, proto) = kfn
  sym = kernel_suffix(kfn)
  fn_id = primary_to_index[kfn]
  cudart, arch = required_cuda(*kfn)
  s = "DEFINE_ncclDevKernel({sym}, ncclFunc{coll}, {redop_cxx}, {ty_cxx}, NCCL_ALGO_{algo}, NCCL_PROTO_{proto}, {fn_id})\n"
  # ...
  out(s.format(...))

DEFINE_ncclDevKernel巨集展開後是:

📎 src/device/common.h:507-509

c
#define DEFINE_ncclDevKernel(suffix, coll, redop, ty, algo, proto, specializedFnId) \
  __global__ void ncclDevKernel_##suffix(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) { \
    ncclKernelMain<specializedFnId, RunWorkBatch<coll, ty, redop<ty>, algo, proto>>(&args4K.args); \
  }

所以每個 kernel 都是一個__global__函式,呼叫ncclKernelMain,模板參數是specializedFnId和RunWorkBatch<coll, ty, redop<ty>, algo, proto>。

設計思考與生產踩坑

〔設計推斷與架構權衡〕

為什麼用「代表性 kernel」而不是每個組合一個 kernel?編譯時間和二進位大小的權衡。完整的組合空間是 7 × 5 × 12 × 7 × 3 ≈ 8820 個 kernel,每個 kernel 編譯需要幾秒鐘,總共需要幾個小時。而且二進位大小會達到幾百 MB。透過映射到代表性 kernel,實際生成的 kernel 數量減少到幾十個。

踩坑點 1:NCCL_EXACT_KERNEL_NAMES導致編譯爆炸。如果設置了這個環境變數,best_kernel會回傳原始函式,每個組合都會生成一個 kernel。這在開發時有用(可以精確控制哪個 kernel 被編譯),但在生產環境會導致編譯時間過長。

踩坑點 2:required_cuda的版本檢查。某些 kernel 需要特定的 CUDA 版本或架構:

📎 src/device/generate.py:130-154

至此,kernel 已經在 GPU 上啟動,裝置側也拿到了工作描述。但真正決定效能的,是裝置內部如何搬運資料。下一章將深入 src/device 下的三種協定原語:LL、LL128 和 Simple,看看同一份 AllReduce 邏輯為什麼需要三套搬運原語,以及它們在同步方式、緩衝區佈局和 flag 語意上的差異。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 09

第 9 章:裝置端通訊原語:LL、LL128、Simple 三種協定的資料搬運實作

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 9 章 / 共 25 章

上一章我們追蹤了 host 側如何將一次 AllReduce 翻譯成 __global__ kernel,並看到裝置側入口 ncclKernelMain 根據演算法與協定完成分發。但分發只是選定了工具,真正決定效能的是這些工具如何執行資料搬運。本章深入 src/device 下的三套搬運原語:LL、LL128 和 Simple,逐一剖析它們的資料搬運實作,理解不同協定在延遲與頻寬之間的取捨。

為什麼同一份 AllReduce 需要三套搬運原語

先建立一個直覺模型。想像一條流水線工廠:原料(使用者資料)從一端進,成品從另一端出,中間有若干工位(rank)要互相交換半成品。搬運半成品的方式有三種:

  • LL(Low Latency):像兩個人面對面遞紙條,遞過去的同時對方就知道「這是給你的」,幾乎零握手開銷。但紙條很小,一次只能遞 8 位元組有效資料。適合小訊息。
  • LL128:把紙條換成 128 位元組的便條紙,一次遞 120 位元組有效資料,但要求便條紙必須 16 位元組對齊擺放,否則要先在共享記憶體裡「重新排版」。適合中等訊息。
  • Simple:像快遞櫃,先把包裹放進櫃子(FIFO 緩衝區),再發一條「第 N 號櫃有貨」的通知。握手開銷大,但一次能搬很多。適合大訊息。
〔設計推斷與架構權衡〕

如果只有一套原語會怎樣?只用 LL,大訊息會因為「每條訊息都要等對方確認 flag」而把頻寬壓死;只用 Simple,小訊息會因為「寫 FIFO + 發通知 + 等通知」的固定開銷而延遲爆炸。NCCL 的效能曲線之所以在 8KB、128KB 附近有明顯的拐點,根源就在這裡。

三套原語共享同一個模板骨架Primitives<T, RedOp, Fan, Direct, Proto, P2p, isNetOffload>,透過Proto這個模板參數特化出三個版本📎 src/device/primitives.h:117-117。ProtoLL、ProtoLL128、ProtoSimple三個結構體各自攜帶協定相關的常數與計算方法📎 src/device/primitives.h:25-75,演算法程式碼只呼叫prims.send()、prims.recvReduceSend()這類統一介面,不關心底層是哪種協定。

mermaid
flowchart TD
    algo["算法层 all_reduce.h<br/>调用 prims.recvReduceSend()"] --> dispatch{"Proto 模板参数?"}
    dispatch -->|ProtoLL| ll["Primitives&lt;..., ProtoLL, ...&gt;<br/>prims_ll.h"]
    dispatch -->|ProtoLL128| ll128["Primitives&lt;..., ProtoLL128, ...&gt;<br/>prims_ll128.h"]
    dispatch -->|ProtoSimple| simple["Primitives&lt;..., ProtoSimple&lt;...&gt;, ...&gt;<br/>prims_simple.h"]
    ll --> llop["LLGenericOp&lt;RECV,SEND,SrcBuf,DstBuf&gt;"]
    ll128 --> ll128op["GenericOp -&gt; recvReduceSendCopy"]
    simple --> simpleop["genericOp -&gt; waitPeer / reduceCopy / postPeer"]

這張圖說明了「同一份 AllReduce 邏輯為什麼需要三套搬運原語」:演算法層是協定無關的,協定差異被封裝在Primitives的三個特化裡。

LL:用 flag 內嵌在資料行裡的零握手搬運

直覺模型

LL 的核心思想是:把「資料」和「資料是否就緒」的標記塞進同一個 16 位元組的讀寫單元。接收方不需要額外的「通知訊息」,只要輪詢資料行裡的 flag 欄位,flag 匹配就說明資料到了。這就像寄信時把「收件人簽名」直接印在信封上,郵遞員一看簽名就知道該不該投遞,不需要另發一張簽收單。

如果沒有這個設計,接收方就得先等一個「資料已寫入」的通知,再回頭讀資料,兩次記憶體往返,延遲翻倍。

資料結構與記憶體佈局

LL 的搬運單元是union ncclLLFifoLine,從storeLL的組譯可以看出它的佈局📎 src/device/prims_ll.h:154-158:

code
st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};
// 写入 4 个 u32:data1, flag, data2, flag

一個ncclLLFifoLine是 16 位元組,排布為[data1(4B) | flag(4B) | data2(4B) | flag(4B)]。有效資料只有 8 位元組(data1 + data2),另外 8 位元組全是 flag。這就是ProtoLL::calcBytePerGrain()回傳sizeof(uint64_t)的原因——「One 16-byte line has 8-bytes of data」📎 src/device/primitives.h:55-57。

關鍵欄位(Primitives的 LL 特化)📎 src/device/prims_ll.h:20-42:

欄位類型作用
recvStep[i] / sendStep[i]uint64_t[MaxRecv/MaxSend]每個 peer 的步進計數,決定緩衝區偏移和 flag 值
recvBuff[i] / sendBuff[i]ncclLLFifoLine*指向每個 peer 的 FIFO 緩衝區基址
recvConnHeadPtrvolatile uint64_t*接收側「已消費到第幾步」的全域指標
sendConnHeadPtrvolatile uint64_t*發送側「對端已消費到第幾步」的全域指標
sendConnHeadCacheuint64_t快取上次讀到的 head 值,避免每次都讀全域記憶體

緩衝區偏移由recvOffset(i) = (recvStep[i] % NCCL_STEPS) * stepLines計算📎 src/device/prims_ll.h:44-46,NCCL_STEPS是環形緩衝區的槽位數,stepLines是每槽的行數。flag 值由recvFlag(i) = NCCL_LL_FLAG(recvStep[i] + 1)計算📎 src/device/prims_ll.h:56-58,注意+1——因為 flag 初值是 0,第一步的 flag 必須是 1 才能和「未寫入」區分開。

場景驅動 Walkthrough:一次 recvReduceSend

假設 rank 0 在 Ring AllReduce 中執行recvReduceSend:從上一個 rank 收資料、和本地資料做 reduce、再發給下一個 rank。呼叫鏈是recvReduceSend(inpIx, eltN) → LLGenericOp<1, 1, Input, -1>(inpIx, -1, eltN, false) 📎 src/device/prims_ll.h:403-405。

第一步:等待發送緩衝區可用。 waitSend檢查sendConnHeadCache + NCCL_STEPS < sendConnHead + 1 📎 src/device/prims_ll.h:73-89。含義是:如果對端消費進度(head)落後我太多,說明環形緩衝區快滿了,必須等。NCCL_STEPS是緩衝區總槽數,sendConnHead + 1是我即將佔用的槽位。等待時輪詢*sendConnHeadPtr更新快取,並週期性呼叫checkAbort檢查是否被 abort📎 src/device/prims_ll.h:73-89。

第二步:載入本地資料。 DataLoader::loadBegin處理對齊問題📎 src/device/prims_ll.h:200-216。當sizeof(T) <= 2(比如 half 或 int8),來源位址可能不是 4 位元組對齊,所以先按 4 位元組對齊讀入u4[0..2],記錄misalign,然後在loadFinish裡用__funnelshift_r做位元組級移位拼出正確的 64 位元值📎 src/device/prims_ll.h:218-225。這是一個典型的「對齊讀 + 移位重組」技巧,避免了非對齊存取的效能懲罰。

第三步:讀對端資料並等 flag。 readLL是核心📎 src/device/prims_ll.h:108-122:

cpp
do {
  asm volatile("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" ...);
  if (checkAbort(abort, 1, spins)) break;
} while ((flag1 != flag) || (flag2 != flag));

它用ld.volatile.global.v4.u32一次性讀 16 位元組(4 個 u32),然後檢查兩個 flag 欄位是否都等於期望值。volatile關鍵字確保編譯器不會把這個讀取最佳化掉或快取到暫存器——因為對端可能隨時寫入新資料。兩個 flag 都要匹配,是因為寫入方storeLL一次寫 4 個 u32,理論上可能被拆成兩次 8 位元組寫入,兩個 flag 都匹配才能保證 16 位元組完整。

第四步:reduce 並發送。收到 peerData 後,applyReduce(redOp, peerData, data)做歸約📎 src/device/prims_ll.h:279。然後storeLL(sendPtr(i) + offset, data, sendFlag(i))把結果寫入發送緩衝區📎 src/device/prims_ll.h:295-296。注意發送順序:先發i=1..MaxSend(通常是網路 peer),最後發i=0(通常是本地 peer)📎 src/device/prims_ll.h:291-297。註解寫得很清楚:「Send : inter-node, then intra-node, then local」——先發慢的(網路),讓它在背景飛,再發快的(本地),這樣本地 peer 不會等網路。

第五步:推進 step 並 post。 incRecv(i)遞增接收步進📎 src/device/prims_ll.h:91-93,postRecv()把recvConnHead寫回全域指標📎 src/device/prims_ll.h:94-97,通知對端「我已經消費了這一步」。發送側incSend有個特殊邏輯📎 src/device/prims_ll.h:99-106:

cpp
if ((sendStep[i] & NCCL_LL_CLEAN_MASK) == NCCL_LL_CLEAN_MASK) {
  for (int o = offset; o < stepLines; o += nthreads) storeLL(sendPtr(i) + o, 0, sendFlag(i));
}

當 step 到達NCCL_LL_CLEAN_MASK邊界時,要把整個 slice 的所有行都用當前 flag 寫一遍(資料填 0)。為什麼?因為 flag 是循環複用的,如果某一行上次的 flag 恰好等於這次的期望值,接收方會誤以為資料已就緒。這個「cleanup」操作把所有行的 flag 統一刷成新值,消除歧義。

並發控制與硬體互動

LL 的同步完全靠volatile讀寫 + flag 輪詢,沒有鎖。barrier()用__syncwarp()(單 warp 時)或barrier_sync(15 - group, nthreads)(多 warp 時)📎 src/device/prims_ll.h:63-69。15 - group是 barrier 編號,NCCL 用不同的 barrier 編號隔離不同的 group,避免相互干擾。

checkAbort是防死循環的關鍵📎 src/device/primitives.h:154-164:每NCCL_SPINS_BEFORE_CHECK_ABORT(10000)次自旋才讀一次abortFlag,避免頻繁讀全域記憶體拖慢熱路徑。一旦發現 abort,設定ncclShmem.aborted並快取,後續所有等待迴圈都會快速退出。

生產踩坑

坑 1:flag 回繞導致的假就緒。如果NCCL_LL_CLEAN_MASK的 cleanup 邏輯被去掉,在長時間執行(step 超過 mask 週期)後,接收方可能讀到上一輪殘留的 flag,誤判資料就緒,讀到髒資料。這類 bug 極難重現,因為它依賴 step 恰好回繞到特定值。

坑 2:MaxRecv == 0的編譯陷阱。程式碼裡MaxRecv = Fan::MaxRecv > 1 ? Fan::MaxRecv : 1 📎 src/device/prims_ll.h:13,因為即使只發不收,也會分配一個長度為 MaxRecv 的接收緩衝區,如果 MaxRecv 是 0 會導致零長度陣列編譯失敗。Windows 上MaxSend也有同樣處理📎 src/device/prims_ll.h:14-19。

LL128:用 128 位元組對齊換取更高有效載荷

直覺模型

LL 的痛點是有效載荷只有 50%(16 位元組裡 8 位元組是 flag)。LL128 的思路是:把 flag 集中到每 128 位元組的最後 8 位元組,前面 120 位元組全是資料。這樣有效載荷從 50% 提升到 93.75%。代價是必須保證 128 位元組對齊,否則要做「共享記憶體重排版」。

資料結構與記憶體佈局

LL128 的搬運單元是uint64_t(8 位元組),但組織成 128 位元組的「line」。NCCL_LL128_LINEELEMS是每 line 的 64 位元元素數(16 個),NCCL_LL128_DATAELEMS是其中資料元素數(15 個),最後一個元素放 flag。

關鍵常數📎 src/device/prims_ll128.h:292-294:

cpp
static constexpr int WireWordPerSlice = WARP_SIZE * NCCL_LL128_SHMEM_ELEMS_PER_THREAD;
static constexpr int DataEltPerSlice =
  (WireWordPerSlice - WireWordPerSlice / NCCL_LL128_LINEELEMS) * (sizeof(uint64_t) / sizeof(T));

WireWordPerSlice是一個 warp 一次搬運的 64 位元字數,DataEltPerSlice是其中有效資料元素數(減去每 line 一個 flag 元素)。

LL128 的 flag 機制和 LL 不同:只有每 8 個執行緒中的第 7 個(flagThread)負責檢查 flag 📎 src/device/prims_ll128.h:373。flagThread = ((tid % 8) == 7)。為什麼?因為 flag 是每 128 位元組一個,而一個 warp 有 32 個執行緒,每 8 個執行緒處理 128 位元組(8 執行緒 × 16 位元組 = 128 位元組),所以每 8 個執行緒裡只有 1 個需要讀 flag。

場景驅動 Walkthrough:一次 recvReduceSendCopy

呼叫鏈:recvReduceSend(inpIx, eltN) → GenericOp<1, 1, Input, -1> → recvReduceSendCopy<NCCL_LL128_SHMEM_ELEMS_PER_THREAD, RECV, SEND, SrcBuf, DstBuf> 📎 src/device/prims_ll128.h:422-423, 296-333。

第一步:載入本地資料到暫存器。 loadRegsBegin分兩種情況📎 src/device/prims_ll128.h:99-142:

  • 16 位元組對齊:直接load128到暫存器,無共享記憶體中轉。注意flagThread只載入一半資料(g % 2 == 0),因為它的另一半暫存器要留給 flag📎 src/device/prims_ll128.h:109-114。
  • 非對齊:先把對齊區域載入到共享記憶體ncclScratchForWarp(warpInBlock),__syncwarp()後從共享記憶體按正確偏移讀回暫存器📎 src/device/prims_ll128.h:115-141。

第二步:等待並讀取對端資料。 recvReduceSendCopy裡的等待迴圈📎 src/device/prims_ll128.h:190-207:

cpp
do {
  needReload = false;
  for (int u = 0; u < ELEMS_PER_THREAD; u += 2) {
    load128(ptr + u * WARP_SIZE, vr[u], vr[u + 1]);
    needReload |= flagThread && (vr[u + 1] != flag);
  }
  needReload &= (0 == checkAbort(abort, 1, spins));
} while (__any_sync(WARP_MASK, needReload));

關鍵點:只有flagThread檢查 flag,然後用__any_sync做 warp 級投票——只要有一個 flagThread 發現 flag 不匹配,整個 warp 繼續自旋。這比每個執行緒都檢查 flag 更省指令。

第三步:暫存器重排。 loadRegsFinish把 flagThread 的 flag 暫存器移到空閒暫存器📎 src/device/prims_ll128.h:145-151。註解解釋了這個設計:「By deferring register shuffle here we've overlapped spinning on first peer's data with memory loads of src data」——把暫存器重排推遲到等待之後,讓等待時間和本地資料載入重疊。

第四步:reduce 並發送。收到資料後做applyReduce 📎 src/device/prims_ll128.h:227-230,然後store128寫入發送緩衝區📎 src/device/prims_ll128.h:274-287。注意發送時flagThread ? flag : v[u+1]——flagThread 寫 flag,其他執行緒寫資料。

第五步:推進 step。和 LL 不同,LL128 的 step 推進在GenericOp末尾統一做📎 src/device/prims_ll128.h:324-332,而不是在recvReduceSendCopy裡。而且postSend用了__threadfence_system()(SM90+)或__threadfence() 📎 src/device/prims_ll128.h:87-96,確保資料對其他 GPU/網卡可見後才更新 tail 指標。

並發控制與硬體互動

LL128 的barrier()總是用barrier_sync(15 - group, nthreads) 📎 src/device/prims_ll128.h:64-66,不像 LL 有單 warp 最佳化。因為 LL128 的資料搬運是 warp 級的,需要跨 warp 同步。

loadRegsBegin裡的共享記憶體重排版用__syncwarp()同步📎 src/device/prims_ll128.h:129,確保所有執行緒寫完共享記憶體後再讀。

生產踩坑

坑 1:非對齊存取的效能懸崖。如果使用者緩衝區不是 16 位元組對齊,每次搬運都要走共享記憶體中轉,效能可能下降 30% 以上。生產環境應確保輸入輸出緩衝區按 16 位元組對齊分配。

坑 2:flagThread的暫存器壓力。flagThread 只載入一半資料,意味著它的暫存器利用率和其他執行緒不同。如果編譯器沒有正確分配暫存器,可能導致暫存器溢出到本地記憶體,效能驟降。

Simple:用 FIFO + 通知實現大訊息高吞吐

直覺模型

Simple 協定像快遞櫃:發送方把資料放進 FIFO 緩衝區(櫃子),然後更新一個「已放入第 N 號櫃」的 step 指標(發通知);接收方輪詢 step 指標,看到新值就去對應櫃子取貨。握手開銷大(要寫指標 + 讀指標),但一次能搬很多資料,適合大訊息。

資料結構與記憶體佈局

Simple 的欄位比 LL/LL128 複雜得多📎 src/device/prims_simple.h:28-46:

欄位類型作用
flagsint位標誌,編碼角色(WaitRecv/WaitSend/PostRecv/PostSend)、Direct 模式、NetReg 等
stepuint64_t當前步進
connStepPtruint64_t*指向連接的對端 step 指標
connStepCacheuint64_t快取上次讀到的 step 值
connEltsFifoT*FIFO 緩衝區基底位址
connStepSizeint每步的位元組數
directBuffT*Direct 模式下的直接緩衝區指標

flags的位定義📎 src/device/prims_simple.h:23-27:

cpp
RoleInput = 0x01, RoleOutput = 0x02, RoleWaitRecv = 0x04, RoleWaitSend = 0x08,
RolePostSend = 0x10, RolePostRecv = 0x20, Aborted = 0x40, NetRegMode = 0x80,
ConnFifoEnabled = 0x100, DirectWrite = 0x200, DirectRead = 0x400, PatMode = 0x800,
NvlsMinPolling = 0x1000, NetDeviceUnpack = 0x2000, AnyNetDeviceUnpack = 0x4000,
RoleWaitPatNvls = 0x8000, RolePostPatNvls = 0x10000;

這是一個典型的「用位元運算代替多個 bool 欄位」的設計,節省暫存器。每個執行緒根據自己的tid被分配一個角色📎 src/device/prims_simple.h:651-666:前nrecv個執行緒是 WaitRecv,接下來nsend個是 WaitSend,最後nrecv個是 PostRecv,倒數nsend個是 PostSend。

場景驅動 Walkthrough:一次 recvReduceSend

呼叫鏈:recvReduceSend(inpIx, eltN) → genericOp<0, 0, 1, 1, Input, -1> 📎 src/device/prims_simple.h:994-996。

第一步:計算 slice 大小。 sliceSize = max(divUp(nelem, 16 * SlicePerChunk) * 16, sliceSize / 32) 📎 src/device/prims_simple.h:185-186。這個公式保證 slice 至少是 16 位元組對齊,且不會太小。

第二步:worker 迴圈。只有tid < nworkers的執行緒進入主迴圈📎 src/device/prims_simple.h:190。nworkers = nthreads - (MaxSend > 0 && nthreads >= NCCL_SIMPLE_EXTRA_GROUP_IF_NTHREADS_GE ? WARP_SIZE : 0) 📎 src/device/prims_simple.h:626——預留一個 warp 做 threadfence 和 copy 的重疊。

第三步:等待對端。 waitPeer是核心📎 src/device/prims_simple.h:103-164:

cpp
while (connStepCache + (isSendNotRecv ? NCCL_STEPS : 0) < step + StepPerSlice) {
  connStepCache = loadStepValue(connStepPtr);
  if (checkAbort(flags, Aborted, spins)) break;
}

isSendNotRecv區分發送和接收:發送時等的是「對端已消費」(head),接收時等的是「對端已生產」(tail)。NCCL_STEPS是緩衝區槽數,StepPerSlice是每 slice 的步進數。

等待完成後,根據 Direct 模式設定ptrs[index] 📎 src/device/prims_simple.h:123-158。Direct 模式允許直接讀寫對端緩衝區,繞過 FIFO,減少一次拷貝。

第四步:reduceCopy。根據 Direct 組合選擇不同的reduceCopy呼叫📎 src/device/prims_simple.h:241-277。最複雜的分支是srcs[0] && dsts[0]都存在時📎 src/device/prims_simple.h:258-271,呼叫reduceCopy<Unroll, RedOp, T, MultimemSrcs, Recv+Src, Recv*MaxRecv+Src, MultimemDsts, Send+Dst, Send*MaxSend+Dst, PreOpSrcs>,參數含義是:從Recv*MaxRecv+Src個來源讀,歸約後寫到Send*MaxSend+Dst個目的地。

第五步:postPeer。 postPeer更新 step 指標📎 src/device/prims_simple.h:167-175:

cpp
if (Send && (flags & RolePostSend) && (dataStored || (flags & ConnFifoEnabled))) {
  fence_acq_rel_sys();
}
st_relaxed_sys_global(connStepPtr, step);

發送側在更新 step 前要fence_acq_rel_sys(),確保資料寫入對其他 GPU/網卡可見。接收側不需要 fence,因為接收方只是通知「我已消費」,不涉及資料可見性。

並發控制與硬體互動

Simple 的同步用st_relaxed_sys_global寫 step 指標📎 src/device/prims_simple.h:167-175,用loadStepValue讀📎 src/device/prims_simple.h:86-100。loadStepValue在 SM90+ 且啟用NvlsMinPolling時用multimem.ld_reduce.acquire.sys.global.min.u64指令📎 src/device/prims_simple.h:86-100,這是 NVLink SHARP 的硬體加速輪詢。

barrier()和subBarrier()的區別📎 src/device/prims_simple.h:49-55:barrier()同步所有nthreads個執行緒,subBarrier()只同步nworkers個 worker 執行緒。subBarrier的 barrier 編號是15 - group - (nworkers != nthreads ? 1 : 0),當 worker 數不等於總執行緒數時用不同的 barrier,避免和barrier()衝突。

生產踩坑

坑 1:NetRegMode 下的解構等待。解構函式裡有一段特殊邏輯📎 src/device/prims_simple.h:794-804:

cpp
if ((flags & NetRegMode) && (flags & RoleWaitSend)) {
  uint64_t prevStep = step - StepPerSlice;
  volatile ssize_t* ptr = &(connFifo[prevStep % NCCL_STEPS].size);
  while (*ptr != -1) { ... }
}

在 NetRegMode 下,發送緩衝區被網卡直接存取,必須等 proxy 執行緒確認已發送(size 被設為 -1)才能返回,否則下一個 kernel 可能覆蓋正在被網卡讀取的資料。

坑 2:DirectRead 的 sendrecv 死鎖。解構函式裡還有一段📎 src/device/prims_simple.h:814-824:

cpp
if ((flags & DirectRead) && (flags & RoleWaitSend) && P2p) {
  while (*tail > *head) { ... }
}

在 sendrecv 的 DirectRead 模式下,發送方必須等接收方讀完資料才能返回。如果接收方因為某種原因沒有推進 tail,發送方會死鎖。這個等待必須在barrier()之後做,否則可能和 post 執行緒競爭。

坑 3:roundUp導致的 step 跳躍。 loadRecvConn和loadSendConn裡都有step = roundUp(step, SlicePerChunk * StepPerSlice) 📎 src/device/prims_simple.h:486, 533。這會把 step 對齊到 slice 邊界,但如果上一步的 step 不是對齊的,會導致跳過的槽位沒有被正確初始化。程式碼在loadRecvConn裡補了一句*connStepPtr = step來歸還 credit📎 src/device/prims_simple.h:489。

三套原語的對比與選型

mermaid
flowchart LR
    subgraph LL["LL 协议"]
        ll_data["ncclLLFifoLine 16B<br/>data1(4B)+flag(4B)+data2(4B)+flag(4B)"]
        ll_sync["flag 内嵌数据行<br/>轮询 flag 匹配"]
    end
    subgraph LL128["LL128 协议"]
        ll128_data["128B line<br/>15×8B data + 1×8B flag"]
        ll128_sync["flagThread 每8线程1个<br/>__any_sync 投票"]
    end
    subgraph Simple["Simple 协议"]
        simple_data["FIFO 缓冲区<br/>connEltsFifo + step*connStepSize"]
        simple_sync["step 指针 + fence<br/>loadStepValue 轮询"]
    end
    ll_data --> ll_sync
    ll128_data --> ll128_sync
    simple_data --> simple_sync
維度LLLL128Simple
有效載荷率50%93.75%~100%
同步方式flag 內嵌,輪詢flagThread + warp 投票step 指標 + fence
對齊要求無(有移位重組)16 位元組無
適用訊息大小小(< 8KB)中(8KB ~ 128KB)大(> 128KB)
緩衝區佈局ncclLLFifoLine[]uint64_t[]按 128B lineT[] FIFO
Direct 支援無(PrimitivesWithoutDirect降級)無(同左)完整支援

LL 和 LL128 都繼承PrimitivesWithoutDirect 📎 src/device/prims_ll.h:9-10, src/device/prims_ll128.h:13-14,因為它們的緩衝區佈局不支援直接讀寫對端記憶體。Simple 則完整實作了 Direct 模式,支援 P2P 直連和 NVLS。

設計思考

〔設計推斷與架構權衡〕

為什麼 LL 的 flag 要重複兩次?因為 GPU 的全域記憶體寫入不保證原子性。storeLL寫 16 位元組,硬體可能拆成兩次 8 位元組寫。如果只放一個 flag,接收方可能在資料只寫了一半時就認為就緒。兩個 flag 分別位於 16 位元組的前半和後半,只有兩次寫都完成,兩個 flag 才會都匹配。

為什麼 Simple 要預留一個 warp? 📎 src/device/prims_simple.h:625-626註解說「For send operations, we need an extra warp to overlap the threadfence and the copy」。fence_acq_rel_sys()是一個昂貴的操作,如果所有執行緒都等 fence 完成再繼續,會浪費大量時間。預留一個 warp 專門做 fence,其他 warp 可以繼續搬運下一批資料。

〔設計推斷與架構權衡〕

為什麼 LL128 的 step 推進在 GenericOp 末尾而不是 recvReduceSendCopy 裡?因為 LL128 的搬運是 warp 級的,多個 warp 可能並行處理不同的 slice。如果在recvReduceSendCopy裡推進 step,每個 warp 都會推進一次,導致 step 被推進多次。放在GenericOp末尾統一推進,確保每個 slice 只推進一次。

本章小結

本章深入了三套搬運原語的實作:

1. LL:用 16 位元組的ncclLLFifoLine把 flag 內嵌在資料行裡,接收方輪詢 flag 匹配即可確認資料就緒。有效載荷 50%,適合小訊息。核心是readLL的ld.volatile.global.v4.u32和storeLL的st.volatile.global.v4.u32。

2. LL128:把 flag 集中到每 128 位元組的最後 8 位元組,有效載荷提升到 93.75%。用flagThread(每 8 執行緒 1 個)檢查 flag,__any_sync做 warp 投票。非對齊時走共享記憶體重排版。

3. Simple:用 FIFO 緩衝區 + step 指標通知實作大訊息高吞吐。flags位標誌編碼角色,waitPeer輪詢 step,postPeer更新 step 並 fence。完整支援 Direct 模式。

三套原語共享同一個模板骨架,透過Proto模板參數特化。演算法層只呼叫統一介面,不關心底層協定。這就是「同一份 AllReduce 邏輯為什麼需要三套搬運原語」的答案:不同訊息大小需要不同的同步策略和緩衝區佈局,三套原語分別針對小、中、大訊息最佳化。

本章思考與自測

Q1: 如果把incSend裡的 cleanup 邏輯(📎 src/device/prims_ll.h:99-106)去掉,在什麼場景下會觸發資料損壞?為什麼?

參考解析:cleanup 邏輯在sendStep[i] & NCCL_LL_CLEAN_MASK == NCCL_LL_CLEAN_MASK時,把整個 slice 的所有行都用當前 flag 寫一遍(資料填 0)。如果去掉,當 step 迴繞到NCCL_LL_CLEAN_MASK邊界時,某些行的 flag 可能還是上一輪的值。如果上一輪的 flag 恰好等於這一輪接收方期望的 flag,接收方會誤以為資料已就緒,讀到上一輪的殘留資料。這是一個典型的 ABA 問題。觸發條件是長時間執行(step 超過NCCL_LL_CLEAN_MASK週期)且 flag 恰好迴繞到相同值。這類 bug 極難復現,因為需要精確的 step 對齊。

Q2:Simple 協定的解構函式中,NetRegMode 下的等待(📎 src/device/prims_simple.h:794-804)和 DirectRead 下的等待(📎 src/device/prims_simple.h:814-824)分別在防什麼?如果去掉其中一個,在高併發場景下會發生什麼?

參考解析:NetRegMode 等待的是 proxy 執行緒把connFifo[prevStep].size設為 -1,表示網卡已完成發送。如果去掉,下一個 kernel 可能覆蓋正在被網卡 DMA 讀取的發送緩衝區,導致網卡讀到髒資料。DirectRead 等待的是接收方推進 tail(*tail > *head),表示接收方已讀完直接緩衝區。如果去掉,發送方可能在接收方還沒讀完時就覆蓋了緩衝區,導致接收方讀到新資料而非舊資料。在高併發場景下,這兩個等待都是必須的,去掉任何一個都會導致資料競爭。區別是 NetRegMode 防的是「網卡讀」,DirectRead 防的是「對端 GPU 讀」。

Q3:LL128 的loadRegsBegin在非對齊時走共享記憶體重排版(📎 src/device/prims_ll128.h:115-141),這個路徑比對齊路徑慢多少?為什麼 NCCL 不直接要求使用者緩衝區必須 16 位元組對齊?

參考解析:非對齊路徑多了三步:寫共享記憶體、__syncwarp()、從共享記憶體讀。共享記憶體的頻寬雖然高,但__syncwarp()是一個同步點,會阻塞 warp 直到所有執行緒完成寫入。粗略估計,非對齊路徑比對齊路徑慢 20-40%,具體取決於共享記憶體 bank 衝突情況。NCCL 不強制要求對齊,是因為使用者可能傳入任意偏移的緩衝區(比如 tensor 切片),強制對齊會限制 API 的靈活性。NCCL 的策略是「對齊時走快路徑,非對齊時走慢路徑但保證正確性」。生產環境建議使用者盡量按 16 位元組對齊分配緩衝區,以走快路徑。

至此,我們已經掌握了 LL、LL128、Simple 三種原語的資料搬運機制,它們為上層演算法提供了靈活的效能調節手段。下一章將深入集合通訊演算法核心,看 AllReduce、AllGather、ReduceScatter 等如何呼叫這些原語,以及 Ring、Tree、CollNet 等演算法如何組織資料流,最終完成端到端的集合通訊。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 10

第 10 章:集體通訊演算法核心:AllReduce、AllGather、ReduceScatter 的裝置端實作

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 10 章 / 共 25 章

上一章拆解了 LL、LL128、Simple 三種協定原語,它們是資料搬運的「發動機」,但發動機本身不知道要搬什麼、往哪搬、按什麼順序搬。本章要看的 src/device 下這一組演算法核心檔案,就是「變速箱」——它們把 AllReduce、AllGather、ReduceScatter 這些集合通訊語義,翻譯成一連串 prims.directSend、prims.directRecvReduceDirectSend 這樣的原語呼叫。一句話概括本章的核心矛盾:同一個 AllReduce,為什麼需要 Ring、Tree、CollNet、NVLS 四套完全不同的裝置側實作?答案藏在「資料流拓撲」與「硬體能力」的匹配裡。Ring 用最少的網路頻寬做兩階段流水,Tree 用樹形歸約把延遲壓到 log(n),CollNet/NVLS 則把歸約卸載到網卡或 NVLink 交換機上。本章逐個拆開看。

10.1 Ring AllReduce:兩階段流水如何在 kernel 內落地

直覺模型:環形流水線上的「接力賽」

想像 n 個工人站成一圈,每人手裡有一箱原料。AllReduce 的目標是讓每個人最終都拿到「所有原料混合後的成品」。Ring 演算法的做法分兩階段:第一階段(reduce-scatter)每人把箱子沿環傳遞,每傳一站就混入自己的原料,轉 n-1 站後每個人手裡恰好有一份「完整混合」的成品,但只有 1/n 的份額;第二階段(all-gather)這些成品份額再沿環傳一圈,每人補齊所有份額。

若沒有 Ring,最樸素的做法是每個 rank 把資料發給 root,root 歸約後再廣播——root 的網路頻寬成為瓶頸,n 越大越慢。Ring 的精妙在於:每個 rank 的發送量和接收量都是 2(n-1)/n 倍資料量,與 n 無關地攤平到所有鏈路。

資料結構與記憶體佈局

Ring 演算法的核心狀態在ncclRing結構裡(定義在 device.h,本章不展開),runRing只取其中兩個欄位:

  • ring->index:本 rank 在環中的邏輯位置,用於計算「第 j 步該處理哪個 chunk」。
  • ring->prev / ring->next:前驅和後繼 rank 編號,作為Primitives建構函式的 recv/send peer 參數。

關鍵的分塊參數由ncclCollCbdPart計算(📎 src/device/all_reduce.h:21-22):

code
ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), (ssize_t*)nullptr, &gridOffset, &channelCount, &chunkCount);

這個函式把整個通訊域的資料按 channel 切分,輸出三個值:gridOffset(本 channel 負責的資料在整個 buffer 中的起始偏移)、channelCount(本 channel 負責的元素總數)、chunkCount(每個 rank 分到的 chunk 元素數)。chunkCount是 Ring 演算法的粒度——每一步搬運一個 chunk。

loopCount = nranks * chunkCount(📎 src/device/all_reduce.h:23)表示「轉一整圈」處理的資料量。外層迴圈for (elemOffset = 0; elemOffset < channelCount; elemOffset += loopCount)(📎 src/device/all_reduce.h:34)意味著:如果 channel 資料量超過一圈能處理的量,就分多圈跑。

Step-by-Step Walkthrough:一次 Ring AllReduce 的完整呼叫流

代入場景:4 個 rank(nranks=4),本 rank 的ringIx=0,chunkCount=100,channelCount=400(正好一圈)。

第 0 步:把「自己的 chunk」推給下一個 GPU(📎 src/device/all_reduce.h:42-47)

code
chunk = modRanks(ringIx + nranks - 1);   // = 3
chunkOffset = chunk * chunkCount;         // = 300
offset = gridOffset + elemOffset + chunkOffset;
nelem = min(chunkCount, remCount - chunkOffset);
prims.directSend(offset, offset, nelem);

modRanks是個 lambda,做模 nranks 的減法(📎 src/device/all_reduce.h:40)。ringIx + nranks - 1表示「本 rank 的前一個 chunk 編號」。為什麼第 0 步發的是 chunk 3?因為 Ring 的 reduce-scatter 階段,每個 rank 先把自己「不該保留」的那份資料(即前驅 rank 的 chunk)發出去。directSend只發不接,因為此時還沒收到任何資料。

第 1 到 nranks-2 步:邊收邊歸約邊轉發(📎 src/device/all_reduce.h:50-56)

code
for (int j = 2; j < nranks; ++j) {
  chunk = modRanks(ringIx + nranks - j);
  ...
  prims.directRecvReduceDirectSend(offset, offset, nelem);
}

directRecvReduceDirectSend是 Ring 的核心原語:從prev收一個 chunk,與本地資料做歸約(比如加法),再把結果發給next。注意offset和nelem在每次迭代都重新計算——因為每步處理的 chunk 不同。j 從 2 到 nranks-1,共 nranks-2 步。

第 nranks-1 步:收下最後一個 chunk 並歸約,產生最終結果(📎 src/device/all_reduce.h:58-64)

code
chunk = ringIx + 0;
...
prims.directRecvReduceCopyDirectSend(offset, offset, nelem, /*postOp=*/true);

這一步的postOp=true是關鍵:歸約完成後要執行後置操作(比如求平均時的除法)。directRecvReduceCopyDirectSend比上一步多了個Copy——把歸約結果同時寫入本地 recvbuff 和發往 next。至此 reduce-scatter 階段結束,每個 rank 手裡有一個「完整歸約」的 chunk。

all-gather 階段:nranks-2 步純轉發(📎 src/device/all_reduce.h:66-73)

code
for (int j = 1; j < nranks - 1; ++j) {
  chunk = modRanks(ringIx + nranks - j);
  ...
  prims.directRecvCopyDirectSend(offset, offset, nelem);
}

注意這裡用的是directRecvCopyDirectSend,沒有Reduce——因為資料已經歸約完了,只需複製轉發。

最後一步:收下最後一個 chunk(📎 src/device/all_reduce.h:75-81)

code
chunk = modRanks(ringIx + 1);
...
prims.directRecv(offset, nelem);

只收不發,補齊最後一塊。

整個流程可以用下面的控制流圖概括:

mermaid
flowchart TD
    start["runRing 入口<br/>计算 chunkCount/loopCount"] --> loop{"elemOffset < channelCount?"}
    loop -->|否| done["返回"]
    loop -->|是| s0["step 0: directSend<br/>chunk = ringIx-1"]
    s0 --> mid{"j 从 2 到 nranks-1?"}
    mid -->|是| s1["directRecvReduceDirectSend<br/>chunk = ringIx-j"]
    s1 --> mid
    mid -->|否| s2["step nranks-1<br/>directRecvReduceCopyDirectSend<br/>postOp=true"]
    s2 --> ag{"j 从 1 到 nranks-2?"}
    ag -->|是| s3["directRecvCopyDirectSend<br/>纯转发"]
    s3 --> ag
    ag -->|否| s4["directRecv<br/>收最后一块"]
    s4 --> loop

設計思考:為什麼 Ring 的 chunk 順序是「倒著走」的

注意 chunk 編號的規律:第 0 步發ringIx-1,第 j 步處理ringIx-j,最後一步處理ringIx+0。這是逆時針推進。為什麼?因為 Ring 的每個 rank 只保留「自己負責歸約的那個 chunk」(即ringIx+0),其餘 chunk 都是路過。逆時針推進保證:當某個 chunk 轉完一圈回到起點時,恰好完成了 nranks 次歸約,產生最終結果。如果順時針推進,chunk 會在錯誤的 rank 上完成歸約。

生產踩坑:remCount < loopCount時的對齊陷阱

📎 src/device/all_reduce.h:38有一行容易被忽略的程式碼:

code
if (remCount < loopCount) chunkCount = alignUp(divUp(remCount, nranks), 16 / sizeof(T));

當剩餘資料不足一圈時,chunkCount 要重新計算,並且alignUp(..., 16/sizeof(T))強制按 16 位元組對齊。為什麼?因為 LL128 協議要求 128 位元組對齊,Simple 協議也有向量化存取的對齊需求。如果去掉這個對齊,非對齊的 chunk 會走慢路徑,效能下降 20-40%。生產環境中如果發現 Ring AllReduce 在小訊息尾部效能抖動,往往就是這個對齊沒生效——檢查channelCount是否是nranks * 16/sizeof(T)的整數倍。

10.2 Tree AllReduce:用樹形歸約把延遲壓到 log(n)

直覺模型:公司裡的「逐級匯報」

Ring 的延遲是 O(n)——資料要轉一圈。當 n 很大(比如 1024 個 GPU)時,即使頻寬攤平了,延遲也受不了。Tree 演算法換了個思路:像公司組織架構一樣,每個 rank 只跟「父節點」和「子節點」通訊。歸約階段,葉子節點把資料往上匯報,父節點合併子節點的資料;廣播階段反過來,根節點把結果往下發。延遲從 O(n) 降到 O(log n)。

若沒有 Tree,大規模叢集的 AllReduce 延遲會隨 rank 數線性增長,訓練迭代時間被通訊拖垮。

資料結構與記憶體佈局

Tree 的狀態在ncclTree裡:

  • tree->up:父節點 rank(-1 表示本 rank 是根)。
  • tree->down[]:子節點陣列,最多NCCL_MAX_TREE_ARITY個(典型是 3,即二元+本地)。

runTreeUpDown和runTreeSplit是兩個變體。前者用「先全部歸約再全部廣播」的兩階段模式,後者把執行緒拆成兩半,一半做歸約一半做廣播,實現流水重疊。

Step-by-Step Walkthrough:runTreeUpDown 的三分支

runTreeUpDown的第一個程式碼區塊是歸約階段(📎 src/device/all_reduce.h:96-118),根據本 rank 在樹中的位置分三種情況:

情況 A:本 rank 是根(tree->up == -1)(📎 src/device/all_reduce.h:99-104)

code
prims.directRecvReduceCopy(offset, offset, nelem, /*postOp=*/true);

根節點只收不發,從所有子節點收資料、歸約、寫入 recvbuff。postOp=true執行後置操作。

情況 B:本 rank 是葉子(tree->down[0] == -1)(📎 src/device/all_reduce.h:105-110)

code
prims.directSend(offset, offset, nelem);

葉子節點只發不收,把自己的資料發給父節點。

情況 C:中間節點(📎 src/device/all_reduce.h:111-117)

code
prims.directRecvReduceDirectSend(offset, offset, nelem);

從子節點收、歸約、發給父節點。

廣播階段(📎 src/device/all_reduce.h:120-142)邏輯對稱:根節點directSendFromOutput(從 recvbuff 發),葉子節點directRecv,中間節點directRecvCopyDirectSend。

runTreeSplit:用執行緒拆分實現歸約-廣播流水

runTreeUpDown的問題是:歸約階段和廣播階段串行,中間有個全域同步點。runTreeSplit把執行緒分成兩組(📎 src/device/all_reduce.h:155-164):

code
if (Proto::Id == NCCL_PROTO_SIMPLE) {
  nthreadsSplit = nthreads / 2;
  if (nthreadsSplit >= 256) nthreadsSplit += 64;
} else {
  nthreadsSplit = (nthreads * 7 / (10 * WARP_SIZE)) * WARP_SIZE;
}

Simple 協議對半分;LL/LL128 協議按 7:3 分,因為「從 3 個來源收資料做歸約」比「發給 3 個目標」計算密集,所以歸約組多分執行緒。

然後tid < nthreadsSplit的執行緒做歸約上推(📎 src/device/all_reduce.h:175-202),其餘執行緒做廣播下推(📎 src/device/all_reduce.h:203-224)。兩組透過Proto::MaxGroupWidth偏移量區分各自的通訊組(📎 src/device/all_reduce.h:189的0 * Proto::MaxGroupWidth和📎 src/device/all_reduce.h:210的1 * Proto::MaxGroupWidth)。

設計思考:為什麼 Tree 的根節點要特殊處理

樹形歸約的根節點是「匯聚點」,它的接收量是子節點數倍,發送量為零(歸約階段)。如果根節點也走通用的directRecvReduceDirectSend,會嘗試往tree->up(-1)發送,導致越界。所以必須用if (tree->up == -1)分支單獨處理。同理葉子節點的tree->down[0] == -1判斷。

生產踩坑:Tree 演算法的「熱點根」問題

Tree 的根節點承擔了所有歸約流量,如果根節點所在 GPU 恰好是慢節點(比如 PCIe 頻寬受限),整個 AllReduce 會被拖慢。NCCL 的應對是:每個 channel 選不同的根,把根節點的負載分散到多個 rank。這就是為什麼runTreeSplit裡根節點分支用FanSymmetric<NCCL_MAX_TREE_ARITY_TOP>(📎 src/device/all_reduce.h:168)——它要同時處理多個子節點的歸約。生產環境如果發現 Tree AllReduce 效能不均,檢查 channel 的根節點分佈是否均勻。

10.3 AllGather 與 ReduceScatter:Ring 的「半程」變體

直覺模型:AllReduce 拆成兩半

AllGather 和 ReduceScatter 本質上是 AllReduce 的兩個階段各自獨立成 API。AllGather 只做「收集」——每個 rank 貢獻一份資料,最終所有人拿到全部資料。ReduceScatter 只做「歸約+分散」——所有人貢獻資料,歸約後每人拿到一份。

若沒有這兩個獨立 API,使用者做「先歸約再收集」或「先收集再歸約」時只能調 AllReduce 再手動切片,浪費一半頻寬。

AllGather 的 Ring 實現

all_gather.h的runRing(📎 src/device/all_gather.h:14-88)比 AllReduce 簡單:沒有歸約,只有複製轉發。

第 0 步:把自己的資料推給下一個 GPU(📎 src/device/all_gather.h:51-60)

code
rankDest = ringRanks[0];
offset = dataOffset + rankDest * count;
if ((inputBuf + dataOffset == outputBuf + offset) || isNetOffload) {
  prims.directSend(dataOffset, offset, nelem);
} else {
  prims.directCopySend(dataOffset, offset, nelem);
}

這裡有個 in-place 判斷:如果inputBuf + dataOffset == outputBuf + offset,說明輸入輸出是同一塊記憶體(in-place AllGather),直接directSend;否則要directCopySend(先拷貝到輸出再發)。

中間 nranks-2 步:純轉發(📎 src/device/all_gather.h:62-67)

code
prims.directRecvCopyDirectSend(offset, offset, nelem);

最後一步:收下最後一塊(📎 src/device/all_gather.h:69-74)

code
prims.directRecv(offset, nelem);

isNetOffload:單 warp 驅動網路 + 多 warp 並行拷貝

📎 src/device/all_gather.h:28-36有個特殊分支:

code
if (isNetOffload) {
  workNthreads = WARP_SIZE;
  chunkCount = NCCL_MAX_NET_SIZE;
} else {
  workNthreads = nthreads;
}

當isNetOffload=true(單 RPN + 網路註冊模式)時,只用 1 個 warp 驅動 Ring 通訊,其餘 warp 並行做「源資料拷貝到目標 buffer」(📎 src/device/all_gather.h:76-82)。這是為了在非 in-place AllGather 時,把拷貝開銷和通訊開銷重疊。

最後有個barrier_sync(14, nthreads)(📎 src/device/all_gather.h:87),註解解釋得很清楚:必須等所有 warp 完成,否則下一個 work 可能複用 outputBuf 導致競爭。用 barrier 14 是為了避開 prims 自己的 barrier 和__syncthreads()。

ReduceScatter 的 Ring 實現

reduce_scatter.h的runRing(📎 src/device/reduce_scatter.h:14-56)是 AllReduce 的 reduce-scatter 階段單獨抽出:

第 0 步:把自己的資料推給下一個 GPU(📎 src/device/reduce_scatter.h:39-42)

code
rankDest = ringRanks[nranks - 1];
offset = dataOffset + rankDest * count;
prims.send(offset, nelem);

中間 nranks-2 步:邊收邊歸約邊轉發(📎 src/device/reduce_scatter.h:44-49)

code
prims.recvReduceSend(offset, nelem);

最後一步:收下並歸約,產生最終結果(📎 src/device/reduce_scatter.h:61-64)

code
prims.recvReduceCopy(offset, dataOffset, nelem, /*postOp=*/true);

注意最後一步的recvReduceCopy有兩個 offset:offset(接收源)和dataOffset(本地輸入),歸約結果寫入dataOffset。

資料流對比圖

mermaid
flowchart LR
    subgraph AllReduce["AllReduce (两阶段)"]
        A1["reduce-scatter<br/>n-1 步"] --> A2["all-gather<br/>n-1 步"]
    end
    subgraph AG["AllGather (单阶段)"]
        B1["directSend<br/>step 0"] --> B2["directRecvCopyDirectSend<br/>n-2 步"] --> B3["directRecv<br/>step n-1"]
    end
    subgraph RS["ReduceScatter (单阶段)"]
        C1["send<br/>step 0"] --> C2["recvReduceSend<br/>n-2 步"] --> C3["recvReduceCopy<br/>step n-1"]
    end
    AllReduce -.->|"拆解"| AG
    AllReduce -.->|"拆解"| RS

生產踩坑:in-place 判斷的邊界

📎 src/device/all_gather.h:55的 in-place 判斷inputBuf + dataOffset == outputBuf + offset依賴指標精確相等。如果使用者傳入的 sendbuff 和 recvbuff 有偏移但邏輯上是同一塊記憶體,這個判斷會失效,導致走directCopySend路徑——雖然正確但多一次拷貝。生產環境建議 in-place AllGather 時確保 sendbuff 和 recvbuff 完全一致。

10.4 CollNet 與 NVLS:把歸約卸載到硬體

直覺模型:讓「交換機」幫忙算

Ring 和 Tree 都是「GPU 自己算歸約」。CollNet 和 NVLS 換了個思路:把歸約操作卸載到網卡(CollNet)或 NVLink 交換機(NVLS)上。GPU 只負責把資料發出去,硬體完成歸約後再廣播回來。這就像從「每個工人自己混合原料」變成「把原料送到中央攪拌機,攪拌機混好再分發」。

若沒有硬體卸載,歸約操作會佔用 GPU 的 SM 資源,且歸約延遲無法隱藏。

CollNet Direct 的執行緒分工

RunWorkColl<ncclFuncAllReduce, ..., NCCL_ALGO_COLLNET_DIRECT, ...>的run(📎 src/device/all_reduce.h:249-386)把執行緒分成四組:

code
const int nThreadsScatter = WARP_SIZE + ((hasUp && hasDn) ? COLLNET_COPY_THREADS : ...);
const int nThreadsGather = ((hasUp && hasDn) ? COLLNET_COPY_THREADS : ...);
const int nThreadsBcast = WARP_SIZE + ((hasUp && hasDn) ? COLLNET_COPY_THREADS : ...);
const int nThreadsReduce = work->nWarps * WARP_SIZE - nThreadsScatter - nThreadsGather - nThreadsBcast;

四組執行緒分別負責:Scatter(把資料分散到各 rail)、Reduce(歸約後發給網路)、Gather(從各 rail 收集)、Bcast(從網路收到後廣播)。COLLNET_COPY_THREADS = 96(📎 src/device/all_reduce.h:250)是固定的拷貝執行緒數。

netRegUsed:網路註冊模式下的緩衝區佈局

📎 src/device/all_reduce.h:280-288有個關鍵分支:

code
if (work->netRegUsed) {
  offsetBase = bid * chunkSize;
  maxNelems = size;
  peerOffset = nChannels * chunkSize;
} else {
  offsetBase = bid * direct->nHeads * chunkSize;
  maxNelems = direct->nHeads * chunkSize;
  peerOffset = chunkSize;
}

netRegUsed模式下,緩衝區按 channel 連續排列(bid * chunkSize),peer 偏移是nChannels * chunkSize;非註冊模式下,按 head 排列(bid * nHeads * chunkSize),peer 偏移是chunkSize。這個差異源於網路註冊模式要求緩衝區連續,以便網卡 DMA。

NVLS 的 warp 分配

RunWorkColl<ncclFuncAllReduce, ..., NCCL_ALGO_NVLS, ...>的run(📎 src/device/all_reduce.h:391-523)用更精細的 warp 分配:

code
const int bcastWarps = hasOut ? (work->regUsed ? ((totalWarps - 2) >> 1) - 1 : 2) : 0;
const int reduceWarps = work->regUsed ? (totalWarps - bcastWarps - 2) : (hasOut ? 3 : nranks <= 6 ? 7 : 5);
const int scatterWarps = work->regUsed ? 1 : (totalWarps - reduceWarps - bcastWarps + 1) >> 1;
const int gatherWarps = work->regUsed ? 1 : (totalWarps - reduceWarps - bcastWarps) >> 1;

regUsed模式下,scatter/gather 各只佔 1 warp(因為 NVLS 硬體直接操作註冊記憶體),reduce 佔大頭;非註冊模式下,scatter/gather 各佔約一半,reduce 根據 rank 數調整(≤6 用 7 warp,否則 5 warp)。

時序互動圖

mermaid
sequenceDiagram
    participant App as 应用层
    participant Scatter as Scatter Warps
    participant NVLS as NVLS 硬件
    participant Reduce as Reduce Warps
    participant Bcast as Bcast Warps

    App->>Scatter: prims.scatter(offset, nelem, chunkSize)
    Scatter->>NVLS: 写入 NVLink SHARP 缓冲区
    NVLS->>NVLS: 硬件归约 (multimem)
    NVLS->>Reduce: prims.directRecvDirectSend(offset, nelem)
    Reduce->>NVLS: 归约结果写回
    NVLS->>Bcast: prims.directRecvDirectSend(offset, nelem)
    Bcast->>App: 广播到所有 rank

生產踩坑:CollNet 的direct->out == -1陷阱

📎 src/device/reduce_scatter.h:521有一行:

code
if (direct->out == -1) __trap();

如果 CollNet 的 out 連接未建立(-1),直接__trap()讓 kernel 崩潰。這是防禦性編程——CollNet 依賴網卡,如果網卡初始化失敗,out 會是 -1,此時繼續執行會導致未定義行為。生產環境如果看到 kernel trap,檢查 CollNet 網卡是否正常初始化。

10.5 Broadcast 與 Reduce:最簡單的兩個集合操作

Broadcast:從 root 扇出

broadcast.h的runRing(📎 src/device/broadcast.h:14-64)邏輯很直接:root 節點發資料,其他節點轉發,最後一個節點只收。

code
if (rank == root) {
  if (inputBuf == outputBuf || isNetOffload) {
    prims.directSend(offset, offset, nelem);
  } else {
    prims.directCopySend(offset, offset, nelem);
  }
} else if (nextRank == root) {
  prims.directRecv(offset, nelem);
} else {
  prims.directRecvCopyDirectSend(offset, offset, nelem);
}

三個分支:root 發、root 的前驅收、中間節點轉發。注意nextRank == root判斷的是「本節點的下一個是 root」,即本節點是環上最後一個——它只收不發。

Reduce:向 root 匯聚

reduce.h的runRing(📎 src/device/reduce.h:14-53)是 Broadcast 的逆操作:

code
if (prevRank == root) {
  prims.send(offset, nelem);
} else if (rank == root) {
  prims.recvReduceCopy(offset, offset, nelem, /*postOp=*/true);
} else {
  prims.recvReduceSend(offset, nelem);
}

prevRank == root的節點只發(它是 root 的前驅),root 只收並歸約,中間節點邊收邊歸約邊轉發。

設計思考:為什麼 Broadcast/Reduce 也用 Ring

Broadcast 和 Reduce 理論上可以用 Tree 實現更低延遲,但 NCCL 選擇 Ring 是因為:這兩個操作的資料量通常較小,Ring 的實現更簡單,且能複用 AllReduce 的 Ring 程式碼路徑。Tree 的複雜度(根節點選擇、執行緒拆分)在小訊息場景下收益不明顯。

生產踩坑:Broadcast 的 root 節點頻寬瓶頸

Broadcast 的 root 節點要發送全部資料,如果 root 是慢節點,整個 Broadcast 被拖慢。NCCL 的應對是:Broadcast 也支援多 channel,每個 channel 的 root 可以不同。但注意work->root是全局的,所有 channel 共享同一個 root——這是 Broadcast 的語義決定的(只有一個源)。生產環境如果 Broadcast 慢,檢查 root 節點的網路頻寬。

10.6 演算法選擇矩陣:RunWorkColl 模板特化

所有演算法核心透過RunWorkColl模板特化註冊(📎 src/device/all_reduce.h:228-788)。每個特化對應「函數 × 演算法 × 協議」的組合:

函數演算法協議特化位置
AllReduceRINGSIMPLE📎 src/device/all_reduce.h:230-233
AllReduceTREESIMPLE📎 src/device/all_reduce.h:238-244
AllReduceCOLLNET_DIRECTSIMPLE📎 src/device/all_reduce.h:249-386
AllReduceNVLSSIMPLE📎 src/device/all_reduce.h:391-523
AllReduceNVLS_TREESIMPLE📎 src/device/all_reduce.h:528-634
AllReduceCOLLNET_CHAINSIMPLE📎 src/device/all_reduce.h:639-759
AllReduceRINGLL📎 src/device/all_reduce.h:764-766
AllReduceTREELL📎 src/device/all_reduce.h:771-773
AllReduceRINGLL128📎 src/device/all_reduce.h:778-780
AllReduceTREELL128📎 src/device/all_reduce.h:785-787

注意:CollNet 和 NVLS 只支援 SIMPLE 協定。因為這兩種演算法依賴硬體卸載,而 LL/LL128 的低延遲同步機制與硬體卸載不相容——硬體歸約的延遲遠大於 LL 的 flag 輪詢,用 LL 反而增加開銷。

協定選擇的內在邏輯

  • LL:小訊息(< 8KB),低延遲優先。Ring 和 Tree 都支援。
  • LL128:中等訊息(8KB - 1MB),128 位元組對齊。Ring 和 Tree 都支援。
  • SIMPLE:大訊息(> 1MB),頻寬優先。所有演算法都支援。

生產踩坑:協定與演算法的組合限制

如果使用者強制指定NCCL_PROTO=LL但演算法是 CollNet,NCCL 會在 tuning 階段回退到 SIMPLE。生產環境如果發現協定設定不生效,檢查演算法是否支援該協定。

設計思考:為什麼同一份 AllReduce 邏輯需要這麼多實作

回顧本章,AllReduce 有 Ring、Tree、CollNet Direct、CollNet Chain、NVLS、NVLS Tree 六種演算法實作。這不是冗餘,而是針對不同硬體拓撲和訊息大小的最佳解:

  • Ring:通用,適合大訊息,頻寬利用率最高。
  • Tree:適合大規模叢集,延遲 O(log n)。
  • CollNet:適合有支援歸約的網卡的叢集,卸載 GPU 計算。
  • NVLS:適合單節點 NVLink 全連接,硬體多播歸約。

NCCL 的 tuning 模組(第 5 章)會根據訊息大小、rank 數、拓撲自動選擇。裝置側的實作只需要保證「每種組合都正確」,選擇邏輯在 host 側。

本章小結

本章拆解了src/device下的六個演算法核心檔案:

1. Ring AllReduce(📎 src/device/all_reduce.h:14-83):兩階段流水,reduce-scatter + all-gather,每階段 n-1 步。

2. Tree AllReduce(📎 src/device/all_reduce.h:86-225):樹形歸約,延遲 O(log n),runTreeSplit用執行緒拆分實現歸約-廣播流水。

3. AllGather(📎 src/device/all_gather.h:14-88):Ring 單階段,支援 in-place 和 netOffload。

4. ReduceScatter(📎 src/device/reduce_scatter.h:14-56):Ring 單階段,是 AllReduce 的 reduce-scatter 階段。

5. Broadcast/Reduce(📎 src/device/broadcast.h:14-64、📎 src/device/reduce.h:14-53):最簡單的 Ring 變體。

6. CollNet/NVLS(📎 src/device/all_reduce.h:247-635):硬體卸載,只支援 SIMPLE 協定。

本章思考與自測

Q1: 在 Ring AllReduce 的 reduce-scatter 階段,第 0 步用directSend,中間步用directRecvReduceDirectSend,最後一步用directRecvReduceCopyDirectSend。如果去掉最後一步的postOp=true,在什麼場景下會產生錯誤結果?

參考解析:postOp=true觸發後置操作(如求平均時的除法)。以ncclAvg為例,歸約是求和,postOp 是除以 nranks。如果去掉postOp,最後一步只做歸約不做除法,recvbuff 裡存的是「和」而非「平均」。在 reduce-scatter 階段,每個 rank 只保留一個 chunk 的最終結果,這個 chunk 恰好是ringIx+0(📎 src/device/all_reduce.h:60)。如果 postOp 缺失,這個 chunk 的和沒有除以 nranks,後續 all-gather 階段會把這個錯誤的「和」傳播給所有 rank。注意:只有最後一步需要 postOp,因為只有這一步產生「完整歸約」的結果;中間步的歸約是部分和,不需要 postOp。生產環境如果發現 AllReduce 結果偏大 nranks 倍,檢查 postOp 是否正確傳遞。

Q2: runTreeSplit在 LL/LL128 協定下把執行緒按 7:3 拆分(📎 src/device/all_reduce.h:163),而 Simple 協定下按 1:1 拆分(📎 src/device/all_reduce.h:157)。如果強行把 LL 協定也改成 1:1,會發生什麼?

參考解析:LL/LL128 的歸約組要從最多 3 個子節點收資料並做歸約(📎 src/device/all_reduce.h:187的FanAsymmetric<NCCL_MAX_TREE_ARITY, 1>),計算密集;廣播組只做複製轉發(📎 src/device/all_reduce.h:208的FanAsymmetric<1, NCCL_MAX_TREE_ARITY>),計算輕。7:3 拆分讓歸約組有足夠執行緒處理 3 路歸約,廣播組執行緒少但夠用。如果改成 1:1,歸約組執行緒不足,歸約成為瓶頸;廣播組執行緒過剩,浪費。更嚴重的是,LL 協定的 flag 輪詢是忙等待,執行緒多了會增加 flag 競爭。生產環境如果發現 Tree AllReduce 在 LL 協定下效能異常,檢查nthreadsSplit的計算是否被修改。

Q3: AllGather 的isNetOffload模式下,只用 1 個 warp 驅動 Ring 通訊(📎 src/device/all_gather.h:32),其餘 warp 並行拷貝(📎 src/device/all_gather.h:76-82)。如果去掉最後的barrier_sync(14, nthreads)(📎 src/device/all_gather.h:87),在什麼場景下會導致資料競爭?

參考解析:barrier_sync保證所有 warp(包括通訊 warp 和拷貝 warp)都完成本 work 後才進入下一個 work。如果去掉,通訊 warp 可能在拷貝 warp 還沒寫完 outputBuf 時就開始下一個 work 的通訊,而下一個 work 可能復用同一塊 outputBuf。具體場景:連續兩次 AllGather,第一次的拷貝 warp 還在寫 outputBuf 的尾部,第二次的通訊 warp 已經開始往 outputBuf 寫新數據,導致第一次的數據被覆蓋。註釋裡說得很清楚:「otherwise, we can have contention if next work will use the outputBuf in this work」。用 barrier 14 而非默認 barrier,是為了避開 prims 內部的 barrier 和__syncthreads(),防止死鎖。生產環境如果發現 AllGather 結果偶發錯誤,檢查isNetOffload路徑的 barrier 是否被優化掉。

至此,我們已經看完了設備側算法內核如何組織數據流。每種算法都通過Primitives調用上一章的原語,算法層只關心「誰發給誰、發哪個 chunk、歸約還是複製」。下一章將深入傳輸層抽象,看 P2P、SHM、NET、NVLS 如何統一成一套接口,以及 host 側的 proxy 線程如何與設備側 kernel 協作完成跨機通信。

核心規律:所有算法都通過 Primitives 模板類調用原語,算法只負責「數據流拓撲」,原語負責「數據搬運」。這種分層讓新增算法只需實現拓撲邏輯,無需關心底層同步。但無論拓撲如何變化,數據最終都要通過物理鏈路傳輸。下一章將深入 src/transport 目錄,看 NCCL 如何用統一的 transport 接口屏蔽 P2P、SHM、NET、NVLS 的差異,以及每種 transport 的 setup/connect/send/recv 語義。這是理解跨機通信的基礎。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 11

第 11 章:傳輸層抽象:P2P、SHM、NET、NVLS 如何統一在同一套接口下

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 11 章 / 共 25 章

上一章我們深入算法內核,看到 Ring AllReduce 如何將數據切分後兩階段歸約,Tree AllReduce 又如何借助樹形結構壓低延遲——但這些算法只定義了「誰發給誰、發哪個 chunk」的邏輯視圖。數據最終必須穿過真實的物理鏈路:NVLink、PCIe、共享內存或網卡。本章拆解 src/transport 目錄,看 NCCL 如何用一套統一的 ncclTransport 接口,把 P2P、SHM、NET、NVLS 四種物理通道屏蔽成同一副面孔,完成從算法拓撲到物理傳輸的最後一公里。

一、統一接口:ncclTransport 如何屏蔽四種物理通道

直覺模型

想像一家物流公司:無論客戶寄的是同城快遞(P2P)、樓內傳遞(SHM)、跨省運輸(NET)還是專線直達(NVLS),前台只填一張「運單」。這張運單就是ncclTransport結構體——它規定了每種運輸方式必須提供canConnect、setup、connect、free等固定動作。若沒有這層抽象,上層算法就得寫四套if-else判斷走哪條鏈路,新增一種硬件就要改遍所有算法。

數據結構與內存佈局

NCCL 用一個全局數組登記所有 transport,順序即優先級:

📎 src/transport.cc:15-20

c
struct ncclTransport* ncclTransports[NTRANSPORTS] = {
  &p2pTransport,
  &shmTransport,
  &netTransport,
  &collNetTransport,
};

數組順序決定選擇順序:P2P 優先,其次 SHM,再次 NET,最後 CollNet。每種 transport 由ncclTransport結構體描述,它包含一個canConnect函數指針和兩個ncclTransportComm(send/recv 各一個)。以 P2P 為例:

📎 src/transport/p2p.cc:1493-1498

c
struct ncclTransport p2pTransport = {"P2P",
                                     p2pCanConnect,
                                     {p2pSendSetup, p2pSendConnect, p2pSendFree, NULL, p2pSendProxySetup, NULL,
                                      p2pSendProxyFree, NULL, p2pProxyRegister, p2pProxyDeregister},
                                     {p2pRecvSetup, p2pRecvConnect, p2pRecvFree, NULL, p2pRecvProxySetup, NULL,
                                      p2pRecvProxyFree, NULL, p2pProxyRegister, p2pProxyDeregister}};

ncclTransportComm的字段順序是固定的「生命週期槽位」:setup(準備資源)、connect(交換連接信息)、free(釋放)、proxySharedInit(代理共享初始化)、proxySetup、proxyConnect、proxyFree、proxyProgress、proxyRegister、proxyDeregister。注意 P2P 的proxyProgress槽位是NULL——因為 P2P 走的是 GPU 直接讀寫對端顯存,不需要 host 代理線程搬運數據;而 NET 的proxyProgress是sendProxyProgress/recvProxyProgress,因為網卡 I/O 必須由 host 線程驅動。

場景驅動 Walkthrough:一次連接如何選中 transport

當 NCCL 需要為某個 channel 的某個 peer 建立連接時,調用selectTransport:

📎 src/transport.cc:23-44

c
template <int type>
static ncclResult_t selectTransport(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclConnect* connect,
                                    int channelId, int peer, int connIndex, int* transportType) {
  struct ncclPeerInfo* myInfo = comm->peerInfo + comm->rank;
  struct ncclPeerInfo* peerInfo = comm->peerInfo + peer;
  struct ncclConnector* connector = (type == 1) ? comm->channels[channelId].peers[peer]->send + connIndex :
                                                  comm->channels[channelId].peers[peer]->recv + connIndex;
  for (int t = 0; t < NTRANSPORTS; t++) {
    struct ncclTransport* transport = ncclTransports[t];
    struct ncclTransportComm* transportComm = type == 1 ? &transport->send : &transport->recv;
    int ret = 0;
    NCCLCHECK(transport->canConnect(&ret, comm, graph, myInfo, peerInfo));
    if (ret) {
      connector->transportComm = transportComm;
      NCCLCHECK(transportComm->setup(comm, graph, myInfo, peerInfo, connect, connector, channelId, connIndex));
      if (transportType) *transportType = t;
      return ncclSuccess;
    }
  }
  WARN("No transport found for rank %d[%lx] -> rank %d[%lx]", myInfo->rank, myInfo->busId, peerInfo->rank,
       peerInfo->busId);
  return ncclSystemError;
}

type==1表示 send 方向,type==0表示 recv 方向。迴圈依次詢問每種 transport 的canConnect:回傳ret=1表示「我能幹這活」,立刻把connector->transportComm指向該 transport 的對應方向,並呼叫其setup。若所有 transport 都回傳 0,列印警告並回傳ncclSystemError。

canConnect的判定邏輯體現了各 transport 的「領地邊界」。以 P2P 為例:

📎 src/transport/p2p.cc:129-157

c
ncclResult_t p2pCanConnect(int* ret, struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* info1,
                           struct ncclPeerInfo* info2) {
  initCeOperation();
  int intermediateRank;
  int isCrossClique;
  NCCLCHECK(ncclTopoCheckP2p(comm, comm->topo, info1->rank, info2->rank, ret, NULL, &intermediateRank, NULL,
                             &isCrossClique));
  if (*ret == 0) return ncclSuccess;
  if (intermediateRank != -1) {
    if (useMemcpy) *ret = 0;
    return ncclSuccess;
  }
  if (!isCrossClique) {
    int useNet = 0;
    NCCLCHECK(ncclTopoCheckNet(comm->topo, info1->rank, info2->rank, &useNet));
    if (useNet) {
      *ret = 0;
      return ncclSuccess;
    }
  }
  if (info1->hostHash != comm->peerInfo[comm->rank].hostHash || info1->hostHash != info2->hostHash) {
    return ncclSuccess;
  }
  ...

P2P 的判定鏈:先問拓撲「兩個 rank 之間有沒有 P2P 路徑」;如果有中間跳(intermediateRank != -1)且啟用了 CE memcpy,則放棄 P2P 讓給 SHM/NET;如果拓撲建議走網路(useNet),也放棄;最後檢查是否同主機。SHM 的判定更簡單:

📎 src/transport/shm.cc:61-83

c
static ncclResult_t shmCanConnect(int* ret, struct ncclComm* comm, struct ncclTopoGraph* graph,
                                  struct ncclPeerInfo* info1, struct ncclPeerInfo* info2) {
  *ret = 0;
  initShmLocality();
  if (ncclParamShmDisable() == 1) return ncclSuccess;
  int useNet = 0;
  NCCLCHECK(ncclTopoCheckNet(comm->topo, info1->rank, info2->rank, &useNet));
  if (useNet) return ncclSuccess;
  if (info1->hostHash != info2->hostHash) return ncclSuccess;
  if (info1->shmDev != info2->shmDev) return ncclSuccess;
  *ret = 1;
  return ncclSuccess;
}

SHM 要求同主機(hostHash相同)且共享同一塊/dev/shm(shmDev相同,用於容器間通訊)。NET 則幾乎總是回傳 1,只在同主機時檢查 intra-node net 是否被禁用:

📎 src/transport/net.cc:160-168

c
static ncclResult_t canConnect(int* ret, struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* info1,
                               struct ncclPeerInfo* info2) {
  *ret = 1;
  if (info1->hostHash == info2->hostHash) {
    NCCLCHECK(ncclTopoCheckNet(comm->topo, info1->rank, info2->rank, ret));
  }
  return ncclSuccess;
}

NET 是「兜底」——只要前面沒人接,它就接。NVLS 的canConnect直接回傳 0:

📎 src/transport/nvls.cc:21-26

c
ncclResult_t nvlsCanConnect(int* ret, struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* info1,
                            struct ncclPeerInfo* info2) {
  // This transport cannot be used for p2p
  *ret = 0;
  return ncclSuccess;
}

NVLS 不走常規的 peer-to-peer 連接路徑,它透過ncclNvlsSetup單獨建立多播組,所以canConnect永遠回傳 0。

mermaid
flowchart TD
    start["selectTransport(comm, peer, connIndex)"] --> loop{"遍历 ncclTransports[t]"}
    loop -->|t=0| p2p["p2pCanConnect()"]
    p2p --> p2p_chk{"拓扑有P2P路径<br/>且非中间跳<br/>且同主机?"}
    p2p_chk -->|是| use_p2p["connector->transportComm = p2pTransport<br/>调用 p2pSendSetup/p2pRecvSetup"]
    p2p_chk -->|否| shm["shmCanConnect()"]
    shm --> shm_chk{"同hostHash<br/>且同shmDev?"}
    shm_chk -->|是| use_shm["connector->transportComm = shmTransport<br/>调用 shmSendSetup/shmRecvSetup"]
    shm_chk -->|否| net["canConnect() (NET)"]
    net --> net_chk{"同主机时<br/>intra-node net 启用?"}
    net_chk -->|是/跨机| use_net["connector->transportComm = netTransport<br/>调用 sendSetup/recvSetup"]
    net_chk -->|否| collnet["collNetTransport"]
    collnet --> fail["WARN: No transport found<br/>return ncclSystemError"]
    use_p2p --> done["return ncclSuccess"]
    use_shm --> done
    use_net --> done

設計思考

〔設計推斷與架構權衡〕

為什麼用「陣列順序 + canConnect 投票」而不是顯式路由表? 因為拓撲是動態的:同一台機器可能因NCCL_P2P_DISABLE、容器隔離、CUDA IPC 可用性等因素導致 P2P 不可用,此時自動降級到 SHM 或 NET。投票機制讓每種 transport 自己判斷「我能不能幹」,新增 transport 只需在陣列裡加一項,無需改動選擇邏輯。這正是開閉原則在系統程式設計中的體現。

二、P2P:同機 GPU 直連的四種形態

直覺模型

P2P 是「鄰居之間直接遞東西」——GPU 0 直接讀寫 GPU 1 的顯存,不經過 CPU 或網卡。若沒有 P2P,同機多卡通訊就得繞道 host 記憶體,延遲翻倍、頻寬腰斬。

資料結構與記憶體佈局

P2P 內部有四種形態,由enum p2pType區分:

📎 src/transport/p2p.cc:19-24

c
enum p2pType {
  P2P_DIRECT,
  P2P_INTERMEDIATE,
  P2P_IPC,
  P2P_CUMEM
};
  • P2P_DIRECT:同進程內不同 GPU,直接用指標存取(最快)。
  • P2P_INTERMEDIATE:兩 GPU 之間沒有直連,需經中間 GPU 轉發。
  • P2P_IPC:跨進程,用傳統cudaIpcOpenMemHandle匯入對端顯存。
  • P2P_CUMEM:跨進程,用 cuMem API(cuMemExportToShareableHandle)匯入,支援更細粒度的記憶體管理。

核心資源結構體:

📎 src/transport/p2p.cc:79-94

c
struct p2pResources {
  enum p2pType type;
  union {
    struct ncclSendMem* sendDevMem;
    struct ncclRecvMem* recvDevMem;
  };
  void* sendMemIpc;
  int sendMemSameProc;
  void* recvMemIpc;
  int recvMemSameProc;
  // CE memcpy support
  struct p2pShmProxyInfo proxyInfo;
  struct p2pShm* shm;
  struct p2pShm* devShm;
  ncclShmIpcDesc_t desc;
};

sendDevMem/recvDevMem是 union——發送方只關心sendDevMem,接收方只關心recvDevMem,共用一塊記憶體。sendMemIpc/recvMemIpc保存匯入的對端記憶體句柄,sendMemSameProc/recvMemSameProc標記是否同進程(決定釋放時用ncclCuMemFreeAddr還是cudaIpcCloseMemHandle)。

連接資訊結構體p2pConnectInfo透過 bootstrap 交換:

📎 src/transport/p2p.cc:38-44

c
struct p2pConnectInfo {
  int rank;
  int read;
  struct ncclP2pBuff p2pBuff;
  // Used by CE memcpy
  ncclShmIpcDesc_t desc;
};
static_assert(sizeof(struct p2pConnectInfo) <= CONNECT_SIZE, "p2pConnectInfo is too large");

static_assert保證連接資訊不超過CONNECT_SIZE(bootstrap 單次交換的固定緩衝區大小)。read欄位決定資料流向:read=1表示接收方主動讀發送方顯存(P2P Read),read=0表示發送方主動寫接收方顯存(P2P Write)。

場景驅動 Walkthrough:P2P Send 的建立

當selectTransport選中 P2P 後,呼叫p2pSendSetup:

📎 src/transport/p2p.cc:393-471

c
ncclResult_t p2pSendSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* myInfo,
                          struct ncclPeerInfo* peerInfo, struct ncclConnect* connectInfo, struct ncclConnector* send,
                          int channelId, int connIndex) {
  struct p2pResources* resources;
  struct ncclP2pRequest req;
  NCCLCHECK(ncclCalloc(&resources, 1));
  send->transportResources = resources;
  int useRead, intermediateRank;
  NCCLCHECK(p2pGetInfo(comm, myInfo, peerInfo, &useRead, &intermediateRank));
  if (useMemcpy) useRead = 0;
  ...
  int sendSize = sizeof(struct ncclSendMem);
  if (info->read) sendSize += comm->buffSizes[NCCL_PROTO_SIMPLE];
  ALIGN_SIZE(sendSize, CUDA_IPC_MIN);
  ...

關鍵點:sendSize在 P2P Read 模式下要額外加上 SIMPLE 協定緩衝區大小——因為讀模式下發送方的 SIMPLE buffer 被接收方直接讀取,必須和ncclSendMem一起分配在同一塊可共享記憶體裡。ALIGN_SIZE(sendSize, CUDA_IPC_MIN)保證大小對齊到 CUDA IPC 最小粒度。

接著根據intermediateRank和進程關係選擇形態:

📎 src/transport/p2p.cc:416-437

c
  if (intermediateRank == -1) {
    info->rank = myInfo->rank;
    if (P2P_SAME_PID(myInfo, peerInfo) && ncclParamP2pDirectDisable() == 0 && useMemcpy == 0) {
      resources->type = P2P_DIRECT;
      ...
    } else {
      if (ncclCuMemEnable()) {
        resources->type = P2P_CUMEM;
        ...
      } else {
        resources->type = P2P_IPC;
        ...
      }
    }
    send->conn.flags |= info->read ? NCCL_P2P_READ : NCCL_P2P_WRITE;
  } else {
    resources->type = P2P_INTERMEDIATE;
    info->rank = intermediateRank;
    ...
  }

P2P_SAME_PID巨集判斷同主機同進程:

📎 src/transport/p2p.cc:334-335

c
#define P2P_SAME_PID(MYINFO, PEERINFO) \
  ((MYINFO->hostHash == PEERINFO->hostHash) && (MYINFO->pidHash == PEERINFO->pidHash))

同進程且未禁用 direct 且未啟用 memcpy,就是最快的P2P_DIRECT——直接拿對端指標。否則走 IPC/CUMEM。

隨後透過代理執行緒分配可共享緩衝區:

📎 src/transport/p2p.cc:457-468

c
  NCCLCHECK(ncclProxyConnect(comm, TRANSPORT_P2P, 1, info->rank, &send->proxyConn));
  if (useMemcpy) {
    NCCLCHECK(ncclProxyCallBlocking(comm, &send->proxyConn, ncclProxyMsgSetup, NULL, 0, &resources->proxyInfo,
                                    sizeof(struct p2pShmProxyInfo)));
    memcpy(&info->desc, &resources->proxyInfo.desc, sizeof(ncclShmIpcDesc_t));
  } else {
    NCCLCHECK(ncclProxyCallBlocking(comm, &send->proxyConn, ncclProxyMsgSetup, &req, sizeof(struct ncclP2pRequest),
                                    &info->p2pBuff, sizeof(struct ncclP2pBuff)));
    NCCLCHECK(p2pMap(comm, &send->proxyConn, myInfo, comm->peerInfo + info->rank, &info->p2pBuff,
                     (void**)&resources->sendDevMem, &resources->sendMemIpc));
    resources->sendMemSameProc = P2P_SAME_PID(myInfo, (comm->peerInfo + info->rank));
  }

ncclProxyCallBlocking是同步 RPC:host 執行緒發訊息給代理執行緒,代理執行緒呼叫p2pSendProxySetup分配可共享緩衝區,把ncclP2pBuff(含 IPC 句柄)回傳。然後p2pMap把對端緩衝區映射到本地位址空間。

p2pMap是核心映射函式:

📎 src/transport/p2p.cc:349-390

c
static ncclResult_t p2pMap(struct ncclComm* comm, struct ncclProxyConnector* proxyConn, struct ncclPeerInfo* myInfo,
                           struct ncclPeerInfo* peerInfo, struct ncclP2pBuff* p2pBuff, void** devMem, void** ipcPtr) {
  if (P2P_SAME_PID(myInfo, peerInfo)) {
    if (peerInfo->cudaDev != myInfo->cudaDev) {
      cudaError_t err = cudaDeviceEnablePeerAccess(peerInfo->cudaDev, 0);
      ...
      if (ncclCuMemEnable()) {
        NCCLCHECK(ncclCuMemAllocAddr(devMem, &p2pBuff->ipcDesc.memHandle, p2pBuff->size));
        CUCHECK(cuMemRelease(p2pBuff->ipcDesc.memHandle));
        *ipcPtr = *devMem;
        ...
      } else {
        *devMem = p2pBuff->directPtr;
        *ipcPtr = NULL;
      }
    } else {
      *devMem = p2pBuff->directPtr;
      *ipcPtr = NULL;
    }
  } else {
    NCCLCHECK(ncclP2pImportShareableBuffer(comm, peerInfo->rank, p2pBuff->size, &p2pBuff->ipcDesc, devMem,
                                           p2pBuff->directPtr, ncclMemOffload));
    *ipcPtr = *devMem;
  }
  return ncclSuccess;
}

同進程不同 GPU:先cudaDeviceEnablePeerAccess打開 P2P 通道,然後直接用directPtr(因為同進程位址空間共享)。跨進程:呼叫ncclP2pImportShareableBuffer匯入對端記憶體句柄。

並發控制與硬體互動

P2P 的同步靠ncclSendMem/ncclRecvMem裡的head/tail指標。發送方寫head告訴接收方「我寫到哪了」,接收方寫tail告訴發送方「我讀到哪了」。這是典型的無鎖生產者-消費者:

📎 src/transport/p2p.cc:571-576

c
  } else {
    send->conn.tail = &remDevMem->tail;
    send->conn.head = &resources->sendDevMem->head;
    send->conn.ptrExchange = &resources->sendDevMem->ptrExchange;
    send->conn.redOpArgExchange = resources->sendDevMem->redOpArgExchange;
  }

head指向本地sendDevMem,tail指向對端remDevMem。GPU kernel 透過讀寫這兩個指標實現跨 GPU 同步,無需 CPU 介入。

生產避坑指南

坑 1:P2P Read 與 memcpy 互斥。看p2pSendConnect:

📎 src/transport/p2p.cc:551-559

c
  for (int p = 0; p < NCCL_NUM_PROTOCOLS; p++) {
    if (info->read && p == NCCL_PROTO_SIMPLE) {
      /* For P2P Read the SIMPLE buffer is local (ncclSendMem) */
      if (resources->sendDevMem == NULL) return ncclInternalError; // We should not use read + memcpy
      send->conn.buffs[p] = (char*)(resources->sendDevMem + 1);
    } else {
      send->conn.buffs[p] = buff;
      buff += comm->buffSizes[p];
    }
  }

如果read=1但sendDevMem==NULL,直接回傳ncclInternalError。生產環境若看到這個錯誤,檢查是否同時設定了NCCL_P2P_READ_ENABLE=1和NCCL_P2P_USE_CUDA_MEMCPY=1——這兩者語義衝突。

坑 2:跨進程釋放順序。 p2pSendFree根據sendMemSameProc決定釋放方式:

📎 src/transport/p2p.cc:624-651

c
ncclResult_t p2pSendFree(struct ncclComm* comm, struct ncclConnector* send) {
  struct p2pResources* resources = (struct p2pResources*)send->transportResources;
  if (resources) {
    if (ncclCuMemEnable()) {
      if (resources->sendMemIpc) {
        if (resources->sendMemSameProc) {
          NCCLCHECK(ncclCuMemFreeAddr(resources->sendMemIpc, comm->memManager));
        } else {
          NCCLCHECK(ncclCudaFree(resources->sendMemIpc, comm->memManager));
        }
      }
      ...

同進程用ncclCuMemFreeAddr(只釋放位址映射,不釋放實體記憶體),跨進程用ncclCudaFree(釋放實體記憶體)。搞反會導致記憶體洩漏或 use-after-free。

三、SHM:共享記憶體的「誰出記憶體」之爭

直覺模型

SHM 是「兩個行程共用一塊白板」——發送方寫,接收方讀。但白板放誰家?放發送方家(sender-side),接收方跑過來讀;還是放接收方家(receiver-side),發送方跑過去寫?這就是NCCL_SHM_LOCALITY參數要解決的問題。

資料結構與記憶體佈局

📎 src/transport/shm.cc:28-34

c
struct shmSendResources {
  struct ncclRecvMem* remHostMem;
  struct ncclRecvMem* devRemHostMem;
  ncclShmIpcDesc_t remDesc;
  struct ncclSendMem* hostMem;
  struct ncclSendMem* devHostMem;
};

struct shmRecvResources {
  struct ncclSendMem* remHostMem;
  struct ncclSendMem* devRemHostMem;
  ncclShmIpcDesc_t remDesc;
  struct ncclRecvMem* hostMem;
  struct ncclRecvMem* devHostMem;
};

注意hostMem和devHostMem成對出現:hostMem是 host 側指標,devHostMem是裝置側指標(透過 UVA 或 cuMem 映射)。remHostMem/devRemHostMem是對端共享記憶體的本地映射。

場景驅動 Walkthrough:SHM 的 locality 選擇

shmSendSetup根據 locality 決定分配多大記憶體:

📎 src/transport/shm.cc:88-119

c
static ncclResult_t shmSendSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* myInfo,
                                 struct ncclPeerInfo* peerInfo, struct ncclConnect* connectInfo,
                                 struct ncclConnector* send, int channelId, int connIndex) {
  struct shmSendResources* resources;
  struct shmConnectInfo* info = (struct shmConnectInfo*)connectInfo;
  size_t shmSize = sizeof(struct ncclSendMem);
  struct shmRequest req;

  NCCLCHECK(ncclCalloc(&resources, 1));
  send->transportResources = resources;

  if (shmLocality == SHM_SEND_SIDE) {
    for (int p = 0; p < NCCL_NUM_PROTOCOLS; p++) shmSize += comm->buffSizes[p];
  }
  req.size = shmSize;
  if (myInfo->hostHash == peerInfo->hostHash && myInfo->pidHash == peerInfo->pidHash) req.legacy = true;
  else req.legacy = false;

  NCCLCHECK(ncclProxyConnect(comm, TRANSPORT_SHM, 1, myInfo->rank, &send->proxyConn));
  NCCLCHECK(ncclProxyCallBlocking(comm, &send->proxyConn, ncclProxyMsgSetup, (void*)&req, sizeof(struct shmRequest),
                                  (void*)info, sizeof(struct shmConnectInfo)));

  info->rank = comm->rank;
  resources->hostMem = (struct ncclSendMem*)info->buf.hptr;
  resources->devHostMem = (struct ncclSendMem*)info->buf.dptr;
  ...

shmLocality == SHM_SEND_SIDE時,發送方分配資料緩衝區(shmSize加上所有協定緩衝區);否則只分配ncclSendMem控制結構。req.legacy標記是否同行程——同行程可以用傳統mmap,跨行程需要 cuMem 或/dev/shm檔案。

shmSendConnect裡根據 locality 決定buffs指向本地還是對端:

📎 src/transport/shm.cc:153-176

c
static ncclResult_t shmSendConnect(struct ncclComm* comm, struct ncclConnect* connectInfo, int nranks, int rank,
                                   struct ncclConnector* send) {
  struct shmConnectInfo* info = (struct shmConnectInfo*)connectInfo;
  struct shmSendResources* resources = (struct shmSendResources*)send->transportResources;
  char* buff;

  NCCLCHECK(ncclShmImportShareableBuffer(comm, info->rank, &info->desc, (void**)&resources->remHostMem,
                                         (void**)&resources->devRemHostMem, &resources->remDesc));

  buff = shmLocality == SHM_SEND_SIDE ? (char*)(resources->devHostMem + 1) : (char*)(resources->devRemHostMem + 1);
  for (int p = 0; p < NCCL_NUM_PROTOCOLS; p++) {
    send->conn.buffs[p] = buff;
    buff += comm->buffSizes[p];
  }
  send->conn.tail = &resources->devRemHostMem->tail;
  send->conn.head = &resources->devHostMem->head;
  send->conn.stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE] / NCCL_STEPS;
  ...

SHM_SEND_SIDE:buffs指向本地devHostMem(發送方寫自己的記憶體);SHM_RECV_SIDE:buffs指向對端devRemHostMem(發送方寫接收方的記憶體)。head永遠指向本地,tail永遠指向對端——因為發送方更新head,接收方更新tail。

設計思考

〔設計推斷與架構權衡〕

為什麼預設SHM_RECV_SIDE? 因為接收方通常需要把資料從共享記憶體拷貝到自己的 GPU 顯存,如果共享記憶體在接收方本地,拷貝路徑更短(本地記憶體 → 本地 GPU),避免跨 NUMA 存取。發送方寫遠端記憶體雖然多一次跨節點寫,但發送方通常是計算密集的 GPU,寫操作可以非同步進行。

生產避坑指南

坑:容器間/dev/shm不共享。 shmCanConnect檢查info1->shmDev != info2->shmDev:

📎 src/transport/shm.cc:76-78

c
  TRACE(NCCL_INIT | NCCL_SHM, "peer1 shmDev %lx peer2 shmDev %lx", info1->shmDev, info2->shmDev);
  if (info1->shmDev != info2->shmDev) return ncclSuccess;

如果兩個容器掛載了不同的/dev/shm,shmDev不同,SHM 自動降級到 NET。生產環境若發現同主機通訊卻走了網路,檢查容器的/dev/shm掛載是否一致。

四、NET:網路傳輸的映射表與代理進度

直覺模型

NET 是「跨城快遞」——資料打包交給網卡,網卡透過光纖送到對端。但網卡不認識 GPU 顯存位址,需要一張「位址映射表」把 GPU 虛擬位址翻譯成網卡能理解的實體位址。這張表就是connectMap。

資料結構與記憶體佈局

📎 src/transport/net.cc:73-86

c
struct connectMapMem {
  char* gpuPtr;
  char* cpuPtr;
  ssize_t size;
  ncclIpcDesc ipcDesc;
  ncclShmIpcDesc_t attachDesc;
  ncclShmIpcDesc_t createDesc;
};

struct connectMap {
  int sameProcess;
  int shared;
  int cudaDev;
  // First 3 bits of offsets determine the mem bank. 001 is host mem, 011 is dev mem, 101 is shared host mem and 111
  // is shared dev mem.
  struct connectMapMem mems[NCCL_NET_MAP_MEMS];
  // Offsets. 3 MSBs indicate mem bank, 111 indicates NULL.
  struct {
    uint32_t sendMem;
    uint32_t recvMem;
    uint32_t buffs[NCCL_NUM_PROTOCOLS];
  } offsets;
};

connectMap是一個「記憶體銀行」系統:mems陣列有 5 個槽位(NCCL_NET_MAP_MEMS=5),分別對應 host mem、dev mem、shared host mem、shared dev mem、GDC mem。offsets裡的每個欄位是一個 32 位元整數,高 3 位元編碼「哪個銀行」,低 29 位元編碼「銀行內偏移」。

解碼巨集:

📎 src/transport/net.cc:36-46

c
#define NCCL_NET_MAP_OFFSET_BANK(mapStruct, offsetName) ((mapStruct)->offsets.offsetName >> 30)

#define NCCL_NET_MAP_OFFSET_NULL(mapStruct, offsetName) (((mapStruct)->offsets.offsetName >> 29) == 0)

#define NCCL_NET_MAP_GET_POINTER(mapStruct, cpuOrGpu, offsetName) \
  (NCCL_NET_MAP_OFFSET_NULL(mapStruct, offsetName) ? \
     NULL : \
     (mapStruct)->mems[NCCL_NET_MAP_OFFSET_BANK(mapStruct, offsetName)].cpuOrGpu##Ptr + \
       ((mapStruct)->offsets.offsetName & NCCL_NET_MAP_MASK_OFFSET))

#define NCCL_NET_MAP_DEV_MEM(mapStruct, offsetName) (((mapStruct)->offsets.offsetName & NCCL_NET_MAP_MASK_DEVMEM) != 0)

NCCL_NET_MAP_GET_POINTER(map, gpu, sendMem)展開後:取offsets.sendMem的高 2 位元作為 bank 索引,從mems[bank].gpuPtr加上低 29 位元偏移,得到實際指標。這套編碼把「哪個記憶體區域 + 區域內偏移」壓縮進一個 32 位元整數,節省了connectMap的傳輸大小。

場景驅動 Walkthrough:sendProxyConnect 的映射建立

sendProxyConnect是 NET 最複雜的函式,負責建立網卡連線、分配緩衝區、註冊記憶體:

📎 src/transport/net.cc:858-1041

c
static ncclResult_t sendProxyConnect(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState,
                                     void* reqBuff, int reqSize, void* respBuff, int respSize, int* done) {
  struct sendNetResources* resources = (struct sendNetResources*)(connection->transportResources);
  ...
  if (resources->shared) {
    // Shared buffers
    ...
    if (resources->maxRecvs > 1 && ncclParamNetSharedComms()) {
      // Connect or reuse connection for a netdev/remote rank.
      ...
      if (comms->sendComm[resources->channelId] == NULL &&
          comms->activeConnect[resources->channelId] == (resources->tpLocalRank + 1)) {
        ret = proxyState->ncclNet->connect(proxyState->netContext, resources->netDev, req->handle,
                                           comms->sendComm + resources->channelId, &resources->netDeviceHandle);
      }
      ...

maxRecvs > 1時啟用「共享連線」:多個 channel 復用同一個網卡連線,減少連線數。activeConnect陣列保證只有一個 local rank 發起連線,避免重複。

接著分配緩衝區並註冊:

📎 src/transport/net.cc:933-956

c
  if (resources->shared == 0) {
    // Only allocate dedicated buffers for ring/tree, not for p2p
    for (int p = 0; p < NCCL_NUM_PROTOCOLS; p++) {
      NCCL_NET_MAP_ADD_POINTER(map, 0, p != NCCL_PROTO_LL && resources->useGdr ? 1 : 0, proxyState->buffSizes[p],
                               buffs[p]);
      resources->buffSizes[p] = proxyState->buffSizes[p];
    }
  } else {
    // Get shared buffers
    int bank = resources->useGdr ? NCCL_NET_MAP_SHARED_DEVMEM : NCCL_NET_MAP_SHARED_HOSTMEM;
    struct connectMapMem* mapMem = map->mems + bank;
    NCCLCHECK(sharedNetBuffersInit(proxyState, resources->useGdr, resources->tpLocalRank, 0, map->sameProcess,
                                   proxyState->p2pnChannels, &mapMem->gpuPtr, &mapMem->cpuPtr, &mapMem->size,
                                   &mapMem->ipcDesc));
    resources->buffSizes[NCCL_PROTO_SIMPLE] = mapMem->size;
    ...

NCCL_NET_MAP_ADD_POINTER巨集把緩衝區登記到connectMap:

📎 src/transport/net.cc:48-62

c
#define NCCL_NET_MAP_ADD_POINTER(mapStruct, shared, dev, memSize, offsetName) \
  do { \
    int bank = NCCL_NET_MAP_MASK_USED + (dev) * NCCL_NET_MAP_MASK_DEVMEM + (shared) * NCCL_NET_MAP_MASK_SHARED; \
    if ((shared) == 0) { \
      if (dev) { \
        (mapStruct)->offsets.offsetName = bank + (mapStruct)->mems[NCCL_NET_MAP_DEVMEM].size; \
        (mapStruct)->mems[NCCL_NET_MAP_DEVMEM].size += memSize; \
      } else { \
        (mapStruct)->offsets.offsetName = bank + (mapStruct)->mems[NCCL_NET_MAP_HOSTMEM].size; \
        (mapStruct)->mems[NCCL_NET_MAP_HOSTMEM].size += memSize; \
      } \
    } else { \
      (mapStruct)->offsets.offsetName = bank; \
    } \
  } while (0);

非共享緩衝區:把當前 bank 的size作為偏移寫入offsets,然後size += memSize——這是 bump allocator。共享緩衝區:直接寫 bank 編號,偏移為 0(因為共享緩衝區整塊就是一個 bank)。

最後註冊記憶體給網卡:

📎 src/transport/net.cc:1004-1035

c
  for (int p = 0; p < NCCL_NUM_PROTOCOLS; p++) {
    resources->buffers[p] = NCCL_NET_MAP_GET_POINTER(map, cpu, buffs[p]);
    if (resources->buffers[p]) {
#if CUDA_VERSION >= 11070
      int type = NCCL_NET_MAP_DEV_MEM(map, buffs[p]) ? NCCL_PTR_CUDA : NCCL_PTR_HOST;
      if (type == NCCL_PTR_CUDA && resources->useDmaBuf) {
        int dmabuf_fd;
        size_t dmaBufSize = resources->buffSizes[p];
        ALIGN_SIZE(dmaBufSize, ncclOsGetPageSize());
        CUCHECK(cuMemGetHandleForAddressRange((void*)&dmabuf_fd, (CUdeviceptr)resources->buffers[p], dmaBufSize,
                                              CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD,
                                              getHandleForAddressRangeFlags(resources->useGdr)));
        NCCLCHECK(proxyState->ncclNet->regMrDmaBuf(resources->netSendComm, resources->buffers[p],
                                                   resources->buffSizes[p], type, 0ULL, dmabuf_fd,
                                                   &resources->mhandles[p]));
        (void)close(dmabuf_fd);
      } else
#endif
      {
        NCCLCHECK(proxyState->ncclNet->regMr(resources->netSendComm, resources->buffers[p], resources->buffSizes[p],
                                             NCCL_NET_MAP_DEV_MEM(map, buffs[p]) ? NCCL_PTR_CUDA : NCCL_PTR_HOST,
                                             &resources->mhandles[p]));
      }
      ...

優先走 DMA-BUF 路徑(cuMemGetHandleForAddressRange拿到 fd,傳給網卡外掛),失敗則回退到regMr(傳統 nv_peermem GDR)。

並發控制與硬體互動:sendProxyProgress 的三段式流水線

sendProxyProgress是 NET 的資料搬運引擎,採用「post → transmit → done」三段式:

📎 src/transport/net.cc:1324-1491

c
static ncclResult_t sendProxyProgress(struct ncclProxyState* proxyState, struct ncclProxyArgs* args) {
  ...
  if (args->state == ncclProxyOpProgress) {
    int p = args->protocol;
    int maxDepth = std::min(NCCL_STEPS, NCCL_SHARED_STEPS / args->nsubs);
    for (int s = 0; s < args->nsubs; s++) {
      struct ncclProxySubArgs* sub = args->subs + s;
      ...
      // Post buffers to the GPU
      if (sub->posted < sub->nsteps && sub->posted < sub->done + maxDepth) {
        ...
        if (resources->shared) {
          ...
          volatile uint64_t* sendHead = resources->gdcSync ? resources->gdcSync : &resources->sendMem->head;
          sub->posted += args->sliceSteps;
          *sendHead = sub->base + sub->posted - NCCL_STEPS;
          if (resources->gdcSync) wc_store_fence(); // Flush out WC write
        } else {
          sub->posted += args->sliceSteps;
        }
        ...
        continue;
      }
      // Check whether we received data from the GPU and send it to the network
      if (sub->transmitted < sub->posted && sub->transmitted < sub->done + NCCL_STEPS) {
        ...
        if (connFifo[buffSlot].size != -1 && (*recvTail > tail || p == NCCL_PROTO_LL)) {
          ...
          if (ready) {
            ...
            NCCLCHECK(proxyState->ncclNet->isend(resources->netSendComm, buff, size, resources->tpRank,
                                                 sub->sendMhandle, phandle, sub->requests + buffSlot));
            ...
  • post:代理執行緒更新sendMem->head,告訴 GPU「緩衝區已就緒,可以寫資料」。
  • transmit:檢查recvMem->tail是否推進(GPU 已寫完),檢查connFifo[buffSlot].size != -1(資料大小已填),然後呼叫ncclNet->isend發起非同步發送。
  • done:呼叫ncclNet->test檢查發送完成,更新sendMem->head歸還緩衝區。

wc_store_fence()是寫合併屏障——GDRCopy 場景下,CPU 寫gdcSync後必須重新整理寫合併緩衝區,否則 GPU 看不到更新。

生產避坑指南

坑 1:LL128 協定的 flag 校驗。當資料在 sysmem(非 GDR)時,代理執行緒必須逐行檢查 LL128 flag:

📎 src/transport/net.cc:1388-1403

c
          if (p == NCCL_PROTO_LL128) {
            ready = resources->useGdr;
            if (!ready) {
              uint64_t flag = sub->base + sub->transmitted + 1;
              int nFifoLines = DIVUP(connFifo[buffSlot].size, sizeof(uint64_t) * NCCL_LL128_LINEELEMS);
              volatile uint64_t* lines = (volatile uint64_t*)buff;
              ready = 1;
              for (int i = 0; i < nFifoLines; i++) {
                if (lines[i * NCCL_LL128_LINEELEMS + NCCL_LL128_DATAELEMS] != flag) {
                  ready = 0;
                  break;
                }
              }
            }
          }

因為 GPU 只呼叫了threadfence(),資料可能還在 L2 快取裡沒落到 sysmem。代理執行緒必須確認每一行的 flag 都正確,才能發送。生產環境若發現 LL128 資料損壞,檢查useGdr是否正確——GDR 路徑下資料直接落顯存,不需要逐行校驗。

坑 2:GDRCopy flush 的記憶體序。接收側在recvProxyProgress裡有一段精妙的內聯組譯:

📎 src/transport/net.cc:1664-1682

c
          if (totalSize > 0 && p == NCCL_PROTO_SIMPLE && needFlush) {
            struct recvNetResources* resources = (struct recvNetResources*)(subGroup->connection->transportResources);
            if (resources->gdcFlush) {
#if defined(__x86_64__)
              asm volatile("mfence" ::: "memory");
              asm volatile("mov (%0), %%eax" ::"l"(resources->gdcFlush) : "%eax", "memory");
#else
              std::atomic_thread_fence(std::memory_order_seq_cst);
              uint64_t dummy;
              NCCLCHECK(ncclGdrCudaRead(resources->gdrDesc, &dummy, resources->gdcFlush, sizeof(dummy)));
#endif
            }

mfence保證 CQE poll 的讀不會重排到 flush 讀之前;mov (%0), %%eax強制發起一次 PCIe 讀,讓 CPU 停頓直到所有先前的 PCIe posted write(包括網卡 DMA)都提交。這是 GDRCopy 場景下防止「網卡說寫完了但資料還在 PCIe 緩衝區」的關鍵。去掉這段,接收方可能讀到舊資料。

mermaid
sequenceDiagram
    participant GPU as GPU Kernel
    participant SM as ncclSendMem
    participant Proxy as sendProxyProgress
    participant NIC as ncclNet->isend
    participant Peer as 对端网卡

    GPU->>SM: 写数据到 buffs[p]
    GPU->>SM: 更新 recvMem->tail
    Proxy->>SM: 读 recvTail, connFifo[buffSlot].size
    Proxy->>Proxy: 检查 ready (LL128 flag / GDR)
    Proxy->>NIC: isend(comm, buff, size, mhandle)
    NIC->>Peer: DMA 发送
    Proxy->>NIC: test(request, &done)
    NIC-->>Proxy: done=1
    Proxy->>SM: 更新 sendMem->head (归还缓冲区)
    Proxy->>GPU: 下一轮 post

五、NVLS:多播組與 UC/MC 記憶體綁定

直覺模型

NVLS 是「廣播電台」——一個 rank 把資料寫到多播組,硬體自動複製給所有訂閱者。傳統 AllReduce 需要 N-1 次點對點傳輸,NVLS 只需 1 次多播寫 + 1 次多播讀。若沒有 NVLS,大規模 AllReduce 的延遲隨 rank 數線性增長。

資料結構與記憶體佈局

NVLS 的核心是「UC(單播)記憶體」和「MC(多播)記憶體」的綁定。nvlsAllocBindUc分配 UC 記憶體並綁定到 MC 組:

📎 src/transport/nvls.cc:225-277

c
static ncclResult_t nvlsAllocBindUc(struct ncclComm* comm, const struct ncclMcPartition* partition, size_t size,
                                    struct ncclNvlsUcSegment* outUc) {
  CUmemAllocationProp ucprop;
  ...
  ucprop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
  ucprop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
  ucprop.location.id = comm->cudaDev;
  ucprop.requestedHandleTypes = ncclCuMemHandleType;
  CUCHECKGOTO(cuMemGetAllocationGranularity(&ucgran, &ucprop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED), ret, fail);
  ALIGN_SIZE(ucsize, ucgran);
  CUCHECKGOTO(cuMemAddressReserve((CUdeviceptr*)&ucptr, ucsize, ucgran, 0U, 0), ret, fail);
  CUCHECKGOTO(cuMemCreate(&ucHandle, ucsize, &ucprop, 0), ret, fail1);
  CUCHECKGOTO(cuMemMap((CUdeviceptr)ucptr, ucsize, 0, ucHandle, 0), ret, fail2);
  CUCHECKGOTO(cuMemSetAccess((CUdeviceptr)ucptr, ucsize, &comm->nvlsResources->accessDesc, 1), ret, fail3);
  CUDACHECKGOTO(cudaMemset(ucptr, 0, ucsize), ret, fail3);
  NCCLCHECKGOTO(ncclMemTrack(comm->memManager, ucptr, ucsize, ucHandle, ncclCuMemHandleType, ncclMemPersist), ret,
                fail3);
  NCCLCHECKGOTO(bootstrapIntraNodeBarrier(comm->bootstrap, comm->localRankToRank, comm->localRank, comm->localRanks,
                                          comm->localRankToRank[0]),
                ret, fail3);
  NCCLCHECKGOTO(ncclMcPartitionBindMem(partition, 0 /*offsetInPartition*/, ucHandle, 0 /*memOffset*/, ucsize), ret,
                fail3);
  ...

流程:cuMemCreate分配實體記憶體 →cuMemMap映射到虛擬位址 →cuMemSetAccess設定 GPU 存取權限 →ncclMcPartitionBindMem把 UC 實體記憶體綁定到 MC 組的指定偏移。綁定後,任何 rank 寫 MC 位址,硬體會把資料複製到所有綁定的 UC 記憶體。

注意bootstrapIntraNodeBarrier在cuMulticastBindMem之前——註解說這是為了「mitigate the possible hang in cuMulticastBindMem during abort」。這是硬體層面的防禦:如果某個 rank 在綁定過程中 abort,其他 rank 可能在cuMulticastBindMem裡掛起。

場景驅動 Walkthrough:ncclNvlsBufferSetup 的緩衝區佈局

📎 src/transport/nvls.cc:279-368

c
ncclResult_t ncclNvlsBufferSetup(struct ncclComm* comm) {
  ...
  nvlsStepSize = comm->nvlsChunkSize;
  buffSize = nvlsStepSize * NCCL_STEPS;
  nvlsPerRankSize = nChannels * 2 * buffSize;
  nvlsTotalSize = nvlsPerRankSize * nHeads;
  ...
  if (resources->dataUc.ptr == NULL) {
    NCCLCHECKGOTO(nvlsAllocBindUc(comm, &resources->dataPartition, nvlsTotalSize, &resources->dataUc), res, fail);
  }
  ...
  for (int h = 0; h < nHeads; h++) {
    int nvlsPeer = comm->nRanks + 1 + h;
    for (int c = 0; c < nChannels; c++) {
      struct ncclChannel* channel = comm->channels + c;
      struct ncclChannelPeer* peer = channel->peers[nvlsPeer];

      // Reduce UC -> MC
      peer->send[1].conn.buffs[NCCL_PROTO_SIMPLE] = (char*)resources->dataUc.ptr + (h * 2 * nChannels + c) * buffSize;
      peer->recv[0].conn.buffs[NCCL_PROTO_SIMPLE] =
        (char*)resources->dataPartition.ptr + (h * 2 * nChannels + c) * buffSize;

      // Broadcast MC -> UC
      peer->recv[1].conn.buffs[NCCL_PROTO_SIMPLE] =
        (char*)resources->dataUc.ptr + ((h * 2 + 1) * nChannels + c) * buffSize;
      peer->send[0].conn.buffs[NCCL_PROTO_SIMPLE] =
        (char*)resources->dataPartition.ptr + ((h * 2 + 1) * nChannels + c) * buffSize;
      ...

緩衝區佈局:每個 head 有2 * nChannels個 buffer(一半用於 reduce,一半用於 broadcast)。send[1]和recv[0]是 reduce 方向(UC → MC),recv[1]和send[0]是 broadcast 方向(MC → UC)。dataUc.ptr是本地 UC 記憶體,dataPartition.ptr是 MC 組映射位址。

設計思考

〔設計推斷與架構權衡〕

為什麼 NVLS 的canConnect返回 0? 因為 NVLS 不是點對點傳輸——它是「一對多」的多播模型。selectTransport的迴圈是為點對點連接設計的,NVLS 的連接建立走ncclNvlsSetup獨立路徑。把 NVLS 放進ncclTransports陣列只是為了統一free介面(nvlsSendFree/nvlsRecvFree),實際連接邏輯完全獨立。

生產避坑指南

坑:MNNVL 不支援 NVLS buffer 註冊。看ncclNvlsSetup:

至此,NCCL 透過 ncclTransport 抽象層,成功將 P2P、SHM、NET、NVLS 四種異構通道統一為一致的介面,算法內核無需關心底層是 NVLink 還是網卡。但傳輸層只解決了「通道如何抽象」,尚未回答「資料如何被異步驅動」。下一章我們將聚焦 src/proxy.cc 與 src/include/proxy.h,看 proxy 執行緒如何在 host 側異步推進網路收發,與 GPU kernel 形成生產者-消費者關係,揭開 NCCL 異步性的關鍵機制。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 12

第 12 章:代理執行緒異步調度:proxy.cc 如何解耦 I/O 與 kernel 執行

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 12 章 / 共 25 章

上一章拆解了 transport 抽象層,看到 NCCL 如何用統一介面屏蔽 P2P/SHM/NET/NVLS 的差異。但傳輸層只回答了「資料走哪條通道」,尚未回答「資料如何被異步驅動」。GPU kernel 若直接阻塞在網路等待上,計算單元就會被 I/O 拖死。本章聚焦src/proxy.cc與src/include/proxy.h,看 NCCL 如何用獨立的 host 執行緒把網路 I/O 從 kernel 執行路徑中剝離出來,與 GPU 形成生產者-消費者關係。

12.1 為什麼需要代理執行緒:從「誰等網路」說起

直覺模型

想像一家餐廳:廚房(GPU kernel)只負責做菜,傳菜員(proxy 執行緒)負責把菜端給客人(網路對端)。如果讓廚師親自端菜,他每端一趟就得停下炒菜,出餐速度暴跌。NCCL 的 proxy 就是那個專職傳菜員——kernel 只管往共享緩衝區裡寫資料、從緩衝區裡讀資料,網路收發的髒活累活全交給 host 側的 proxy 執行緒。

〔設計推斷與架構權衡〕

若沒有 proxy,系統會面臨什麼災難? GPU kernel 是 SIMT 大規模並行的,一個 warp 阻塞在網路輪詢上會浪費整個 SM 的算力;更致命的是,網路收發涉及 socket 系統呼叫、verbs 輪詢、DMA 描述符提交,這些操作根本無法在 device 程式碼裡執行。因此 NCCL 必須把網路 I/O 搬到 host,讓 kernel 與 proxy 透過共享記憶體中的 FIFO 交換「資料就緒」訊號。

兩類執行緒的分工

NCCL 在 host 側啟動了兩類 proxy 執行緒,職責截然不同:

  • Service 執行緒(ncclProxyService):處理控制面請求——連線建立、記憶體註冊、FD 查詢。它監聽一個 socket,接收來自本地 rank 的 RPC 請求,非同步推進 setup/connect 等操作。
  • Progress 執行緒(ncclProxyProgress):處理資料面——真正驅動網路收發。它從共享記憶體池裡取 proxy op,呼叫 transport 的proxyProgress回呼推進資料搬運。

📎 src/include/proxy.h:343-345顯示ncclProxyState同時持有thread(Service)和threadUDS(UDS 服務),而 Progress 執行緒的句柄藏在progressState.thread裡📎 src/include/proxy.h:261-261。

生產者-消費者關係的建立

📎 src/proxy.cc:2130-2166的ncclProxyCreate是執行緒誕生的地方:當refCount == 1(首個 comm 建立)時,它把 comm 的關鍵欄位拷貝進proxyState,然後啟動 Service 執行緒和 UDS 執行緒。注意 Progress 執行緒不在這裡啟動——它由proxyProgressInit在首次需要 proxy progress 的連線建立時才懶啟動📎 src/proxy.cc:1523-1524。

mermaid
flowchart TD
    create["ncclProxyCreate(comm)"] --> check_ref{"proxyState->refCount == 1?"}
    check_ref -->|否| skip["复用已有线程,直接返回"]
    check_ref -->|是| copy["拷贝 comm 字段到 proxyState"]
    copy --> start_svc["std::thread(ncclProxyService)"]
    start_svc --> start_uds["std::thread(ncclProxyServiceUDS)"]
    start_uds --> wait["等待连接建立请求"]
    wait --> conn_init{"proxyConnInit 发现<br/>tcomm->proxyProgress != NULL?"}
    conn_init -->|是| prog_init["proxyProgressInit()"]
    conn_init -->|否| no_prog["不启动 Progress 线程"]
    prog_init --> shm["ncclShmOpen 创建 opsPool 共享内存"]
    shm --> start_prog["std::thread(ncclProxyProgress)"]

這張圖錨定了執行緒啟動的真實分支:只有tcomm->proxyProgress非空(即該 transport 需要資料面推進)時,Progress 執行緒才會被建立。

12.2 資料結構與記憶體佈局:共享記憶體池與 op 池

核心結構體全景

proxy 的並發模型建立在兩塊共享記憶體之上,理解它們的記憶體佈局是理解整個機制的前提。

第一塊:ncclProxyOpsPool(📎 src/include/proxy.h:218-226)。這是主執行緒與 Progress 執行緒之間的「任務投遞箱」,透過/dev/shm跨行程共享。

欄位類型作用
ops[]ncclProxyOp[]預分配的 op 陣列,大小MAX_OPS_PER_PEER * NCCL_MAX_LOCAL_RANKS
nextOpsvolatile int待處理 op 鏈結串列頭索引,-1 表示空
nextOpsEndvolatile int待處理 op 鏈結串列尾索引
freeOps[]volatile int[]每個 local rank 的空閒 op 鏈結串列頭
syncObjectsInitializedint標記 mutex/cond 是否已初始化
mutex / condstd::mutex / std::condition_variable跨行程同步原語

MAX_OPS_PER_PEER的定義📎 src/include/proxy.h:218-226是2 * MAXCHANNELS * 2 * NCCL_MAX_DEV_WORK_P2P_PER_BATCH。註解解釋了為什麼是 2 倍:每個 p2p work 包含一個 send 和一個 recv proxy op,所以要乘 2;再乘 2 是為了能存兩輪完整操作,否則無法「投遞一半、釋放一半」。

第二塊:ncclProxyArgs(📎 src/include/proxy.h:174-209)。這是 Progress 執行緒內部使用的「執行時 op 描述」,從ncclProxyPool裡分配,不跨行程共享。

關鍵欄位:

  • subs[NCCL_PROXY_MAX_SUBS]:子操作陣列,NCCL_PROXY_MAX_SUBS = MAXCHANNELS 📎 src/include/proxy.h:55-55。多個 channel 的同類操作會被聚合到一個 args 的多個 sub 裡。
  • progress:函式指標,指向 transport 的proxyProgress回呼📎 src/include/proxy.h:176-176。
  • next / nextPeer / proxyAppendPtr:三根鏈結串列指標,構成複雜的 op 組織關係。
  • state:ncclProxyOpNone / ncclProxyOpReady / ncclProxyOpProgress三態📎 src/include/proxy.h:48-52。

記憶體池的分層設計

ncclProxyPool 📎 src/proxy.cc:50-53是一個批次分配單元,每個 pool 含PROXYARGS_ALLOCATE_SIZE(即NCCL_MAX_OPS)個ncclProxyArgs。allocateArgs 📎 src/proxy.cc:207-231的分配邏輯值得細看:

c
if (state->pool == NULL) {
    struct ncclProxyPool* newPool;
    NCCLCHECK(ncclCalloc(&newPool, 1));
    struct ncclProxyArgs* newElems = newPool->elems;
    for (int i = 0; i < PROXYARGS_ALLOCATE_SIZE; i++) {
      if (i + 1 < PROXYARGS_ALLOCATE_SIZE) newElems[i].next = newElems + i + 1;
    }
    state->pool = newElems;
    newPool->next = state->pools;
    state->pools = newPool;
}
elem = state->pool;
state->pool = state->pool->next;

📎 src/proxy.cc:207-231

〔設計推斷與架構權衡〕

這裡的設計動機是 :ncclProxyArgs結構體很大(含subs[MAXCHANNELS]陣列,每個 sub 又有requests[NCCL_STEPS]),如果每個 op 單獨 malloc,會造成嚴重的記憶體碎片和分配開銷。批次分配 + 空閒鏈結串列複用,把分配成本攤薄到幾乎為零。註解「Make sure we allocate the memory close to the network thread」暗示這是為了 NUMA 親和性——pool 在 Progress 執行緒首次分配時建立,天然靠近該執行緒運行的 CPU。

偽共享與原子變數

ncclProxyOpsPool裡的nextOps、nextOpsEnd、freeOps[]都是volatile int。它們被主執行緒和 Progress 執行緒同時讀寫,但 NCCL 沒有用鎖保護所有存取——而是用原子操作 + 記憶體序來保證正確性。

看ncclLocalOpAppend裡從 freeOps 取空閒 op 的邏輯📎 src/proxy.cc:503-513:

c
int freeOp = -1;
while (freeOp == -1) {
  freeOp = COMPILER_ATOMIC_EXCHANGE(&pool->freeOps[tpLocalRank], -1, std::memory_order_acquire);
  if (freeOp == -1) std::this_thread::yield();
}

主執行緒用atomic_exchange把freeOps[tpLocalRank]置為 -1 並取回舊值——這是一個「搶佔式取用」:誰先 exchange 成功誰拿到整條空閒鏈結串列。Progress 執行緒歸還 op 時用 CAS 迴圈📎 src/proxy.cc:898-907:

c
oldFree = COMPILER_ATOMIC_LOAD(&pool->freeOps[i], std::memory_order_acquire);
do {
  pool->ops[freeOpEnd[i]].next = oldFree;
} while (!COMPILER_ATOMIC_COMPARE_EXCHANGE(&pool->freeOps[i], &oldFree, newFree,
                                           std::memory_order_release,
                                           std::memory_order_acquire));
〔設計推斷與架構權衡〕

這裡用 acquire/release 而非 seq_cst,是因為只需要保證「鏈結串列節點的 next 指標寫入」對取用方可見,不需要全域順序。freeOps[]陣列每個元素對應一個 local rank,天然分散在不同快取行附近,減少了偽共享。

12.3 控制面:連線建立與 RPC 機制

直覺模型

〔設計推斷與架構權衡〕

Service 執行緒像一個「前台接待」:本地 rank 要建立網路連線時,不是自己直接去連,而是發一個 RPC 請求給 Service 執行緒,由它代為執行 setup/connect。為什麼要這樣? 因為網路連線建立(尤其是 verbs 的 QP 建立、記憶體註冊)可能阻塞,而且某些資源(如 listen socket)必須由單一執行緒持有。把控制面集中到 Service 執行緒,主執行緒就能非阻塞地繼續做別的事。

RPC 請求的編碼

ncclProxyCallAsync 📎 src/proxy.cc:1369-1394是 RPC 的發送端。它透過 socket 依次發送:type、connection 指標、reqSize、respSize、reqBuff、opId。

c
NCCLCHECKGOTO(ncclSocketSend(sock, &type, sizeof(int)), ret, error);
NCCLCHECKGOTO(ncclSocketSend(sock, &proxyConn->connection, sizeof(void*)), ret, error);
NCCLCHECKGOTO(ncclSocketSend(sock, &reqSize, sizeof(int)), ret, error);
NCCLCHECKGOTO(ncclSocketSend(sock, &respSize, sizeof(int)), ret, error);
if (reqSize) NCCLCHECKGOTO(ncclSocketSend(sock, reqBuff, reqSize), ret, error);
NCCLCHECKGOTO(ncclSocketSend(sock, &opId, sizeof(opId)), ret, error);
NCCLCHECK(expectedProxyResponseEnqueue(sharedProxyState, opId, respSize));

📎 src/proxy.cc:1369-1394

注意最後一步:發送完請求後,立刻把 opId 登記到expectedResponses佇列。這是非同步 RPC 的關鍵——呼叫方不等回覆,而是先登記「我期待這個 opId 的回應」,之後用ncclPollProxyResponse輪詢。

回應佇列的鏈結串列實作

expectedProxyResponseEnqueue 📎 src/proxy.cc:97-117用單向鏈結串列儲存待回應的 op。expectedProxyResponseStore 📎 src/proxy.cc:67-95在收到回應時按 opId 匹配,把回應資料 memcpy 進預先分配的respBuff,標記done = true。expectedProxyResponseDequeue 📎 src/proxy.cc:119-141在輪詢時查找已完成的回應並摘除。

這裡有個細節:expectedProxyResponseStore檢查respSize是否匹配📎 src/proxy.cc:72-75,不匹配就報ncclInternalError。這是防禦性編程——如果請求方和回應方對回應大小的理解不一致,說明協議錯亂,必須立即失敗而非靜默繼續。

Service 執行緒的主迴圈

ncclProxyService 📎 src/proxy.cc:1789-2016的核心是一個 poll 迴圈。它用pollfds陣列管理所有連線,包括 listen socket 和每個 peer 的 socket。

c
while (stop == PROXY_RUNNING || npeers > 0) {
    if (COMPILER_ATOMIC_LOAD(proxyState->abortFlag, std::memory_order_acquire) != 0) stop = PROXY_ABORT;
    int ret = 0;
    const int timeout = asyncOpCount ? 0 : 500;
    ...
    ret = poll(activePollfds, nfds_to_poll, timeout);

📎 src/proxy.cc:1842-1863

timeout的選擇很講究:如果有非同步 op 在推進(asyncOpCount > 0),timeout 設為 0(非阻塞輪詢),因為需要頻繁呼叫proxyProgressAsync推進它們;否則設 500ms,避免空轉燒 CPU。註解「never let proxy service thread blocks in poll, or it cannot receive abortFlag」📎 src/proxy.cc:1847-1847點明了為什麼不能無限阻塞——必須週期性醒來檢查 abortFlag。

非同步 op 的推進

proxyProgressAsync 📎 src/proxy.cc:1626-1700是 Service 執行緒推進非同步操作的核心。它根據 op 類型分發到不同的 transport 回呼:

c
if (op->type == ncclProxyMsgSetup) {
    res = op->connection->tcomm->proxySetup(op->connection, proxyState, op->reqBuff, op->reqSize, op->respBuff,
                                            op->respSize, &done);
} else if (op->type == ncclProxyMsgConnect) {
    res = op->connection->tcomm->proxyConnect(...);
} else if (op->type == ncclProxyMsgInit) {
    res = proxyConnInit(peer, connectionPool, proxyState, ...);
}

📎 src/proxy.cc:1631-1664

每個回呼都帶一個done輸出參數。如果done == 0,說明操作還沒完成(比如網路連線還在三次握手),返回ncclInProgress,下次迴圈繼續推進。如果done == 1,則發送回應標頭 + 回應主體給請求方📎 src/proxy.cc:1681-1689。

mermaid
sequenceDiagram
    participant Main as 主线程 (ncclSend)
    participant Svc as Service 线程
    participant Net as 网络插件 (ncclNet)
    Main->>Svc: ncclProxyCallAsync(ncclProxyMsgConnect)
    Note over Main: expectedProxyResponseEnqueue(opId)
    Svc->>Svc: proxyServiceInitOp 读取请求
    Svc->>Net: proxyConnect() 调用 ncclNet->connect
    alt connect 未完成
        Net-->>Svc: netSendComm == NULL, done=0
        Svc->>Svc: 返回 ncclInProgress,下次 poll 重试
    else connect 完成
        Net-->>Svc: netSendComm != NULL, done=1
        Svc->>Main: ncclSocketSend(resp header + connectMap)
    end
    Main->>Main: ncclPollProxyResponse 轮询
    Main->>Main: expectedProxyResponseDequeue 取回结果

這張時序圖錨定了sendProxyConnect裡*done = 0; return ncclInProgress的真實分支📎 src/transport/net.cc:913-916。

12.4 資料面:Progress 執行緒如何驅動網路收發

直覺模型

Progress 執行緒是「傳送帶操作員」:它盯著共享緩衝區裡的 FIFO,一旦 GPU 寫好了資料(FIFO 裡 size != -1),就立刻呼叫isend把資料發出去;一旦網路收完了資料,就更新 recvTail 通知 GPU 可以讀了。整個過程 GPU 和 proxy 透過 FIFO 裡的 head/tail 指標同步,不需要任何鎖。

op 的投遞:從主執行緒到 Progress 執行緒

主執行緒在ncclProxySaveOp 📎 src/proxy.cc:591-761裡根據 pattern 決定需要哪些 proxy op,然後透過SaveProxy → ncclLocalOpAppend把 op 寫入共享記憶體池。

ncclLocalOpAppend 📎 src/proxy.cc:488-554的流程:

1. 從proxyOps->freeOp或pool->freeOps[tpLocalRank]取一個空閒 op 槽位。

2. memcpy(op, proxyOp, sizeof(struct ncclProxyOp))把 op 內容拷進共享記憶體📎 src/proxy.cc:515-515。

3. 把 op 掛到proxyOps->nextOps鏈結串列尾部。

4. 如果累積的 op 數達到MAX_OPS_PER_PEER,觸發一次批次投遞📎 src/proxy.cc:525-551。

批次投遞的邏輯很微妙:它不能簡單地把所有 op 都發出去,因為「同一個 opCount 的多個 op 必須一起投遞,否則會破壞 proxyArgs 的 sub 聚合」。所以它找到最後一個 opCount 變化的邊界,只投遞到那裡📎 src/proxy.cc:529-548。

投遞透過ncclProxyPost 📎 src/proxy.cc:476-486完成,它加鎖、更新pool->nextOps、notify_one喚醒 Progress 執行緒。

Progress 執行緒的主迴圈

ncclProxyProgress 📎 src/proxy.cc:951-1011的結構:

c
do {
    int idle = 1;
    ncclResult_t ret = progressOps(proxyState, state, state->active, &idle);
    ...
    if (idle || !state->active || (++proxyOpAppendCounter == ncclParamProgressAppendOpFreq())) {
      int added = 0;
      proxyOpAppendCounter = 0;
      ret = ncclProxyGetPostedOps(proxyState, &added);
      ...
    }
    lastIdle = idle;
    stopv = state->stop.load(std::memory_order_acquire);
} while ((stopv == 0 || (stopv == 1 && state->active)) &&
         COMPILER_ATOMIC_LOAD(proxyState->abortFlag, std::memory_order_acquire) == 0);

📎 src/proxy.cc:976-1009

這裡有個效能優化值得注意:proxyOpAppendCounter計數器📎 src/proxy.cc:974-974。註解解釋📎 src/proxy.cc:969-973:太頻繁呼叫ncclProxyGetPostedOps會導致小訊息通訊效能回退,所以每推進ProgressAppendOpFreq(預設 8)次才去取一次新 op。

op 的聚合:ProxyAppend

ProxyAppend 📎 src/proxy.cc:437-474決定一個 op 是「追加到已有 args 的 sub 裡」還是「新建一個 args」。判斷依據是connection->shared && args->opCount == op->opCount 📎 src/proxy.cc:443-443——同一連線、同一 opCount 的多個 channel 操作會被聚合。

〔設計推斷與架構權衡〕

聚合的價值 :多個 channel 的同類操作合併成一個 args,Progress 執行緒一次迴圈就能推進所有 channel,減少了函式呼叫開銷和快取失效。ncclProxyOpToArgs 📎 src/proxy.cc:368-435在追加 sub 時會校驗sliceSteps、chunkSteps、protocol、dtype、redOp、coll是否一致📎 src/proxy.cc:401-406,不一致就報錯——這是防止錯誤聚合的防線。

sendProxyProgress:發送側的四階段狀態機

sendProxyProgress 📎 src/transport/net.cc:1324-1491是發送側的核心。它按 sub 逐個推進,每個 sub 有四個計數器:posted、transmitted、done。

階段一:Ready 初始化 📎 src/transport/net.cc:1326-1339

c
sub->base = ROUNDUP(resources->step, args->chunkSteps);
resources->step = sub->base + sub->nsteps;
sub->posted = sub->transmitted = sub->done = 0;

base是 step 的起始編號,ROUNDUP保證對齊到chunkSteps。resources->step累加,為下一個 op 預留空間。

階段二:Post 緩衝區給 GPU 📎 src/transport/net.cc:1355-1376

c
if (sub->posted < sub->nsteps && sub->posted < sub->done + maxDepth) {
    int buffSlot = (sub->base + sub->posted) % NCCL_STEPS;
    if (resources->shared) {
        ...
        *sendHead = sub->base + sub->posted - NCCL_STEPS;
    } else {
        sub->posted += args->sliceSteps;
    }
}

maxDepth是流水線深度📎 src/transport/net.cc:1343-1343,限制同時 in-flight 的 step 數。shared 模式下,proxy 透過更新sendHead告訴 GPU「這個 slot 可以寫了」。

階段三:檢查 GPU 是否寫好,發起 isend 📎 src/transport/net.cc:1378-1452

c
if (sub->transmitted < sub->posted && sub->transmitted < sub->done + NCCL_STEPS) {
    int buffSlot = (sub->base + sub->transmitted) % NCCL_STEPS;
    volatile uint64_t* recvTail = &resources->recvMem->tail;
    uint64_t tail = sub->base + sub->transmitted;
    if (connFifo[buffSlot].size != -1 && (*recvTail > tail || p == NCCL_PROTO_LL)) {
        int size = connFifo[buffSlot].size;
        ...
        NCCLCHECK(proxyState->ncclNet->isend(resources->netSendComm, buff, size, resources->tpRank,
                                             sub->sendMhandle, phandle, sub->requests + buffSlot));
        if (sub->requests[buffSlot] != NULL) {
            sub->transmitted += args->sliceSteps;
        }
    }
}

這裡的關鍵判斷是connFifo[buffSlot].size != -1 && *recvTail > tail——GPU 寫好資料後會更新 FIFO 的 size 和 recvTail,proxy 看到這兩個條件滿足才發起 isend。對於 LL 協議,因為它是「零拷貝」語意,不需要等 recvTail。

階段四:檢查發送完成,更新 sendHead 📎 src/transport/net.cc:1455-1481

c
if (sub->done < sub->transmitted) {
    int buffSlot = (sub->base + sub->done) % NCCL_STEPS;
    NCCLCHECK(proxyState->ncclNet->test(sub->requests[buffSlot], &done, &size));
    if (done) {
        connFifo[buffSlot].size = -1;
        std::atomic_thread_fence(std::memory_order_seq_cst);
        sub->done += args->sliceSteps;
        if (resources->shared == 0) {
            volatile uint64_t* sendHead = resources->gdcSync ? resources->gdcSync : &resources->sendMem->head;
            *sendHead = sub->base + sub->done;
        }
    }
}

test返回 done 後,先把 FIFO size 重置為 -1,插入一個 seq_cst fence,再更新 sendHead 通知 GPU「這個 slot 可以複用了」。fence 的作用是防止 size 重置和 head 更新的重排序——如果 head 先更新,GPU 可能在 size 還是舊值時就開始寫。

recvProxyProgress:接收側的四階段

recvProxyProgress 📎 src/transport/net.cc:1493-1788更複雜,因為它涉及 sub 分組(多個 sub 共享同一個 recvComm 時用 multirecv)。

階段一:Ready 時按 recvComm 分組 📎 src/transport/net.cc:1495-1538

c
for (int s = 0; s < args->nsubs; s++) {
    ...
    if (groupSize == maxRecvs) {
        groupSize = 0;
    } else if (s > 0) {
        int next;
        for (next = s; next < args->nsubs; next++) {
            struct recvNetResources* nextRes = ...;
            if (nextRes->netRecvComm == recvComm) break;
        }
        if (next == args->nsubs) {
            groupSize = 0;
        } else if (s != next) {
            // swap subs
        }
    }
    groupSize++;
    ...
    for (int i = 0; i < groupSize; i++) sub[-i].groupSize = groupSize;
}
〔設計推斷與架構權衡〕

這段程式碼把使用同一recvComm的 sub 排到一起,並記錄groupSize。為什麼要分組? 因為irecv支援一次接收多個 buffer(multirecv),把同 comm 的請求合併成一次呼叫能顯著降低外掛開銷。

階段二:發起 irecv 📎 src/transport/net.cc:1543-1631

c
if (subCount) {
    uint64_t step = subGroup->posted;
    void** requestPtr = subGroup->requests + (step % NCCL_STEPS);
    bool ignoreCompletion = ncclParamNetOptionalRecvCompletion() &&
                            ((args->protocol == NCCL_PROTO_LL128) || (args->protocol == NCCL_PROTO_LL)) &&
                            (subCount == 1);
    if (ignoreCompletion) *requestPtr = (void*)NCCL_NET_OPTIONAL_RECV_COMPLETION;
    NCCLCHECK(proxyState->ncclNet->irecv(resources->netRecvComm, subCount, ptrs, sizes, tags, mhandles, phandles,
                                         requestPtr));
    if (*requestPtr) {
        subGroup->recvRequestsCache[step % NCCL_STEPS] = *requestPtr;
        subGroup->recvRequestsSubCount = subCount;
        for (int i = 0; i < subGroup->groupSize; i++) {
            sub->posted += args->sliceSteps;
        }
    }
}

ignoreCompletion優化📎 src/transport/net.cc:1608-1610:對於 LL/LL128 協議的單 buffer 接收,完成通知是可選的(因為資料本身帶 flag),可以跳過 completion 檢查。

階段三:檢查接收完成,更新 recvTail 📎 src/transport/net.cc:1634-1743

c
NCCLCHECK(proxyState->ncclNet->test(subGroup->requests[step % NCCL_STEPS], &done, sizes));
if (done) {
    for (int i = 0; i < subGroup->groupSize; i++) {
        struct ncclProxySubArgs* sub = subGroup + i;
        int buffSlot = (sub->base + sub->received) % NCCL_STEPS;
        connFifo[buffSlot].size = -1;
        sub->received += args->sliceSteps;
    }
    ...
}

接收完成後,重置 FIFO size,然後進入 flush 階段(GDRDMA 場景需要 flush 保證資料可見性)。

階段四:等待 GPU 消費,更新 done 📎 src/transport/net.cc:1745-1779

c
if (sub->transmitted > sub->done) {
    volatile uint64_t* sendHead = &resources->sendMem->head;
    uint64_t done = *sendHead;
    while (done > sub->base + sub->done && sub->transmitted > sub->done) {
        if (subGroup->recvRequestsCache[sub->done % NCCL_STEPS]) {
            if (proxyState->ncclNet->irecvConsumed) {
                NCCLCHECK(proxyState->ncclNet->irecvConsumed(resources->netRecvComm, subGroup->recvRequestsSubCount,
                                                             subGroup->recvRequestsCache[sub->done % NCCL_STEPS]));
            }
            subGroup->recvRequestsCache[sub->done % NCCL_STEPS] = NULL;
        }
        sub->done += args->sliceSteps;
    }
}

這裡透過讀sendHead判斷 GPU 是否已經消費了資料。irecvConsumed是給外掛的回呼,告訴它「這個接收請求的 buffer 已經被消費,可以複用了」。

資料流全景

mermaid
flowchart LR
    subgraph GPU["GPU Kernel"]
        gpu_write["写入数据到 buff"]
        gpu_fifo["更新 connFifo.size<br/>和 recvTail"]
    end
    subgraph SHM["共享内存 FIFO"]
        fifo["ncclConnFifo<br/>size / offset"]
        head["sendMem->head"]
        tail["recvMem->tail"]
    end
    subgraph PROXY["Progress 线程"]
        check["检查 size != -1<br/>且 recvTail > tail"]
        isend["ncclNet->isend()"]
        test["ncclNet->test()"]
        update["更新 sendHead"]
    end
    gpu_write --> gpu_fifo
    gpu_fifo --> fifo
    gpu_fifo --> tail
    fifo --> check
    tail --> check
    check -->|数据就绪| isend
    isend --> test
    test -->|发送完成| update
    update --> head
    head -->|GPU 可复用 slot| gpu_write

這張資料流圖展示了 GPU 與 proxy 透過 FIFO 和 head/tail 指標形成的閉環:GPU 寫資料 → 更新 tail → proxy 檢測到並發 isend → test 確認完成 → 更新 head → GPU 複用 slot。

12.5 並發控制、記憶體屏障與硬體互動

無鎖 FIFO 的記憶體序

proxy 與 GPU 之間的同步完全依賴ncclConnFifo和 head/tail 指標,沒有任何鎖。這要求極其謹慎的記憶體序控制。

發送側,proxy 在test返回 done 後📎 src/transport/net.cc:1460-1473:

c
connFifo[buffSlot].size = -1;
std::atomic_thread_fence(std::memory_order_seq_cst);
...
*sendHead = sub->base + sub->done;

seq_cst fence 保證 size 重置對 GPU 可見後,head 更新才可見。如果順序反了,GPU 可能看到新 head 但舊 size,誤以為 slot 裡有資料。

接收側,proxy 在更新 recvTail 前📎 src/transport/net.cc:1731-1736:

c
if (step < sub->nsteps) {
    std::atomic_thread_fence(std::memory_order_seq_cst);
    volatile uint64_t* recvTail = resources->gdcSync ? resources->gdcSync : &resources->recvMem->tail;
    *recvTail = sub->base + sub->transmitted;
}

同樣的道理:先 fence 保證資料寫入可見,再更新 tail 通知 GPU 可以讀。

GDRCOPY 的 flush 機制

當使用 GDRDMA 時,NIC 直接寫 GPU 顯存,但寫操作可能還在 PCIe 總線上未提交。proxy 需要主動 flush 才能保證資料可見。看recvProxyProgress裡的 flush 邏輯📎 src/transport/net.cc:1664-1709:

c
if (totalSize > 0 && p == NCCL_PROTO_SIMPLE && needFlush) {
    if (resources->gdcFlush) {
#if defined(__x86_64__)
        asm volatile("mfence" ::: "memory");
        asm volatile("mov (%0), %%eax" ::"l"(resources->gdcFlush) : "%eax", "memory");
#else
        std::atomic_thread_fence(std::memory_order_seq_cst);
        uint64_t dummy;
        NCCLCHECK(ncclGdrCudaRead(resources->gdrDesc, &dummy, resources->gdcFlush, sizeof(dummy)));
#endif
    } else {
        // iflush 路径
        NCCLCHECK(proxyState->ncclNet->iflush(resources->netRecvComm, subCount, ptrs, sizes, mhandles,
                                              subGroup->requests + (step % NCCL_STEPS)));
    }
}

x86 路徑的註解非常精彩📎 src/transport/net.cc:1668-1674:mfence阻止 CQE-poll 的 load 被重排到 flush load 之前;mov (%0), %%eax強制一次 PCIe 讀,讓 CPU 停頓直到所有先前的 PCIe posted write(包括 NIC DMA)提交到端點。這是硬體級別的記憶體序控制,比任何軟體 fence 都硬核。

原子變數與 stop/abort 的協作

Progress 執行緒的退出條件📎 src/proxy.cc:1007-1009:

c
stopv = state->stop.load(std::memory_order_acquire);
} while ((stopv == 0 || (stopv == 1 && state->active)) &&
         COMPILER_ATOMIC_LOAD(proxyState->abortFlag, std::memory_order_acquire) == 0);

stop == 1但state->active != NULL時繼續運行——這是為了「優雅停止」:已經投遞的 op 必須推進完,否則 GPU 會永遠等不到資料。只有stop == 2(abort)或abortFlag != 0才強制退出。

ncclProxyProgressDestroy 📎 src/proxy.cc:1039-1065的停止流程:

c
std::lock_guard<std::mutex> lock(state->opsPool->mutex);
state->stop.store(1, std::memory_order_release);
state->opsPool->cond.notify_one();
state->thread.join();

先加鎖再 store stop,然後 notify——這是防止 lost wakeup 的標準模式。Progress 執行緒在pool->cond.wait時持有鎖並檢查謂詞📎 src/proxy.cc:850-851,保證不會錯過喚醒。

12.6 生產避坑指南與故障恢復鏈

坑一:連線洩漏導致 Service 執行緒無法退出

ncclProxyService的主迴圈條件是stop == PROXY_RUNNING || npeers > 0 📎 src/proxy.cc:1842-1842。註解解釋📎 src/proxy.cc:1843-1845:即使本地 comm abort,只要還有 peer 連線,proxy 執行緒就不能退出,否則可能段錯誤。

排查場景:如果某個 rank 崩潰但沒通知對端,對端的 Service 執行緒會一直卡在npeers > 0的迴圈裡。此時需要依賴abortFlag或逾時機制。生產環境中如果看到行程 hang 在ncclProxyService,先檢查是否有對端 rank 異常退出。

坑二:回應佇列不匹配導致記憶體洩漏

expectedProxyResponseStore在 opId 不匹配時返回ncclInternalError 📎 src/proxy.cc:93-94。但如果回應到達時請求方已經放棄(比如逾時),這個回應會永遠留在佇列裡,respBuff洩漏。

防禦措施:expectedProxyResponseFree 📎 src/proxy.cc:55-65在ncclProxyDestroy時清理整個佇列📎 src/proxy.cc:2226-2226。但這是最後兜底,正常運行中不應該有殘留。

坑三:shared 模式下 head 初始化為負值

sendProxyConnect裡📎 src/transport/net.cc:999-1000:

c
// Don't give credits yet in shared mode.
(resources->gdcSync ? *resources->gdcSync : resources->sendMem->head) = (map->shared ? -NCCL_STEPS : 0);

shared 模式下 head 初始化為-NCCL_STEPS,意味著 GPU 一開始沒有 credit 可寫。proxy 需要在 post 階段逐步增加 head 來「發放 credit」。如果忘記這個初始化,GPU 會誤以為有 credit 而寫入未就緒的 slot,導致資料錯亂。

坑四:LL128 協議的 flag 校驗

sendProxyProgress裡 LL128 的 ready 判斷📎 src/transport/net.cc:1388-1403:

c
if (p == NCCL_PROTO_LL128) {
    ready = resources->useGdr;
    if (!ready) {
        uint64_t flag = sub->base + sub->transmitted + 1;
        int nFifoLines = DIVUP(connFifo[buffSlot].size, sizeof(uint64_t) * NCCL_LL128_LINEELEMS);
        volatile uint64_t* lines = (volatile uint64_t*)buff;
        ready = 1;
        for (int i = 0; i < nFifoLines; i++) {
            if (lines[i * NCCL_LL128_LINEELEMS + NCCL_LL128_DATAELEMS] != flag) {
                ready = 0;
                break;
            }
        }
    }
}

當資料在 sysmem(非 GDR)時,GPU 只調用了threadfence(),proxy 必須逐行檢查 flag 才能確認資料完整。如果跳過這個檢查直接 isend,可能發出半截資料。這是 LL128 特有的陷阱。

故障恢復鏈

當proxyProgressAsync返回非ncclSuccess/ncclInProgress時📎 src/proxy.cc:1929-1937,Service 執行緒會關閉連線並清理該 peer 的所有 async op📎 src/proxy.cc:1984-1995。這個清理是「全量 drain」——不只清理失敗的那個 op,而是把整個 peer 的 asyncOps 佇列清空,防止殘留 op 引用已釋放的連線。

Progress 執行緒遇到錯誤時📎 src/proxy.cc:979-983,把錯誤碼寫入proxyState->asyncResult並退出迴圈。主執行緒後續可以透過檢查這個欄位感知錯誤。

本章小結

本章我們拆解了 NCCL 代理執行緒的完整機制:

1. 兩類執行緒分工:Service 執行緒處理控制面 RPC(連線建立、記憶體註冊),Progress 執行緒處理資料面(網路收發推進)。

2. 共享記憶體池:ncclProxyOpsPool跨行程傳遞 op,ncclProxyArgs在 Progress 執行緒內聚合多個 channel 的操作。

3. 無鎖 FIFO 同步:GPU 與 proxy 透過connFifo和 head/tail 指標交換資料就緒訊號,用 seq_cst fence 保證記憶體序。

4. 四階段狀態機:send/recv 各自的 posted → transmitted → received → done 計數器驅動流水線。

5. 硬體級 flush:GDRDMA 場景下用mfence+ PCIe 讀強制提交 posted write。

本章思考與自測

Q1: 如果把sendProxyProgress中sub->done == sub->nsteps時更新sendHead的邏輯去掉(即不通知 GPU slot 已釋放),在什麼場景下會觸發死鎖?為什麼?

參考解析:sendHead是 GPU 判斷「哪些 slot 可以複用」的唯一依據。看📎 src/transport/net.cc:1469-1473:

c
if (resources->shared == 0) {
    volatile uint64_t* sendHead = resources->gdcSync ? resources->gdcSync : &resources->sendMem->head;
    *sendHead = sub->base + sub->done;
}

如果去掉這段,GPU 的 head 永遠停在初始值(shared 模式下是-NCCL_STEPS,非 shared 是 0)。GPU kernel 在waitSend時會檢查head + NCCL_STEPS > step才認為有 credit 可寫。head 不推進,GPU 寫滿NCCL_STEPS個 slot 後就永遠阻塞在等待 credit 上,而 proxy 又在等 GPU 寫新資料才能 isend——經典的生產者-消費者死鎖。在 shared 模式下更嚴重,因為初始 head 是負值,GPU 一開始就沒有 credit。

Q2: ncclLocalOpAppend在累積 op 達到MAX_OPS_PER_PEER時會觸發批量投遞,但程式碼特意「不投遞最後一個 opCount 的所有 op」。如果改成簡單地把所有 op 都投遞,會破壞什麼機制?

參考解析:看📎 src/proxy.cc:525-548的註釋和邏輯:

c
// Do not post last operations as we could have more coming with the same opCount, and posting
// them in different batches would break proxyArgs aggregation with subs.
uint64_t lastOpCount = pool->ops[proxyOps->nextOpsEnd].opCount;
int lastOp = -1;
...
for (int op = proxyOps->nextOps; op != proxyOps->nextOpsEnd; op = pool->ops[op].next) {
    ops++;
    if (pool->ops[op].opCount != lastOpCount) {
        lastOp = op;
        toSend = ops;
    }
}

ProxyAppend的聚合邏輯📎 src/proxy.cc:443-443依賴args->opCount == op->opCount來判斷是否追加 sub。如果同一個 opCount 的多個 channel op 被拆到兩個批次投遞,第一批會建立一個 args,第二批到達時args->opCount已經不等於新 op 的 opCount(因為 args 可能已經被推進),導致本應聚合的 sub 被拆成獨立的 args。這不僅降低效能,還可能破壞ncclProxyOpToArgs裡的nChannels/nPeers取 min 的邏輯📎 src/proxy.cc:399-400,導致錯誤的通道數計算。

Q3: recvProxyProgress的 Ready 階段會按recvComm對 sub 重新排序分組。如果去掉這個分組邏輯,讓每個 sub 獨立呼叫irecv,在maxRecvs > 1的網卡上會有什麼後果?

參考解析:看📎 src/transport/net.cc:1495-1538的分組邏輯和📎 src/transport/net.cc:1613-1614的 multirecv 呼叫:

c
NCCLCHECK(proxyState->ncclNet->irecv(resources->netRecvComm, subCount, ptrs, sizes, tags, mhandles, phandles,
                                     requestPtr));

maxRecvs是網卡插件宣告的「單次 irecv 能接收的最大 buffer 數」📎 src/transport/net.cc:1525-1525。當maxRecvs > 1時,插件(如 IB)支援一次 WQE 接收多個 buffer,能顯著降低 doorbell 開銷和 CQE 處理成本。如果去掉分組,每個 sub 單獨 irecv,subCount永遠是 1,插件退化為單 buffer 模式,吞吐量會下降。更關鍵的是,recvRequestsCache和irecvConsumed機制📎 src/transport/net.cc:1616-1617是為 multirecv 設計的——單 buffer 模式下這些快取邏輯會失效,可能導致請求洩漏。

至此,我們理解了 proxy 執行緒如何將網路 I/O 與 kernel 執行解耦,讓 GPU 計算與通訊真正並行。但 proxy 只是驅動者,底層網路傳輸的具體實現仍待揭曉。下一章我們將深入net_ib,看 NCCL 如何封裝 verbs API 實現 InfiniBand 傳輸,以及 GPUDirect RDMA 如何讓網卡直接讀寫 GPU 顯存。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 13

第 13 章:InfiniBand 網路傳輸:net_ib 如何封裝 verbs 與 GPUDirect RDMA

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 13 章 / 共 25 章

上一章我們看到 proxy 執行緒如何把網路 I/O 從 GPU kernel 中剝離出來,讓計算與通訊真正並行。但 proxy 只是一個「驅動者」——它呼叫 ncclNet->isend/irecv 這些抽象介面,卻不知道底下到底是 TCP、InfiniBand 還是別的什麼。本章我們掀開這層抽象,進入 src/transport/net_ib 與 src/misc/ibvwrap.cc,看 NCCL 如何把 libibverbs 這套 C 庫封裝成可插拔的符號表,如何建立 Queue Pair(QP),以及 GPUDirect RDMA 如何讓網卡繞過 host 記憶體直接讀寫 GPU 顯存。

13.1 為什麼 NCCL 不直接呼叫 libibverbs

直覺模型:符號表就是「可插拔的電源插座」

想像你買了一台進口電器,插頭形狀和家裡插座不匹配。你有兩個選擇:要麼把電器拆開改線(直接#include <infiniband/verbs.h>並連結-libverbs),要麼買一個萬能轉換插頭(執行時動態載入符號)。NCCL 選擇了後者。

〔設計推斷與架構權衡〕

這個選擇的核心動機是部署靈活性:NCCL 作為一個庫被 PyTorch、TensorFlow 等上層框架載入,它無法假設執行環境一定裝了libibverbs.so。如果編譯期硬連結,那麼在沒有 InfiniBand 驅動的機器上,整個 NCCL 庫都無法載入——哪怕你只想用 NVLink 做單機通訊。透過執行時dlopen+ 符號解析,NCCL 可以在沒有 IB 的機器上優雅降級。

如果缺少這一層封裝,系統會面臨的災難是:一個純 NVLink 的單機訓練任務,因為機器上沒裝 IB 驅動而直接崩潰。這在雲環境、開發機上極其常見。

資料結構與記憶體佈局:符號表容器

核心資料結構是ncclIbvSymbols,定義在ibvsymbols.h中(本章材料未包含該檔案,但從使用方式可推斷其結構)。它是一個純函式指標容器,每個欄位對應一個 libibverbs 函式:

c
struct ncclIbvSymbols {
  int (*ibv_internal_fork_init)(void);
  struct ibv_device** (*ibv_internal_get_device_list)(int* num_devices);
  int (*ibv_internal_modify_qp)(struct ibv_qp*, struct ibv_qp_attr*, int);
  // ... 数十个函数指针
};

全域只有一個實例,配合std::once_flag保證執行緒安全初始化:

📎 src/misc/ibvwrap.cc:26-29

c
static std::once_flag initOnceFlag;
static ncclResult_t initResult;
struct ncclIbvSymbols ibvSymbols;

這裡的設計非常克制:initOnceFlag是std::once_flag,initResult快取初始化結果,ibvSymbols是全域符號表。三者都是靜態儲存期,生命週期貫穿整個行程。

〔設計推斷與架構權衡〕

為什麼用std::once_flag而不是pthread_once?因為 NCCL 的 C++ 程式碼已經依賴<mutex>和<thread>,用標準函式庫更一致。call_once的語意是:無論多少執行緒同時呼叫wrap_ibv_symbols(),lambda 只執行一次,其餘執行緒阻塞等待,然後都拿到同一個initResult。這比手寫雙重檢查鎖定(DCLP)安全得多——DCLP 在 C++ 記憶體模型下有著名的重排序陷阱。

Step-by-Step:符號解析的完整流程

當 NCCL 第一次需要 IB 傳輸時,會呼叫wrap_ibv_symbols():

📎 src/misc/ibvwrap.cc:26-29

c
ncclResult_t wrap_ibv_symbols(void) {
  std::call_once(initOnceFlag, []() { initResult = buildIbvSymbols(&ibvSymbols); });
  return initResult;
}

buildIbvSymbols定義在ibvsymbols.cc(本章未包含),它的工作是用dlopen("libibverbs.so")開啟函式庫,然後對每個函式名稱呼叫dlsym填充指標。如果某個符號找不到,對應欄位保持 NULL。

這個「允許 NULL」的設計貫穿整個封裝層。看CHECK_NOT_NULL巨集:

📎 src/misc/ibvwrap.cc:26-29

c
#define CHECK_NOT_NULL(container, internal_name) \
  if (container.internal_name == NULL) { \
    WARN("lib wrapper not initialized."); \
    return ncclInternalError; \
  }

每個封裝函式在呼叫前都會檢查對應符號是否非空。這意味著:如果某個舊版本 libibverbs 缺少某個新函式,NCCL 不會在載入時崩潰,而是在真正用到該函式時才報錯。這是漸進式降級的關鍵。

設計思考:巨集封裝的三重職責

ibvwrap.cc裡定義了 7 個巨集,它們不是簡單的語法糖,而是承擔了三重職責:

1. 空指標防護:CHECK_NOT_NULL攔截未初始化

2. 錯誤碼正規化:把 libibverbs 的多種錯誤約定(回傳 -1、回傳 errno、回傳 NULL 指標)統一翻譯成ncclResult_t

3. 日誌埋點:失敗時WARN列印函式名稱和 errno

看IBV_PTR_CHECK_ERRNO這個最複雜的巨集:

📎 src/misc/ibvwrap.cc:38-45

c
#define IBV_PTR_CHECK_ERRNO(container, internal_name, call, retval, error_retval, name) \
  CHECK_NOT_NULL(container, internal_name); \
  retval = container.call; \
  if (retval == error_retval) { \
    WARN("Call to " name " failed with error %s", strerror(errno)); \
    return ncclSystemError; \
  } \
  return ncclSuccess;

它展開後做四件事:檢查符號非空、執行呼叫、把回傳值寫入retval(通常是透過指標參數回傳的ibv_pd*等)、判斷是否等於錯誤值。注意strerror(errno)——libibverbs 的指標回傳型函式(如ibv_alloc_pd)失敗時回傳 NULL 並設定errno,所以這裡讀errno是對的。

而IBV_INT_CHECK用於回傳 int 的函式:

📎 src/misc/ibvwrap.cc:84-91

c
#define IBV_INT_CHECK(container, internal_name, call, error_retval, name) \
  CHECK_NOT_NULL(container, internal_name); \
  int ret = container.call; \
  if (ret == error_retval) { \
    WARN("Call to " name " failed"); \
    return ncclSystemError; \
  } \
  return ncclSuccess;

這裡不讀errno,因為這類函式(如ibv_fork_init)直接回傳 -1 表示失敗,錯誤資訊已經遺失。

〔設計推斷與架構權衡〕

這種「每個函式用不同巨集」的做法看起來繁瑣,但它是必要的:libibverbs 的 API 錯誤約定極不統一,有的回傳 0/-1,有的回傳 errno 值,有的回傳指標。如果強行統一,反而會遺失錯誤資訊。NCCL 選擇「如實翻譯」,把複雜性留在封裝層,讓上層net_ib.cc只需判斷ncclSuccess。

13.2 ibvcore.h:不依賴標頭檔的 ABI 契約

直覺模型:自帶字典的翻譯官

ibvcore.h是一個奇特的檔案——它把 libibverbs 的核心結構體、列舉、常數重新定義了一遍。為什麼?因為 NCCL 要在不#include <infiniband/verbs.h>的前提下使用這些型別。

〔設計推斷與架構權衡〕

這解決了一個真實的工程問題:infiniband/verbs.h在不同發行版、不同驅動版本下內容不同。如果 NCCL 直接包含它,編譯期就綁定了某個版本。而透過自己定義一份「最小必要子集」,NCCL 可以在編譯時不需要 IB 標頭檔,執行時透過dlopen載入任意版本的函式庫。

如果缺少這層,災難是:在沒裝libibverbs-dev的機器上無法編譯 NCCL。而實際上執行時可能透過rdma-core提供了函式庫檔案。

關鍵結構體的記憶體佈局

我們挑幾個對理解 RDMA 最關鍵的結構體剖析。

ibv_gid:全域識別碼

📎 src/include/ibvcore.h:58-64

c
union ibv_gid {
	uint8_t			raw[16];
	struct {
		uint64_t	subnet_prefix;
		uint64_t	interface_id;
	} global;
};

GID 是 InfiniBand 的「IP 位址」,16 位元組。它既可作為 16 位元組陣列存取,也可作為兩個 64 位元整數存取。RoCE(RDMA over Converged Ethernet)場景下,GID 實際上就是 IPv6 位址——這也是為什麼ibvGetGidStr用inet_ntop(AF_INET6, ...)來格式化:

📎 src/include/ibvwrap.h:102-108

c
static inline const char* ibvGetGidStr(union ibv_gid* gid, char* gidStr, size_t strLen) {
  static_assert(sizeof(union ibv_gid) == sizeof(struct in6_addr),
                "the sizeof struct ibv_gid must be the size of struct in6_addr");
  return inet_ntop(AF_INET6, gid->raw, gidStr, strLen);
}

static_assert在編譯期保證ibv_gid和in6_addr大小一致,這樣inet_ntop才能正確解釋這 16 位元組。

ibv_mr:記憶體註冊句柄

📎 src/include/ibvcore.h:402-410

c
struct ibv_mr {
	struct ibv_context     *context;
	struct ibv_pd	       *pd;
	void		       *addr;
	size_t			length;
	uint32_t		handle;
	uint32_t		lkey;
	uint32_t		rkey;
};

這是 GPUDirect RDMA 的核心。addr是註冊的記憶體起始位址(可以是 host 記憶體,也可以是 GPU 顯示記憶體映射到 host 的位址),length是長度。lkey(local key)和rkey(remote key)是網卡用來驗證存取權限的「鑰匙」——傳送方在 WQE 裡帶上lkey,接收方用rkey校驗。

〔設計推斷與架構權衡〕

為什麼需要註冊?因為網卡做 DMA 時用的是實體位址,而addr是虛擬位址。註冊過程讓驅動把這段虛擬位址的分頁表「釘住」(pin),建立 IOMMU 映射,並回傳lkey/rkey作為後續引用的句柄。註冊是昂貴的(涉及分頁表遍歷和 IOMMU 程式設計),所以 NCCL 會快取 MR,避免每次傳輸都註冊。

ibv_send_wr:傳送工作請求

📎 src/include/ibvcore.h:704-738

c
struct ibv_send_wr {
	uint64_t		wr_id;
	struct ibv_send_wr     *next;
	struct ibv_sge	       *sg_list;
	int			num_sge;
	enum ibv_wr_opcode	opcode;
	int			send_flags;
	uint32_t		imm_data;
	union {
		struct {
			uint64_t	remote_addr;
			uint32_t	rkey;
		} rdma;
		// ...
	} wr;
};

這是「我要網卡做什麼」的描述。wr_id是使用者自訂的標籤(完成時會原樣回傳),sg_list是散列表(scatter-gather list),opcode決定操作類型(RDMA_WRITE、SEND 等),wr.rdma.remote_addr和wr.rdma.rkey指定對端的目標位址和存取密鑰。

ibv_sge描述一段本地記憶體:

📎 src/include/ibvcore.h:698-702

c
struct ibv_sge {
	uint64_t		addr;
	uint32_t		length;
	uint32_t		lkey;
};

注意addr是uint64_t而非指標——因為 WQE 會被網卡硬體讀取,必須是固定的 64 位格式。

內聯函式:繞過符號表的快路徑

有些函式 NCCL 選擇內聯實作,而不是走符號表。比如ibv_post_send:

📎 src/include/ibvcore.h:1099-1101

c
static inline int ibv_post_send(struct ibv_qp *qp, struct ibv_send_wr *wr, struct ibv_send_wr **bad_wr) {
  return qp->context->ops.post_send(qp, wr, bad_wr);
}

它直接透過qp->context->ops.post_send函式指標呼叫。這是 libibverbs 的經典設計:ibv_context裡有一個ops結構體,包含所有操作函式指標,由具體驅動填充。

〔設計推斷與架構權衡〕

為什麼post_send走ops而不走符號表?因為post_send是資料路徑上的熱函式,每次發送都要呼叫。如果走dlsym解析的全域符號表,會多一次間接定址。而透過qp->context->ops,編譯器可以做更好的最佳化,且這個指標在 QP 建立時就固定了。相比之下,ibv_modify_qp是控制路徑函式,呼叫頻率低,走符號表無所謂。

NCCL 的封裝wrap_ibv_post_send也是內聯的:

📎 src/include/ibvwrap.h:77-85

c
static inline ncclResult_t wrap_ibv_post_send(struct ibv_qp* qp, struct ibv_send_wr* wr, struct ibv_send_wr** bad_wr) {
  int ret = qp->context->ops.post_send(
    qp, wr, bad_wr);
  if (ret != IBV_SUCCESS) {
    WARN("ibv_post_send() failed with error %s, Bad WR %p, First WR %p", strerror(ret), wr, *bad_wr);
    return ncclSystemError;
  }
  return ncclSuccess;
}

注意IBV_SUCCESS定義為 0:

📎 src/include/ibvwrap.h:23-25

c
typedef enum ibv_return_enum {
  IBV_SUCCESS = 0,
} ibv_return_t;

設計思考:ABI 相容性的「版本探測」

ibvcore.h裡有一段精妙的 ABI 版本探測程式碼:

📎 src/include/ibvcore.h:81

c
static void *__VERBS_ABI_IS_EXTENDED = ((uint8_t *)NULL) - 1;

這是一個「魔法指標」——值為(uint8_t*)0 - 1,即0xFFFFFFFFFFFFFFFF。它被用作ibv_context.abi_compat欄位的標記值:

📎 src/include/ibvcore.h:1072-1081

c
static inline struct verbs_context *verbs_get_ctx(struct ibv_context *ctx)
{
	if (ctx->abi_compat != __VERBS_ABI_IS_EXTENDED)
		return NULL;
	return (struct verbs_context *)(((uintptr_t)ctx) -
					offsetof(struct verbs_context,
						 context));
}

如果abi_compat等於這個魔法值,說明底層函式庫支援擴充 ABI,此時可以透過container_of技巧從ibv_context反推出外層的verbs_context。verbs_context的最後一個欄位就是ibv_context:

📎 src/include/ibvcore.h:1068-1069

c
	size_t   sz;			/* Must be immediately before struct ibv_context */
	struct ibv_context context;	/* Must be last field in the struct */
〔設計推斷與架構權衡〕

這是 C 語言實作「繼承」的經典手法:verbs_context「繼承」了ibv_context,透過把基類放在末尾,可以用container_of從基類指標反推衍生類指標。sz欄位記錄結構體大小,用於版本相容——新版本函式庫可以擴充結構體,老版本程式碼透過檢查sz判斷某個欄位是否存在。

verbs_get_ctx_op巨集進一步封裝了這個檢查:

📎 src/include/ibvcore.h:1083-1086

c
#define verbs_get_ctx_op(ctx, op) ({ \
	struct verbs_context *__vctx = verbs_get_ctx(ctx); \
	(!__vctx || (__vctx->sz < sizeof(*__vctx) - offsetof(struct verbs_context, op)) || \
	 !__vctx->op) ? NULL : __vctx; })

它檢查三件事:是否是擴充 ABI、結構體是否足夠大包含該欄位、該欄位是否非空。只有全部滿足才回傳有效指標。這就是ibv_query_port_ex能安全呼叫的基礎:

📎 src/include/ibvcore.h:1121-1132

c
static inline int ibv_query_port_ex(struct ibv_context *context,
				    uint8_t port_num,
				    struct ibv_port_attr *port_attr)
{
	struct verbs_context *vctx = verbs_get_ctx_op(context, query_port);
        if (vctx) {
          return vctx->query_port(context, port_num, port_attr, sizeof(*port_attr));
        }
        return -1;
}

如果底層函式庫不支援擴充query_port,回傳 -1,呼叫方wrap_ibv_query_port會回退到老 API:

📎 src/misc/ibvwrap.cc:156-171

c
ncclResult_t wrap_ibv_query_port(struct ibv_context* context, uint8_t port_num, struct ibv_port_attr* port_attr) {
#ifndef NCCL_BUILD_RDMA_CORE
  // First try and query the extended port attributes (e.g. active_speed_ex)
  if (ibv_query_port_ex(context, port_num, port_attr) != 0) {
    // Fall back to the original attribute API call, but zero all members first
    memset(port_attr, 0, sizeof(*port_attr));
    IBV_INT_CHECK_RET_ERRNO(ibvSymbols, ibv_internal_query_port, ibv_internal_query_port(context, port_num, port_attr),
                            0, "ibv_query_port");
  }
#else
  IBV_INT_CHECK_RET_ERRNO(ibvSymbols, ibv_internal_query_port, ibv_internal_query_port(context, port_num, port_attr), 0,
                          "ibv_query_port");
#endif
  return ncclSuccess;
}

注意memset(port_attr, 0, sizeof(*port_attr))——回退前先清零,因為老 API 不會填充active_speed_ex等新欄位,如果不清零會讀到堆疊上的垃圾值。

13.3 QP 狀態機與 modify_qp 的重試藝術

直覺模型:QP 是「打電話」的完整流程

Queue Pair(QP)是 RDMA 通訊的基本單位,它包含發送佇列(SQ)和接收佇列(RQ)。建立一條 QP 就像打電話:先撥號(RESET→INIT),等對方接聽(INIT→RTR),確認雙方都能聽見(RTR→RTS),然後才能通話。

如果 QP 狀態機出錯,災難是:網卡無法建立連線,所有跨機通訊失敗,訓練任務卡死或崩潰。而 QP 狀態轉換恰恰是最容易出問題的地方——網路抖動、GID 變化、跨 rail 連線錯誤都會導致ibv_modify_qp失敗。

狀態列舉與轉換

📎 src/include/ibvcore.h:636-645

c
enum ibv_qp_state {
	IBV_QPS_RESET,
	IBV_QPS_INIT,
	IBV_QPS_RTR,
	IBV_QPS_RTS,
	IBV_QPS_SQD,
	IBV_QPS_SQE,
	IBV_QPS_ERR,
	IBV_QPS_UNKNOWN
};

這是標準的 RDMA QP 狀態機。NCCL 的ibvQpStateName把列舉翻譯成可讀字串用於日誌:

📎 src/misc/ibvwrap.cc:263-293

c
static void ibvQpStateName(enum ibv_qp_state state, char* msg, const size_t len) {
  switch (state) {
  case (IBV_QPS_RESET):
    snprintf(msg, len, "RESET");
    break;
  case (IBV_QPS_INIT):
    snprintf(msg, len, "INIT");
    break;
  // ...
  }
}

下面這張狀態圖精確對應原始碼中的列舉與轉換語意:

mermaid
stateDiagram-v2
    [*] --> RESET : ibv_create_qp()
    RESET --> INIT : modify_qp(IBV_QPS_INIT) [设置 pkey_index, port]
    INIT --> RTR : modify_qp(IBV_QPS_RTR) [设置 ah_attr, dest_qp_num, rq_psn]
    RTR --> RTS : modify_qp(IBV_QPS_RTS) [设置 sq_psn, timeout, retry_cnt]
    RTS --> SQD : modify_qp(IBV_QPS_SQD) [SQ Drain]
    SQD --> RTS : modify_qp(IBV_QPS_RTS)
    RTS --> ERR : 硬件错误 / WC 错误
    RTR --> ERR : 硬件错误
    ERR --> RESET : modify_qp(IBV_QPS_RESET) [错误恢复]
〔設計推斷與架構權衡〕

注意IBV_QPS_SQD(SQ Drained)和IBV_QPS_SQE(SQ Error)這兩個狀態。SQD 用於優雅關閉——排空發送佇列後再轉換。SQE 表示發送佇列出錯。NCCL 在正常路徑上不會主動進入這兩個狀態,但錯誤處理時需要識別它們。

Step-by-Step:modify_qp 的重試邏輯

wrap_ibv_modify_qp是本章最複雜的函式,它實作了一套完整的重試機制:

📎 src/misc/ibvwrap.cc:360-385

c
ncclResult_t wrap_ibv_modify_qp(struct ibv_qp* qp, struct ibv_qp_attr* attr, int attr_mask) {
  char qpMsg[1024];
  int ret = 0, attempts = 0;
  int maxCnt = (int)ncclParamIbMQpRetryCnt() + 1; // number of attempts = number of retry + 1
  int timeOut = (int)ncclParamIbMQpRetryTimeout();
  CHECK_NOT_NULL(ibvSymbols, ibv_internal_modify_qp);
  do {
    if (attempts > 0) {
      unsigned int sleepTime = timeOut * attempts;
      ibvModifyQpLog(qp, attr->qp_state, attr, attr_mask, qpMsg, sizeof(qpMsg));
      INFO(NCCL_NET, "Call to ibv_modify_qp failed with %d %s, %s, retrying %d/%d after %u msec of sleep", ret,
           strerror(ret), qpMsg, attempts, maxCnt, sleepTime);
      // sleep before retrying
      std::this_thread::sleep_for(std::chrono::milliseconds(sleepTime));
    }
    ret = ibvSymbols.ibv_internal_modify_qp(qp, attr, attr_mask);
    attempts++;
  } while (IBV_MQP_RETRY_ERRNO_ALL(ret) && attempts < maxCnt);
  if (ret != 0) {
    ibvModifyQpLog(qp, attr->qp_state, attr, attr_mask, qpMsg, sizeof(qpMsg));
    WARN("Call to ibv_modify_qp failed with %d %s, %s", ret, strerror(ret), qpMsg);
    printIbModifyQpHint(ret);
    return ncclSystemError;
  }
  return ncclSuccess;
}

逐步拆解:

第一步:讀取參數。maxCnt = IbMQpRetryCnt() + 1,預設重試 34 次,所以最多嘗試 35 次。timeOut預設 100 毫秒。

第二步:進入重試迴圈。第一次attempts == 0,不 sleep,直接呼叫。之後每次失敗,sleepTime = timeOut * attempts——這是線性退避,第 1 次重試等 100ms,第 2 次等 200ms,第 34 次等 3400ms。

第三步:判斷是否重試。IBV_MQP_RETRY_ERRNO_ALL(ret)決定是否繼續:

📎 src/misc/ibvwrap.cc:107-109

c
#define IBV_ERR_EQ(e, code) (e == code || e == (-code))
#define IBV_MQP_RETRY_ERRNO(e) (IBV_ERR_EQ(e, ETIMEDOUT))
#define IBV_MQP_RETRY_ERRNO_ALL(e) (ncclParamIbMQpRetryAll() ? (e != 0) : IBV_MQP_RETRY_ERRNO(e))

預設只對ETIMEDOUT重試。IBV_ERR_EQ同時匹配正負值,因為不同驅動可能回傳ETIMEDOUT或-ETIMEDOUT。如果設定了NCCL_IB_MQP_RETRY_ALL=1,則對任何非零錯誤都重試。

第四步:失敗時列印診斷資訊。ibvModifyQpLog收集裝置名、埠號、當前狀態、目標狀態、本地/遠端 GID:

📎 src/misc/ibvwrap.cc:297-339

c
static void ibvModifyQpLog(struct ibv_qp* qp, enum ibv_qp_state qpState, struct ibv_qp_attr* userAttr, int userFlag,
                           char* msg, size_t msgLen) {
  // ...
  char nextState[32], currState[32];
  ibvQpStateName(qp->state, currState, sizeof(currState));
  ibvQpStateName(qpState, nextState, sizeof(nextState));
  char devName[IBV_SYSFS_NAME_MAX] = "";
  snprintf(devName, sizeof(devName), "%s",
           (qp->pd->context) ? wrap_ibv_get_device_name(qp->pd->context->device) : "N/A");
  // ...
}

注意QP_ATTR巨集的巧妙設計:

📎 src/misc/ibvwrap.cc:295

c
#define QP_ATTR(attr, userAttr, userFlag, mask) ((userFlag & mask) ? (userAttr) : (attr))

它優先使用使用者傳入的屬性(如果attr_mask裡設定了對應位),否則回退到query_qp查到的當前屬性。這樣即使query_qp失敗,也能從使用者參數裡拿到部分資訊。

第五步:失敗時給出提示。printIbModifyQpHint針對常見錯誤碼給出排查建議:

📎 src/misc/ibvwrap.cc:341-358

c
static void printIbModifyQpHint(int status) {
  switch (status) {
  case ETIMEDOUT:
    INFO(NCCL_NET, "HINT: In many cases this error indicates that the NICs are not cross-rail connected.");
    INFO(NCCL_NET, "HINT: To confirm, set NCCL_CROSS_NIC=0 to disable cross-rail communication ...");
    return;
  case EINVAL:
    INFO(NCCL_NET, "HINT: In many cases this error indicates that an incorrect GID index is forced by "
                   "NCCL_IB_GID_INDEX, or that a NIC's GID changed mid-run.");
    // ...
  }
}
〔設計推斷與架構權衡〕

這段提示是生產經驗的結晶。ETIMEDOUT最常見的原因是跨 rail 連線問題——在多 rail 網路裡,如果 rank A 的 NIC 0 試圖連接 rank B 的 NIC 1,而它們不在同一 rail,就會逾時。EINVAL通常是 GID 索引配置錯誤,或者執行中 GID 發生變化(比如網卡重置)。

並發控制與硬體互動

wrap_ibv_modify_qp本身沒有加鎖——它假設呼叫者保證同一個 QP 不會被多執行緒同時修改。這在 NCCL 裡是成立的:QP 建立發生在初始化階段,由單個執行緒完成。

〔設計推斷與架構權衡〕

但重試迴圈裡的std::this_thread::sleep_for值得注意。它會讓出 CPU,但不釋放任何鎖(因為本來就沒持鎖)。 在 proxy 執行緒裡呼叫這個函式時,sleep 會阻塞 proxy 的進度推進——如果 QP 建立卡住,整個通訊會停滯。這就是為什麼預設重試次數是 34 次、總時間約 60 秒——足夠覆蓋短暫的網路抖動,但不會無限等待。

13.4 記憶體註冊:GPUDirect RDMA 的入口

直覺模型:給網卡發一張「門禁卡」

網卡要直接讀寫記憶體,必須先「認識」這塊記憶體。記憶體註冊(ibv_reg_mr)就是給網卡發一張門禁卡——告訴它這塊記憶體的實體位址範圍,並回傳一個lkey(本地鑰匙)和rkey(遠端鑰匙)。之後網卡做 DMA 時,就憑這把鑰匙存取。

如果缺少記憶體註冊,災難是:網卡無法存取任何記憶體,RDMA 完全無法工作。更隱蔽的問題是:如果註冊了 host 記憶體但想存取 GPU 顯存,網卡會讀到錯誤的資料或觸發保護錯誤。

三種註冊路徑

NCCL 封裝了三種記憶體註冊函式,對應不同的使用場景:

路徑一:普通註冊

📎 src/misc/ibvwrap.cc:198-201

c
ncclResult_t wrap_ibv_reg_mr(struct ibv_mr** ret, struct ibv_pd* pd, void* addr, size_t length, int access) {
  IBV_PTR_CHECK_ERRNO(ibvSymbols, ibv_internal_reg_mr, ibv_internal_reg_mr(pd, addr, length, access), *ret, NULL,
                      "ibv_reg_mr");
}

這是標準路徑,addr是虛擬位址,access是存取權限標誌(IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE等)。

路徑二:指定 IOVA 註冊

📎 src/misc/ibvwrap.cc:211-219

c
ncclResult_t wrap_ibv_reg_mr_iova2(struct ibv_mr** ret, struct ibv_pd* pd, void* addr, size_t length, uint64_t iova,
                                   int access) {
  if (ibvSymbols.ibv_internal_reg_mr_iova2 == NULL) {
    return ncclInternalError;
  }
  if (ret == NULL) return ncclSuccess; // Assume dummy call
  IBV_PTR_CHECK_ERRNO(ibvSymbols, ibv_internal_reg_mr_iova2, ibv_internal_reg_mr_iova2(pd, addr, length, iova, access),
                      *ret, NULL, "ibv_reg_mr_iova2");
}

iova(I/O Virtual Address)允許指定網卡看到的位址。這在需要固定位址映射的場景有用。注意ret == NULL時直接回傳成功——這是「探測呼叫」,只檢查函式是否存在,不真正註冊。

路徑三:DMA-BUF 註冊(GPUDirect RDMA 的關鍵)

📎 src/misc/ibvwrap.cc:222-227

c
ncclResult_t wrap_ibv_reg_dmabuf_mr(struct ibv_mr** ret, struct ibv_pd* pd, uint64_t offset, size_t length,
                                    uint64_t iova, int fd, int access) {
  IBV_PTR_CHECK_ERRNO(ibvSymbols, ibv_internal_reg_dmabuf_mr,
                      ibv_internal_reg_dmabuf_mr(pd, offset, length, iova, fd, access), *ret, NULL,
                      "ibv_reg_dmabuf_mr");
}

這是 GPUDirect RDMA 的核心。fd是一個 DMA-BUF 檔案描述符——它代表一塊 GPU 顯存。NCCL 透過cuMemGetHandleForAddressRange之類的 CUDA API 拿到這個 fd,然後傳給ibv_reg_dmabuf_mr。網卡驅動透過 DMA-BUF 機制直接映射 GPU 顯存,無需經過 host 記憶體拷貝。

〔設計推斷與架構權衡〕

DMA-BUF 是 Linux 核心的緩衝區共享框架。GPU 驅動(如 NVIDIA 的 nvidia.ko)把顯存匯出為 DMA-BUF,網卡驅動(如 mlx5)匯入它,建立 IOMMU 映射。整個過程在核心完成,使用者態只傳遞一個 fd。這就是「網卡直接讀寫 GPU 顯存」的底層機制。

直接註冊 vs 封裝註冊

注意有兩個「direct」版本:

📎 src/misc/ibvwrap.cc:203-209

c
struct ibv_mr* wrap_direct_ibv_reg_mr(struct ibv_pd* pd, void* addr, size_t length, int access) {
  if (ibvSymbols.ibv_internal_reg_mr == NULL) {
    WARN("lib wrapper not initialized.");
    return NULL;
  }
  return ibvSymbols.ibv_internal_reg_mr(pd, addr, length, access);
}

📎 src/misc/ibvwrap.cc:229-236

c
struct ibv_mr* wrap_direct_ibv_reg_dmabuf_mr(struct ibv_pd* pd, uint64_t offset, size_t length, uint64_t iova, int fd,
                                             int access) {
  if (ibvSymbols.ibv_internal_reg_dmabuf_mr == NULL) {
    errno = EOPNOTSUPP; // ncclIbDmaBufSupport() requires this errno being set
    return NULL;
  }
  return ibvSymbols.ibv_internal_reg_dmabuf_mr(pd, offset, length, iova, fd, access);
}

它們直接回傳ibv_mr*而非ncclResult_t,且不列印 WARN 日誌。為什麼?

〔設計推斷與架構權衡〕

因為這兩個函式被用於能力探測。ncclIbDmaBufSupport()會呼叫wrap_direct_ibv_reg_dmabuf_mr試探網卡是否支援 DMA-BUF。如果失敗,它期望拿到errno == EOPNOTSUPP來判斷「不支援」而非「出錯」。如果這裡列印 WARN,會在不支援 DMA-BUF 的機器上刷屏。所以 direct 版本把錯誤處理的責任交給呼叫者。

存取權限標誌

📎 src/include/ibvcore.h:365-372

c
enum ibv_access_flags {
	IBV_ACCESS_LOCAL_WRITE		= 1,
	IBV_ACCESS_REMOTE_WRITE		= (1<<1),
	IBV_ACCESS_REMOTE_READ		= (1<<2),
	IBV_ACCESS_REMOTE_ATOMIC	= (1<<3),
	IBV_ACCESS_MW_BIND		= (1<<4),
	IBV_ACCESS_RELAXED_ORDERING     = (1<<20),
};

這些標誌是位元遮罩,可以組合。LOCAL_WRITE允許本地寫(接收資料時需要),REMOTE_WRITE允許遠端寫(RDMA WRITE 的目標需要),REMOTE_READ允許遠端讀(RDMA READ 的目標需要)。

IBV_ACCESS_RELAXED_ORDERING是一個效能最佳化標誌——它允許網卡用更寬鬆的記憶體序存取,可能提升吞吐,但需要應用層保證正確性。

資料流:從 GPU 顯存到網卡的完整路徑

下面這張圖展示一次跨機 RDMA 寫的資料流,錨定本章涉及的結構體:

mermaid
flowchart LR
    subgraph GPU["GPU 显存"]
        buf["ncclSendBuff<br/>(device ptr)"]
    end
    subgraph Host["Host 进程"]
        dmabuf["DMA-BUF fd<br/>(cuMemGetHandleForAddressRange)"]
        mr["ibv_mr<br/>{addr, lkey, rkey}"]
        wr["ibv_send_wr<br/>{opcode=RDMA_WRITE,<br/>sg_list, wr.rdma.remote_addr, rkey}"]
    end
    subgraph NIC["网卡 mlx5"]
        qp["ibv_qp<br/>(SQ + RQ)"]
        wqe["WQE<br/>(硬件工作队列元素)"]
    end
    buf -->|导出| dmabuf
    dmabuf -->|ibv_reg_dmabuf_mr| mr
    mr -->|填充 sge.lkey| wr
    wr -->|ibv_post_send| qp
    qp -->|DMA 读取| wqe
    wqe -->|PCIe P2P| buf
    wqe -->|网络| remote["对端 GPU 显存<br/>(remote_addr + rkey)"]

圖中每個節點都對應原始碼中的真實型別:ibv_mr來自📎 src/include/ibvcore.h:402-410,ibv_send_wr來自📎 src/include/ibvcore.h:704-738,ibv_qp來自📎 src/include/ibvcore.h:787-802。

13.5 工作完成與錯誤診斷

直覺模型:快遞簽收單

RDMA 是非同步的——你post_send之後不會立即知道結果。網卡完成操作後,會在 Completion Queue(CQ)裡放一個 Work Completion(WC),就像快遞員把簽收單放進你的信箱。你需要主動poll_cq去取。

如果缺少 WC 診斷,災難是:通訊失敗時你只知道「失敗了」,不知道「為什麼失敗」。RDMA 的錯誤碼有 20 多種,每種對應不同的根因。

WC 結構體

📎 src/include/ibvcore.h:349-363

c
struct ibv_wc {
	uint64_t		wr_id;
	enum ibv_wc_status	status;
	enum ibv_wc_opcode	opcode;
	uint32_t		vendor_err;
	uint32_t		byte_len;
	uint32_t		imm_data;	/* in network byte order */
	uint32_t		qp_num;
	uint32_t		src_qp;
	int			wc_flags;
	uint16_t		pkey_index;
	uint16_t		slid;
	uint8_t			sl;
	uint8_t			dlid_path_bits;
};

wr_id是你 post 時填的標籤,status是完成狀態,opcode是操作型別,byte_len是實際傳輸位元組數。qp_num和src_qp用於多 QP 場景下識別是哪個 QP 完成的。

狀態碼翻譯

ibvWcStatusStr把狀態列舉翻譯成字串:

📎 src/misc/ibvwrap.cc:415-464

c
const char* ibvWcStatusStr(enum ibv_wc_status status) {
  switch (status) {
  case IBV_WC_SUCCESS:
    return "IBV_WC_SUCCESS";
  case IBV_WC_LOC_LEN_ERR:
    return "IBV_WC_LOC_LEN_ERR";
  // ... 20 多个 case
  default:
    return "UNKNOWN_STATUS";
  }
}

這些狀態碼的含義:

狀態碼含義常見根因
IBV_WC_SUCCESS成功—
IBV_WC_LOC_LEN_ERR本地長度錯誤SGE 長度超過 MR 範圍
IBV_WC_LOC_ACCESS_ERR本地存取錯誤lkey 無效或權限不足
IBV_WC_REM_ACCESS_ERR遠端存取錯誤rkey 無效或對端 MR 已註銷
IBV_WC_RETRY_EXC_ERR重試耗盡網路不通或對端 QP 未就緒
IBV_WC_RNR_RETRY_EXC_ERRRNR 重試耗盡對端沒有 post recv
IBV_WC_RESP_TIMEOUT_ERR回應逾時對端無回應
〔設計推斷與架構權衡〕

IBV_WC_RNR_RETRY_EXC_ERR(Receiver Not Ready)是生產環境最常見的問題之一。它意味著發送方發了資料,但接收方沒有預先 post 足夠的 recv buffer。在 NCCL 裡,這通常發生在連線建立階段——雙方 QP 狀態不同步,一方已經開始發送,另一方還沒準備好接收。

opcode 翻譯

ibvWcOpcodeStr和ibvWrOpcodeStr分別翻譯完成 opcode 和請求 opcode:

📎 src/misc/ibvwrap.cc:467-488

c
const char* ibvWcOpcodeStr(enum ibv_wc_opcode opcode) {
  switch (opcode) {
  case IBV_WC_SEND:
    return "IBV_WC_SEND";
  case IBV_WC_RDMA_WRITE:
    return "IBV_WC_RDMA_WRITE";
  case IBV_WC_RDMA_READ:
    return "IBV_WC_RDMA_READ";
  // ...
  }
}

注意IBV_WC_RECV的值是1 << 7:

📎 src/include/ibvcore.h:329-342

c
enum ibv_wc_opcode {
	IBV_WC_SEND,
	IBV_WC_RDMA_WRITE,
	IBV_WC_RDMA_READ,
	IBV_WC_COMP_SWAP,
	IBV_WC_FETCH_ADD,
	IBV_WC_BIND_MW,
	IBV_WC_RECV			= 1 << 7,
	IBV_WC_RECV_RDMA_WITH_IMM
};
〔設計推斷與架構權衡〕

為什麼IBV_WC_RECV是1 << 7而不是順序值?因為接收完成和發送完成是兩類不同的操作,用高位區分可以讓程式碼用opcode & IBV_WC_RECV快速判斷「這是不是一個接收完成」。這是 libibverbs 的 API 設計約定。

輪詢 CQ

wrap_ibv_poll_cq是內聯的:

📎 src/include/ibvwrap.h:60-69

c
static inline ncclResult_t wrap_ibv_poll_cq(struct ibv_cq* cq, int num_entries, struct ibv_wc* wc, int* num_done) {
  int done = cq->context->ops.poll_cq(cq, num_entries,
                                      wc);
  if (done < 0) {
    WARN("Call to ibv_poll_cq() returned %d", done);
    return ncclSystemError;
  }
  *num_done = done;
  return ncclSuccess;
}

它透過cq->context->ops.poll_cq呼叫,和post_send一樣走ops快路徑。回傳值done是本次輪詢到的 WC 數量,0 表示沒有新完成,負數表示錯誤。

〔設計推斷與架構權衡〕

poll_cq是忙輪詢——它不阻塞,立即回傳。NCCL 的 proxy 執行緒會在迴圈裡反覆呼叫它,直到拿到完成事件。這是低延遲的關鍵:相比中斷驅動,忙輪詢避免了中斷上下文切換的開銷。代價是 CPU 佔用高,但在高效能運算場景下這是可接受的。

13.6 生產避坑指南

坑一:跨 rail 連線逾時

現象:ibv_modify_qp回傳ETIMEDOUT,重試 34 次後失敗。

根因:在多 rail 網路裡,每個 GPU 通常綁定到特定的 NIC。如果 rank A 的 GPU 0 綁定了 NIC 0,rank B 的 GPU 0 綁定了 NIC 1,而 NIC 0 和 NIC 1 不在同一 rail(即它們連接不同的交換器),那麼 QP 建立會逾時。

排查:原始碼已經給出了提示:

📎 src/misc/ibvwrap.cc:343-347

c
  case ETIMEDOUT:
    INFO(NCCL_NET, "HINT: In many cases this error indicates that the NICs are not cross-rail connected.");
    INFO(NCCL_NET, "HINT: To confirm, set NCCL_CROSS_NIC=0 to disable cross-rail communication ...");
    return;

設定NCCL_CROSS_NIC=0可以強制同 rail 通訊。如果這樣能解決,說明確實是跨 rail 問題。

恢復鏈:NCCL 的重試機制(34 次、線性退避)給了網路足夠時間恢復。但如果根因是拓撲配置錯誤,重試無用,必須修正NCCL_IB_HCA或NCCL_CROSS_NIC配置。

坑二:GID 索引錯誤

現象:ibv_modify_qp回傳EINVAL。

根因:NCCL_IB_GID_INDEX強制指定了一個不存在的 GID 索引,或者執行中網卡的 GID 發生了變化(比如 RoCE 網卡重新取得 IP)。

排查:

📎 src/misc/ibvwrap.cc:341-358

c
  case EINVAL:
    INFO(NCCL_NET, "HINT: In many cases this error indicates an incorrect GID index is forced by "
                   "NCCL_IB_GID_INDEX, or that a NIC's GID changed mid-run.");
    INFO(NCCL_NET, "HINT: To confirm, set NCCL_IB_GID_INDEX=-1 to enable automatic detection and check "
                   "'dmesg | grep -i gid' for GID changes ...");
    return;

設定NCCL_IB_GID_INDEX=-1啟用自動偵測。同時檢查dmesg裡是否有 GID 變化事件。

坑三:DMA-BUF 不支援導致回退到 host 複製

現象:GPUDirect RDMA 沒有生效,效能低於預期。

根因:網卡驅動或核心不支援 DMA-BUF,wrap_direct_ibv_reg_dmabuf_mr回傳 NULL 並設定errno = EOPNOTSUPP:

📎 src/misc/ibvwrap.cc:229-236

c
struct ibv_mr* wrap_direct_ibv_reg_dmabuf_mr(struct ibv_pd* pd, uint64_t offset, size_t length, uint64_t iova, int fd,
                                             int access) {
  if (ibvSymbols.ibv_internal_reg_dmabuf_mr == NULL) {
    errno = EOPNOTSUPP; // ncclIbDmaBufSupport() requires this errno being set
    return NULL;
  }
  return ibvSymbols.ibv_internal_reg_dmabuf_mr(pd, offset, length, iova, fd, access);
}

注意註解:ncclIbDmaBufSupport()依賴這個errno來判斷是否支援。如果這裡不設定EOPNOTSUPP,上層會誤判為「出錯」而非「不支援」。

排查:檢查核心版本(需要 5.12+)、網卡驅動版本、以及nvidia-peermem模組是否載入。如果確實不支援,NCCL 會回退到 host 記憶體中轉,效能會下降但功能正常。

坑四:MR 快取與記憶體洩漏

〔設計推斷與架構權衡〕

記憶體註冊是昂貴的操作(涉及 IOMMU 編程),NCCL 會快取ibv_mr。但如果快取策略不當,會導致兩個問題:一是記憶體洩漏(MR 一直不註銷),二是快取失效(記憶體被釋放但 MR 還指向舊位址)。

wrap_ibv_dereg_mr是註銷入口:

📎 src/misc/ibvwrap.cc:238-241

c
ncclResult_t wrap_ibv_dereg_mr(
  struct ibv_mr* mr) {
  IBV_INT_CHECK_RET_ERRNO(ibvSymbols, ibv_internal_dereg_mr, ibv_internal_dereg_mr(mr), 0, "ibv_dereg_mr");
}
〔設計推斷與架構權衡〕

生產環境中,如果訓練任務頻繁建立/銷毀通訊域,而 MR 沒有正確註銷,會導致 IOMMU 映射表膨脹,最終觸發ibv_reg_mr失敗(回傳ENOMEM)。排查方法是監控/sys/kernel/debug/iommu下的映射數量。

設計思考:為什麼封裝層如此「厚」

回顧本章,ibvwrap.cc有 509 行,ibvcore.h有 1134 行。對於一個「只是呼叫 libibverbs」的封裝層,這個體量相當大。為什麼?

〔設計推斷與架構權衡〕

三個原因:

第一,錯誤處理的複雜性。libibverbs 的 API 錯誤約定極不統一,NCCL 需要為每種約定寫一個巨集,並在每個函式裡正確使用。這不是過度設計,而是「如實翻譯」的必要成本。

第二,ABI 相容性的負擔。ibvcore.h重新定義了所有結構體,還要處理verbs_context的版本探測。這是為了在編譯期不依賴 IB 標頭檔,執行時相容任意版本。

第三,診斷資訊的價值。ibvModifyQpLog、printIbModifyQpHint、ibvWcStatusStr這些函式在正常路徑上不會被呼叫,但在故障排查時價值巨大。NCCL 選擇把診斷資訊「預埋」在封裝層,而不是等到出錯時再臨時收集。

這種「厚封裝」的代價是程式碼量大、維護成本高。但收益是:上層net_ib.cc可以用統一的ncclResult_t介面編寫,不必關心 libibverbs 的各種怪癖。這是典型的「複雜性隔離」設計。

本章小結

本章我們深入了 NCCL 的 InfiniBand 傳輸封裝層,核心要點:

1. 符號表封裝:ncclIbvSymbols透過dlopen + dlsym執行時載入 libibverbs,配合std::once_flag保證執行緒安全初始化。這讓 NCCL 在沒有 IB 驅動的機器上也能載入。

2. ABI 契約:ibvcore.h重新定義了 libibverbs 的核心類型,透過__VERBS_ABI_IS_EXTENDED魔法指標和verbs_context的container_of技巧實現版本探測。

3. QP 狀態機:wrap_ibv_modify_qp實現了 34 次線性退避重試,針對ETIMEDOUT和EINVAL給出診斷提示。

4. GPUDirect RDMA:wrap_ibv_reg_dmabuf_mr透過 DMA-BUF 機制讓網卡直接映射 GPU 顯存,wrap_direct_ibv_reg_dmabuf_mr用於能力探測。

5. 錯誤診斷:ibvWcStatusStr、ibvWcOpcodeStr、ibvWrOpcodeStr把硬體錯誤碼翻譯成可讀字串,是生產排查的關鍵工具。

本章思考與自測

Q1: 如果把wrap_ibv_symbols裡的std::call_once換成普通的if (initResult == ncclSuccess) return initResult;雙檢鎖,在什麼並發場景下會出問題?

參考解析:看📎 src/misc/ibvwrap.cc:26-29:

c
ncclResult_t wrap_ibv_symbols(void) {
  std::call_once(initOnceFlag, []() { initResult = buildIbvSymbols(&ibvSymbols); });
  return initResult;
}

如果換成樸素的雙檢鎖,問題在於記憶體重排序。buildIbvSymbols會填充ibvSymbols的各個欄位,然後寫入initResult。在沒有記憶體屏障的情況下,CPU 或編譯器可能把initResult = ncclSuccess重排到 `

至此,我們看清了 NCCL 如何透過 net_ib 將 libibverbs 封裝為可插拔的傳輸層,並利用 GPUDirect RDMA 實現網卡對 GPU 顯存的直接存取。這套機制解決了跨機通訊的延遲與頻寬瓶頸。但機內通訊同樣關鍵——下一章我們將進入對稱記憶體與 NVLS,看 NCCL 如何利用 NVLink 多播實現硬體加速的集合通訊。屆時你會發現,本章的 RDMA 機制與 NVLS 形成互補:前者負責跨機,後者負責機內。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 14

第 14 章:對稱記憶體與 NVLS:多播加速與 LSA 裝置端直接定址

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 14 章 / 共 25 章

上一章我們跟隨一次跨機 AllReduce,看資料如何從 GPU 顯存經網卡到達對端 GPU,那條路徑解決的是機器之間的通訊。但現代 AI 叢集裡,同一台機器甚至同一個 NVLink 域內部的 GPU 間通訊量同樣巨大——資料並行訓練中的梯度同步、張量並行中的激活值交換,絕大多數都發生在機內。如果機內通訊仍走 GPU→顯存→網卡→對端網卡→顯存→GPU 這套跨機流程,就相當於同城寄快遞非要走航空件,延遲白白浪費。本章要拆解的,正是 NCCL 為機內通訊準備的兩把利器:對稱記憶體與 NVLS。前者讓每個 rank 用同一套虛擬位址存取所有 rank 的緩衝區,後者利用 NVSwitch 硬體的多播能力做歸約。兩者結合,能把小訊息集合通訊的延遲壓到接近硬體極限。

14.1 對稱記憶體:讓「第 3 排第 5 座」在每個人家里都指同一個位置

直覺模型

想像一個班級要交換作業本。傳統做法是:每個人把自己的本子編號,然後喊「張三,我的第 5 本給你;李四,我的第 8 本給你」——每個人都要記住「誰的本子放在哪、第幾本」。這就是普通通訊:位址是相對的、私有的,你要存取對端資料,得先知道對端的位址映射。

對稱記憶體換了個思路:全班約定「第 3 排第 5 座」這個座標,在每個人家里都指向同一個實體位置。於是張三要拿李四的第 5 本,直接說「李四家第 3 排第 5 座」就行,不需要任何位址轉譯。這就是對稱記憶體的核心:每個 rank 的緩衝區在所有 rank 的位址空間裡映射到相同的虛擬位址。

〔設計推斷與架構權衡〕

如果沒有對稱記憶體,機內集合通訊會面臨什麼災難? 每個 rank 存取對端緩衝區時,都要經過一次「位址轉譯」——查表、計算偏移、可能還要跨行程通訊確認映射關係。對於小訊息(幾 KB),這次轉譯的開銷可能比資料本身傳輸還大。對稱記憶體把這個開銷徹底消除,這正是它「顯著降低小訊息延遲」的根本原因。

資料結構與記憶體佈局

對稱記憶體的註冊類型由ncclSymRegType_t描述,ncclGetSymRegType根據 send/recv 視窗是否帶NCCL_WIN_COLL_SYMMETRIC標誌,把註冊狀態分成四類。

📎 src/sym_kernels.cc:395-412

c
ncclResult_t ncclGetSymRegType(struct ncclDevrWindow* sendWin, struct ncclDevrWindow* recvWin,
                               ncclSymRegType_t* winRegType) {
  bool isSendSymmReg = false;
  bool isRecvSymmReg = false;
  if (sendWin && (sendWin->winFlags & NCCL_WIN_COLL_SYMMETRIC)) isSendSymmReg = true;
  if (recvWin && (recvWin->winFlags & NCCL_WIN_COLL_SYMMETRIC)) isRecvSymmReg = true;
  // determine the registration type
  if (!isSendSymmReg && !isRecvSymmReg) {
    *winRegType = ncclSymSendNonregRecvNonreg;
  } else if (isSendSymmReg && !isRecvSymmReg) {
    *winRegType = ncclSymSendRegRecvNonreg;
  } else if (!isSendSymmReg && isRecvSymmReg) {
    *winRegType = ncclSymSendNonregRecvReg;
  } else if (isSendSymmReg && is isRecvSymmReg) {
    *winRegType = ncclSymSendRegRecvReg;
  }
  return ncclSuccess;
}

這四個狀態決定了後續 kernel 走哪條路徑:全對稱註冊(SendRegRecvReg)走最快的 LSA 路徑,全非註冊(SendNonregRecvNonreg)走普通路徑,混合狀態則要特殊處理。winFlags裡的NCCL_WIN_COLL_SYMMETRIC位就是「這個視窗是否已做對稱註冊」的標記。

對稱記憶體的初始化入口是ncclSymkInitOnce,它做了一件關鍵的事:判斷當前通訊域是否支援 LSA 多播(hasLsaMultimem)。

📎 src/sym_kernels.cc:185-196

c
ncclResult_t ncclSymkInitOnce(struct ncclComm* comm) {
  // ncclTeamLsa() below calls this internally but drops the error code so we do it here.
  NCCLCHECK(ncclDevrInitOnce(comm));

  struct ncclSymkState* symk = &comm->symkState;
  if (!symk->initialized) {
    symk->initialized = true;
    struct ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER;
    // Disable LSA multicast for cross-clique since NVLS isn't available across cliques
    symk->hasLsaMultimem =
      ncclNvlsSymmetricMultimemEnabled(comm) && ncclTeamLsa(comm).nRanks > 2 && !comm->p2pCrossClique;
    reqs.lsaMultimem = symk->hasLsaMultimem;

hasLsaMultimem的三個條件缺一不可:NVLS 對稱多播已啟用、LSA 團隊 rank 數大於 2(兩個 rank 直接點對點更快,不需要多播)、且不跨 clique(跨 clique 時 NVSwitch 多播不可用)。這個判斷直接決定了reqs.lsaMultimem是否置位,進而影響裝置側通訊器的資源分配。

場景驅動的 Step-by-Step Walkthrough

假設我們發起一次 AllReduce,訊息大小 4KB,8 個 rank 在同一 NVLink 域內。ncclSymkMask會決定哪些 kernel 可用。

📎 src/sym_kernels.cc:304-352

c
uint32_t ncclSymkMask(struct ncclComm* comm, ncclFunc_t coll, int /*ncclDevRedOp_t*/ red, ncclDataType_t ty,
                      size_t nElts, bool symAligned16B) {
  uint32_t kmask = kernelMask_coll(coll);

  bool hasSTMC = comm->symkState.hasLsaMultimem;
  bool hasLDMC = false;
  if (comm->symkState.hasLsaMultimem) {
    switch (ty) {
    case ncclInt32:
    ...
      hasLDMC = red == ncclDevSum || red == ncclDevMinMax || red == ncclDevSumPostDiv;
      break;
    ...
    }
  }
  if (!hasSTMC) kmask &= ~kernelMask_STMC;
  if (!hasLDMC) kmask &= ~kernelMask_LDMC;

第一步:kernelMask_coll根據集合類型(AllReduce)取出候選 kernel 集合kernelMask_AR。第二步:檢查hasLsaMultimem,如果支援多播,則進一步判斷資料類型和歸約操作是否支援 LDMC(Load-Multicast)。第三步:用位元遮罩清除不支援的特性——kmask &= ~kernelMask_STMC把不支援 STMC 的 kernel 全部剔除。

接著是大小限制:

📎 src/sym_kernels.cc:336-342

c
  size_t nBytes = alignUp(nElts * ncclTypeSize(ty), NCCL_SYM_KERNEL_CELL_SIZE);
  size_t nBusBytes = (coll == ncclFuncAllReduce ? 1 : comm->nRanks) * nBytes;
  // LL kernels use 32-bit ints to track element counts and indices.
  if (nBusBytes >= (size_t(2) << 30)) kmask &= ~kernelMask_LL;
  // Any kernel might use 32-bit int to track unrolled loop chunks (which are going
  // to be at least 32 bytes per chunk)
  if (nBusBytes >= 32 * (size_t(2) << 30)) kmask = 0;

這裡有兩個硬邊界:LL 系列 kernel 用 32 位元整數追蹤元素計數,所以當匯流排位元組數超過 2GB 時,LL kernel 被剔除;當超過 64GB 時,所有 kernel 都被剔除(kmask = 0)。這是典型的「用位寬換效能」——32 位元索引比 64 位元省暫存器、省指令,但代價是訊息大小上限。

最後是 TMA 和 GIN 的可用性檢查:

📎 src/sym_kernels.cc:344-350

c
  if (!ncclSymkTmaAvailable(comm)) kmask &= ~kernelMask_Tma;
  if (!symAligned16B) kmask &= ~kernelMask_Tma;

  bool hasGin = ncclParamSymGinKernelsEnable() != 0;
  if (!hasGin) kmask &= ~kernelMask_Gin;
  bool needGin = ncclTeamLsa(comm).nRanks < comm->nRanks;
  kmask &= needGin ? kernelMask_Gin : ~kernelMask_Gin;
  return kmask;

TMA 需要 SMEM 容量達標(ncclSymkTmaAvailable檢查maxSharedMemOptin)且 16 位元組對齊。GIN 則只在「LSA 團隊 rank 數小於總 rank 數」時才需要——也就是說,只有當通訊域跨越了 LSA 邊界(需要走網路)時,GIN 才有意義。如果整個通訊域都在 LSA 內,GIN kernel 被剔除。

並發控制與硬體互動

對稱記憶體的位址解析最終落到裝置側。ncclSymkMakeDevWork把 host 側的任務描述翻譯成裝置側可讀的工作項。

📎 src/sym_kernels.cc:380-393

c
ncclResult_t ncclSymkMakeDevWork(struct ncclComm* comm, struct ncclTaskColl* task, struct ncclSymkDevWork* outDevWork) {
  outDevWork->rootRank = task->root;
  outDevWork->redOpArg = task->opDev.scalarArg;
  outDevWork->nElts = task->count;
  outDevWork->inputWin = task->sendWin ? task->sendWin->vidmem : nullptr;
  outDevWork->inputOff =
    task->sendWin ? (uint8_t*)task->sendbuff - (uint8_t*)task->sendWin->userPtr : (size_t)task->sendbuff;
  outDevWork->outputWin = task->recvWin ? task->recvWin->vidmem : nullptr;
  outDevWork->outputOff =
    task->recvWin ? (uint8_t*)task->recvbuff - (uint8_t*)task->recvWin->userPtr : (size_t)task->recvbuff;
  outDevWork->sChannelId = 0xffff;
  outDevWork->nChannels = 0;
  return ncclSuccess;
}

注意inputOff的計算:如果 sendWin 存在(對稱註冊視窗),偏移是sendbuff - sendWin->userPtr——這是視窗內偏移,裝置側拿到inputWin(視窗基址)加上inputOff就能算出實際位址。如果 sendWin 不存在,偏移直接是sendbuff的絕對位址。這個設計讓裝置側 kernel 用同一套邏輯處理註冊和非註冊緩衝區。

ncclSymkInitOnce裡還初始化了 GIN 相關的資源需求,包括 inbox、outbox、accumulation buffer 和 rail signal。

📎 src/sym_kernels.cc:208-251

c
    struct ncclDevResourceRequirements ginInboxRailReq = {};
    struct ncclDevResourceRequirements ginOutboxReq = {};
    struct ncclDevResourceRequirements rsGinAccumReq = {};
    struct ncclDevResourceRequirements railSignalReq = {};
    if (ncclParamSymGinKernelsEnable() && ncclTeamLsa(comm).nRanks < comm->nRanks) {
      int maxBlocks;
      size_t bufSize;
      getRequirements_gin(comm, &maxBlocks, &bufSize);

      maxBlocks = std::max(maxBlocks, comm->config.minCTAs);
      maxBlocks = std::min(maxBlocks, comm->config.maxCTAs);
      if (ncclParamSymCTAs() >= 1) maxBlocks = ncclParamSymCTAs();
      maxBlocks = std::min(maxBlocks, ncclSymkMaxBlocks);
      symk->maxGinInboxBlocks = maxBlocks;
      symk->kcomm.rsGinAccumBytesPerBlock = ncclSymkRsGinAccumBytesPerBlock();

      rsGinAccumReq.bufferSize = (size_t)maxBlocks * symk->kcomm.rsGinAccumBytesPerBlock;
      rsGinAccumReq.bufferAlign = 128;
      rsGinAccumReq.outBufferHandle = &symk->kcomm.rsGinAccumBuf;
      ...
      uint32_t railSignalCount = ncclTeamRail(comm).nRanks * ncclSymkMaxBlocks;
      ...
      reqs.barrierCount = ncclSymkMaxBlocks;
      reqs.ginConnectionType = NCCL_GIN_CONNECTION_RAIL;
      reqs.ginStrongSignalsRequired = true;
      reqs.ginVaSignalsRequired = true;
    }

getRequirements_gin用調優模型算出需要的 block 數和緩衝區大小,然後被 clamp 到[minCTAs, maxCTAs]區間。rsGinAccumBytesPerBlock是每個 block 的累加緩衝區大小,對齊到 128 位元組——這是快取行大小,避免偽共享。

mermaid
flowchart TD
    start["ncclSymkMask(comm, coll, red, ty, nElts)"] --> coll{"集合类型?"}
    coll -->|AllGather| mask_ag["kmask = kernelMask_AG"]
    coll -->|AllReduce| mask_ar["kmask = kernelMask_AR"]
    coll -->|ReduceScatter| mask_rs["kmask = kernelMask_RS"]
    mask_ag --> check_stmc{"hasLsaMultimem?"}
    mask_ar --> check_stmc
    mask_rs --> check_stmc
    check_stmc -->|否| clear_stmc["kmask &= ~kernelMask_STMC"]
    check_stmc -->|是| check_ldmc{"数据类型+归约支持LDMC?"}
    clear_stmc --> size_check
    check_ldmc -->|否| clear_ldmc["kmask &= ~kernelMask_LDMC"]
    check_ldmc -->|是| size_check
    clear_ldmc --> size_check
    size_check{"nBusBytes >= 2GB?"} -->|是| clear_ll["kmask &= ~kernelMask_LL"]
    size_check -->|否| tma_check
    clear_ll --> tma_check{"TMA可用且16B对齐?"}
    tma_check -->|否| clear_tma["kmask &= ~kernelMask_Tma"]
    tma_check -->|是| gin_check
    clear_tma --> gin_check{"需要GIN? LSA rank < 总rank"}
    gin_check -->|否| clear_gin["kmask &= ~kernelMask_Gin"]
    gin_check -->|是| done
    clear_gin --> done["返回 kmask"]

這張圖完整刻畫了ncclSymkMask的決策鏈:從集合類型出發,依次經過多播支援、資料類型、大小邊界、TMA 可用性、GIN 需求五道過濾,最終返回一個位元遮罩。每一道過濾都可能把一批 kernel 剔除,這正是 NCCL「按場景選最優 kernel」的體現。

生產避坑指南

坑 1:跨 clique 時多播靜默失效。 hasLsaMultimem的第三個條件是!comm->p2pCrossClique。如果你的叢集配置了 MNNVL(Multi-Node NVLink),但某些 rank 跨了 clique,多播會被禁用,效能悄悄退化到普通路徑。排查時看ncclNvlsSymmetricMultimemEnabled的日誌輸出。

坑 2:16 位元組對齊的隱性要求。 ncclSymkMask裡if (!symAligned16B) kmask &= ~kernelMask_Tma;——如果使用者緩衝區不是 16 位元組對齊,TMA kernel 被剔除。TMA 是 Hopper/Blackwell 上最快的拷貝引擎,失去它意味著效能下降。生產環境裡,使用者傳入的 buffer 往往來自cudaMalloc,天然對齊;但如果來自自訂 allocator 或切片,就可能踩坑。

坑 3:2GB 邊界。LL kernel 用 32 位元索引,超過 2GB 匯流排位元組數就被剔除。對於大模型訓練,單次 AllReduce 的梯度可能超過這個值,此時 NCCL 會自動切到 STMC 或 Simple 協議。這不是 bug,但如果你手動指定了 LL 協議,會得到ncclInvalidArgument。

---

14.2 NVLS:讓 NVSwitch 硬體替你做歸約

直覺模型

傳統 AllReduce 是「軟體歸約」:每個 GPU 把資料發給鄰居,鄰居做加法,再轉發——資料在 GPU 之間來回搬運,加法在 SM 上執行。這就像 8 個人傳紙條算總和,每個人都要讀一遍、加一遍、再傳出去。

NVLS 換了個思路:NVSwitch 晶片內建了多播(multicast)和歸約(reduction)能力。你把資料往多播位址一寫,NVSwitch 自動把它廣播給所有成員,並在硬體裡完成加法。這就像 8 個人把數字寫在同一塊白板上,白板自動顯示總和——GPU 只寫一次、讀一次,中間的搬運和加法全由交換器硬體完成。

如果沒有 NVLS,機內 AllReduce 的頻寬會被 GPU 之間的點對點鏈路限制,且 SM 要花大量週期做加法。NVLS 把這兩件事都卸載到硬體,SM 可以去做別的計算。

資料結構與記憶體佈局

NVLS 的核心是多播組(MC group)。ncclMcGroup結構體描述了一個多播組的全部狀態。

📎 src/transport/multicast.cc:72-77

c
struct ncclMcGroup {
  CUmemGenericAllocationHandle handle;  // the MC object
  char* base;                          // mapped MC VA base
  size_t capacity;                      // total mapped VA size
  int dev;                           // local device, for unbind
};

四個欄位:handle是 CUDA 多播物件的句柄,base是多播虛擬位址的基址,capacity是總映射大小,dev是本地裝置號(用於解綁)。注意這裡沒有鎖——多播組的建立和銷毀都在初始化/銷毀階段,不在熱路徑上。

多播組被切分成多個分區(partition),每個分區是一個不可變的切片。ncclMcPartition描述一個分區。

📎 src/transport/multicast.cc:162-170

c
  // A partition is self-sufficient for binds: it carries the group's handle, device and
  // bind granularity alongside its own extent.
  for (int i = 0; i < nRequests; i++) {
    if (outPartitions[i].size == 0) continue;
    outPartitions[i].ptr = group->base + outPartitions[i].offset;
    outPartitions[i].mcHandle = mcHandle;
    outPartitions[i].minGranularity = minGran;
    outPartitions[i].dev = comm->cudaDev;
  }

每個分區攜帶自己的offset、size、ptr,以及所屬組的mcHandle、minGranularity、dev。這種「自給自足」的設計讓分區可以獨立傳遞給綁定函式,不需要再查組資訊。

場景驅動的 Step-by-Step Walkthrough

假設 8 個 rank 要建立一個 NVLS 域。ncclMcGroupBuildPartitions負責建立多播組並切分分區。

📎 src/transport/multicast.cc:79-121

c
ncclResult_t ncclMcGroupBuildPartitions(struct ncclComm* comm, const struct ncclMcRequest* requests, int nRequests,
                                        struct ncclMcGroup** outGroup, struct ncclMcPartition* outPartitions) {
  ...
  mcprop.numDevices = comm->localRanks;
  mcprop.handleTypes = ncclCuMemHandleType;
  mcprop.flags = 0;
  mcprop.size = 0;
  for (int i = 0; i < nRequests; i++) mcprop.size += requests[i].size;
  CUCHECKGOTO(cuMulticastGetGranularity(&recGran, &mcprop, CU_MULTICAST_GRANULARITY_RECOMMENDED), ret, fail);
  CUCHECKGOTO(cuMulticastGetGranularity(&minGran, &mcprop, CU_MULTICAST_GRANULARITY_MINIMUM), ret, fail);

  // Bump-allocate an immutable slice per request. Offsets and sizes are rounded
  // to the recommended granularity (a multiple of the MC minimum) so every slice
  // boundary is a valid bind offset.
  for (int i = 0; i < nRequests; i++) {
    outPartitions[i] = {};
    if (requests[i].size == 0) continue;
    size_t align = requests[i].alignment > recGran ? requests[i].alignment : recGran;
    ALIGN_SIZE(capacity, align);
    size_t slice = requests[i].size;
    ALIGN_SIZE(slice, recGran);
    outPartitions[i].offset = capacity;
    outPartitions[i].size = slice;
    capacity += slice;
  }

第一步:累加所有請求的大小,得到多播組總大小。第二步:查詢 CUDA 的推薦粒度和最小粒度——這是硬體約束,多播物件的位址和大小必須是粒度的整數倍。第三步:bump 分配——每個請求切一塊,偏移和大小都對齊到推薦粒度。ALIGN_SIZE(capacity, align)確保每個切片的起始偏移是合法的綁定偏移。

接下來是跨 rank 的建立與匯入:

📎 src/transport/multicast.cc:125-146

c
  if (comm->localRank == 0) {
    NCCLCHECKGOTO(ncclMcCreate(comm, &mcprop, comm->localRank, comm->localRanks, &mcHandle, shareableHandle), ret,
                  fail);
    mcCreated = 1;
    NCCLCHECKGOTO(bootstrapIntraNodeBroadcast(comm->bootstrap, comm->localRankToRank, comm->localRank, comm->localRanks,
                                              0, shareableHandle, NVLS_HANDLE_SIZE),
                  ret, fail);
  } else {
    NCCLCHECKGOTO(bootstrapIntraNodeBroadcast(comm->bootstrap, comm->localRankToRank, comm->localRank, comm->localRanks,
                                              0, shareableHandle, NVLS_HANDLE_SIZE),
                  ret, fail);
    NCCLCHECKGOTO(ncclMcImport(comm, shareableHandle, comm->localRankToRank[0], &mcHandle), ret, fail);
    mcCreated = 1;
  }
  CUCHECKGOTO(cuMulticastAddDevice(mcHandle, comm->cudaDev), ret, fail);

  // cuMemMap of an MC object blocks until every device has been added. This
  // abort-aware barrier makes a peer failing before cuMulticastAddDevice trip the
  // abort flag here instead of stranding survivors in the blocking cuMemMap.
  NCCLCHECKGOTO(bootstrapIntraNodeBarrier(comm->bootstrap, comm->localRankToRank, comm->localRank, comm->localRanks,
                                          comm->localRankToRank[0]),
                ret, fail);

localRank 0 建立多播物件,然後透過 bootstrap 廣播 shareable handle;其他 rank 接收 handle 並匯入。cuMulticastAddDevice把本地裝置加入多播組。注意那個 barrier——註解說得很清楚:cuMemMap會阻塞直到所有裝置都加入,如果某個 peer 在cuMulticastAddDevice之前失敗,倖存者會卡死在cuMemMap裡。這個 barrier 讓失敗在阻塞前就被 abort 標誌捕獲。

最後是映射和存取權限設定:

📎 src/transport/multicast.cc:148-155

c
  // Reserve and map the whole MC VA once; each consumer slice is a view into it.
  CUCHECKGOTO(cuMemAddressReserve(&base, capacity, recGran, 0U, 0), ret, fail);
  CUCHECKGOTO(cuMemMap(base, capacity, 0, mcHandle, 0), ret, fail);
  mapped = 1;
  desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
  desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
  desc.location.id = comm->cudaDev;
  CUCHECKGOTO(cuMemSetAccess(base, capacity, &desc, 1), ret, fail);

整個多播 VA 只保留和映射一次,每個消費者切片是這個 VA 的一個視圖。這是「一次映射、多次切片」的設計——比每個消費者單獨建立多播物件省資源。

並行控制與硬體互動

綁定是 NVLS 最關鍵的操作。ncclMcPartitionBindMem把一個 UC(單播)記憶體句柄綁定到多播組的某個偏移。

📎 src/transport/multicast.cc:200-225

c
ncclResult_t ncclMcPartitionBindMem(const struct ncclMcPartition* partition, size_t offsetInPartition,
                                    CUmemGenericAllocationHandle mem, size_t memOffset, size_t bindSize) {
  // A bind overrunning its partition would corrupt the next consumer's partition; fail
  // cleanly instead (possible when UC rounding exceeds the MC-rounded partition).
  if (offsetInPartition + bindSize > partition->size) {
    WARN("NVLS MC bind of size %zu at slice offset %zu exceeds slice size %zu (UC/MC granularity mismatch)", bindSize,
         offsetInPartition, partition->size);
    return ncclInternalError;
  }
  size_t mcOffset = partition->offset + offsetInPartition;
  ...
  CUresult err = CUPFN(cuMulticastBindMem(partition->mcHandle, mcOffset, mem, memOffset, bindSize, 0 /*flags*/));
  if (err != CUDA_SUCCESS) {
    ...
    WARN("Failed to bind NVLink SHARP (NVLS) Multicast memory of size %zu at MC group %llx offset %zu : CUDA error %d "
         "'%s'.\nThis is usually caused by a system or configuration error in the Fabric Manager or NVSwitches.\n"
         "Disable NVLS (NCCL_NVLS_ENABLE=0) if you wish to avoid this error in the future.",
         bindSize, partition->mcHandle, mcOffset, err, errStr);
    return ncclUnhandledCudaError;
  }
  return ncclSuccess;
}

第一道防線是邊界檢查:offsetInPartition + bindSize > partition->size就報錯。註解解釋了原因——UC 記憶體的粒度可能比 MC 分區大,如果 UC 對齊後超出了 MC 分區的邊界,會踩到下一個消費者的分區。這是典型的「兩種粒度不匹配」陷阱。

cuMulticastBindMem是硬體呼叫,註解說它「blocks until all ranks have been added to the group」——這是 NVLS 最容易出問題的地方。如果 Fabric Manager 配置錯誤或 NVSwitch 韌體有問題,這裡會掛起或返回錯誤。錯誤訊息裡直接建議使用者NCCL_NVLS_ENABLE=0,這是生產環境的標準逃生艙。

還有一個「嘗試綁定」的變體,用於使用者緩衝區註冊:

📎 src/transport/multicast.cc:237-268

c
ncclResult_t ncclMcPartitionTryBindAddr(const struct ncclMcPartition* partition, size_t offsetInPartition,
                                        CUdeviceptr address, size_t bindSize, enum ncclMcBindStatus* outStatus) {
  const char* errStr = NULL;

  *outStatus = ncclMcBindStatusTransient;
  if (offsetInPartition + bindSize > partition->size) {
    ...
    return ncclInternalError;
  }
  size_t mcOffset = partition->offset + offsetInPartition;
  CUresult err = CUPFN(cuMulticastBindAddr(partition->mcHandle, mcOffset, address, bindSize, 0 /*flags*/));
  if (err == CUDA_SUCCESS) {
    *outStatus = ncclMcBindStatusOk;
    return ncclSuccess;
  }

  (void)pfn_cuGetErrorString(err, &errStr);
  // Only an outright rejection of the input is a property of the buffer. Anything else,
  // notably OUT_OF_MEMORY, may succeed later, so it must not be reported as permanent.
  if (err == CUDA_ERROR_INVALID_VALUE || err == CUDA_ERROR_NOT_SUPPORTED || err == CUDA_ERROR_NOT_PERMITTED) {
    *outStatus = ncclMcBindStatusNoSupport;
    ...
  } else {
    WARN("NVLS Multicast bind of size %zu at MC group %llx offset %zu dev %d failed transiently: CUDA error %d '%s'.\n"
         "The buffer is left unregistered for this operation and will be retried; repeated occurrences indicate "
         "sustained resource pressure.",
         bindSize, partition->mcHandle, mcOffset, partition->dev, err, errStr);
  }
  return ncclSuccess;
}

這裡有個精妙的錯誤分類:CUDA_ERROR_INVALID_VALUE、NOT_SUPPORTED、NOT_PERMITTED被歸類為ncclMcBindStatusNoSupport——這是永久性失敗,說明這個 buffer 本身不支援多播綁定。而其他錯誤(尤其是OUT_OF_MEMORY)被歸類為ncclMcBindStatusTransient——這是臨時性失敗,可以重試。這個區分至關重要:如果把 OOM 當成永久失敗,會錯誤地放棄一個本可以成功的註冊;如果把參數錯誤當成臨時失敗,會無限重試。

生產避坑指南

坑 1:Fabric Manager 配置錯誤導致cuMulticastBindMem掛起。這是 NVLS 最經典的生產故障。錯誤訊息裡明確指向 Fabric Manager 或 NVSwitch。排查步驟:先NCCL_NVLS_ENABLE=0確認問題消失,然後檢查 Fabric Manager 日誌和 NVSwitch 韌體版本。

坑 2:UC/MC 粒度不匹配。 ncclMcPartitionBindMem的邊界檢查會捕獲這個問題,但如果你看到 "UC/MC granularity mismatch" 警告,說明某個請求的 UC 大小對齊後超出了 MC 分區。這通常發生在請求大小接近粒度邊界時。

坑 3:多播組建立失敗後的資源洩漏。 ncclMcGroupBuildPartitions的 fail 路徑用了CUCALL(best-effort)而不是CUCHECK:

📎 src/transport/multicast.cc:179-184

c
fail:
  // Best-effort (CUCALL) so a failing cleanup op cannot skip releasing the MC handle.
  if (mapped) CUCALL(cuMemUnmap(base, capacity));
  if (base) CUCALL(cuMemAddressFree(base, capacity));
  if (mcCreated) CUCALL(cuMemRelease(mcHandle));
  return ret;

註解解釋了原因:如果 cleanup 操作本身失敗,不能因此跳過釋放 MC handle——MC slot 是稀缺資源,洩漏會導致後續建立失敗。這是「清理路徑必須盡力而為」的典型設計。

mermaid
sequenceDiagram
    participant R0 as "Rank 0 (localRank=0)"
    participant R1 as "Rank 1..N-1"
    participant BS as "bootstrapIntraNode"
    participant CU as "CUDA Driver"

    R0->>CU: "cuMulticastCreate(mcHandle, prop)"
    CU-->>R0: "mcHandle"
    R0->>BS: "bootstrapIntraNodeBroadcast(shareableHandle)"
    BS-->>R1: "shareableHandle"
    R1->>CU: "cuMemImportFromShareableHandle(mcHandle)"
    CU-->>R1: "mcHandle"
    R0->>CU: "cuMulticastAddDevice(mcHandle, cudaDev)"
    R1->>CU: "cuMulticastAddDevice(mcHandle, cudaDev)"
    R0->>BS: "bootstrapIntraNodeBarrier()"
    R1->>BS: "bootstrapIntraNodeBarrier()"
    Note over R0,R1: "barrier 防止 cuMemMap 阻塞时 peer 失败"
    R0->>CU: "cuMemAddressReserve(base, capacity)"
    R0->>CU: "cuMemMap(base, capacity, mcHandle)"
    R0->>CU: "cuMemSetAccess(base, capacity, desc)"
    R0->>CU: "cuMulticastBindMem(mcHandle, mcOffset, ucHandle)"
    CU-->>R0: "绑定完成,硬件多播就绪"

這張時序圖刻畫了多播群組從建立到綁定的完整流程。關鍵點是那個 barrier——它把「peer 失敗」和「cuMemMap 阻塞」解耦,避免倖存者卡死。

---

14.3 對稱記憶體與 NVLS 的合體:LSA 指標如何在裝置側解析

直覺模型

對稱記憶體解決了「位址一致」問題,NVLS 解決了「硬體歸約」問題。但兩者要真正協同,還需要一個關鍵機制:裝置側如何知道某個位址是對稱的、可以走多播路徑?

答案在 LSA(Load-Store Accessible)指標。LSA 是「可載入-儲存存取」的縮寫,意思是這個指標指向的記憶體,GPU 可以直接用普通的 load/store 指令存取——不管它實體上在本地還是遠端。如果位址落在多播群組內,load/store 會被 NVSwitch 硬體攔截並廣播。

資料結構與記憶體佈局

ncclSymkDevWork是裝置側的工作描述符,它攜帶了對稱記憶體的關鍵資訊。

📎 src/sym_kernels.cc:380-393

c
ncclResult_t ncclSymkMakeDevWork(struct ncclComm* comm, struct ncclTaskColl* task, struct ncclSymkDevWork* outDevWork) {
  outDevWork->rootRank = task->root;
  outDevWork->redOpArg = task->opDev.scalarArg;
  outDevWork->nElts = task->count;
  outDevWork->inputWin = task->sendWin ? task->sendWin->vidmem : nullptr;
  outDevWork->inputOff =
    task->sendWin ? (uint8_t*)task->sendbuff - (uint8_t*)task->sendWin->userPtr : (size_t)task->sendbuff;
  outDevWork->outputWin = task->recvWin ? task->recvWin->vidmem : nullptr;
  outDevWork->outputOff =
    task->recvWin ? (uint8_t*)task->recvbuff - (uint8_t*)task->recvWin->userPtr : (size_t)task->recvbuff;
  outDevWork->sChannelId = 0xffff;
  outDevWork->nChannels = 0;
  return ncclSuccess;
}

inputWin是視窗的裝置側虛擬位址(vidmem),inputOff是緩衝區在視窗內的偏移。裝置側 kernel 拿到這兩個值後,計算inputWin + inputOff就得到實際位址。如果這個位址落在多播群組內,硬體會自動處理廣播。

ncclSymkInitOnce裡還設定了 LSA barrier 和 LLA2A(Low-Latency All-to-All)資源。

📎 src/sym_kernels.cc:197-206

c
    reqs.lsaBarrierCount = ncclSymkMaxBlocks;
    reqs.ginStrongSignalsRequired = false;
    reqs.ginVaSignalsRequired = false;

    struct ncclDevResourceRequirements lla2aReq;
    ncclLLA2ACreateRequirement(ncclSymkMaxBlocks,
                               ncclLLA2ACalcSlots(ncclTeamLsa(comm).nRanks * ncclSymkMaxThreads, ncclSymkLLMaxEltSize),
                               &symk->kcomm.lsaLLA2A, &lla2aReq);
    lla2aReq.next = reqs.resourceRequirementsList;
    reqs.resourceRequirementsList = &lla2aReq;

lsaBarrierCount設為ncclSymkMaxBlocks——每個 block 一個 barrier 槽位。LLA2A 是低延遲 all-to-all 的縮寫,用於在 LSA 域內做快速資料交換。ncclLLA2ACalcSlots根據 rank 數、執行緒數和最大元素大小算出需要的槽位數。

場景驅動的 Step-by-Step Walkthrough

假設一次 AllReduce 使用AllReduce_AGxLLMC_Rkernel(AllGather + LL + MC + Reduce)。這個 kernel 的工作流程是:

1. AllGather 階段:每個 rank 把自己的資料寫入多播群組,NVSwitch 硬體廣播給所有 rank。

2. Reduce 階段:每個 rank 從多播群組讀取所有 rank 的資料,在本地做歸約。

ncclSymkMask會檢查這個 kernel 是否可用。kernelMask_LL包含AllReduce_AGxLLMC_R,但前提是hasLsaMultimem為真(否則kernelMask_STMC被清除,而AllReduce_AGxLLMC_R屬於 STMC 集合)。

等等,這裡有個細節:kernelMask_STMC包含AllReduce_AGxLLMC_R嗎?看原始碼:

📎 src/sym_kernels.cc:17-21

c
constexpr uint32_t kernelMask_STMC =
  1 << ncclSymkKernelId_AllGather_LLMC | 1 << ncclSymkKernelId_AllGather_STMC |
  1 << ncclSymkKernelId_AllGather_TmaSTMC | 1 << ncclSymkKernelId_AllReduce_AGxLLMC_R |
  1 << ncclSymkKernelId_AllReduce_RSxLDMC_AGxSTMC | 1 << ncclSymkKernelId_ReduceScatter_LDMC |
  1 << ncclSymkKernelId_AllGather_RailRing_LsaSTMC;

是的,AllReduce_AGxLLMC_R在kernelMask_STMC裡。所以如果hasLsaMultimem為假,這個 kernel 會被剔除。這解釋了為什麼對稱記憶體和 NVLS 必須協同工作——沒有多播,MC 系列 kernel 全部不可用。

裝置側拿到ncclSymkDevWork後,會根據inputWin和inputOff計算位址。如果位址在多播群組內,load/store 指令會被 NVSwitch 攔截。這就是 LSA 指標的解析過程:不需要軟體轉譯,硬體根據位址範圍自動判斷。

並發控制與硬體互動

NVLS 的同步機制依賴credit(信用)。ncclNvlsSetup裡初始化了 credit 分區。

📎 src/transport/nvls.cc:407-447

c
    int nChannels = comm->nvlsChannels;
    size_t creditSize = nChannels * 2 * memSize * nHeads;
    int nvlsStepSize = comm->nvlsChunkSize;

    NCCLCHECKGOTO(ncclCalloc(&comm->nvlsResources, 1), res, fail);
    comm->nvlsResources->inited = false;
    comm->nvlsResources->refCount = 1;
    comm->nvlsResources->nChannels = nChannels;
    comm->nvlsResources->nHeads = nHeads;
    comm->nvlsResources->chunkSize = comm->nvlsChunkSize;
    comm->nvlsResources->treeMaxChunkSize = comm->nvlsTreeMaxChunkSize;
    resources = comm->nvlsResources;

    for (int c = 0; c < nChannels; c++) {
      NCCLCHECKGOTO(initNvlsChannel(comm, c, NULL, false), res, fail);
    }

    memset(&resources->accessDesc, 0, sizeof(resources->accessDesc));
    resources->accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
    resources->accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
    resources->accessDesc.location.id = comm->cudaDev;
    resources->dev = comm->cudaDev;

    // Build the single shared MC group for this NVLS domain. The data slice is
    // reserved here but bound later by ncclNvlsBufferSetup.
    {
      size_t buffSize = nvlsStepSize * NCCL_STEPS;
      size_t dataSize = nChannels * 2 * buffSize * nHeads;
      size_t ubSize = ncclNvlsUbSize(comm);
      struct ncclMcRequest requests[3] = {{creditSize, 0}, {dataSize, 0}, {ubSize, 0}};
      struct ncclMcPartition partitions[3];
      NCCLCHECKGOTO(ncclMcGroupBuildPartitions(comm, requests, 3, &resources->mcGroup, partitions), res, fail);
      resources->creditPartition = partitions[0];
      resources->dataPartition = partitions[1];
      if (ubSize) {
        resources->ubPartition = partitions[2];
        NCCLCHECKGOTO(ncclMcArenaInit(comm, &resources->ubArena, &resources->ubPartition), res, fail);
        resources->ubEnabled = true;
      }
      NCCLCHECKGOTO(nvlsAllocBindUc(comm, &resources->creditPartition, creditSize, &resources->creditUc), res, fail);
    }

多播群組被切成三個分區:creditPartition(信用)、dataPartition(資料)、ubPartition(使用者緩衝區)。credit 分區用於同步——每個 channel 有獨立的 head/tail 指標,透過多播群組共享。

credit 的初始化在後面的迴圈裡:

📎 src/transport/nvls.cc:456-491

c
    for (int h = 0; h < nHeads; h++) {
      int nvlsPeer = comm->nRanks + 1 + h;
      for (int c = 0; c < nChannels; c++) {
        struct ncclChannel* channel = comm->channels + c;
        char* mem = NULL;
        struct ncclChannelPeer* peer = channel->peers[nvlsPeer];

        // Reduce UC -> MC
        mem = (char*)resources->creditUc.ptr + (h * 2 * nChannels + c) * memSize;
        peer->send[1].transportComm = &nvlsTransport.send;
        peer->send[1].conn.buffs[NCCL_PROTO_SIMPLE] = NULL;
        peer->send[1].conn.head = (uint64_t*)mem;
        peer->send[1].conn.tail = (uint64_t*)(mem + memSize / 2);
        peer->send[1].conn.stepSize = nvlsStepSize;
        mem = (char*)resources->creditPartition.ptr + (h * 2 * nChannels + c) * memSize;
        peer->recv[0].transportComm = &nvlsTransport.recv;
        peer->recv[0].conn.buffs[NCCL_PROTO_SIMPLE] = NULL;
        peer->recv[0].conn.head = (uint64_t*)mem;
        peer->recv[0].conn.tail = (uint64_t*)(mem + memSize / 2);
        peer->recv[0].conn.stepSize = nvlsStepSize;
        peer->recv[0].conn.flags |= NCCL_NVLS_MIN_POLL;

每個 head 和 channel 組合都有獨立的 credit 區域。head和tail是 64 位元指標,memSize是 64 位元組(size_t memSize = 64;),所以 head 和 tail 各佔 32 位元組——正好半個快取行。NCCL_NVLS_MIN_POLL標誌讓接收方用最小輪詢模式,減少 CPU 開銷。

生產避坑指南

坑 1:credit 分區的 head/tail 競爭。多個 channel 共享同一個多播群組,但每個 channel 有獨立的 credit 區域。如果 channel 數配置不當(比如nvlsCTAs設得太大),credit 區域會膨脹,佔用寶貴的多播位址空間。ncclNvlsChannels會根據 GPU 架構和節點數自動調整 channel 數:

📎 src/transport/nvls.cc:100-133

c
  if (comm->config.nvlsCTAs != NCCL_CONFIG_UNDEF_INT) {
    channels = comm->config.nvlsCTAs;
  } else if (channels == 0 && comm->compCap >= 100) {
    // Use a reduced number of channels for single node/MNNVL domain on Blackwell and above.
    // comm->nNodes is not yet initialized at this point so we need to use local information.
    bool multiNode = false;
    if (comm->MNNVL) {
      multiNode = (comm->clique.size < comm->nRanks);
    } else {
      int i;
      for (i = 1; i < comm->nRanks; i++) {
        if (comm->peerInfo[i].hostHash != comm->peerInfo[0].hostHash) break;
      }
      multiNode = (i < comm->nRanks);
    }
    if (multiNode) {
      channels = RUBIN_AND_LATER(comm->compCap) ? /*RUBIN=*/64 : /*SM100=*/32;
    } else {
      channels = RUBIN_AND_LATER(comm->compCap) ? /*RUBIN=*/48 : /*SM100=*/24;
    }
  } else if (channels == 0) {
    channels = /*SM90=*/16;
  }

注意comm->nNodes在這個階段還沒初始化,所以程式碼用peerInfo[i].hostHash手動判斷是否多節點。這是初始化順序的經典陷阱——你不能依賴還沒算出來的欄位。

坑 2:MNNVL 不支援 NVLS buffer 註冊。 📎 src/transport/nvls.cc:516-517

c
  // MNNVL does not support NVLS buffer registration
  if (!comm->MNNVL && comm->nvlsResources->nvlsShmemHandle == NULL) {

MNNVL(Multi-Node NVLink)環境下,使用者緩衝區註冊被跳過。如果你的叢集是 MNNVL 且依賴 UB 註冊來提升效能,會發現註冊沒生效。這是硬體限制,不是 bug。

坑 3:共享資源的引用計數。 ncclNvlsSetup支援父子通訊域共享 NVLS 資源:

📎 src/transport/nvls.cc:380-392

c
  if (nvlsShare) {
    /* reuse NVLS resources */
    comm->nvlsChannels = std::min(comm->nvlsChannels, parent->nvlsResources->nChannels);
    /* Inherit chunk sizes from the shared resource since we're reusing the parent's
     * NVLS buffers, which were allocated and laid out based on these values. */
    comm->nvlsChunkSize = parent->nvlsResources->chunkSize;
    comm->nvlsTreeMaxChunkSize = parent->nvlsResources->treeMaxChunkSize;
    for (int c = 0; c < comm->nvlsChannels; c++) {
      NCCLCHECKGOTO(initNvlsChannel(comm, c, parent, true), res, fail);
    }

    comm->nvlsResources = parent->nvlsResources;
    ncclAtomicRefCountIncrement(&parent->nvlsResources->refCount);
  }

子通訊域復用父通訊域的資源,引用計數加一。ncclNvlsFree裡引用計數減到零才真正釋放。如果引用計數管理出錯,會導致資源提前釋放或洩漏。注意nvlsChunkSize和nvlsTreeMaxChunkSize必須繼承父通訊域的值——因為緩衝區是按這些值佈局的,改了會導致位址計算錯誤。

mermaid
flowchart LR
    subgraph host["Host 侧"]
        task["ncclTaskColl<br/>sendbuff/recvbuff"]
        devwork["ncclSymkDevWork<br/>inputWin + inputOff"]
        task -->|"ncclSymkMakeDevWork"| devwork
    end
    subgraph device["Device 侧"]
        kernel["SymKernel<br/>load/store"]
        lsa{"地址在多播组内?"}
        devwork --> kernel
        kernel --> lsa
    end
    subgraph hw["NVSwitch 硬件"]
        mc["多播组<br/>MC group"]
        reduce["硬件归约<br/>Reduction"]
        lsa -->|"是"| mc
        lsa -->|"否"| local["本地显存<br/>UC memory"]
        mc --> reduce
        reduce -->|"广播结果"| kernel
    end

這張資料流圖展示了從 host 側任務到裝置側執行的完整鏈路。關鍵分支是lsa{"地址在多播组内?"}——如果是,走 NVSwitch 硬體多播和歸約;如果否,走本地顯存。這個判斷由硬體根據位址範圍自動完成,不需要軟體干預。

---

14.4 設計思考:為什麼對稱記憶體能降低小訊息延遲

回到本章開頭的核心問題:為什麼對稱記憶體能顯著降低小訊息延遲?

第一,消除了位址轉譯開銷。傳統通訊裡,每個 rank 存取對端緩衝區都要查表、計算偏移。對稱記憶體讓所有 rank 用同一套位址,裝置側 kernel 直接算base + offset就行。對於小訊息,這次轉譯的開銷佔比很高。

第二,消除了控制訊息往返。傳統通訊需要交換「我要寫你的哪個緩衝區」這類控制資訊。對稱記憶體下,位址是預先約定好的,不需要執行時協商。

第三,讓硬體多播成為可能。只有當位址對稱時,NVSwitch 才能用同一套位址做多播。如果每個 rank 的位址不同,硬體無法知道該廣播到哪裡。

第四,減少了 SM 的歸約負擔。NVLS 把加法卸載到 NVSwitch,SM 只需要發起一次寫、一次讀。對於小訊息,SM 的指令開銷是延遲的主要來源。

這四個因素疊加,讓小訊息延遲從「微秒級」降到「亞微秒級」。

〔設計推斷與架構權衡〕

從工程角度看,對稱記憶體的設計體現了 NCCL 的一個核心哲學:把複雜性推到初始化階段,讓熱路徑盡可能簡單。位址協商、多播組建立、credit 分配都在初始化時完成,執行時 kernel 只需要做最簡單的位址計算和 load/store。這種「初始化重、執行時輕」的設計,是高效能通訊庫的通用模式。

---

本章小結

本章拆解了 NCCL 機內通訊的兩大支柱:

1. 對稱記憶體:透過ncclSymkInitOnce和ncclSymkMask建立位址一致的緩衝區,讓每個 rank 用同一套位址存取所有 rank 的資料。ncclSymkMakeDevWork把 host 側任務轉譯成裝置側工作項,inputWin + inputOff是位址解析的核心公式。

2. NVLS 多播:透過ncclMcGroupBuildPartitions建立多播組,ncclMcPartitionBindMem把 UC 記憶體綁定到多播組,cuMulticastBindMem是硬體呼叫。多播組被切成 credit、data、ub 三個分區,分別用於同步、資料傳輸和使用者緩衝區註冊。

3. LSA 指標解析:裝置側根據位址範圍自動判斷是否走多播路徑,不需要軟體轉譯。NCCL_NVLS_MIN_POLL標誌優化輪詢開銷。

4. 錯誤處理:ncclMcPartitionTryBindAddr區分永久性失敗和臨時性失敗,ncclMcGroupBuildPartitions的 fail 路徑用CUCALL確保資源釋放。

本章思考與自測

Q1: 如果把ncclMcPartitionBindMem裡的邊界檢查if (offsetInPartition + bindSize > partition->size)去掉,在什麼場景下會觸發記憶體越界?為什麼這個檢查不能用「UC 和 MC 粒度相同」來替代?

參考解析:看📎 src/transport/multicast.cc:200-208:

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 15

第 15 章:RMA 與 GIN:遠端記憶體存取與 GPU 直連通訊的演進

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 15 章 / 共 25 章

上一章我們看到,對稱記憶體讓每個 rank 用同一套位址存取所有 rank 的緩衝區,NVLS 則藉助 NVSwitch 的多播能力把硬體加速歸約推向極致。但集合通訊並非全部——當應用需要點對點遠端記憶體操作,或希望 GPU kernel 直接發起網路請求時,就需要 RMA 與 GIN 登場。RMA 提供 put/get 語意的遠端記憶體存取,GIN 則讓 GPU 繞過 host proxy 執行緒直接與網路互動。本章按「先 RMA 後 GIN」的順序,逐層拆解這兩套機制的資料結構、排程邏輯、並發控制與生產陷阱。

RMA 的雙通道模型:CE 與 Proxy 的分工

直覺模型

想像一個跨國快遞系統:同城快遞(LSA 可達的 rank)可以直接由本地配送車送達,而跨城快遞(非 LSA 可達的 rank)必須交給航空貨運代理。NCCL 的 RMA 正是這個模型——同一個 put 操作,根據目標 rank 是否在 LSA(Load-Store Accessible)團隊內,被路由到兩條完全不同的執行路徑:CE(Copy Engine,拷貝引擎)路徑和 Proxy(代理執行緒)路徑。

如果沒有這個分流機制,所有 RMA 操作都走 proxy 執行緒,那麼同機內的 put 也要經過 host 執行緒中轉,白白增加一次 host-device 往返延遲。反之,如果所有操作都走 CE,跨機操作就無法利用網路外掛的異步能力。

資料結構與記憶體佈局

RMA 的核心調度結構是ncclRmaArgs,它記錄了一個 plan 中 RMA 任務的分流結果。關鍵欄位包括:

欄位含義
func操作類型(PutSignal / Signal / WaitSignal)
nRmaTasks總任務數
nRmaTasksProxy走 proxy 路徑的任務數
nRmaTasksCe走 CE 路徑的任務數

每個 plan 內部維護兩個侵入式佇列:rmaTaskQueueCe和rmaTaskQueueProxy,分別存放兩條路徑的任務。📎 src/rma/rma.cc:166-171

判斷一個 rank 是否 LSA 可達的邏輯很直接——遍歷lsaRankList陣列做線性查找。📎 src/rma/rma.cc:34-41這個查找在任務調度時對每個 peer 執行一次,複雜度 O(lsaSize),對於典型的小規模 LSA 團隊(通常 2-8 個 rank)開銷可忽略。

Step-by-Step 調度流程

當應用呼叫一次 RMA put 操作後,任務進入planner->rmaTaskQueues[ctx]。scheduleRmaTasksToPlan負責把佇列中的任務分配到 plan 中。📎 src/rma/rma.cc:141-296

第一步:找到第一個非空的 context 佇列。NCCL 支援多個 RMA context(由numRmaCtx配置),每個 context 有獨立的佇列。📎 src/rma/rma.cc:148-155

第二步:取出第一個任務,判斷操作類型。如果是 WaitSignal,走特殊的分裂邏輯;如果是 Put/Signal,走批量合併邏輯。📎 src/rma/rma.cc:163-168

對於 WaitSignal 任務,調度器需要把 peers 列表按 LSA 可達性拆分成兩組:CE 組和 Proxy 組。📎 src/rma/rma.cc:187-204拆分後分別建立兩個新的ncclTaskRma結構,各自持有對應組的 peers 陣列。📎 src/rma/rma.cc:207-246原始任務被釋放。📎 src/rma/rma.cc:251

對於 Put/Signal 任務,邏輯更複雜——調度器會遍歷所有 context 的佇列,把連續的 put/signal 任務全部拉入同一個 plan,直到遇到 WaitSignal 才停止。📎 src/rma/rma.cc:279-295這個設計的目的在註解中寫得很清楚:讓一次 kernel launch 覆蓋所有 context 的 put/signal,proxy 可以在任何阻塞操作之前一次性發起所有異步請求,CE 路徑則把所有 context 的拷貝和信號批量提交。📎 src/rma/rma.cc:270-278

mermaid
flowchart TD
    start["scheduleRmaTasksToPlan(comm, plan)"]
    find_ctx{"找到非空 ctx 队列?"}
    no_task["返回 ncclSuccess"]
    dequeue["取出 firstTask"]
    check_func{"firstTask->func == WaitSignal?"}
    ws_split["按 isLsaAccessible 拆分 peers"]
    ws_ce{"npeersCe > 0?"}
    ws_proxy{"npeersProxy > 0?"}
    ws_ce_task["创建 CE WaitSignal 任务"]
    ws_proxy_task["创建 Proxy WaitSignal 任务"]
    ws_free["释放原始 firstTask"]
    put_check{"firstTask 的 peer LSA 可达?"}
    put_ce["入队 rmaTaskQueueCe"]
    put_proxy["入队 rmaTaskQueueProxy"]
    batch_loop["遍历所有 ctx 队列, 拉取连续 put/signal"]
    batch_check{"isRmaPutOrSignal(task->func)?"}
    batch_route{"isLsaAccessible(comm, task->peer)?"}
    batch_ce["入队 CE, nRmaTasksCe++"]
    batch_proxy["入队 Proxy, nRmaTasksProxy++"]
    done["记录 INFO 日志, 返回"]

    start --> find_ctx
    find_ctx -->|否| no_task
    find_ctx -->|是| dequeue
    dequeue --> check_func
    check_func -->|是| ws_split
    ws_split --> ws_ce
    ws_ce -->|是| ws_ce_task
    ws_ce -->|否| ws_proxy
    ws_ce_task --> ws_proxy
    ws_proxy -->|是| ws_proxy_task
    ws_proxy -->|否| ws_free
    ws_proxy_task --> ws_free
    ws_free --> done
    check_func -->|否| put_check
    put_check -->|是| put_ce
    put_check -->|否| put_proxy
    put_ce --> batch_loop
    put_proxy --> batch_loop
    batch_loop --> batch_check
    batch_check -->|否, 遇到 WaitSignal| done
    batch_check -->|是| batch_route
    batch_route -->|是| batch_ce
    batch_route -->|否| batch_proxy
    batch_ce --> batch_loop
    batch_proxy --> batch_loop

並行執行與流同步

調度完成後,ncclLaunchRma根據func欄位分發到ncclRmaPut或ncclRmaWaitSignal。📎 src/rma/rma.cc:109-131

以ncclRmaPut為例,當 plan 中同時存在 proxy 和 CE 任務時,兩條路徑需要並行執行。NCCL 的做法是:在輸入流上記錄一個 event,讓 CE 流等待這個 event,然後同時在兩條流上啟動操作,最後在 CE 流上再記錄一個 event,讓輸入流等待它。📎 src/rma/rma.cc:80-96這個 event 鏈確保了:CE 操作不會在輸入流的依賴就緒前開始,輸入流的後續操作也不會在 CE 完成前開始。

如果只有 proxy 任務或只有 CE 任務,則直接在輸入流上啟動對應操作,無需額外的流同步。📎 src/rma/rma.cc:97-101

設計思考與生產陷阱

陷阱一:LSA 可達性判斷的靜態性。 isLsaAccessible在調度時查詢comm->devrState.lsaRankList,這個列表在通訊域初始化後就不再變化。如果運行過程中拓撲發生變化(比如 NVLink 故障降級),LSA 列表不會自動更新,可能導致本應走 proxy 的操作仍然走 CE 路徑,觸發不可恢復的錯誤。

陷阱二:批量合併的 FIFO 保證。批量合併邏輯只拉取連續的 put/signal 任務,遇到 WaitSignal 就停止。📎 src/rma/rma.cc:283這保證了每個 context 內的 FIFO 順序,但跨 context 的任務可能被合併到同一個 plan 中。如果應用依賴跨 context 的操作順序,需要顯式使用 WaitSignal 來建立屏障。

陷阱三:記憶體洩漏路徑。在 WaitSignal 分支中,如果npeersProxy == 0,程式碼會釋放peersProxy、nsignalsProxy、signalIdxsProxy三個陣列。📎 src/rma/rma.cc:239-244但如果npeersCe == 0且npeersProxy > 0,peersCe等陣列是透過ncclMemoryStackAlloc分配的,不需要手動釋放(棧式分配器統一回收)。📎 src/rma/rma.cc:176-178這個不對稱性容易讓讀者困惑,但實際上是正確的——棧分配的記憶體由comm->memScoped統一管理。

RMA Proxy 上下文:訊號、佇列與無鎖環形緩衝

直覺模型

Proxy 上下文就像一個「郵局分揀中心」:GPU 把要發送的包裹(put 請求)放進收件箱(環形緩衝),proxy 執行緒從收件箱取出包裹,交給快遞公司(網路外掛),快遞公司送達後在回執單(訊號)上蓋章。整個過程中,GPU 和 proxy 執行緒透過無鎖資料結構通訊,避免昂貴的鎖競爭。

資料結構與記憶體佈局

ncclRmaProxyCtx是 proxy 上下文的宿主結構,其核心欄位包括:

訊號區(signalsDev):在 GPU 上分配的一塊記憶體,大小為nRanks * numRmaSig * sizeof(uint64_t)。📎 src/rma/rma_proxy.cc:120-123每個 rank 有numRmaSig個訊號槽,用於接收來自該 rank 的訊號。這塊記憶體註冊到網路外掛時帶有NCCL_NET_MR_FLAG_FORCE_SO(強制強序)和NCCL_NET_MR_FLAG_SIGNAL_NEVER_RESET(訊號永不重置)標誌。📎 src/rma/rma_proxy.cc:125-127強序標誌確保 put 和 signal 之間的順序關係——如果 put 先於 signal 發出,網路必須保證 signal 在 put 資料到達後才寫入。

序列號區(opSeqs/readySeqs/doneSeqs):每個 rank 一組,透過allocMemCPUAccessible分配,可能是 GDR(GPU Direct RDMA)記憶體或普通 host 記憶體。📎 src/rma/rma_proxy.cc:132-137這三個序列號分別追蹤:已提交的操作序號、已就緒的操作序號、已完成的操作序號。

無鎖環形緩衝(circularBuffers):大小為nRanks * queueSize的指標陣列,每個 rank 一個獨立的環形佇列。📎 src/rma/rma_proxy.cc:163-164配套的pis(Producer Index)和cis(Consumer Index)陣列各nRanks個元素。📎 src/rma/rma_proxy.cc:165-166佇列大小必須是 2 的冪,這樣索引回繞可以用位與運算& (queueSize - 1)代替取模。📎 src/rma/rma_proxy.cc:156-160

InProgress 佇列:每個 peer 一個侵入式鏈結串列,存放已提交給網路外掛但尚未完成的描述符。📎 src/rma/rma_proxy.cc:170-175這是單消費者佇列,只有 proxy 執行緒存取,無需原子操作。

Step-by-Step:從上下文建立到進度推進

上下文建立:ncclRmaProxyCreateContext首先透過 RMA 外掛建立網路上下文。📎 src/rma/rma_proxy.cc:229然後呼叫ncclRmaProxyCtxAlloc分配訊號、序列號、環形緩衝等資源。📎 src/rma/rma_proxy.cc:231接著呼叫ncclRmaProxyCtxAllocGraph分配圖捕獲模式所需的資源——CPU 可存取的訊號、flush 緩衝、持久化佇列。📎 src/rma/rma_proxy.cc:232

圖捕獲模式的存在是因為 CUDA Graph 要求所有操作可重放。在普通模式下,訊號在 GPU 記憶體中,proxy 透過 GDR 讀取;在圖捕獲模式下,訊號在 CPU 可存取記憶體中,proxy 可以直接讀寫,避免 GDR 的不確定性。📎 src/rma/rma_proxy.cc:184-190

進度執行緒:ncclRmaProxyProgressThread是 proxy 的主迴圈。📎 src/rma/rma_proxy.cc:354-389它根據rmaProgress狀態字決定行為:

  • rmaProgress == 1:正常推進模式,遍歷所有 proxy 上下文呼叫ncclRmaProxyProgress。📎 src/rma/rma_proxy.cc:361-372
  • rmaProgress == 2:暫停模式,用於資源回收。執行緒確認暫停後等待條件變數。📎 src/rma/rma_proxy.cc:373-378
  • rmaProgress == -1:退出訊號,執行緒返回。📎 src/rma/rma_proxy.cc:379-380
  • rmaProgress == 0:空閒等待。📎 src/rma/rma_proxy.cc:381-382

如果ncclRmaProxyProgress返回錯誤,執行緒把錯誤碼寫入asyncResult,設定rmaProgress = -2,然後退出。📎 src/rma/rma_proxy.cc:365-369這個錯誤碼會被主執行緒在後續的ncclCommGetAsyncError呼叫中讀取。

並發控制與記憶體序

RMA proxy 的並發模型是「單生產者-單消費者」:GPU kernel 是生產者,proxy 執行緒是消費者。環形緩衝的 PI 由 GPU 更新,CI 由 proxy 更新。由於是單生產者單消費者,不需要 CAS 操作,只需要正確的記憶體序。

訊號區的強序標誌NCCL_NET_MR_FLAG_FORCE_SO是關鍵。📎 src/rma/rma_proxy.cc:127沒有這個標誌,網路外掛可能重排 put 和 signal 的順序,導致接收方在資料到達前就看到訊號,讀取到髒資料。

NCCL_NET_MR_FLAG_SIGNAL_NEVER_RESET標誌告訴網路外掛:訊號一旦寫入就不會被重置。📎 src/rma/rma_proxy.cc:127這允許外掛優化訊號的寫入路徑——不需要每次寫入前清零。

生產陷阱

陷阱一:佇列大小不是 2 的冪。如果使用者透過NCCL_RMA_PROXY_QUEUE_SIZE設定了一個非 2 的冪的值,程式碼會回退到預設值並列印 INFO 日誌。📎 src/rma/rma_proxy.cc:156-159這個回退是靜默的(只有 INFO 級別),在生產環境中容易被忽略。如果使用者期望更大的佇列來吸收突發流量,實際使用的卻是預設值,可能導致背壓。

陷阱二:DMA-BUF 註冊失敗的回退鏈。 ncclRmaProxyRegMrSym對 CUDA 記憶體的註冊有三層回退:先嘗試 DataDirect 模式的 DMA-BUF,失敗後嘗試非 DataDirect 的 DMA-BUF,再失敗才回退到普通regMrSym。📎 src/rma/rma_proxy.cc:76-108註解中特別警告:如果一個 MR 進入了非 DataDirect 路徑,所有其他 MR 也必須如此,混合使用會破壞 GIN 的順序保證。📎 src/gin/gin_host_proxy.cc:429-430這個約束在 RMA 路徑中沒有顯式檢查,是一個潛在的隱患。

陷阱三:進度執行緒的錯誤傳播延遲。當ncclRmaProxyProgress返回錯誤時,執行緒設定asyncResult並退出。📎 src/rma/rma_proxy.cc:366-369但主執行緒可能正在執行一個長時間的 kernel,不會立即檢查asyncResult。在這段時間內,後續的 RMA 操作會繼續入隊但不會被處理,直到主執行緒發現錯誤。這是非同步錯誤傳播的固有延遲,應用需要定期呼叫ncclCommGetAsyncError來縮短這個窗口。

GIN 架構:GPU 直接發起網路請求

直覺模型

傳統模式下,GPU 要發送網路資料,必須經過「GPU → host 記憶體 → proxy 執行緒 → 網卡」的路徑。GIN(GPU-Initiated Networking)的目標是讓 GPU 直接寫網卡的發送佇列,就像 CPU 直接寫網卡的 MMIO 暫存器一樣。這需要網卡支援 GPU 發起的 doorbell 寫入,以及一套 GPU 和 proxy 執行緒之間的通訊協定。

資料結構與記憶體佈局

GIN 的核心資料結構是ginProxyHostGpuCtx,它代表一個 GPU-host 通訊上下文:

欄位類型含義
queuesncclGinProxyGfd_t*GFD 佇列,大小nRanks * queueSize
pisuint32_t*生產者索引(GPU 寫)
cisuint32_t*消費者索引(proxy 寫)
cisShadowuint32_t*CI 的影子副本(proxy 本地)
sisuint32_t*已見索引(proxy 本地)
statesginProxyGfdState*每個 GFD 槽的狀態
inlinesuint64_t*內聯資料緩衝區

GFD(GIN Forwarding Descriptor)是 GPU 寫給 proxy 的請求描述符。每個 GFD 由多個 qword 組成,包含操作類型、來源位址、目標位址、大小、信號資訊等。📎 src/gin/gin_host_proxy.cc:158-163

queues陣列的記憶體分配有一個關鍵細節:它透過allocMemCPUAccessible分配,但傳入了forceHost=true參數。📎 src/gin/gin_host_proxy.cc:564這意味著佇列本身在 host 記憶體中,GPU 透過 PCIe 寫入。而cis陣列則分配在 GPU 可存取記憶體中(可能是 GDR),因為 proxy 需要頻繁更新它。📎 src/gin/gin_host_proxy.cc:565-566

cisShadow和sis是 proxy 執行緒的本地副本,避免每次都讀取可能位於 GPU 記憶體的cis。📎 src/gin/gin_host_proxy.cc:44-47只有當cisShadow前進時,才批量更新cis。

Step-by-Step:GFD 的輪詢與處理

ncclGinProxyProgress是 GIN proxy 的主迴圈。📎 src/gin/gin_host_proxy.cc:648-669

第一步:對每個 context,先呼叫proxyGinPollCompletions檢查已提交請求的完成狀態。📎 src/gin/gin_host_proxy.cc:653

第二步:對每個 target rank,批量輪詢 GFD。pollBatch控制每次最多處理多少個 GFD。📎 src/gin/gin_host_proxy.cc:654-655

第三步:proxyGinPollGfd檢查佇列頭部是否有新的 GFD。判斷依據是 GFD 頭部的 flag 位是否非零。📎 src/gin/gin_host_proxy.cc:176-182如果有,先拷貝第一個 qword(頭部),然後等待其餘 qword 就緒。📎 src/gin/gin_host_proxy.cc:194-202拷貝完成後,把佇列中的 GFD 清零,防止重複處理。📎 src/gin/gin_host_proxy.cc:206-208

第四步:proxyGinProcessGfd根據操作類型分發到不同的處理路徑。📎 src/gin/gin_host_proxy.cc:246-340

mermaid
flowchart TD
    poll_start["proxyGinPollGfd(ctx, hostGpuCtx, targetRank)"]
    check_avail{"isGfdAvailable?"}
    no_gfd["返回 0, 跳出批量循环"]
    copy_header["拷贝 GFD header qword"]
    copy_rest["循环等待并拷贝其余 qword"]
    reset_gfd["清零队列中的 GFD"]
    set_state["设置 state->op, counterId, done=0"]
    inc_sis["sis[targetRank]++"]
    process["proxyGinProcessGfd(ctx, hostGpuCtx, targetRank, gfd, state, isLastInBatch)"]
    check_va{"op & ncclGinProxyOpVASignal?"}
    check_get{"op & ncclGinProxyOpGet?"}
    check_flush{"op & ncclGinProxyOpFlush?"}
    check_inline{"op & ncclGinProxyOpWithInline?"}
    va_signal["rmaBackend->iputSignal(...)"]
    get_op["rmaBackend->iget(...)"]
    flush_op["rmaBackend->iflush(...)"]
    inline_src["从 inlines 缓冲区取源地址"]
    normal_src["从 GFD 取源地址"]
    put_signal["rmaBackend->iputSignal(...)"]
    put_only["rmaBackend->iput(...)"]

    poll_start --> check_avail
    check_avail -->|否| no_gfd
    check_avail -->|是| copy_header
    copy_header --> copy_rest
    copy_rest --> reset_gfd
    reset_gfd --> set_state
    set_state --> inc_sis
    inc_sis --> process
    process --> check_va
    check_va -->|是| va_signal
    check_va -->|否| check_get
    check_get -->|是| get_op
    check_get -->|否| check_flush
    check_flush -->|是| flush_op
    check_flush -->|否| check_inline
    check_inline -->|是| inline_src
    check_inline -->|否| normal_src
    inline_src --> put_signal
    normal_src --> put_signal
    put_signal --> put_only

完成輪詢與計數器更新

proxyGinPollCompletions負責檢查已提交請求的完成狀態。📎 src/gin/gin_host_proxy.cc:113-156

對每個 target rank,從cisShadow到sis遍歷所有已見但未消費的 GFD 狀態。📎 src/gin/gin_host_proxy.cc:117如果狀態未完成,呼叫rmaBackend->test檢查。📎 src/gin/gin_host_proxy.cc:122如果完成且操作帶有計數器標誌,更新計數器值。📎 src/gin/gin_host_proxy.cc:132-141

計數器更新使用原子載入和原子存儲,但註釋解釋了為什麼不需要原子加法:GPU kernel 不允許在有未完成操作時重置計數器,因此不存在競爭。📎 src/gin/gin_host_proxy.cc:133-135

CI 的更新有一個「允許空洞」的機制:只有當state->done && i == cisShadow[targetRank]時才推進 CI。📎 src/gin/gin_host_proxy.cc:145-151這確保了 CI 是單調遞增的,即使某些 GFD 先完成,也不會跳過未完成的 GFD。

並發控制與記憶體屏障

GIN proxy 的並發模型比 RMA proxy 更複雜,因為存在多個 proxy 執行緒(由GIN_PROXY_NTHREADS控制)。📎 src/gin/gin_host.cc:90

ncclGinProgress中,每個執行緒負責一組連接:執行緒 t 處理連接 t, t+proxyNthreads, t+2*proxyNthreads, ...。📎 src/gin/gin_host.cc:72這個分配方式確保了每個連接只被一個執行緒處理,避免了連接級別的競爭。

devComms 鏈表的修改需要寫鎖保護。ginProgressWriteLock先設置writePending標誌,然後獲取寫鎖。📎 src/gin/gin_host.cc:43-47進度執行緒在每次迴圈開始時檢查writePending,如果為真則讓出 CPU。📎 src/gin/gin_host.cc:63-66這個設計避免了進度執行緒在持有讀鎖時被寫鎖阻塞。

writePending使用std::atomic<bool>,但註釋指出這個邏輯假設只有一個寫者。📎 src/gin/gin_host.cc:43-47在 NCCL 的使用場景中,只有主執行緒會修改 devComms 鏈表,所以這個假設成立。

生產陷阱

陷阱一:GFD 佇列的記憶體位置。 queues被強制分配在 host 記憶體中(forceHost=true),📎 src/gin/gin_host_proxy.cc:564這意味著 GPU 寫入 GFD 需要經過 PCIe 總線。如果 GFD 寫入頻率很高(小訊息場景),PCIe 頻寬可能成為瓶頸。相比之下,cis分配在 GPU 可存取記憶體中,因為 proxy 需要頻繁更新它。📎 src/gin/gin_host_proxy.cc:565-566

陷阱二:內聯資料的重建。當 GFD 帶有內聯資料時,proxy 需要從多個 qword 中重建內聯值。📎 src/gin/gin_host_proxy.cc:298-305重建邏輯根據 size 決定讀取哪些 qword:size ≤ 4 只讀低 32 位,size > 4 讀低 64 位,size > 6 再讀高 16 位。這個分段邏輯與 GPU 側的寫入邏輯必須嚴格對應,任何不一致都會導致資料損毀。

陷阱三:多執行緒進度與連線分配。如果不同 rank 設定了不同的GIN_PROXY_NTHREADS,經過 AllGather 取最小值後,某些執行緒可能沒有分配到任何連線。📎 src/gin/gin_host.cc:181-183註解指出這些執行緒會在 stride 迴圈中空轉,不會造成正確性問題,但會浪費 CPU 資源。

GIN 後端選擇與版本相容

直覺模型

GIN 支援多種後端:Proxy(基於 RMA 外掛的軟體模擬)、GDAKI(GPU Direct Async Kernel Initiated)、GPI(GPU-Initiated)、EFA GDA(AWS EFA 的 GPU Direct Async)。這就像同一個 API 可以有多種實作——軟體模擬版相容性最好但效能一般,硬體卸載版效能最好但需要特定網卡支援。

後端版本矩陣

每種後端有一個版本相容陣列,索引是後端版本號,值是該版本要求的最低 NCCL 版本。📎 src/gin/gin_host.cc:27-33

後端版本 0版本 1版本 2版本 3
Proxy02.30.32.30.52.32.0
GDAKI02.30.32.30.5-
GPI02.30.5--
EFA GDA02.31.02.32.0-

版本選擇邏輯:遍歷版本陣列,找到第一個要求版本高於當前裝置程式碼版本的條目,前一個版本即為可用版本。📎 src/gin/gin_host.cc:300-304

後端選擇流程

ncclGinDevCommSetup遍歷所有活躍後端,嘗試用每個後端建立 DevComm。📎 src/gin/gin_host.cc:427-442選擇條件包括:請求的 GIN 類型匹配(或未指定)、訊號能力滿足要求。📎 src/gin/gin_host.cc:430-435

ncclGinValidateSignalRequest檢查兩個能力:強訊號(supportsStrongSignals)和 VA 訊號(supportsVASignals)。📎 src/gin/gin_host.cc:230-243如果請求要求強訊號但後端不支援,跳過該後端。

連線建立與 stride 計算

ncclGinConnectOnce建立 GIN 連線。📎 src/gin/gin_host.cc:92-228

連線類型決定 stride:FULL 模式下 stride 為 1(連接所有 rank),RAIL 模式下 stride 為contiguousRanksPerHost(只連接同一 rail 的 rank)。📎 src/gin/gin_host.cc:139-145

在ginDevCommSetupWithBackend中,stride 的校驗邏輯很嚴格:

  • 請求的 stride 不能為 0。📎 src/gin/gin_host.cc:318-323
  • 請求的 stride 不能大於 rail team 的 stride。📎 src/gin/gin_host.cc:324-330
  • 請求的 stride 必須是已連接 stride 的倍數。📎 src/gin/gin_host.cc:331-337

這些約束的動機是:分層屏障假設 GIN 至少是 RAIL 連接的。📎 src/gin/gin_host.cc:325如果 stride 不滿足這些條件,某些 rank 之間的通訊路徑可能不存在。

生產陷阱

陷阱一:後端版本不匹配。如果裝置程式碼版本低於後端要求的最低版本,backendVersion會停留在較低值。📎 src/gin/gin_host.cc:301-303這可能導致某些新特性不可用(比如訊號永不重置),但不會導致錯誤。然而,如果裝置程式碼版本高於所有已知版本,backendVersion會取最大值,可能觸發未定義行為。

陷阱二:stride 校驗的邊界。如果requestedStride % connectedStride != 0,建立失敗。📎 src/gin/gin_host.cc:331-337這個檢查假設 connectedStride 是 2 的冪(FULL 模式為 1,RAIL 模式為contiguousRanksPerHost)。如果contiguousRanksPerHost不是 2 的冪(比如 3),倍數檢查可能拒絕合法的 stride。

本章思考與自測

Q1: 在scheduleRmaTasksToPlan的 WaitSignal 分支中,如果去掉plan->rmaArgs->nRmaTasks = (npeersCe > 0 ? 1 : 0) + (npeersProxy > 0 ? 1 : 0)這一行,改為直接設為 1,在什麼場景下會導致問題?

參考解析:看📎 src/rma/rma.cc:248。nRmaTasks記錄的是實際入列的任務數。如果所有 peers 都是 LSA 可達的(npeersProxy == 0),實際只有 1 個 CE 任務入列,nRmaTasks應該為 1。如果所有 peers 都不可達(npeersCe == 0),實際只有 1 個 Proxy 任務入列,nRmaTasks也應該為 1。但如果 peers 混合分布,兩個任務都入列,nRmaTasks應該為 2。

如果把這一行改為plan->rmaArgs->nRmaTasks = 1,在混合分布場景下,nRmaTasks會低估實際任務數。後續ncclRmaWaitSignal中的判斷plan->rmaArgs->nRmaTasksProxy > 0 && plan->rmaArgs->nRmaTasksCe > 0仍然能正確工作(因為用的是nRmaTasksProxy和nRmaTasksCe),📎 src/rma/rma.cc:47但任何依賴nRmaTasks做資源估算或日誌統計的程式碼會得到錯誤結果。更嚴重的是,如果後續程式碼用nRmaTasks來分配陣列或計算迴圈次數,可能導致緩衝區溢位或任務遺漏。

Q2: 在proxyGinPollGfd中,如果把hostGpuCtx->sis[targetRank]++移到proxyGinProcessGfd呼叫之後,在什麼並發場景下會導致 GFD 被重複處理?

參考解析:看📎 src/gin/gin_host_proxy.cc:228。sis是「已見索引」,表示 proxy 已經看到並開始處理的 GFD 數量。proxyGinPollGfd在拷貝完 GFD 後立即遞增sis,然後返回 1 表示成功。呼叫者ncclGinProxyProgress在迴圈中呼叫proxyGinPollGfd,如果返回 1 則繼續處理下一個 GFD。📎 src/gin/gin_host_proxy.cc:648-669

如果把sis++移到proxyGinProcessGfd之後,那麼在proxyGinProcessGfd執行期間(可能涉及網路外掛的非同步呼叫),sis仍然指向當前 GFD。如果此時 GPU 寫入了一個新的 GFD 到同一個槽位(因為佇列是環形的,pis可能已經回繞),proxyGinPollGfd會再次看到這個槽位,但sis沒有前進,導致重複處理同一個槽位。

更危險的是,proxyGinPollGfd在拷貝 GFD 後會清零隊列中的 GFD。📎 src/gin/gin_host_proxy.cc:206-208如果sis沒有前進,下一次輪詢會看到清零後的 GFD(flag 為 0),isGfdAvailable返回 false,導致 GFD 丟失。這會造成 GPU 側等待一個永遠不會被處理的請求,最終死鎖。

Q3: 在ncclRmaProxyProgressThread中,如果rmaProgress == 2分支中忘記調用rmaProxyState->cond.notify_one(),在什麼場景下會導致主線程永久阻塞?

參考解析:看📎 src/rma/rma_proxy.cc:373-378。rmaProgress == 2是「暫停請求」狀態,用於資源回收。主線程設置rmaProgress = 2後,會等待進度線程確認暫停。進度線程在cond.wait(lock)中等待,主線程需要調用cond.notify_one()來喚醒它。📎 src/rma/rma_proxy.cc:377

如果進度線程在設置rmaProgress = 0後忘記notify_one(),主線程會一直等待條件變量。但更關鍵的是,進度線程在cond.wait(lock)中等待時,主線程需要先獲取鎖才能設置rmaProgress = 2。如果進度線程在wait之前沒有釋放鎖,主線程無法獲取鎖,形成死鎖。

正確的順序是:進度線程設置rmaProgress = 0,調用notify_one()喚醒主線程,然後調用cond.wait(lock)釋放鎖並等待。主線程被喚醒後獲取鎖,設置rmaProgress = 2,調用notify_one()喚醒進度線程,然後等待進度線程確認。進度線程被喚醒後,設置rmaProgress = 0,再次notify_one(),然後wait。這個握手協議中任何一步的notify_one()缺失都會導致永久阻塞。

從 RMA 的 put/get 語義到 GIN 的 GPU 發起網絡通信,我們走完了 NCCL 向通用遠程內存訪問引擎演進的關鍵一步。但無論機制多麼精巧,最終都要通過插件體系與外部網絡後端、調優策略和性能採集器對接。下一章將進入插件世界,看 NCCL 如何在不修改核心代碼的前提下,動態加載 net、tuner、profiler、env 等擴展,並以 google-fastsocket 和 google-CoMMA 為例揭示生態擴展性的實現要點。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 16

第 16 章:插件生態與環境變量:net、tuner、profiler、env 如何擴展 NCCL 行為

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 16 章 / 共 25 章

上一章我們看到 NCCL 如何通過 RMA 與 GIN 將通信能力從集合操作延伸到點對點遠程訪問,甚至讓 GPU 直接發起網絡請求。這種向新硬件與低延遲場景的演進,對通信引擎的靈活性提出了更高要求:如果每次適配新網絡、新調優策略或新採集工具都要重新編譯核心代碼,NCCL 將難以跟上生態變化。本章拆解 src/plugin 與 plugins 目錄,回答一個核心問題:NCCL 如何在不重新編譯核心代碼的前提下,替換網絡後端、調優策略、性能採集器與配置來源。

16.1 插件加載器:plugin_open.cc 如何把 .so 變成可用的後端

直覺模型

把plugin_open.cc想像成 NCCL 的「招聘中介」:它手裡有一份崗位清單(NET、GIN、RMA、TUNER、PROFILER、ENV),每個崗位對應一個候選庫名。當 NCCL 需要某個崗位的人時,中介按固定順序去人才市場(動態鏈接器)找人,找到就簽合同(dlopen),找不到就記錄「這個人不存在」,最後交回一個句柄。若沒有這層中介,NCCL 就只能把網絡後端硬編碼進二進制,任何網卡廠商想接入都得改 NCCL 源碼——這正是插件體系要消滅的災難。

數據結構與內存佈局

加載器的全部狀態就是六個並行數組,索引即插件類型枚舉:

code
static char* libNames[NUM_LIBS];              // 已加载库的名字
char* ncclPluginLibPaths[NUM_LIBS];           // 库的绝对路径
static void* libHandles[NUM_LIBS];            // dlopen 返回的句柄
static const char* pluginNames[NUM_LIBS];     // 日志用的人类可读名
static const char* pluginPrefix[NUM_LIBS];    // 库名前缀
static const char* pluginFallback[NUM_LIBS];  // 找不到时的提示
static unsigned long subsys[NUM_LIBS];        // 日志子系统位掩码

這七個數組的下標必須嚴格對齊,pluginNames[type]、pluginPrefix[type]、subsys[type]描述的是同一個插件類型。📎 src/plugin/plugin_open.cc:18-29定義了NUM_LIBS = 6,類型順序為{"NET", "GIN", "RMA", "TUNER", "PROFILER", "ENV"},前綴為{"libnccl-net", "libnccl-gin", "libnccl-rma", "libnccl-tuner", "libnccl-profiler", "libnccl-env"}。

〔設計推斷與架構權衡〕

這裡用並行數組而非結構體數組,是為了讓openPluginLib這個單一函數能同時服務六種插件——類型只作為下標,邏輯完全復用。代價是新增插件類型時必須同步修改六個數組,編譯器無法幫你檢查漏改。

subsys數組決定日誌歸屬:NET/GIN/RMA 都掛NCCL_INIT | NCCL_NET,TUNER 掛NCCL_INIT | NCCL_TUNING,PROFILER 只掛NCCL_INIT,ENV 掛NCCL_INIT | NCCL_ENV。📎 src/plugin/plugin_open.cc:26-29這樣NCCL_DEBUG_SUBSYS=NET時只會看到網絡插件的日誌,不會淹沒在調優日誌裡。

Step-by-Step Walkthrough:一次ncclOpenNetPluginLib("mlx5")的完整旅程

假設用戶設置NCCL_NET_PLUGIN=mlx5,NCCL 初始化時調用ncclOpenNetPluginLib("mlx5"),它直接轉發到openPluginLib(ncclPluginTypeNet, "mlx5")。📎 src/plugin/plugin_open.cc:132-134

第一步:構造候選庫名。因為傳入了非空libName,走snprintf(libName_, MAX_STR_LEN, "%s", libName)分支,libName_變成"mlx5"。📎 src/plugin/plugin_open.cc:85-89注意此時它還不是一個合法的庫文件名——沒有前綴也沒有.so後綴。

第二步:第一次嘗試打開。 tryOpenLib("mlx5", ...)被調用。📎 src/plugin/plugin_open.cc:91進入tryOpenLib後,先檢查name是否為空或長度為零,然後有一個特殊分支:如果名字以STATIC_PLUGIN開頭,就把name置為nullptr。📎 src/plugin/plugin_open.cc:37-39這是給靜態連結進 NCCL 的插件用的哨兵——dlopen(nullptr)在 Linux 上返回主程式句柄,從而讓dlsym能在主程式符號表裡找到插件符號。

接著呼叫ncclOsDlopen(name)。📎 src/plugin/plugin_open.cc:41因為"mlx5"既不是路徑也不是合法庫名,dlopen會失敗。失敗後程式碼取ncclOsDlerror()的錯誤串,並做一個精細判斷:如果錯誤串裡同時包含name和"No such file or directory",就把*err設為ENOENT。📎 src/plugin/plugin_open.cc:42-55這個判斷的意義在於區分「檔案根本不存在」和「檔案存在但載入失敗」——前者只是候選名不對,應該靜默嘗試下一個候選名;後者是真實錯誤,應該打日誌。

第三步:第一次失敗後的處理。回到openPluginLib,libHandles[type]為空,且openErr == ENOENT,於是把"mlx5"追加到eNoEntNameList。📎 src/plugin/plugin_open.cc:97-101這個列表最終會拼成一句「Could not find: mlx5 libnccl-net-mlx5.so」的日誌。

第四步:第二次嘗試——加前綴。程式碼檢查libName是否既不是路徑(不含/)也不是庫名(不以lib開頭、不以.so結尾)。📎 src/plugin/plugin_open.cc:105-107 "mlx5"滿足條件,於是拼出"libnccl-net-mlx5.so"再次嘗試。📎 src/plugin/plugin_open.cc:108這一次dlopen成功,libHandles[type]被賦值,libNames[type]記錄庫名,ncclPluginLibPaths[type]透過getLibPath拿到絕對路徑,函式返回句柄。📎 src/plugin/plugin_open.cc:110-115

第五步:拿到絕對路徑。 getLibPath在 Linux 上用dlinfo(handle, RTLD_DI_LINKMAP, &lm)取出link_map,再strdup(lm->l_name)。📎 src/plugin/plugin_open.cc:65-69這個路徑會出現在後續所有日誌裡,讓使用者一眼看出到底載入了哪個檔案——生產環境排查「為什麼載入了錯誤的插件」時,這行日誌是第一現場。

整個決策流如下:

mermaid
flowchart TD
    start["openPluginLib(type, libName)"] --> build{"libName 非空?"}
    build -->|是| use_name["libName_ = libName"]
    build -->|否| use_prefix["libName_ = pluginPrefix[type] + .so"]
    use_name --> try1["tryOpenLib(libName_)"]
    use_prefix --> try1
    try1 --> ok1{"handle 非空?"}
    ok1 -->|是| success["记录 libNames/libPaths, 返回 handle"]
    ok1 -->|否| enoent{"openErr == ENOENT?"}
    enoent -->|是| append1["appendNameToList(eNoEntNameList)"]
    enoent -->|否| log1["INFO 打印 dlopen 错误"]
    append1 --> shape{"非路径且非库名?"}
    log1 --> shape
    shape -->|是| try2["tryOpenLib(prefix-libName.so)"]
    shape -->|否| report["打印 Could not find 列表"]
    try2 --> ok2{"handle 非空?"}
    ok2 -->|是| success
    ok2 -->|否| report
    report --> retnull["返回 nullptr"]

設計思考與生產踩坑

〔設計推斷與架構權衡〕

候選名順序即優先級。先試使用者給的裸名,再試加前綴的名字。這意味著如果當前目錄恰好有一個叫mlx5的檔案,它會被優先載入—— 這是一個潛在的安全面,生產環境應避免在LD_LIBRARY_PATH裡放入與插件同名的可執行檔。

STATIC_PLUGIN的語義。當NCCL_NET_PLUGIN=STATIC_PLUGIN時,tryOpenLib把名字置空,dlopen(nullptr)打開主程式,dlsym從主程式符號表找ncclNet_v12等符號。📎 src/plugin/plugin_open.cc:37-39這允許把插件靜態連結進 NCCL 二進位檔,省去部署.so的麻煩,代價是失去執行時替換能力。

引用計數與卸載。 ncclClosePluginLib只在libHandles[type] == handle時才真正dlclose,並清空路徑和名字。📎 src/plugin/plugin_open.cc:176-186這個相等判斷防止誤關一個已經被替換的句柄。GIN 和 RMA 插件透過ncclGetGinPluginLib/ncclGetNetPluginLib復用 NET 庫的句柄,實現方式是再次dlopen同一個庫名來增加引用計數。📎 src/plugin/plugin_open.cc:156-164這是dlopen的引用計數語義——同一個庫被打開兩次,需要dlclose兩次才真正卸載。

16.2 net.cc:網路插件的狀態機與生命週期

直覺模型

net.cc是網路插件的「調度中心」。它維護一個插件庫陣列,每個庫有自己的狀態(未載入、載入失敗、待載入、待初始化、已啟用)。當一個新的通訊域(communicator)誕生時,調度中心遍歷所有候選插件,逐個嘗試初始化,第一個成功的就被「分配」給這個通訊域,其餘外部插件全部禁用。若沒有這層狀態機,NCCL 就無法處理「插件載入了但裝置不可用」「多個插件共存時選哪個」「通訊域銷毀時如何安全卸載」這些現實問題。

資料結構與記憶體佈局

核心結構是netPluginLib_t:

欄位類型含義
namechar[255]插件庫名
dlHandlevoid*dlopen 句柄
ncclNetncclNet_t*網路函式表
ncclNetVerint網路 API 版本號
ncclCollNetncclCollNet_t*集合通訊卸載函式表
ncclNetPluginState列舉網路插件狀態
ncclCollNetPluginState列舉CollNet 插件狀態
ncclNetPluginRefCountint引用計數
netPhysDevs/netVirtDevsint物理/虛擬裝置數
collNetPhysDevs/collNetVirtDevsintCollNet 裝置數

📎 src/plugin/net.cc:63-76定義了這些欄位。注意ncclNet和ncclCollNet是分開的兩個函式表,狀態也是分開的兩個列舉——一個插件可以提供網路功能但不提供 CollNet 卸載。

狀態列舉有五個值:Disabled = -2(初始化失敗)、LoadFailed = -1(載入失敗)、LoadReady = 0(待載入)、InitReady = 1(已載入待初始化)、Enabled = 2(已啟用)。📎 src/plugin/net.cc:54-60用負數表示失敗態,使得「狀態 >= InitReady」這樣的比較能自然表達「至少已載入」。

全域狀態是三個變數:pluginCount記錄插件總數,netPluginLibs[NCCL_NET_MAX_PLUGINS]是插件陣列,netPluginMutex保護並發存取,initPluginLibsOnceFlag保證初始化只做一次。📎 src/plugin/net.cc:78-81

Step-by-Step Walkthrough:一次ncclNetInit(comm)的完整旅程

第一步:一次性初始化。 std::call_once(initPluginLibsOnceFlag, initPluginLibsOnceFunc)保證插件列表只構建一次。📎 src/plugin/net.cc:360 initPluginLibsOnceFunc讀取NCCL_NET_PLUGIN環境變數,若未設定則預設加入"libnccl-net.so",然後註冊兩個內建插件ncclNetIb和ncclNetSocket。📎 src/plugin/net.cc:288-340

環境變數解析用strtok_r按逗號切分,支援多個插件名。📎 src/plugin/net.cc:303-324有一個容量檢查:外部插件數量不能超過NCCL_NET_MAX_PLUGINS - NCCL_NET_NUM_INTERNAL_PLUGINS,超出部分被忽略並打日誌。📎 src/plugin/net.cc:307-311內建插件固定為 2 個(IB 和 Socket),所以外部插件最多NCCL_NET_MAX_PLUGINS - 2個。

第二步:加鎖遍歷。 std::lock_guard<std::mutex> lock(netPluginMutex)保護整個遍歷過程。📎 src/plugin/net.cc:361對每個插件索引,先判斷它是否是外部插件且處於LoadReady狀態,若是則呼叫ncclNetPluginLoad。📎 src/plugin/net.cc:364-367

第三步:載入插件。 ncclNetPluginLoad呼叫ncclOpenNetPluginLib拿到句柄,然後從高版本到低版本依次嘗試getNcclNet_v12到getNcclNet_v6,第一個返回非空的版本被採用。📎 src/plugin/net.cc:103-112版本陣列ncclNetVersion和函式指標陣列getNcclNet按降序排列,保證優先使用最新 API。📎 src/plugin/net.cc:41-43

如果所有版本都拿不到ncclNet,說明這個庫不是合法的網路插件。此時檢查NCCL_NET_PLUGIN是否被顯式設定:若設定了,用ATTN級別告警(使用者明確要求卻失敗);若沒設定,用INFO級別(只是預設嘗試失敗)。📎 src/plugin/net.cc:115-125這個區分很重要——使用者顯式配置失敗必須讓他看見。

第四步:初始化外掛。回到ncclNetInit,對狀態>= InitReady且名字匹配comm->config.netName的外掛呼叫ncclNetPluginInit。📎 src/plugin/net.cc:369-372 ncclNetPluginInit做兩件事:呼叫外掛的init函式建立通訊域上下文,以及首次初始化時呼叫devices探測裝置數。📎 src/plugin/net.cc:186-236

注意init的呼叫條件:pluginLib->ncclNetPluginState >= ncclNetPluginStateInitReady。📎 src/plugin/net.cc:190註解明確說明「每個新通訊域都必須呼叫 init 來設定正確的上下文」。📎 src/plugin/net.cc:189但裝置探測只在== InitReady時做一次。📎 src/plugin/net.cc:201這個「init 每次呼叫,devices 只調一次」的區分是效能最佳化——裝置探測可能很慢,但上下文必須每個通訊域獨立。

第五步:分配與停用。初始化成功後呼叫ncclNetPluginAssignToComm,它把外掛的ncclNet賦給comm->ncclNet,遞增引用計數,設定comm->netPluginIndex。📎 src/plugin/net.cc:238-255分配成功後立即呼叫ncclNetPluginDisableOtherExternal停用其他所有外部外掛。📎 src/plugin/net.cc:377-380

〔設計推斷與架構權衡〕

停用邏輯有個關鍵判斷:只有當被分配的外掛是外部外掛(pluginIndex >= pluginCount - NCCL_NET_NUM_INTERNAL_PLUGINS)時才停用其他外部外掛。📎 src/plugin/net.cc:257-259如果分配的是內建 IB 外掛,外部外掛保持原狀——這為後續通訊域留了選擇空間。

mermaid
flowchart TD
    init["ncclNetInit(comm)"] --> once["call_once(initPluginLibsOnceFunc)"]
    once --> lock["lock(netPluginMutex)"]
    lock --> loop{"遍历 pluginIndex"}
    loop -->|外部且 LoadReady| load["ncclNetPluginLoad()"]
    loop -->|状态 >= InitReady| namechk{"netName 匹配?"}
    load --> namechk
    namechk -->|否| loop
    namechk -->|是| plugininit["ncclNetPluginInit()"]
    plugininit --> enabled{"状态 == Enabled?"}
    enabled -->|否| loop
    enabled -->|是| assign["ncclNetPluginAssignToComm()"]
    assign --> assigned{"isAssigned?"}
    assigned -->|否| finalize["ncclNetPluginFinalize()"]
    finalize --> loop
    assigned -->|是| disable["ncclNetPluginDisableOtherExternal()"]
    disable --> ok["返回 ncclSuccess"]
    loop -->|遍历结束| fail["WARN 无可用插件, 返回 ncclInvalidUsage"]

並行控制與硬體互動

netPluginMutex保護所有對netPluginLibs的讀寫。ncclNetInit、ncclNetFinalize都加鎖。📎 src/plugin/net.cc:361📎 src/plugin/net.cc:411-416但ncclNetGetDevCount等函式註解說「不需要鎖,因為呼叫者已在ncclTopoGetSystem的鎖內」。📎 src/plugin/net.cc:418-429這是一種「鎖由上层持有」的約定,減少了巢狀鎖的開銷,代價是呼叫者必須遵守約定。

ncclGpuGdrSupport展示了外掛與硬體的直接互動:它分配 2MB GPU 緩衝,透過外掛的listen/connect/accept建立回環連線,然後嘗試regMr註冊 GPU 記憶體。📎 src/plugin/net.cc:464-535如果註冊成功,說明網卡支援 GPUDirect RDMA。這個探測結果快取在gdrSupportMatrix[32]裡,按 CUDA 裝置號索引。📎 src/plugin/net.cc:478-480

〔設計推斷與架構權衡〕

注意gdrSupportMatrix是static的,跨通訊域共享。📎 src/plugin/net.cc:478這意味著同一行程內多個通訊域會複用探測結果,避免重複的昂貴探測。但陣列大小硬編碼為 32,超過 32 個 GPU 的機器會越界——這是一個隱含的上限假設。

生產避坑指南

坑一:外掛載入成功但裝置數為零。 ncclNetPluginInit檢查devices(&ndev) != ncclSuccess || ndev <= 0就跳轉到失敗分支。📎 src/plugin/net.cc:202失敗後呼叫finalize清理已建立的上下文,把裝置數重置為NCCL_UNDEF_DEV_COUNT,狀態設為Disabled。📎 src/plugin/net.cc:229-234如果不做這個清理,後續通訊域會看到一個「已初始化但無裝置」的外掛,導致難以診斷的錯誤。

〔設計推斷與架構權衡〕

坑二:init成功但devices失敗。程式碼用initCompleted標誌追蹤init是否成功。📎 src/plugin/net.cc:178-184📎 src/plugin/net.cc:198失敗分支裡只有initCompleted為真才呼叫finalize。📎 src/plugin/net.cc:230這防止對未初始化的上下文呼叫finalize——很多外掛的finalize不檢查空指標,誤呼叫會崩潰。

坑三:通訊域銷毀時的引用計數。 ncclNetPluginFinalize先呼叫外掛的finalize,再遞減引用計數,最後在引用計數歸零且是外部外掛時卸載庫。📎 src/plugin/net.cc:342-355 ncclNetPluginUnload檢查dlHandle非空且引用計數為零才真正dlclose。📎 src/plugin/net.cc:84-101卸載後重置欄位但保留name,以便重新載入時複用。📎 src/plugin/net.cc:84-101

16.3 tuner.cc 與 profiler.cc:策略外掛與觀測外掛的不同契約

直覺模型

Tuner 外掛像「導航軟體的路線偏好設定」——它不改變車怎麼開,只改變選哪條路。Profiler 外掛像「行車記錄器」——它不干預駕駛,只記錄發生了什麼。兩者的共同點是都透過函式表接入,區別在於 Tuner 是「每個通訊域一個實例」的輕量策略物件,而 Profiler 需要一個獨立執行緒來非同步消費 GPU 產生的事件。

tuner.cc:極簡的全域單例

Tuner 的狀態極其簡單:一個互斥鎖、一個引用計數、一個庫控制代碼、一個符號指標、一個狀態變數。📎 src/plugin/tuner.cc:24-37沒有外掛陣列,沒有多外掛共存——全域只有一個 tuner。

ncclTunerPluginLoad的邏輯是「首次載入,後續複用」:如果狀態是LoadSuccess,直接把符號賦給comm->tuner並遞增引用計數。📎 src/plugin/tuner.cc:53-57否則讀取NCCL_TUNER_PLUGIN環境變數,若為"none"則直接失敗。📎 src/plugin/tuner.cc:59-63

〔設計推斷與架構權衡〕

版本協商從 v6 降到 v2,逐個嘗試。📎 src/plugin/tuner.cc:75-87注意這裡沒有 v1——tuner API 從 v2 開始才有穩定的函式表結構。

〔設計推斷與架構權衡〕

一個有趣的細節:如果ncclOpenTunerPluginLib傳回空,程式碼嘗試ncclGetNetPluginLib(ncclPluginTypeTuner)。📎 src/plugin/tuner.cc:65-70這意味著 tuner 可以打包在 net 外掛庫裡——這降低了部署複雜度,一個.so同時提供網路和調優功能。

profiler.cc:非同步事件消費執行緒

Profiler 是本章最複雜的外掛,因為它需要處理 GPU 非同步產生的事件。核心結構是ncclProfilerThread:

欄位類型作用
threadstd::thread消費執行緒
mutexstd::mutex保護佇列
condcondition_variable有新工作時喚醒
condIterationInactivecondition_variable等待迭代結束
stopint停止標誌
refCountint通訊域引用計數
cudaDevint綁定的 CUDA 裝置
abortFlagvolatile uint32_t*中止標誌
iterationActivebool是否正在迭代
pending/pendingTail鏈結串列待處理工作
active/activeTail鏈結串列處理中工作
opStack/opPool記憶體池工作物件分配
inflight/maxInflightSeen/maxInflightsize_t背壓觀測
droppedOpsuint64_t分配失敗計數

📎 src/plugin/profiler.cc:38-69定義了這個結構。注意pending和active是兩個獨立鏈結串列:生產者往pending追加,消費執行緒在鎖內把pending拼接到active,然後在鎖外遍歷active。📎 src/plugin/profiler.cc:56-59

iterationActive標誌是並行正確性的關鍵:消費執行緒在鎖內設為true後釋放鎖去呼叫外掛回呼,銷毀執行緒必須等這個標誌變回false才能拆除通訊域狀態。📎 src/plugin/profiler.cc:52-55

Step-by-Step Walkthrough:一次 KernelCh 事件的產生與消費

第一步:主機側入隊。當內核計劃(kernel plan)被提交時,ncclProfilerPostPlanWork遍歷計劃裡的集合任務,對每個啟用了ncclProfileKernelCh的任務,按通道範圍調用profilerPostWorkInternal。📎 src/plugin/profiler.cc:1315-1331

profilerPostWorkInternal先遞增comm->profiler.workCounter[channelId],然後調用profilerEnqueueOp。📎 src/plugin/profiler.cc:1259-1266註釋強調這個遞增必須「每次調用恰好一次,即使分配失敗」,以保持與設備內核的同步。📎 src/plugin/profiler.cc:1259-1266

第二步:分配工作對象。 profilerEnqueueOp在鎖內從內存池分配ncclProfilerWorkOp,填充通道號、工作計數器、激活掩碼、任務事件句柄、通訊域上下文等字段。📎 src/plugin/profiler.cc:1199-1223分配失敗時遞增droppedOps並記錄日誌,但不回退workCounter——這是保持與設備同步的關鍵。📎 src/plugin/profiler.cc:1202-1207

分配成功後把對象追加到pending鏈表尾部,遞增inflight,更新maxInflightSeen,喚醒消費線程。📎 src/plugin/profiler.cc:1225-1239

第三步:消費線程等待。 ncclProfilerThreadFunc循環調用waitForAction。📎 src/plugin/profiler.cc:1074-1077 waitForAction在鎖內等待條件變量,直到pending或active非空,或收到停止/中止信號。📎 src/plugin/profiler.cc:1017-1031

被喚醒後,它調用appendWorkToActiveQueue把pending拼接到active尾部,設置iterationActive = true,返回NCCL_PROFILER_THREAD_PROGRESS。📎 src/plugin/profiler.cc:1017-1031

第四步:處理工作。 profilerProgressOps在鎖外遍歷active鏈表。📎 src/plugin/profiler.cc:958-999對每個工作對象,檢查設備是否已經寫入了啟動時間戳:wc <= op->workStarted[ch].data[slot].counter。📎 src/plugin/profiler.cc:972注意用的是<=而非==,因為設備會環繞MAX_PROFILER_EVENTS_PER_CHANNEL個槽位,主機落後時設備可能已經覆蓋了該槽位。📎 src/plugin/profiler.cc:969-971

如果啟動條件滿足,調用ncclProfilerStartKernelChEvent通知插件。📎 src/plugin/profiler.cc:973然後檢查完成條件,若滿足則先觸發階段事件,再調用ncclProfilerStopKernelChEvent。📎 src/plugin/profiler.cc:978-985

完成的工作對象被摘出鏈表,收集到recycled列表。📎 src/plugin/profiler.cc:987-991

第五步:回收與發布。 cleanupAndStop在鎖內回收recycled列表,發布新的activeTail,清除iterationActive並通知等待者。📎 src/plugin/profiler.cc:1036-1050

mermaid
sequenceDiagram
    participant Host as 主机线程
    participant PT as Profiler 线程
    participant Plugin as Profiler 插件
    participant Dev as GPU 内核

    Host->>Host: profilerPostWorkInternal() 递增 workCounter
    Host->>PT: profilerEnqueueOp() 追加到 pending
    Host->>PT: cond.notify_one()
    PT->>PT: waitForAction() 返回 PROGRESS
    PT->>PT: appendWorkToActiveQueue() 拼接 pending 到 active
    Dev->>Dev: 内核写入 workStarted/workCompleted 时间戳
    PT->>PT: profilerProgressOps() 检查 wc <= counter
    PT->>Plugin: startEvent(ncclProfileKernelCh)
    PT->>Plugin: recordEventState(ncclProfilerKernelChStop)
    PT->>Plugin: stopEvent()
    PT->>PT: cleanupAndStop() 回收对象, 清除 iterationActive

並發控制與背壓

NCCL_PROFILER_DEFAULT_MAX_INFLIGHT定義為MAXCHANNELS * MAX_PROFILER_EVENTS_PER_CHANNEL * 4。📎 src/plugin/profiler.cc:32-32這是一個「軟上限」——超過它不會阻止入隊,只會打日誌。📎 src/plugin/profiler.cc:1233-1238註釋說明保持入隊是為了讓 KernelCh 事件與其父任務事件配對。📎 src/plugin/profiler.cc:32-32

日誌用 2 的冪次觸發:(pt->inflight & (pt->inflight - 1)) == 0。📎 src/plugin/profiler.cc:1233這保證只在 inflight 為 1、2、4、8... 時打日誌,避免刷屏。

消費線程的退避策略在updateProgressInterval裡:有進展時立即重試,無進展時從 1 微秒開始翻倍,上限 10 微秒。📎 src/plugin/profiler.cc:1054-1057這個設計平衡了延遲和 CPU 佔用。

生產避坑指南

坑一:銷毀時的工作洩漏。 ncclProfilerThreadDestroy先等待iterationActive變假,然後調用profilerPurgeByContext清除所有引用該通訊域上下文的待處理工作。📎 src/plugin/profiler.cc:1162-1169如果不做這個清除,插件回調會拿到已銷毀的上下文指針,導致 use-after-free。

坑二:停止時的排空。當收到停止信號但active非空時,返回NCCL_PROFILER_THREAD_CLEANUP_AND_STOP,cleanupAndStop的drainStuck參數為真,直接回收所有剩餘工作。📎 src/plugin/profiler.cc:1029📎 src/plugin/profiler.cc:1036-1050註釋說這些工作的內核永遠不會運行,所以直接丟棄。📎 src/plugin/profiler.cc:1034-1035

坑三:CUDA 設備綁定。消費線程啟動時調用cudaSetDevice(pt->cudaDev)。📎 src/plugin/profiler.cc:1054-1057註釋解釋:線程本身只讀主機固定內存,但插件可能做依賴上下文的驅動調用,所以防禦性綁定。📎 src/plugin/profiler.cc:1054-1057綁定失敗只打日誌不中止,因為線程本身不依賴 CUDA。📎 src/plugin/profiler.cc:1065-1070

16.4 官方示例:google-fastsocket 與 google-CoMMA 的實現要點

直覺模型

官方示例是插件 API 的「參考實現」。google-fastsocket展示如何用用戶態網絡棧替換內核 TCP;google-CoMMA展示如何實現一個 profiler 插件來採集通信性能。它們的存在證明插件 API 足夠表達真實需求。

google-fastsocket:替換網絡後端

〔設計推斷與架構權衡〕

FastSocket 是 Google 開源的用戶態網絡棧,通過AF_FABRIC地址族繞過內核 TCP/IP 棧。作為 NCCL net 插件,它需要實現ncclNet_t的全部函數:init、devices、getProperties、listen、connect、accept、regMr、isend、irecv、test、closeSend等。

關鍵實現點在於getProperties返回的ptrSupport:如果 FastSocket 支持 GPUDirect RDMA,應設為NCCL_PTR_HOST|NCCL_PTR_CUDA;否則只能設為NCCL_PTR_HOST,NCCL 會在發送前把 GPU 數據拷到主機內存。📎 plugins/net/README.md:245-245

connect和accept的「非阻塞」契約是插件實現的核心難點:它們必須立即返回,把sendComm/recvComm設為NULL,讓 NCCL 反覆調用直到成功。📎 plugins/net/README.md:299-311這要求插件內部維護連接狀態機,把耗時的握手放在後台。

google-CoMMA:實現 profiler 插件

〔設計推斷與架構權衡〕

CoMMA(Collective Memory Monitoring Agent)是 Google 的通信性能採集器。作為 profiler 插件,它實現ncclProfiler_t函數表:init、finalize、startEvent、stopEvent、recordEventState。

init接收ncclProfilerEventMask指針,插件通過寫入這個掩碼來選擇訂閱哪些事件。📎 src/plugin/profiler.cc:341NCCL 支持的事件類型包括 Group、Coll、P2p、ProxyOp、ProxyStep、ProxyCtrl、KernelCh、KernelPhase、NetPlugin 等。📎 src/plugin/profiler.cc:285-307

startEvent返回一個事件句柄,後續stopEvent和recordEventState用這個句柄關聯事件。📎 src/plugin/profiler.cc:392📎 src/plugin/profiler.cc:400-407插件可以用句柄存儲自己的狀態,實現事件配對和耗時統計。

設計思考

為什麼 net 插件有版本協商而 tuner/profiler 沒有?因為 net API 涉及裝置側程式碼(ncclNetDeviceHandle),版本不匹配會導致核心崩潰;而 tuner/profiler 是純主機側,版本不匹配最多是功能缺失。📎 src/plugin/net.cc:153-176展示了ncclNetCheckDeviceVersion如何檢查裝置類型和版本,不匹配時返回ncclInternalError。

為什麼 profiler 需要獨立執行緒?因為 profiler 回呼可能阻塞(比如寫檔案、發網路請求),如果在主機執行緒呼叫會拖慢通訊。📎 src/plugin/profiler.cc:950-952註解明確說「外掛回呼可能阻塞,所以不能在持鎖時呼叫」。

16.5 生產避坑指南與故障恢復鏈

坑一:外掛版本不匹配導致核心崩潰

ncclNetCheckDeviceVersion檢查props.netDeviceType和props.netDeviceVersion。📎 src/plugin/net.cc:153-176如果外掛報告的NCCL_NET_DEVICE_UNPACK版本與 NCCL 編譯時的NCCL_NET_DEVICE_UNPACK_VERSION不一致,返回ncclInternalError並告警。📎 src/plugin/net.cc:153-176這個檢查在ncclNetPluginAssignToComm裡被呼叫,失敗時外掛不會被分配給通訊域。📎 src/plugin/net.cc:241

恢復鏈:版本不匹配 →ncclNetCheckDeviceVersion返回錯誤 →ncclNetPluginAssignToComm返回isAssigned = false → ncclNetInit繼續嘗試下一個外掛 → 最終可能回退到內建 Socket 外掛。

坑二:profiler 執行緒無法退出

如果 profiler 外掛在stopEvent裡阻塞,消費執行緒會卡在profilerProgressOps裡,iterationActive永遠為真,ncclProfilerThreadDestroy會永久等待。📎 src/plugin/profiler.cc:1166這是一個真實的死鎖風險。

〔設計推斷與架構權衡〕

恢復鏈:comm->abortFlag被設定 →waitForAction偵測到中止 → 返回CLEANUP_AND_STOP → cleanupAndStop排空佇列。📎 src/plugin/profiler.cc:1017-1031但如果執行緒已經卡在外掛回呼裡,中止標誌無法打斷它—— 這是外掛實作者的責任,回呼必須有逾時。

坑三:tuner 外掛的引用計數洩漏

ncclTunerPluginLoad在成功時遞增tunerPluginRefCount。📎 src/plugin/tuner.cc:98 ncclTunerPluginUnload在comm->tunerPluginLoaded為真時遞減。📎 src/plugin/tuner.cc:111-123如果某個通訊域載入了 tuner 但銷毀時tunerPluginLoaded被意外清零,引用計數永遠不會歸零,外掛程式庫永遠不會卸載。

本章思考與自測

Q1: 如果把ncclNetPluginLoad裡「從高版本到低版本嘗試」的迴圈改成「只嘗試最高版本」,在什麼場景下會導致原本可用的外掛無法載入?

參考解析:看📎 src/plugin/net.cc:108-112。迴圈遍歷NCCL_NET_VERSION_COUNT個版本,從 v12 降到 v6,第一個返回非空的被採用。如果只嘗試 v12,那麼一個只實作了 v11 的舊外掛會載入失敗。

〔設計推斷與架構權衡〕

這個設計是為了向後相容:NCCL 核心升級到支援 v12 後,仍然能載入只提供 v11 的外掛。外掛作者被鼓勵提供多個版本的符號(見📎 plugins/net/README.md:35-37),這樣同一個.so能服務多個 NCCL 版本。

如果去掉降級嘗試,使用者升級 NCCL 後舊外掛會突然不可用,只能回退到內建 Socket 外掛,效能大幅下降。這正是版本協商存在的意義。

Q2: 在profilerProgressOps裡,如果把wc <= op->workStarted[ch].data[slot].counter改成wc == op->workStarted[ch].data[slot].counter,在什麼高併發場景下會導致事件永遠不觸發?

參考解析:看📎 src/plugin/profiler.cc:969-972。註解明確說明裝置會環繞MAX_PROFILER_EVENTS_PER_CHANNEL個槽位。如果主機消費速度落後於裝置生產速度,裝置可能已經用計數器wc + N覆蓋了槽位wc % MAX_PROFILER_EVENTS_PER_CHANNEL。

此時op->workStarted[ch].data[slot].counter的值是wc + N,而op->workCounter是wc。用==判斷會失敗,事件永遠不會觸發,工作物件永遠留在active鏈結串列裡,inflight只增不減,最終耗盡記憶體池。

用<=則能正確處理這種情況:只要裝置寫入的計數器不小於期望值,就認為事件已就緒。這是一個典型的「生產者-消費者環繞緩衝區」的正確性條件。

Q3: 如果ncclProfilerThreadDestroy裡去掉等待iterationActive變假的迴圈,在什麼時序下會導致 profiler 外掛存取已釋放的通訊域上下文?

參考解析:看📎 src/plugin/profiler.cc:1162-1166。註解說明ncclProfilerPluginFinalize會在ncclProfilerThreadDestroy返回後立即銷毀通訊域的profilerContext。

消費執行緒在profilerProgressOps裡呼叫外掛回呼時,傳入的是op->profilerContext。📎 src/plugin/profiler.cc:938如果銷毀執行緒不等待iterationActive變假就返回,ncclProfilerPluginFinalize會釋放上下文,而消費執行緒可能正在用這個上下文呼叫外掛——use-after-free。

iterationActive的握手協定是:消費執行緒在鎖內置為true後釋放鎖去呼叫外掛,銷毀執行緒在鎖內等待它變回false。📎 src/plugin/profiler.cc:1028📎 src/plugin/profiler.cc:1054-1057這個協定保證外掛回呼期間上下文始終有效。

去掉等待後,銷毀執行緒可能在消費執行緒剛進入外掛回呼時就返回,導致外掛拿到懸空指標。這是一個典型的「生命週期與併發存取」競態。

外掛體系讓 NCCL 從封閉走向開放:網路後端、調優策略、效能採集器、配置來源都可以在不改核心程式碼的前提下替換。但外掛也引入了新的故障面——版本不匹配、生命週期競態、引用計數洩漏。下一章我們將進入 RAS 與診斷子系統,看 NCCL 如何偵測故障、監控進度並在長時間訓練任務中實現自癒。

外掛體系讓 NCCL 的核心通訊路徑與可替換元件之間劃出了清晰邊界,net、tuner、profiler、env 四類外掛各自透過註冊與引用計數機制安全地介入執行時行為。但一個可擴展的通訊引擎不僅要能靈活替換元件,更要在長時間訓練中穩定運行——當網卡或 GPU 出現故障時,NCCL 如何偵測、監控並觸發恢復?下一章我們將進入 RAS 與診斷機制,看生產環境下的可靠性如何被系統性地保障。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 17

第 17 章:RAS 機制與容錯:鏈路故障偵測、心跳與優雅降級

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 17 章 / 共 25 章

上一章我們看到外掛體系如何讓核心通訊路徑與可替換元件劃清邊界,從而在不修改核心程式碼的前提下替換網路後端、調優策略和效能採集器。但可擴展性只是生產可用的一個維度,另一個同樣硬核的問題是:當一次 AllReduce 已經跑了 72 小時,某台機器的網卡悄悄掛了,NCCL 憑什麼能發現、能隔離、能繼續?RAS 子系統正是 NCCL 從「能跑通」走向「生產可用」的分水嶺,本章將拆解故障偵測、進度監控與自癒機制背後的設計。

17.1 RAS 總控:一個行程一個 RAS 執行緒的全域協調者

直覺模型

把 RAS 想像成整個作業的「值班室」。每個 NCCL 行程(每個 rank)在初始化時都會開一間值班室,裡面坐著一個專職執行緒。所有通訊域(communicator)的建立、銷毀、診斷請求,都要先向值班室登記;值班室之間再透過一條獨立的 RAS 網路互相通報「誰還活著、誰已經死了」。

如果沒有這間值班室,NCCL 就只能靠通訊路徑本身的逾時來感知故障——而通訊路徑上的逾時既慢又容易誤判(一次網路抖動就可能被當成節點死亡)。RAS 把「故障感知」從資料面剝離到控制面,用獨立的輕量心跳和診斷通道來判定健康狀態。

資料結構與記憶體佈局

RAS 的核心狀態散落在ras.cc的全域變數裡,我們逐一拆解:

變數類型作用
rasInitMutexstd::mutex保護 RAS 單例初始化
rasInitializedbool是否已初始化
rasInitRefCountint引用計數,等於活躍 comm 數
rasNetListeningSocketstruct ncclSocketRAS 網路監聽套接字
rasNotificationPipe[2]ncclSocketPairDescriptor本地執行緒 → RAS 執行緒的通知管道
rasPfdsstruct pollfd*主事件迴圈的 poll 陣列
ncclCommsstruct ncclComm**所有通訊域指標陣列

📎 src/ras/ras.cc:49-61定義了這些全域狀態。注意rasInitRefCount用ncclAtomicRefCountIncrement增減📎 src/ras/ras.cc:129,而rasInitialized用普通 bool 加雙重檢查鎖保護📎 src/ras/ras.cc:103-105——這是典型的「初始化一次、之後唯讀」模式。

ncclComms陣列的分配策略值得注意:它不是按需增長,而是每次擴容RAS_INCREMENT * 8(即 32 個槽位)📎 src/ras/ras.cc:139-140。陣列裡允許出現nullptr空洞(comm 銷毀時置空),新 comm 會複用第一個空洞📎 src/ras/ras.cc:135-137。

場景驅動 Walkthrough:從 comm 初始化到 RAS 執行緒啟動

第一步:ncclRasCommInit被呼叫。這是每個 comm 初始化時第一個呼叫的 RAS 函式📎 src/ras/ras.cc:101。它先檢查rasInitialized,若未初始化則進入臨界區:

1. 用 bootstrap 網路介面位址初始化rasNetListeningSocket,埠設為 0 讓核心隨機分配📎 src/ras/ras.cc:108-109

2. 監聽該套接字📎 src/ras/ras.cc:113

3. 建立本地通知管道📎 src/ras/ras.cc:118

4. 初始化診斷子系統📎 src/ras/ras.cc:120

5. 啟動rasThreadMain執行緒📎 src/ras/ras.cc:121

6. 註冊atexit(rasTerminate)保證行程退出時清理📎 src/ras/ras.cc:126

第二步:登記 comm。無論是否首次初始化,都會把comm指標寫入ncclComms陣列📎 src/ras/ras.cc:142,並把ncclCommsSorted置 false📎 src/ras/ras.cc:143——因為陣列順序變了,之前的排序失效。

第三步:回填埠。函式最後把rasNetListeningSocket.addr(含核心分配的埠)拷回myRank->addr 📎 src/ras/ras.cc:146,這樣呼叫方就能知道 RAS 網路監聽在哪個埠。

主事件迴圈:poll 驅動的多路複用

rasThreadMain是 RAS 執行緒的心臟📎 src/ras/ras.cc:633。它先註冊三個固定 fd:通知管道、RAS 網路監聽套接字、客戶端監聽套接字📎 src/ras/ras.cc:641-652。然後進入無限迴圈:

code
for (int64_t nextWakeup = 0;;) {
  // 计算超时
  timeoutMs = min(..., 1000);
  nEvents = poll(rasPfds, nRasPfds, timeoutMs);
  // 处理事件
  for (pollIdx...) { ... }
  // 处理各类超时
  rasSocksHandleTimeouts(now, &nextWakeup);
  rasConnsHandleTimeouts(now, &nextWakeup);
  rasNetHandleTimeouts(now, &nextWakeup);
  rasCollsHandleTimeouts(now, &nextWakeup);
}

📎 src/ras/ras.cc:655-728展示了這個迴圈。注意timeoutMs被硬性限制在 1000ms 以內📎 src/ras/ras.cc:664——即使nextWakeup很遠,也要每秒醒一次,保證逾時檢查的及時性。

事件分發邏輯用 fd 值做路由📎 src/ras/ras.cc:684-715:如果是通知管道就調rasLocalHandle;如果是監聽套接字就 accept;否則遍歷rasSocketsHead和rasClientsHead鏈結串列找到對應的 socket 處理。

本地通知機制:管道 + 定長結構

本地 NCCL 執行緒與 RAS 執行緒透過一個 socketpair 通訊。通知結構rasNotification是定長的📎 src/ras/ras.cc:35-46,並用static_assert保證不超過PIPE_BUF 📎 src/ras/ras.cc:47——這是為了確保寫入的原子性(POSIX 保證小於 PIPE_BUF 的寫入是原子的)。

發送端rasLocalNotify用rasNotificationMutex序列化多個使用者執行緒的寫入📎 src/ras/ras.cc:224-237,然後迴圈寫直到全部寫完📎 src/ras/ras.cc:224-237。接收端rasLocalHandle同樣迴圈讀滿整個結構📎 src/ras/ras.cc:247-256,讀到 EOF 返回ncclSystemError 📎 src/ras/ras.cc:251-253。

三種通知類型:RAS_ADD_RANKS(新 rank 加入)、RAS_RUN_DIAG(執行診斷)、RAS_TERMINATE(終止)📎 src/ras/ras.cc:28-32。

訊息收發:長度前綴 + 增量進度

RAS 訊息的線格式是「4 位元組長度 + 訊息體」📎 src/ras/ras_internal.h:110-117。發送時rasConnSendMsg先發長度再發訊息體📎 src/ras/ras.cc:362-390,用meta->offset記錄進度,支援部分發送後下次繼續。接收時rasMsgRecv先收長度、按長度分配緩衝區、再收訊息體📎 src/ras/ras.cc:393-412。

這裡有個細節:rasMsgAlloc分配的是rasMsgMeta結構,msg欄位在結構末尾,透過offsetof計算偏移📎 src/ras/ras.cc:313-319。釋放時反向計算📎 src/ras/ras.cc:323-328。這種「元資料前置」的佈局讓訊息可以攜帶發送進度、入隊時間等本地資訊,而不佔用線格式。

設計思考

〔設計推斷與架構權衡〕

為什麼用 poll 而不是 epoll?poll 的 O(n) 複雜度在 RAS 場景下可接受——RAS 連線數遠小於資料面連線數,且 RAS 執行緒本身不是效能關鍵路徑。poll 的跨平台性也更好(Windows 相容)。

〔設計推斷與架構權衡〕

為什麼通知用管道而不是條件變數?管道可以無縫整合進 poll 迴圈,讓 RAS 執行緒用統一的poll等待所有事件源。如果用條件變數,就需要額外的機制來喚醒 poll。

mermaid
flowchart TD
    start["rasThreadMain 启动"] --> reg_pipe["注册通知管道 fd"]
    reg_pipe --> reg_net["注册 RAS 网络监听 fd"]
    reg_net --> reg_client["注册客户端监听 fd"]
    reg_client --> poll["poll(rasPfds, timeout<=1000ms)"]
    poll --> check{"nEvents == -1?"}
    check -->|"是且非 EINTR"| log_err["记录 poll 错误并继续"]
    check -->|"否"| dispatch["遍历 revents 分发事件"]
    log_err --> dispatch
    dispatch --> is_pipe{"fd == 通知管道?"}
    is_pipe -->|"是"| local_handle["rasLocalHandle()"]
    is_pipe -->|"否"| is_net{"fd == RAS 监听?"}
    is_net -->|"是"| accept_net["rasNetAcceptNewSocket()"]
    is_net -->|"否"| is_client{"fd == 客户端监听?"}
    is_client -->|"是"| accept_client["rasClientAcceptNewSocket()"]
    is_client -->|"否"| find_sock["遍历 rasSocketsHead 找匹配 socket"]
    find_sock --> sock_loop["rasSockEventLoop(sock, pollIdx)"]
    local_handle --> terminate{"terminate?"}
    terminate -->|"是"| cleanup["rasThreadCleanup() 并退出"]
    terminate -->|"否"| timeouts
    sock_loop --> timeouts["rasSocksHandleTimeouts / rasConnsHandleTimeouts / rasNetHandleTimeouts / rasCollsHandleTimeouts"]
    accept_net --> timeouts
    accept_client --> timeouts
    timeouts --> poll

17.2 進度監控:用 DMA 把 GPU 計數器搬到主機

直覺模型

進度監控像汽車儀表板上的「引擎轉速表」。它不參與駕駛(不參與通訊),但持續把 GPU 內部的進度計數器抄到主機記憶體,讓主機能判斷「這個通訊域是不是卡住了」。如果沒有它,一次 AllReduce 卡死時你只能看到「程式不返回」,卻不知道是 GPU 在算、在等網路、還是徹底死鎖。

資料結構與記憶體佈局

每個 CUDA 裝置對應一個ncclGpuProgressCounterMonitor工作執行緒📎 src/ras/progress_monitor.cc:35-52:

欄位類型作用
cudaDevint綁定的 CUDA 裝置號
threadstd::thread工作執行緒
mutex / cvstd::mutex / condition_variable保護可變狀態與喚醒
running / shouldStopbool執行緒生命週期標誌
copyInFlightbool是否有 DMA 拷貝在途
copyStallWarnedbool是否已對本次卡頓告警
copyStartNsuint64_t本次拷貝開始時間
sideStreamcudaStream_t專用非阻塞流
copyDonecudaEvent_t拷貝完成事件
warningMutexstd::mutex保護告警時間戳
lastStaleWarnNs / lastErrorWarnNsuint64_t限流時間戳
destroyRefsint銷毀引用計數
registrations侵入式佇列註冊到本裝置的 comm 列表

📎 src/ras/progress_monitor.cc:59-62明確了鎖順序:gpuProgressCounterMonitorsMu先於ncclGpuProgressCounterMonitor::mutex。這是避免死鎖的關鍵約定。

全域陣列gpuProgressCounterMonitors[kRasMaxCudaDevices]按裝置號索引📎 src/ras/progress_monitor.cc:59-62。

場景驅動 Walkthrough:一次計數器拷貝

第一步:註冊。 ncclProgressCounterMonitorInit被呼叫📎 src/ras/progress_monitor.cc:319。若deviceCountersBlock為空則直接返回(該 comm 不參與監控)📎 src/ras/progress_monitor.cc:323。否則在全域鎖內查找或建立該裝置的 worker📎 src/ras/progress_monitor.cc:328-335,然後把 comm 入隊到registrations 📎 src/ras/progress_monitor.cc:339。

第二步:工作執行緒啟動。 createGpuProgressCounterMonitor建立 worker,設定cudaSetDevice、建立sideStream(cudaStreamNonBlocking)和copyDone事件📎 src/ras/progress_monitor.cc:280-282,啟動執行緒後等待最多 2000ms 確認running變 true📎 src/ras/progress_monitor.cc:287-303。

第三步:迴圈拷貝。 progressCounterMonitorLoop先綁定裝置、設定 relaxed 流捕獲模式(避免干擾應用的 graph capture)📎 src/ras/progress_monitor.cc:97-121,然後進入主迴圈:

1. 等待pollIntervalMs(預設 1000ms)📎 src/ras/progress_monitor.cc:132-136

2. 若上次拷貝還在途,用cudaEventQuery檢查📎 src/ras/progress_monitor.cc:140。若cudaErrorNotReady且超過 stale 閾值(預設 5000ms),發出限流告警📎 src/ras/progress_monitor.cc:141-154

3. 遍歷所有註冊的 comm,對每個呼叫cudaMemcpyAsync把deviceCountersBlock拷到hostCountersBlock 📎 src/ras/progress_monitor.cc:170-185

4. 若有任何拷貝成功,記錄copyDone事件並置copyInFlight 📎 src/ras/progress_monitor.cc:194-202

並發控制與限流

告警限流由progressCounterMonitorShouldWarn實現📎 src/ras/progress_monitor.cc:78-87:在warningMutex保護下檢查距上次告警是否超過warnIntervalNs,超過才更新並返回 true。預設staleWarnSec是 600 秒📎 src/ras/progress_monitor.cc:27,即同一類告警最多每 10 分鐘一條。

參數有下限鉗制:poll 間隔最小 50ms📎 src/ras/progress_monitor.cc:29,stale 閾值最小 1000ms📎 src/ras/progress_monitor.cc:30。這防止使用者配置過激導致 CPU 空轉。

銷毀:引用計數 + 流同步

ncclProgressCounterMonitorDestroy的銷毀邏輯是本章最精妙的並發設計之一📎 src/ras/progress_monitor.cc:352-354:

1. 在全域鎖 + worker 鎖內從registrations刪除 comm📎 src/ras/progress_monitor.cc:368

2. 若刪除成功,destroyRefs++並置haveDestroyRef 📎 src/ras/progress_monitor.cc:371-372

3. 若註冊列表變空,從全域陣列摘除並置shouldStop 📎 src/ras/progress_monitor.cc:373-376

4. 釋放鎖後,cudaStreamSynchronize(g->sideStream)排空可能仍引用該 comm 緩衝區的拷貝📎 src/ras/progress_monitor.cc:393

5. 最後releaseGpuProgressCounterMonitorDestroyRef遞減引用計數,歸零且佇列空時 join 執行緒並刪除📎 src/ras/progress_monitor.cc:219-246

〔設計推斷與架構權衡〕

為什麼需要destroyRefs?因為cudaStreamSynchronize在鎖外執行,期間可能有另一個執行緒也在銷毀同一個 worker。引用計數保證只有最後一個銷毀者才真正 join 和 delete。

mermaid
sequenceDiagram
    participant App as 应用线程
    participant Mon as 监控线程
    participant GPU as CUDA 设备
    App->>Mon: ncclProgressCounterMonitorInit(comm)
    Mon->>Mon: 查找/创建 worker
    Mon->>Mon: registrations 入队 comm
    loop 每 pollIntervalMs
        Mon->>GPU: cudaEventQuery(copyDone)
        GPU-->>Mon: cudaErrorNotReady / cudaSuccess
        Mon->>GPU: cudaMemcpyAsync(hostCounters, deviceCounters, D2H, sideStream)
        Mon->>GPU: cudaEventRecord(copyDone, sideStream)
    end
    App->>Mon: ncclProgressCounterMonitorDestroy(comm)
    Mon->>Mon: registrations 删除 comm, destroyRefs++
    Mon->>GPU: cudaStreamSynchronize(sideStream)
    GPU-->>Mon: 拷贝排空完成
    Mon->>Mon: releaseGpuProgressCounterMonitorDestroyRef
    Mon->>Mon: join 线程, delete worker

生產避坑

坑 1:cudaSetDevice失敗導致監控靜默失效。若執行緒啟動時cudaSetDevice失敗,worker 會置shouldStop並退出📎 src/ras/progress_monitor.cc:97-107,但註冊它的 comm 仍然認為監控在跑。此時計數器鏡像會一直陳舊,直到 Init 階段暴露失敗。排查時要看NCCL_RAS日誌裡是否有 "progress-counter mirrors will remain stale"。

坑 2:graph capture 衝突。監控執行緒呼叫 CUDA API 時若應用正在做 stream capture,會污染捕獲圖。程式碼用cudaThreadExchangeStreamCaptureMode(cudaStreamCaptureModeRelaxed)規避📎 src/ras/progress_monitor.cc:110-111,這是必須的防護。

17.3 診斷框架:表驅動的檢查分發

直覺模型

診斷框架像醫院的「體檢套餐」。每個檢查項(GPU 型號、ECC 狀態、NVLink 健康、XID 錯誤等)是一個獨立的「檢查科室」,框架負責把各 rank 的檢查結果收集起來、彙總成一份報告。沒有它,維運只能靠nvidia-smi逐台機器手工排查,在千卡叢集上完全不可行。

資料結構:檢查分發表

核心是一張靜態分發表rasDiagnosticsChecks 📎 src/ras/diagnostics.cc:63-77,每個條目綁定一個檢查 ID 和兩個回呼:collectLocal(本地採集)和summarize(彙總)。共 11 項檢查:GPU 型號、CUDA 驅動版本、ECC、NVLink、NCCL 環境、RDMA 拓撲、IOMMU 模式、ATS、XID/SXID、NVIDIA 驅動版本、路徑。

rasDiagnosticsGetCheck做三重校驗:ID 範圍、表項 ID 匹配、回調非空📎 src/ras/diagnostics.cc:104-128。這是防禦性編程——防止表項被錯誤修改導致調用空指針。

場景驅動 Walkthrough:一次診斷的完整生命週期

第一步:構建本地 payload。 rasDiagnosticsCollectLocalPeerPayload先寫入 peer 頭📎 src/ras/diagnostics.cc:226-227,然後遍歷分發表,對每項調用rasDiagnosticsAppendCheckPayload 📎 src/ras/diagnostics.cc:229-231。

rasDiagnosticsAppendCheckPayload調用collectLocal拿到rasDiagnosticsLocalData,用ncclUniquePtr接管 records 所有權📎 src/ras/diagnostics.cc:191-192,校驗元數據📎 src/ras/diagnostics.cc:193,若記錄數為 0 則跳過📎 src/ras/diagnostics.cc:194,否則寫入檢查頭 + 記錄數據📎 src/ras/diagnostics.cc:196-201。

第二步:發起集合通信。 rasDiagnosticsStart構造RAS_COLL_DIAG請求📎 src/ras/diagnostics.cc:532-537,通過rasNetSendCollReq發出📎 src/ras/diagnostics.cc:539,客戶端狀態置為RAS_CLIENT_DIAG_FINI 📎 src/ras/diagnostics.cc:541。

第三步:合併響應。 rasCollDiagMerge把各 peer 的 payload 追加到集合緩衝區📎 src/ras/diagnostics.cc:310-337。注意它做了大量溢出檢查:peer 數上限📎 src/ras/diagnostics.cc:320-324、總大小上限📎 src/ras/diagnostics.cc:325-328。

第四步:彙總。 rasDiagnosticsSummarizePeerPayloads是兩遍掃描📎 src/ras/diagnostics.cc:399:

  • 第一遍:校驗每個 peer 頭和檢查頭,累計每類檢查的記錄數和字節數📎 src/ras/diagnostics.cc:418-470
  • 分配每類檢查的合併緩衝區📎 src/ras/diagnostics.cc:472-476
  • 第二遍:把各 peer 的記錄拷貝到對應緩衝區📎 src/ras/diagnostics.cc:479-497
  • 最後對每類檢查調用summarize 📎 src/ras/diagnostics.cc:499-506

客戶端狀態與取消

診斷狀態存在rasDiagnosticsClientState裡📎 src/ras/diagnostics.cc:242-245,掛在rasClient->diagnostics上。rasDiagnosticsCancelTarget在客戶端 socket 關閉時把 reporter 換成 noop📎 src/ras/diagnostics.cc:286-293,防止異步診斷完成後向已關閉的 socket 寫入📎 src/ras/diagnostics.cc:48-52。

設計思考

〔設計推斷與架構權衡〕

為什麼用兩遍掃描?因為 payload 是變長的,第一遍才能算出每類檢查需要多大緩衝區。一遍掃描要麼動態增長(多次 realloc),要麼預分配過大。兩遍掃描用一次精確分配換取確定性。

為什麼檢查頭裡帶recordStride? 📎 src/ras/diagnostics.cc:197因為不同檢查的記錄結構大小不同,彙總時需要知道步長才能正確拷貝和校驗。rasDiagnosticsAccountCheckRecords強制同一檢查的 stride 一致📎 src/ras/diagnostics.cc:381-385。

mermaid
flowchart TD
    start["rasDiagnosticsStart"] --> build_req["构造 RAS_COLL_DIAG 请求"]
    build_req --> send["rasNetSendCollReq"]
    send --> all_done{"allDone?"}
    all_done -->|"是"| fini["client->status = DIAG_FINI"]
    all_done -->|"否"| in_progress["返回 ncclInProgress"]
    fini --> resume["rasDiagnosticsResume"]
    in_progress --> resume
    resume --> summarize["rasDiagnosticsSummarizePeerPayloads"]
    summarize --> pass1["第一遍: 校验头 + 累计每类记录数"]
    pass1 --> valid{"payload 合法?"}
    valid -->|"否"| err["返回 ncclInternalError"]
    valid -->|"是"| alloc["为每类检查分配合并缓冲区"]
    alloc --> pass2["第二遍: 拷贝各 peer 记录"]
    pass2 --> emit["对每类检查调用 summarize"]
    emit --> finish["reporter.finish + rasCollFree"]

17.4 對等體管理:排序數組 + 哈希同步

直覺模型

peers.cc維護的是「全班同學名單」。每個 RAS 線程都保存一份完全相同的名單,記錄每個 NCCL 進程的地址、PID、管理的 GPU。當有新同學加入或有人「失聯」時,通過 RAS 網絡把變更廣播出去。名單用哈希值做版本號,避免每次全量同步。

數據結構與內存佈局

兩個核心數組:

  • rasPeers:所有已知 peer,按地址排序📎 src/ras/peers.cc:18-19。包含已死 peer。
  • rasDeadPeers:已死 peer 地址,單獨存放📎 src/ras/peers.cc:37-38。

為什麼死 peer 單獨存? 📎 src/ras/peers.cc:25-28的註釋解釋得很清楚:rasPeers在大規模下基本靜態且很大,而rasDeadPeers動態且小得多。分開存避免每次同步都傳輸龐大的rasPeers數組。

rasPeerInfo結構📎 src/ras/ras_internal.h:110-117:

字段類型說明
addrncclSocketAddress網絡地址(排序鍵)
pidncclPid_t進程 ID
cudaDevsuint64_tCUDA 設備位掩碼(受 CUDA_VISIBLE_DEVICES 影響)
nvmlDevsuint64_tNVML 設備位掩碼(不受影響)
hostHash / pidHashuint64_t從 comm 提取,減去 commHash 使其與通信域無關

兩個哈希rasPeersHash和rasDeadPeersHash是同步的核心📎 src/ras/peers.cc:21📎 src/ras/peers.cc:37-38。

場景驅動 Walkthrough:新 rank 加入

第一步:轉換。 rasRanksConvertToPeers把rasRankInit數組轉成rasPeerInfo 📎 src/ras/peers.cc:104。先按地址 + cudaDev 排序📎 src/ras/peers.cc:114,跳過空地址📎 src/ras/peers.cc:127-130,合併同地址的多 GPU 進程(位掩碼 OR)📎 src/ras/peers.cc:134-139。

第二步:更新本地數組。 rasPeersUpdate是本章最複雜的合併算法📎 src/ras/peers.cc:197。它先計算新數組大小📎 src/ras/peers.cc:202-229,然後歸併兩個有序數組📎 src/ras/peers.cc:244-361。關鍵點:合併過程中把rankPeers改造成「差異」——只保留真正新增的 GPU 位📎 src/ras/peers.cc:301-308,最後清除無貢獻的條目📎 src/ras/peers.cc:393-402。這樣廣播的數據量最小。

第三步:傳播。 rasNetUpdatePeers沿rasNextLink和rasPrevLink兩個方向傳播📎 src/ras/peers.cc:430-450,然後重建連接📎 src/ras/peers.cc:443-444。

第四步:發送更新。 rasConnSendPeersUpdate先檢查哈希📎 src/ras/peers.cc:500-508:若對端已知當前哈希則跳過。消息裡帶peersHash和deadPeersHash 📎 src/ras/peers.cc:521-524,接收方合併後若哈希仍不匹配則回發📎 src/ras/peers.cc:608-653。

死 peer 的聲明與傳播

rasPeerDeclareDead把地址加入rasDeadPeers,排序後重算哈希📎 src/ras/peers.cc:793-812。rasMsgHandleBCDeadPeer處理廣播的死 peer 消息📎 src/ras/ras.cc:578-591:若本地未知則斷開連接並聲明死亡,否則標記*pDone = true停止重廣播。

rasDeadPeersUpdate用歸併排序合併新舊死 peer 列表📎 src/ras/peers.cc:838-893。注意它用memmove而非memcpy 📎 src/ras/peers.cc:855,因為源和目標可能重疊。

連接重建:避免重複連接競態

rasLinkReinitConns在 peer 更新後重建鏈路連接📎 src/ras/peers.cc:680。核心策略:從地址較小的一方發起連接📎 src/ras/peers.cc:706-711,避免雙方同時發起導致重複。

rasLinkCalculatePeer計算下一個 peer 索引,跳過死 peer📎 src/ras/peers.cc:743-785。對 fallback 還有額外優化:跳過與前一 fallback 同節點的 peer📎 src/ras/peers.cc:743-785,避免整節點宕機時逐個等待。

生產避坑

坑 1:地址比較的字節序陷阱。 ncclSocketsCompare按地址族 → 地址 → 端口排序📎 src/ras/peers.cc:960-990。註釋指出不能簡單memcmp整個結構,因為內存佈局順序與期望排序順序不同📎 src/ras/peers.cc:957-959。IPv4 地址和端口在網絡字節序下可以逐字節比較,但地址族字段不行。

坑 2:myPeerIdx失效。陣列增長時myPeerIdx會變📎 src/ras/peers.cc:22-23。rasPeersUpdate在合併過程中同步更新它📎 src/ras/peers.cc:312📎 src/ras/peers.cc:358,若更新失敗則回退到二分查找📎 src/ras/peers.cc:374-388。

〔設計推斷與架構權衡〕

坑 3:雜湊碰撞導致同步遺漏。雜湊只用於「是否需要同步」的判斷,不用於正確性 。即使雜湊碰撞導致跳過同步,後續 keep-alive 交換仍會帶上雜湊,最終收斂。

mermaid
flowchart LR
    subgraph 输入
        ranks["rasRankInit[]"]
    end
    subgraph 转换
        convert["rasRanksConvertToPeers: 排序+合并同地址"]
        rankPeers["rasPeerInfo[] (rankPeers)"]
    end
    subgraph 合并
        update["rasPeersUpdate: 归并到 rasPeers"]
        diff["rankPeers 改造为差异"]
        hash["重算 rasPeersHash"]
    end
    subgraph 传播
        send["rasConnSendPeersUpdate: 带哈希"]
        recv["rasMsgHandlePeersUpdate: 合并+回发"]
        reinit["rasLinkReinitConns: 重建连接"]
    end
    ranks --> convert --> rankPeers --> update
    update --> diff --> hash
    hash --> send --> recv --> reinit

17.5 設計思考:RAS 與主通訊路徑的邊界

RAS 子系統最核心的設計決策是與資料面完全解耦。RAS 執行緒不參與任何集合通訊的資料搬運,它只做三件事:維護 peer 名單、檢測連線健康、執行診斷。這種解耦帶來幾個好處:

1. 故障隔離:RAS 執行緒崩潰不會直接導致通訊失敗(雖然會失去故障感知能力)

2. 效能無損:RAS 的心跳和同步流量走獨立網路,不佔用資料面頻寬

3. 可觀測性:診斷和監控可以在通訊進行時並行執行

代價是狀態一致性的挑戰:RAS 看到的 comm 狀態可能滯後於資料面。ncclRasCommInit和ncclRasCommFini透過ncclCommsMutex保護📎 src/ras/ras.cc:77-77,但 RAS 執行緒讀取時只做快照,不做強一致保證。

另一個關鍵設計是超時分層。ras_internal.h定義了一整套超時常量📎 src/ras/ras_internal.h:214-249:keep-alive 間隔 1 秒、警告閾值 5 秒、錯誤閾值 20 秒、peer 死亡閾值 60 秒。這種分層讓系統能在不同嚴重程度下採取不同動作——先警告、再嘗試備用連線、最後才宣告死亡。

17.6 本章小結

本章拆解了 NCCL RAS 子系統的四個核心模組:

  • ras.cc:單例 RAS 執行緒 + poll 事件迴圈,透過管道接收本地通知、透過獨立網路與其他 rank 交換訊息
  • progress_monitor.cc:每裝置一個工作執行緒,用 DMA 把 GPU 進度計數器搬到主機,帶限流告警和引用計數銷毀
  • diagnostics.cc:表驅動的檢查分發框架,兩遍掃描彙總各 rank 的診斷 payload
  • peers.cc:排序陣列 + 雜湊同步的 peer 名單管理,死 peer 單獨存放以節省頻寬

本章思考與自測

Q1:rasLocalNotify用rasNotificationMutex串行化寫入,但rasLocalHandle讀取時沒有對應的鎖。為什麼這樣是安全的?如果把static_assert(sizeof(struct rasNotification) <= PIPE_BUF)去掉,在什麼場景下會出問題?

參考解析:安全性來自 POSIX 對管道寫入原子性的保證——小於PIPE_BUF的寫入是原子的📎 src/ras/ras.cc:47。rasLocalNotify的迴圈寫📎 src/ras/ras.cc:224-237在單次寫入就能完成時不會與其他寫入交錯。rasLocalHandle的迴圈讀📎 src/ras/ras.cc:247-256可能讀到部分資料,但由於寫入是原子的,讀到的必然是完整訊息的前綴,下次讀補齊即可。

去掉static_assert後,若rasNotification超過PIPE_BUF,寫入可能被拆成多次非原子寫。兩個執行緒並發寫時,它們的位元組可能交錯,導致 RAS 執行緒讀到拼接了兩次通知的畸形資料。msg.type可能來自執行緒 A 而msg.addRanks.ranks來自執行緒 B,觸發rasLocalHandle的未知類型分支📎 src/ras/ras.cc:267-269或更糟的野指標解引用。

Q2:ncclProgressCounterMonitorDestroy在釋放鎖後才執行cudaStreamSynchronize 📎 src/ras/progress_monitor.cc:381-400。如果在同步期間另一個執行緒也呼叫 Destroy 銷毀同一個 comm,會發生什麼?destroyRefs如何防止問題?

參考解析:destroyRefs是防止 worker 被過早刪除的引用計數。第一個執行緒刪除 comm 後destroyRefs++ 📎 src/ras/progress_monitor.cc:371,此時haveDestroyRef = true。第二個執行緒嘗試刪除同一 comm 時,ncclIntruQueueDelete返回 nullptr(已被刪),haveDestroyRef保持 false📎 src/ras/progress_monitor.cc:368,直接跳過同步和釋放。

第一個執行緒完成cudaStreamSynchronize後呼叫releaseGpuProgressCounterMonitorDestroyRef 📎 src/ras/progress_monitor.cc:402,遞減destroyRefs到 0,且註冊佇列為空,才真正 join 執行緒並 delete📎 src/ras/progress_monitor.cc:225。

若沒有destroyRefs,第一個執行緒可能在同步期間被第二個執行緒的delete g釋放 worker,導致 use-after-free。注意releaseGpuProgressCounterMonitorDestroyRef在全局鎖 + worker 鎖內遞減📎 src/ras/progress_monitor.cc:222-225,保證檢查registrations為空和destroyRefs == 0的原子性。

Q3:rasDiagnosticsSummarizePeerPayloads第一遍掃描時校驗checkHeader->payloadBytes != checkHeader->nRecords * checkHeader->recordStride 📎 src/ras/diagnostics.cc:451-454。如果某個惡意或損壞的 peer 發送recordStride = 0且nRecords = 0,這個校驗會通過嗎?後續會發生什麼?

參考解析:recordStride <= 0會被第一個條件攔截📎 src/ras/diagnostics.cc:451,返回ncclInternalError。所以recordStride = 0不會通過。

但若recordStride > 0且nRecords = 0,則payloadBytes = 0,校驗通過。rasDiagnosticsAccountCheckRecords對nRecords == 0直接返回成功📎 src/ras/diagnostics.cc:378,不更新combined。後續分配時recordsBytes == 0不分配📎 src/ras/diagnostics.cc:473,拷貝時payloadBytes > 0為假跳過📎 src/ras/diagnostics.cc:490。最終summarize收到records = nullptr, recordsBytes = 0,各檢查的 summarize 實作需要處理空輸入。

真正的風險在nRecords > INT_MAX / recordStride的檢查📎 src/ras/diagnostics.cc:453——這防止nRecords * recordStride整數溢位繞過相等校驗。若去掉這個檢查,攻擊者可以構造nRecords = 2^31, recordStride = 2,乘積溢位為 0,與payloadBytes = 0相等,通過校驗後rasDiagnosticsAccountCheckRecords會累計一個巨大的nRecords,導致後續分配或拷貝越界。

RAS 讓 NCCL 在長時間訓練中具備了故障感知與自癒能力,但它依賴的是一套獨立於資料面的控制網路。下一章我們將進入記憶體管理子系統,看 NCCL 如何透過 allocator、註冊快取和使用者緩衝區註冊來優化顯存分配與 RDMA 註冊開銷——這是效能與可靠性之外的第三個支柱。

貫穿全章的設計原則是:控制面與數據面解耦、狀態用哈希做版本、超時分層處理、併發用引用計數保護生命週期。這些原則讓 RAS 能在不拖累通信性能的前提下實現故障發現與自癒。而通信性能的另一個關鍵支撐點——內存管理,同樣需要精細的工程權衡:為什麼 NCCL 通信前需要註冊內存?註冊緩存如何影響性能?下一章我們將深入 allocator、註冊緩存與用戶緩衝區註冊,揭開這些問題的答案。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 18

第 18 章:內存分配與顯存管理:allocator、註冊緩存與用戶註冊內存優化

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 18 章 / 共 25 章

上一章我們看到 RAS 子系統如何在控制面上獨立於數據面運行,用哈希做版本、用引用計數保護生命週期。本章進入 NCCL 的第三個支柱——內存管理。通信性能的上限,往往不取決於算法本身,而取決於「數據能不能被網卡直接讀寫」。NCCL 為此構建了三層機制:底層用ncclSpace和ncclShadowPool管理地址空間與影子對象,中層用ncclMemManager跟蹤動態內存的導入導出與掛起恢復,上層用ncclCommRegister把用戶緩衝區註冊進緩存,避免每次通信都重複 pin 內存。本章將逐層拆解這三套機制,回答「為什麼 NCCL 通信前需要註冊內存」以及「註冊緩存如何影響性能」。

18.1 ncclSpace:把地址空間切成滿/空交替的段

直覺模型

想像一條無限長的停車位編號線,從 0 開始向右延伸。有些車位停了車(已分配),有些空著(未分配)。ncclSpace就是這條編號線的「車位狀態記錄本」——它不記錄每個車位,只記錄「狀態發生翻轉的邊界點」。若沒有它,NCCL 在管理對稱內存的虛擬地址區間時,就得為每個字節維護一個標記位,內存開銷與地址空間成正比,完全不可接受。

數據結構與內存佈局

ncclSpace的定義極簡📎 src/include/allocator.h:20-24:

c
struct ncclSpace {
  int count;        // cuts[] 中有效元素个数
  int capacity;     // cuts[] 已分配容量
  int64_t* cuts;    // 升序排列的边界点数组
};

核心洞察在源碼註釋裡寫得很清楚📎 src/allocator.cc:151-153:cuts[]把非負整數軸切成「滿」和「空」交替的段,切割點升序排列,最後一個切割點之後的段必然是空的(未分配前沿)。由此可以推導出判斷第i段是否已滿的公式:

code
isFull(i) = (i%2 != ncuts%2)

這個公式的含義是:段的滿/空狀態由「段索引奇偶性」和「切割點總數的奇偶性」共同決定。當ncuts為偶數時,第 0 段(cuts[0]之前)是空的;當ncuts為奇數時,第 0 段是滿的。這個不變量貫穿整個模塊。

Step-by-Step Walkthrough:一次分配如何改變 cuts[]

代入場景:初始ncclSpace為空(count=0),調用ncclSpaceTryAlloc(a, limit=1000, size=100, align=1, &outOffset)。

第一步:定位第一個空段 📎 src/allocator.cc:209。i = a->count % 2,此時count=0,所以i=0,從第 0 段開始掃描。

第二步:計算段邊界 📎 src/allocator.cc:212-213。i==0時lo=0;i==a->count時hi=limit=1000。所以空段是[0, 1000)。

第三步:對齊並檢查容量 📎 src/allocator.cc:214-215。off = alignUp(0, 1) = 0,0 + 100 <= 1000成立,分配成功。

第四步:插入切割點 📎 src/allocator.cc:217-223。因為i==0(在頭部插入),走慢路徑insertSegment(a, 0, 0, 100)。insertSegment在index=0處插入兩個切割點lo=0, hi=100 📎 src/allocator.cc:172-174,然後執行「相鄰重複值過濾」📎 src/allocator.cc:185-203。過濾邏輯很精妙:它用讀寫雙游標掃描,遇到重複值就回退寫游標,把成對的重複值刪掉——因為成對重複意味著一個空段被夾在兩個滿段之間,可以合併。但前導零是特例,可以單獨刪除📎 src/allocator.cc:182-184。

分配後cuts = [0, 100],count=2。此時isFull(0) = (0%2 != 2%2) = false,第 0 段([0,0),空)為空;第 1 段([0,100))為滿。正確。

第五步:釋放 📎 src/allocator.cc:239-267。調用ncclSpaceFree(a, 0, 100)。先檢查cuts[count-1] <= offset是否成立📎 src/allocator.cc:231-237,即100 <= 0為假,繼續。定位第一個滿段i = 1 - count%2 = 1 - 0 = 1 📎 src/allocator.cc:246,cuts[1]=100 > 0,所以i=1。lo = cuts[0] = 0,hi = cuts[1] = 100。檢查offset < lo || hi < offset+size 📎 src/allocator.cc:252,0<0假,100<100假,通過。因為lo==offset且offset+size==hi,兩個快速路徑都不滿足(第一個要求offset+size != hi,第二個要求lo != offset),走慢路徑insertSegment(a, 1, 0, 100) 📎 src/allocator.cc:264。插入後cuts = [0, 0, 100, 100],過濾後變成[],count=0。回到初始狀態。

這個「插入後過濾」的設計避免了在分配/釋放時做複雜的段合併邏輯,把複雜度集中在insertSegment一處。

設計思考與生產踩坑

為什麼用 int64_t 而不是 size_t?因為ncclSpace管理的是「偏移量」而非「指針」,偏移量可能為負(雖然實際使用中不會),且需要與 CUDA 的CUdeviceptr寬度一致。用有符號類型便於在調試時發現越界。

性能陷阱:ncclSpaceFree的註釋直言「This could be binary search, but since allocate is linear there's no point」📎 src/allocator.cc:245。這意味著分配和釋放都是 O(n) 掃描。如果某個通信域頻繁分配釋放大量小段,cuts[]會膨脹,每次操作都變慢。生產環境中應盡量復用已註冊的緩衝區,而不是反覆註冊/註銷。

對齊溢出風險:alignUp(lo, align)在lo接近INT64_MAX且align較大時可能溢出。源碼沒有顯式檢查,因為limit由調用方保證在合理範圍內。

18.2 ncclShadowPool:裝置物件與主機影子的配對管理

直覺模型

GPU kernel 執行在裝置上,無法直接存取主機記憶體中的 C++ 物件(比如ncclDevComm裡的元資料)。ncclShadowPool就像一個「翻譯官」:它為每個裝置側物件分配一塊顯存,同時在主機側分配一塊對應的「影子」記憶體,並維護「裝置位址 → 主機位址」的映射表。當 host 需要修改某個裝置物件的配置時,先改主機影子,再拷貝到裝置。若沒有它,每次 kernel 要讀元資料都得透過cudaMemcpy從 host 拉取,延遲高得無法接受。

資料結構與記憶體佈局

兩個核心結構體📎 src/allocator.cc:272-277:

c
struct ncclShadowPage {   // 最多 64 个对象的连续块
  struct ncclShadowPage* next;
  int objSize;
  uint64_t freeMask;      // 位图,1=空闲,0=已占用
  void* devObjs;
};
struct ncclShadowObject {
  struct ncclShadowObject* next;
  void* devObj;
  void* hostObj;
  struct ncclShadowPage* page;  // null 表示直接分配在 CUDA mempool
};

ncclShadowPool本身📎 src/include/allocator.h:42-47:

c
struct ncclShadowPool {
  int count, hbits;                       // 对象数、哈希位数
  struct ncclShadowObject** table;        // 哈希桶数组
  cudaMemPool_t memPool;                  // 可选的 CUDA 内存池
  struct ncclShadowPage* pages;           // 页链表
};

關鍵設計點:freeMask是 uint64_t,所以每頁最多 64 個物件。這不是隨意選的——64 位正好是一個快取行的寬度,popFirstOneBit可以用單條__builtin_ctzll指令找到第一個空閒槽位,無需迴圈。

雜湊表增長策略:原始碼註解「Maintain 2:1 object:bucket ratio」📎 src/allocator.cc:368,即物件數超過桶數兩倍時擴容。初始hbits=4(16 個桶)📎 src/allocator.cc:363,每次翻倍。

Step-by-Step Walkthrough:一次分配如何選擇頁或直連

代入場景:ncclShadowPoolAlloc(pool, size=1024, &devObj, &hostObj, stream)。

第一步:惰性初始化 📎 src/allocator.cc:347-366。若hbits==0,先查詢裝置是否支援記憶體池📎 src/allocator.cc:352,支援則建立cudaMemPool_t,設定maxSize為參數SHADOW_MEMPOOL_MAX_SIZE(預設 1GB)📎 src/allocator.cc:359。然後分配 16 個桶的雜湊表。

第二步:檢查是否需要擴容 📎 src/allocator.cc:369-386。若count+1 > 2<<hbits,分配雙倍桶陣列,遍歷舊表重新插入(hashInsert用ncclHashPointer計算桶索引📎 src/allocator.cc:333-337),釋放舊表。

第三步:決定走頁路徑還是直連路徑 📎 src/allocator.cc:390。判斷條件(64<<10)/size >= 3,即size <= 21845時走頁路徑。對於size=1024,65536/1024=64 >= 3,走頁路徑。

第四步:計算頁內物件大小 📎 src/allocator.cc:391-392。shift = max(0, log2Down(1024)+1-4) = max(0, 10+1-4) = 7。pageObjSize = ((1024 + 127) >> 7) << 7 = 1024。即頁內物件大小按 2 的冪對齊到 128 位元組的倍數。

第五步:查找或建立頁 📎 src/allocator.cc:393-415。遍歷pool->pages鏈表,找objSize == pageObjSize的頁。若沒有,建立新頁:pageSize = min(65536, 64*1024) = 65536,freeMask = uint64_t(-1) >> (64 - 65536/1024) = uint64_t(-1) >> 0 = 全 1(64 個槽位全空)📎 src/allocator.cc:400。用cudaMallocFromPoolAsync或cudaMalloc分配顯存📎 src/allocator.cc:403-404,並cudaMemsetAsync清零📎 src/allocator.cc:405。

第六步:從頁中取槽位 📎 src/allocator.cc:408-412。popFirstOneBit(&page->freeMask)找到第一個空閒位,devObj = page->devObjs + slot * pageObjSize。若freeMask變為 0(頁滿),把頁從空閒鏈表移除📎 src/allocator.cc:411。

第七步:分配主機影子物件 📎 src/allocator.cc:423-428。malloc(sizeof(ncclShadowObject) + alignof(max_align_t)-1 + size),注意這裡多分配了alignof(max_align_t)-1位元組用於對齊填充。hostObj = alignUp((char*)(obj+1), alignof(max_align_t)),即物件頭之後對齊到最大對齊邊界。然後memset(hostObj, 0, size)清零。

第八步:插入雜湊表並更新計數 📎 src/allocator.cc:429-430。

併發控制與硬體互動

ncclShadowPool本身沒有鎖。這意味著它只能在單執行緒上下文中使用,或者由呼叫方保證互斥。從 NCCL 的實際使用看,它主要在通訊域初始化階段被呼叫,此時是單執行緒的。

cudaMallocFromPoolAsync和cudaFreeAsync是非同步操作,依賴stream參數保證順序📎 src/allocator.cc:403,459。ncclShadowPoolDestruct在釋放所有資源後呼叫cudaStreamSynchronize(stream) 📎 src/allocator.cc:333-337,確保所有非同步釋放完成後再銷毀記憶體池。

生產避坑指南

坑 1:頁內物件大小對齊導致的記憶體浪費。pageObjSize按 2 的冪對齊,若size=1000,shift = log2Down(1000)+1-4 = 9+1-4 = 6,pageObjSize = ((1000+63)>>6)<<6 = 1024。每個物件浪費 24 位元組,頁內 64 個物件浪費 1536 位元組。對於大量小物件,這個開銷不可忽視。

坑 2:ncclShadowPoolFree找不到物件時的行為 📎 src/allocator.cc:442-445。它返回ncclInternalError並列印警告,但不釋放任何資源。如果呼叫方忽略返回值,會導致記憶體洩漏。生產程式碼必須檢查返回值。

坑 3:ncclShadowPoolDestruct中freeMask==0的頁被回收 📎 src/allocator.cc:301-306。注意這裡把freeMask設為 1(而非全 1),意味著只標記第一個槽位為空。這是為了把「滿頁」重新放入pool->pages鏈表,但頁內其他槽位仍然被佔用——實際上這些物件即將被釋放,所以這個操作是安全的。但如果解構過程中存在併發存取,會讀到不一致狀態。

18.3 ncclMemManager:動態記憶體的引用計數與掛起恢復

直覺模型

訓練任務可能運行數天,期間 GPU 可能被其他任務搶佔,或者需要做檢查點。ncclMemManager就像一個「記憶體管家」:它記錄所有動態分配的記憶體(scratch/offload),在需要時把 GPU 記憶體「掛起」(unmap 物理頁,保留虛擬位址),把資料備份到 CPU,等恢復時再重新分配物理頁、重新映射、恢復資料。若沒有它,任務被搶佔後只能從頭開始,浪費數小時訓練進度。

資料結構與記憶體佈局

ncclMemManager的核心欄位(從初始化程式碼推斷)📎 src/mem_manager.cc:32-60:

欄位類型含義
entriesncclDynMemEntry*動態記憶體條目鏈表頭
numEntriesint鏈表長度
releasedint0=活躍,1=已掛起
refCountint引用計數(多個 comm 可共享)
totalPersistsize_t持久記憶體總量(原子)
totalScratchsize_tscratch 記憶體總量(原子)
totalOffloadsize_toffload 記憶體總量(原子)
cpuBackupUsagesize_tCPU 備份記憶體總量
lockstd::mutex保護 entries 鏈表
initializedint原子標誌,防止存取已銷毀的 mutex

記憶體佈局的關鍵設計:lock是一個std::mutex,但ncclMemManager是用ncclCalloc分配的(C 風格),所以必須用 placement new 顯式建構📎 src/mem_manager.cc:39,解構時顯式呼叫~mutex() 📎 src/mem_manager.cc:120。這是 C/C++ 混合編程的經典陷阱。

原子變數與鎖的分工:統計欄位(totalPersist等)用原子操作更新,不需要鎖;entries鏈表用lock保護。這樣統計查詢(ncclCommMemStats)可以無鎖讀取📎 src/mem_manager.cc:1117-1130,而鏈表操作必須持鎖。

Step-by-Step Walkthrough:掛起與恢復的完整流程

掛起流程 ncclCommMemSuspend 📎 src/mem_manager.cc:418-540:

第一步:前置檢查 📎 src/mem_manager.cc:419-430。檢查記憶體管理器是否禁用、comm 是否為空、是否已經掛起。

第二步:裝置同步與 barrier 📎 src/mem_manager.cc:440-441。cudaDeviceSynchronize()確保所有 GPU 操作完成,然後bootstrapBarrier確保所有 rank 同步。barrier tag 是0xBEEF。

第三步:第一遍掃描——unmap 所有 peer 匯入的緩衝區 📎 src/mem_manager.cc:444-465。對每個isImportedFromPeer && state==Active的條目,呼叫cuMemUnmap解除映射📎 src/mem_manager.cc:451,釋放 handle📎 src/mem_manager.cc:456,狀態改為Released。

第四步:第二遍掃描——offload 本地記憶體 📎 src/mem_manager.cc:468-526。跳過 peer 匯入和已釋放的條目。對ncclMemOffload類型,先分配 CPU 備份📎 src/mem_manager.cc:484,然後cudaMemcpy從 GPU 拷貝到 CPU📎 src/mem_manager.cc:492。對ncclMemScratch類型,只累加統計。然後關閉 shareable FD📎 src/mem_manager.cc:508-513,cuMemUnmap 📎 src/mem_manager.cc:516,cuMemRelease 📎 src/mem_manager.cc:519,狀態改為Released。

第五步:標記已掛起 📎 src/mem_manager.cc:528。

恢復流程 ncclCommMemResume 📎 src/mem_manager.cc:550-942:

第一步:恢復本地記憶體 📎 src/mem_manager.cc:577-668。對每個!isImportedFromPeer && state==Released的條目,重新cuMemCreate 📎 src/mem_manager.cc:599,ncclCuMemMapAndSetAccess映射到相同虛擬位址📎 src/mem_manager.cc:602,恢復 peer 存取權限📎 src/mem_manager.cc:610-626,對 offload 類型從 CPU 備份恢復資料📎 src/mem_manager.cc:632-643,重新匯出 FABRIC handle📎 src/mem_manager.cc:646-658。

第二步:barrier 同步 📎 src/mem_manager.cc:671-679。tag 仍是0xBEEF。

第三步:交換新 handle 資訊 📎 src/mem_manager.cc:688-816。統計每個 rank 有多少本地緩衝區需要廣播📎 src/mem_manager.cc:689-696,用bootstrapAllGather交換計數📎 src/mem_manager.cc:710,計算偏移📎 src/mem_manager.cc:724-728,然後先bootstrapSend再bootstrapRecv(註解明確「send first, then receive to avoid deadlock」📎 src/mem_manager.cc:783)。

第四步:重新匯入 peer 緩衝區 📎 src/mem_manager.cc:822-911。對每個isImportedFromPeer && state==Released的條目,在交換結果中查找匹配的 handle 資訊📎 src/mem_manager.cc:829-835。POSIX FD 類型需要檢查 hostHash 是否相同📎 src/mem_manager.cc:853-859,然後透過 proxy 取得 FD📎 src/mem_manager.cc:866,cuMemImportFromShareableHandle匯入📎 src/mem_manager.cc:873。FABRIC 類型直接匯入📎 src/mem_manager.cc:878。然後ncclCuMemMapAndSetAccess重新映射📎 src/mem_manager.cc:893。

第五步:最終 barrier 📎 src/mem_manager.cc:916-928。tag 是0xCAFE,與前面的0xBEEF區分。

並發控制與硬體互動

引用計數保護生命週期:ncclMemManagerDestroy先遞減refCount 📎 src/mem_manager.cc:76,若仍大於 0 則只清除當前 comm 的指標📎 src/mem_manager.cc:81,不釋放資源。這允許多個 comm 共享同一個記憶體管理器(比如 split_share 場景)。

原子 initialized 標誌:所有操作前都檢查COMPILER_ATOMIC_LOAD(&manager->initialized, memory_order_acquire) 📎 src/mem_manager.cc:136,242,338,358,防止存取已銷毀的 mutex。銷毀時用memory_order_release儲存 0📎 src/mem_manager.cc:87,確保之前的寫操作對其他執行緒可見。

CUDA VMM API 的使用:cuMemCreate/cuMemMap/cuMemUnmap/cuMemRelease是 CUDA 虛擬記憶體管理 API,允許實體記憶體和虛擬位址分離。這是掛起/恢復的基礎——掛起時 unmap 實體頁但保留虛擬位址,恢復時重新映射到相同虛擬位址,這樣所有已建立的指標關係都不需要修改。

生產避坑指南

坑 1:split_share 通訊域不支援掛起 📎 src/mem_manager.cc:1014-1018。若refCount > 1,直接返回ncclInvalidUsage。因為多個 comm 共享記憶體管理器時,掛起一個 comm 會影響其他 comm 的記憶體。

坑 2:POSIX FD 跨節點失效 📎 src/mem_manager.cc:853-859。POSIX 檔案描述符只在同一節點內有效,跨節點恢復時必須跳過。原始碼用hostHash比較判斷是否同節點。

坑 3:offload 資料恢復失敗時保留備份 📎 src/mem_manager.cc:635。若cudaMemcpy從 CPU 恢復到 GPU 失敗,原始碼列印警告並保留cpuBackup,不釋放。這是為了給呼叫方一個重試的機會,但如果不重試就會洩漏 CPU 記憶體。

坑 4:ncclMemUntrackDynamic中的 use-after-free 風險。原始碼在持鎖狀態下找到條目、儲存必要資訊、釋放條目📎 src/mem_manager.cc:302,然後在鎖外更新統計📎 src/mem_manager.cc:311-327。這個順序是正確的,但如果info指標指向呼叫方的堆疊記憶體,且呼叫方在鎖外讀取,需要確保info的生命週期覆蓋整個函式。

mermaid
flowchart TD
    start["ncclCommMemSuspend(comm)"] --> check{"manager->released?"}
    check -->|"是"| err1["返回 ncclInvalidUsage"]
    check -->|"否"| sync["cudaDeviceSynchronize()"]
    sync --> barrier1["bootstrapBarrier(tag=0xBEEF)"]
    barrier1 --> pass1["第一遍: 遍历 entries"]
    pass1 --> cond1{"isImportedFromPeer && Active?"}
    cond1 -->|"是"| unmap1["cuMemUnmap + cuMemRelease"]
    cond1 -->|"否"| skip1["跳过"]
    unmap1 --> pass2["第二遍: 遍历 entries"]
    skip1 --> pass2
    pass2 --> cond2{"memType == Offload?"}
    cond2 -->|"是"| backup["ncclCudaHostCalloc + cudaMemcpy D2H"]
    cond2 -->|"否"| scratch["累加 releasedScratch"]
    backup --> unmap2["cuMemUnmap + cuMemRelease"]
    scratch --> unmap2
    unmap2 --> mark["manager->released = 1"]
    mark --> done["返回 ncclSuccess"]
    err1 --> done

上圖展示了掛起流程的控制流。注意兩個關鍵分支:第一遍只處理 peer 匯入的緩衝區,第二遍只處理本地緩衝區,順序不能顛倒——必須先解除對 peer 記憶體的引用,再釋放本地記憶體。

18.4 註冊快取:ncclRegister 如何避免重複 pin

直覺模型

網卡要直接讀寫 GPU 顯存(GPUDirect RDMA),必須先「註冊」這塊記憶體——告訴網卡「這塊位址你可以直接存取」。註冊過程涉及 pin 頁、建立 IOMMU 映射,開銷很大(毫秒級)。如果每次 AllReduce 都重新註冊,小訊息通訊的延遲會被註冊開銷完全淹沒。ncclRegister就是一個「註冊快取」:它把已註冊的位址範圍記錄在有序陣列裡,下次遇到相同或包含的緩衝區,直接複用,不重複註冊。

資料結構與記憶體佈局

ncclRegCache的核心是一個有序陣列slots,每個元素是ncclReg*。ncclReg的關鍵欄位(從使用推斷):

欄位類型含義
begAddruintptr_t頁對齊的起始位址
endAddruintptr_t頁對齊的結束位址
localRefsint本地引用計數
graphRefsint圖引用計數
stateint註冊狀態位(NET/NVLS/COLLNET/IPC)
netHandleHeadncclRegNetHandles*網路 handle 鏈結串列
ipcInfosncclIpcInfo**IPC 資訊陣列

頁對齊:begAddr = (uintptr_t)data & -pageSize 📎 src/register/register.cc:31,endAddr = ((uintptr_t)data + size + pageSize - 1) & -pageSize 📎 src/register/register.cc:32。-pageSize是pageSize的二進位補碼,等價於「向下對齊到 pageSize 的倍數」。這樣做的原因是:註冊的最小粒度是頁,即使只註冊 1 位元組,也要註冊整頁。

Step-by-Step Walkthrough:一次註冊如何命中快取

代入場景:ncclCommRegister(comm, buff=0x7f0000001000, size=4096, &handle)。

第一步:參數檢查與頁對齊 📎 src/register/register.cc:18-24。CommCheck驗證 comm 有效性。假設pageSize=4096,begAddr = 0x7f0000001000 & -4096 = 0x7f0000001000,endAddr = (0x7f0000001000 + 4096 + 4095) & -4096 = 0x7f0000002000。

第二步:系統記憶體檢查 📎 src/register/register.cc:36-64。若ncclCuMemEnable(),查詢位址範圍和記憶體類型。若memType == CU_MEMORYTYPE_HOST,說明是 CPU 記憶體,跳過註冊📎 src/register/register.cc:58-61。否則檢查是否有 Sysmem 段📎 src/register/register.cc:50-55。

第三步:遍歷快取尋找插入位置 📎 src/register/register.cc:66-89。迴圈slot從 0 開始:

  • 若slot == population(到達末尾)或begAddr < slots[slot]->begAddr(當前位址在快取條目之前),說明需要新建條目📎 src/register/register.cc:67。
  • 若slots[slot]->begAddr <= begAddr && slots[slot]->endAddr >= endAddr,說明當前緩衝區被已有條目完全包含,直接增加引用計數📎 src/register/register.cc:83-87。

第四步:新建條目 📎 src/register/register.cc:68-82。若快取滿,擴容(初始 32,之後翻倍)📎 src/register/register.cc:70。用memmove在slot位置騰出空間📎 src/register/register.cc:73,ncclCalloc分配新條目📎 src/register/register.cc:74,設定begAddr/endAddr,根據isGraph設定graphRefs或localRefs為 1📎 src/register/register.cc:78-79,population++,返回 handle。

第五步:註銷 📎 src/register/register.cc:172-195。commDeregister先找到 handle 對應的 slot📎 src/register/register.cc:180,遞減引用計數📎 src/register/register.cc:185-186。若仍有引用,直接返回📎 src/register/register.cc:187。否則呼叫regCleanup清理所有底層註冊📎 src/register/register.cc:188,釋放條目,用memmove填補空洞📎 src/register/register.cc:190,population--。

設計思考與生產踩坑

為什麼用有序陣列而不是雜湊表?因為註冊查詢是「範圍包含」查詢,不是精確匹配。有序陣列支援二分查找(雖然原始碼用線性掃描),且記憶體局部性好。雜湊表無法高效處理「這個位址是否被某個更大的範圍包含」這類查詢。

regCleanup的狀態位設計 📎 src/register/register.cc:95-134。state是一個位元遮罩,每個位對應一種註冊類型(NET/NVLS/COLLNET/IPC)。清理時逐位檢查,只清理已完成的註冊。這種設計允許部分註冊成功、部分失敗的情況——比如網路註冊成功但 IPC 註冊失敗,清理時只清理網路部分。

生產陷阱:註冊快取不感知記憶體釋放。如果使用者註冊了一塊緩衝區,然後在未註銷的情況下cudaFree了它,快取中仍然保留著這個條目。下次分配可能復用同一位址,導致快取命中但實際記憶體已失效。NCCL 的約定是:註冊和註銷必須配對,使用者負責保證註冊期間記憶體不被釋放。

ncclCommRegister的跳過條件 📎 src/register/register.cc:150-159。若LocalRegister=0或P2pUsesMemcpy=1,直接返回NULLhandle。這意味著在某些配置下(比如 P2P 走 memcpy 而非 RDMA),註冊被完全跳過。呼叫方必須檢查 handle 是否為 NULL。

18.5 集合通訊註冊:coll_reg 如何為不同演算法選擇註冊策略

直覺模型

不同的集合通訊演算法走不同的傳輸路徑:NVLS 走 NVLink SHARP,Ring 走 P2P 或網路,Tree 走樹形拓撲。每條路徑需要不同的註冊方式:NVLS 需要註冊到 NVLS 硬體,網路需要註冊到網卡,IPC 需要註冊到對端 GPU。coll_reg.cc就是「註冊策略路由器」:它根據演算法、協定、緩衝區類型,決定呼叫哪些註冊函式。若沒有它,每種演算法都得自己實作註冊邏輯,程式碼重複且容易出錯。

Step-by-Step Walkthrough:Ring 演算法的註冊決策

代入場景:ncclRegisterCollBuffers(comm, info, outRegBufSend, outRegBufRecv, cleanupQueue, regNeedConnect),其中info->algorithm == NCCL_ALGO_RING,info->protocol == NCCL_PROTO_SIMPLE。

第一步:前置檢查 📎 src/register/coll_reg.cc:155-157。設定regBufType = NCCL_REGULAR_BUFFER,regNeedConnect = true。若LocalRegister=0且非持久圖註冊,直接退出。

第二步:進入 Ring 分支 📎 src/register/coll_reg.cc:338。初始化recvRegRecord/sendRegRecord為 NULL,分配sendNetConns/sendNetHandles/recvNetConns/recvNetHandles/srecvNetHandles陣列📎 src/register/coll_reg.cc:356-360。

第三步:查找已有註冊記錄 📎 src/register/coll_reg.cc:351-355。ncclRegFind在快取中查找 recv/send 緩衝區。若 recv 未找到且非持久圖註冊,退出📎 src/register/coll_reg.cc:352。若跨節點且 send 未找到且非持久圖註冊,退出📎 src/register/coll_reg.cc:354。

第四步:遍歷所有 channel 收集 peer 📎 src/register/coll_reg.cc:362-393。對每個 channel,檢查ring.prev和ring.next。若連接標誌包含NCCL_DIRECT_NIC,記錄到recvNetConns/sendNetConns 📎 src/register/coll_reg.cc:370-379。若包含NCCL_P2P_READ | NCCL_P2P_WRITE,把 peer 加入peerRanks陣列📎 src/register/coll_reg.cc:382-391。

第五步:IPC 註冊 📎 src/register/coll_reg.cc:394-407。若nPeers > 0 && comm->isAllDirectP2p,先嘗試圖註冊📎 src/register/coll_reg.cc:395-399,失敗則嘗試本地註冊📎 src/register/coll_reg.cc:400-403。若成功,設定regBufType = NCCL_IPC_REG_BUFFER 📎 src/register/coll_reg.cc:406。

第六步:網路註冊 📎 src/register/coll_reg.cc:409-457。檢查!comm->useNetPXN && comm->useGdr && netDeviceType != UNPACK且非 AllReduce 的 PreMulSum/SumPostDiv📎 src/register/coll_reg.cc:415-418。先嘗試圖註冊📎 src/register/coll_reg.cc:419-430,失敗則本地註冊📎 src/register/coll_reg.cc:431-442。若成功,設定regBufType |= NCCL_NET_REG_BUFFER,保存 handle 陣列📎 src/register/coll_reg.cc:445-452。

第七步:調整通道數 📎 src/register/coll_reg.cc:551-554。若只有 IPC 註冊且單節點且通道數在 17-24 之間,降到 16。這是為了匹配 IPC 註冊後的頻寬特性。

設計思考與生產踩坑

為什麼 NVLS 和 Ring 的註冊順序相反?NVLS 分支先嘗試圖註冊再本地註冊📎 src/register/coll_reg.cc:86-94,而 Ring 分支先本地再圖📎 src/register/coll_reg.cc:395-403。這是因為 NVLS 的圖註冊更可能成功(NVLS 硬體對持久緩衝區有優化),而 Ring 的本地註冊更輕量。

isMloPartBufRdmaCapable的全域決策 📎 src/register/coll_reg.cc:14-37。註解強調「Registration decision must be global, using communicator-wide guarantees」📎 src/register/coll_reg.cc:20。這意味著即使某個 rank 的緩衝區支援 RDMA,只要通訊域內有一個 rank 不支援,整個通訊域都不註冊。這是為了避免部分 rank 註冊、部分不註冊導致的不一致。

生產陷阱:註冊失敗時的靜默降級。ncclRegisterCollBuffers在註冊失敗時不會報錯,只是不設定regBufType的對應位。這意味著通訊仍然能工作,只是效能下降。生產環境中如果發現效能不達預期,應該檢查NCCL_REG日誌確認註冊是否成功。

mermaid
flowchart LR
    subgraph input["输入"]
        task["ncclTaskColl<br/>algorithm=RING<br/>protocol=SIMPLE"]
    end
    subgraph ipc["IPC 注册路径"]
        find["ncclRegFind<br/>查找缓存"]
        collect["遍历 channel<br/>收集 peerRanks"]
        ipcReg["ncclIpcLocalRegisterBuffer<br/>或 GraphRegister"]
    end
    subgraph net["网络注册路径"]
        checkGdr{"useGdr &&<br/>!useNetPXN?"}
        netReg["ncclNetLocalRegisterBuffer<br/>或 GraphRegister"]
    end
    subgraph output["输出"]
        regType["info->regBufType<br/>NCCL_IPC_REG_BUFFER<br/>NCCL_NET_REG_BUFFER"]
        handles["info->sendNetHandles<br/>info->recvNetHandles"]
    end
    task --> find
    find --> collect
    collect --> ipcReg
    ipcReg --> regType
    find --> checkGdr
    checkGdr -->|"是"| netReg
    checkGdr -->|"否"| regType
    netReg --> regType
    netReg --> handles

上圖展示了 Ring 演算法下兩條並行的註冊路徑:IPC 路徑處理同節點 P2P 連接,網路路徑處理跨節點 RDMA 連接。兩條路徑獨立執行,最終都彙總到info->regBufType。

18.6 生產避坑與故障恢復鏈

坑 1:註冊快取與記憶體池的互動

當使用ncclMemAlloc分配記憶體時,底層走 CUDA VMM API📎 src/allocator.cc:38-94。這種分配方式建立的實體記憶體帶有gpuDirectRDMACapable標誌📎 src/allocator.cc:54,意味著它天然支援 RDMA。但ncclMemFree釋放時,如果記憶體管理器已銷毀,會走cudaFree回退路徑📎 src/allocator.cc:130-132。這可能導致 VMM 分配的記憶體被錯誤地用cudaFree釋放。生產環境中必須確保ncclMemAlloc/ncclMemFree配對使用,且不要在記憶體管理器銷毀後釋放。

坑 2:掛起期間的通訊請求

ncclCommMemSuspend執行期間,如果有新的通訊請求到達,會怎樣?原始碼在掛起前呼叫cudaDeviceSynchronize() 📎 src/mem_manager.cc:440,確保所有已入隊的 GPU 操作完成。但如果有 host 側的通訊請求正在入隊,沒有顯式保護。生產環境中應該在掛起前停止所有通訊執行緒,或者使用 group 語意確保掛起操作與其他操作串行。

坑 3:FABRIC handle 的相容性

ncclMemAlloc在 CUDA 12.3+ 上會嘗試使用 FABRIC handle📎 src/allocator.cc:60-71。如果cuMemCreate返回CUDA_ERROR_NOT_PERMITTED或CUDA_ERROR_NOT_SUPPORTED,會回退到 POSIX FD📎 src/allocator.cc:63-65。但恢復時,如果 handle 類型是 FABRIC 但匯出失敗,會直接報錯並 unmap📎 src/mem_manager.cc:649-655。這意味著在混合環境中(部分 GPU 支援 FABRIC,部分不支援),掛起/恢復可能失敗。

坑 4:引用計數洩漏

ncclRegister每次命中快取都會增加引用計數📎 src/register/register.cc:84-85。如果呼叫方註冊了 N 次但只註銷了 M 次(M < N),引用計數永遠不會歸零,regCleanup永遠不會被呼叫,底層註冊資源洩漏。生產程式碼必須嚴格配對ncclCommRegister/ncclCommDeregister。

mermaid
sequenceDiagram
    participant App as 应用层
    participant Reg as ncclRegister
    participant Cache as ncclRegCache
    participant Net as ncclNetLocalRegisterBuffer
    participant GPU as CUDA Driver

    App->>Reg: ncclCommRegister(comm, buff, size, &handle)
    Reg->>Reg: begAddr = data & -pageSize
    Reg->>Cache: 遍历 slots 查找包含范围
    alt 缓存命中
        Cache-->>Reg: 返回已有 ncclReg*
        Reg->>Reg: localRefs++
    else 缓存未命中
        Reg->>Cache: memmove 腾出插入位置
        Reg->>Cache: ncclCalloc 新条目
        Reg->>Reg: localRefs = 1
    end
    Reg-->>App: 返回 handle
    App->>Net: 首次注册时调用
    Net->>GPU: cuMemExportToShareableHandle
    GPU-->>Net: 返回 handle
    Net-->>App: 注册完成

本章思考與自測

Q1: 若將ncclSpaceFree中的if (a->count == 0 || a->cuts[a->count - 1] <= offset)檢查📎 src/allocator.cc:231-237去掉,在什麼場景下會觸發越界存取?

參考解析:這個檢查有兩個作用。第一,a->count == 0防止空陣列存取cuts[-1]。第二,a->cuts[a->count-1] <= offset防止offset超出已分配範圍。如果去掉,當count == 0時,a->cuts[a->count - 1]會讀取cuts[-1],這是未定義行為,可能讀到堆元資料或觸發段錯誤。更隱蔽的是,即使count > 0,如果offset大於最後一個切割點,後續的while (a->cuts[i] <= offset) i += 2迴圈📎 src/allocator.cc:247會一直遞增i直到越界,因為cuts[]中不存在大於offset的元素。這在生產中的觸發場景是:呼叫方傳入了一個從未分配過的偏移量(比如緩衝區被外部釋放後再次呼叫 free),或者ncclSpace被並行修改導致狀態不一致。修復方式是保留這個檢查,並在返回錯誤時列印offset和count便於排查。

Q2: ncclMemManagerDestroy中,如果refCount遞減後仍大於 0,只清除當前 comm 的指標而不釋放資源📎 src/mem_manager.cc:78-83。如果此時另一個 comm 正在呼叫ncclMemTrack,會發生什麼?

參考解析:ncclMemTrack首先檢查manager->initialized 📎 src/mem_manager.cc:136。由於refCount > 0時不會設定initialized = 0,所以檢查通過。然後它會取得manager->lock並修改entries鏈結串列📎 src/mem_manager.cc:188-192。這是安全的,因為refCount > 0意味著至少還有一個 comm 持有引用,記憶體管理器不會被銷毀。真正的風險在於:如果最後一個 comm 呼叫ncclMemManagerDestroy時,refCount遞減到 0,它會設定initialized = 0 📎 src/mem_manager.cc:87並釋放所有資源。如果此時另一個執行緒正在ncclMemTrack中已經通過了initialized檢查但還沒取得鎖,它會存取已釋放的manager->lock,導致 use-after-free。原始碼透過memory_order_acquire/release配對來緩解這個問題,但嚴格來說仍存在競態視窗。生產環境中應該確保所有通訊執行緒在銷毀記憶體管理器前已停止。

Q3: 在ncclCommMemResume中,POSIX FD 類型的 peer 緩衝區在跨節點時被跳過📎 src/mem_manager.cc:853-859。如果所有 peer 緩衝區都被跳過,restoredPeerCount為 0,但manager->released仍被設為 0📎 src/mem_manager.cc:913。這會導致什麼後果?

參考解析:manager->released = 0表示記憶體管理器認為恢復已完成。但如果有 peer 緩衝區被跳過,它們的state仍然是ncclDynMemStateReleased,handle仍然是 0。後續通訊如果存取這些緩衝區,會觸發 CUDA 錯誤(存取未映射的虛擬位址)。更嚴重的是,ncclCommMemStats查詢ncclStatGpuMemSuspended會返回 0(活躍)📎 src/mem_manager.cc:1130,但實際有部分記憶體未恢復。這個問題的根源是:跨節點 POSIX FD 本身就不應該被匯入——在掛起前,這些緩衝區就不應該存在於entries中。正確的做法是在掛起時就把跨節點的 POSIX FD 條目標記為不可恢復,或者在恢復時返回錯誤而非靜默跳過。生產環境中,如果使用 POSIX FD 且跨節點,應該改用 FABRIC handle 或確保掛起/恢復只在單節點內進行。

記憶體管理是 NCCL 效能的隱形支柱:ncclSpace用極簡的切割點陣列管理位址空間,ncclShadowPool用 64 位元位圖和雜湊表管理裝置/主機物件配對,ncclMemManager用引用計數和 CUDA VMM API 實現掛起恢復,ncclRegister用有序陣列快取註冊結果避免重複 pin。這四層機制共同支撐起「通訊前不需要重新註冊記憶體」這一關鍵效能保證。下一章我們將進入裝置側通訊器與 ABI 相容,看devcomm如何把這些 host 側的記憶體佈局映射到 GPU kernel 可存取的結構中。

上圖展示了註冊的時序:快取命中時只增加引用計數,不呼叫底層註冊;快取未命中時才建立新條目並觸發底層註冊。至此,host 側的記憶體管理機制已經清晰。但通訊最終發生在 GPU 上,kernel 需要直接存取對端 rank 的位址和連線狀態。下一章將進入裝置側通訊器與 ABI 相容,看 devcomm 如何把 host 側 ncclComm 的元資料映射到裝置側可存取的結構,以及版本化 ABI 如何保證新舊 kernel 與函式庫的相容。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 19

第 19 章:裝置端通訊域與 ABI 相容:devcomm 與 kernel 的通訊契約

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 19 章 / 共 25 章

上一章我們看到,host 側的 ncclMemManager 用引用計數和 CUDA VMM API 管理著通訊緩衝區的生命週期。但通訊真正發生的地方是 GPU kernel——kernel 裡的執行緒需要知道:我是哪個 rank?對端 rank 的緩衝區在哪個虛擬位址?連線是否就緒?這些資訊在 host 側的 ncclComm 結構裡,但 kernel 不能直接解引用 host 指標。如果 NCCL 讓 kernel 每次都透過參數傳遞或全域記憶體查詢來取得這些元資料,那麼每次通訊都要付出額外的延遲和頻寬開銷。更糟糕的是,kernel 程式碼一旦編譯,其存取的欄位偏移就固定了——如果函式庫升級後 ncclComm 的佈局變了,舊 kernel 就會讀到錯誤的資料。這就是 devcomm 要解決的核心問題:把 host 側通訊域的關鍵元資料,以穩定的、版本化的記憶體佈局,映射到裝置側可存取的結構中。src/devcomm 目錄下的 devcomm_v22902.cc、devcomm_v22907.cc、devcomm_v23000.cc、devcomm_v23100.cc 就是這套版本化 ABI 的具體實現。每個檔案對應一個 NCCL 版本區間,定義了該區間內 ncclDevComm 的精確記憶體佈局,以及新舊版本之間的欄位拷貝邏輯。本章將依次拆解:裝置側通訊器的核心資料結構長什麼樣、版本化 ABI 的註冊與匹配機制如何運作、新舊版本之間如何做欄位級轉換、以及這套機制在生產環境中的邊界與陷阱。

一、裝置側通訊器的核心結構:ncclDevComm 的記憶體佈局

直覺模型

把ncclDevComm想像成一張「工位卡」:每個 GPU kernel 啟動時,都會拿到一張卡片,上面印著「你是 3 號 rank,總共 8 個 rank,你的 LSA 組裡有 4 個 rank,對端緩衝區基底位址在 0x7f...」。這張卡片必須足夠小(能塞進 kernel 參數),又必須包含所有關鍵資訊。如果這張卡片不存在,kernel 就只能靠 host 側反覆傳遞參數,每次通訊都要重新組裝——延遲高、易出錯。

資料結構與記憶體佈局

以ncclDevComm_v23000為例,它的完整定義在📎 src/devcomm/devcomm_v23000.cc:25-62:

c
struct ncclDevComm_v23000 {
  unsigned int magic;          // 偏移 0,魔数校验
  unsigned int version;        // 偏移 4,版本号

  int rank, nRanks;            // 偏移 8, 12
  uint32_t nRanks_rcp32;       // 偏移 16,nRanks 的倒数(定点数)
  int lsaRank, lsaSize;        // 偏移 20, 24
  uint32_t lsaSize_rcp32;      // 偏移 28

  ncclDevCommWindowTable_t windowTable;  // 偏移 32
  ncclWindow_t resourceWindow;           // 偏移 40
  ncclResourceWindow_vidmem_v23000_t resourceWindow_inlined;  // 偏移 48
  ncclGinBarrierHandle_t hybridWorldGinBarrier;  // 偏移 112
  ...
};

📎 src/devcomm/devcomm_v23000.cc:64-93用一連串static_assert把每個欄位的偏移釘死。這不是裝飾——它是 ABI 相容性的編譯期契約。如果某個欄位的偏移因為編譯器對齊策略變化而移動,編譯就會失敗,而不是在執行時產生難以除錯的記憶體錯位。

幾個關鍵欄位的設計動機:

〔設計推斷與架構權衡〕

nRanks_rcp32和lsaSize_rcp32:這是nRanks和lsaSize的倒數,用 32 位定點數表示。 kernel 裡做 rank 到 buffer 偏移的除法運算時,GPU 的整數除法很慢,用乘以倒數再移位的方式可以顯著加速。這是典型的「用空間換時間」——多存 4 位元組,省掉每次除法的幾十個時脈週期。

resourceWindow_inlined:這是一個內聯的視窗描述符,類型為ncclResourceWindow_vidmem_v23000_t。注意📎 src/devcomm/devcomm_v23000.cc:11-18中它的定義:

c
typedef struct ncclResourceWindow_vidmem_v23000 {
  char reserved1[8];
  char* lsaFlatBase;
  char reserved2[8];
  uint32_t stride4G;
  uint32_t mcOffset4K;
  char reserved3[32];  // NOTE: shrunk from 40 in 2.30u1 to reclaim 8 bytes
} ncclResourceWindow_vidmem_v23000_t;

這裡的reserved1、reserved2、reserved3是填充欄位,用來佔位。為什麼需要填充?因為ncclDevComm_v23000的佈局必須與某個「基準版本」保持偏移一致,即使某些欄位在當前版本中不再使用,也要保留佔位以保證後續欄位的偏移不變。📎 src/devcomm/devcomm_v23000.cc:11-18的註解明確說明:2.30u1 把reserved3從 40 位元組縮小到 32 位元組,騰出 8 位元組給hybridWorldGinBarrier。這是一次佈局重排——透過縮小填充區,在不改變整體大小的前提下塞入新欄位。

📎 src/devcomm/devcomm_v23000.cc:11-18的static_assert進一步驗證:lsaFlatBase、stride4G、mcOffset4K三個欄位的偏移必須與「當前版本」的ncclWindow_vidmem一致,且整個結構體大小為 64 位元組。這意味著resourceWindow_inlined在 v23000 和當前版本之間是二進位相容的——可以直接 memcpy。

版本化結構體的家族

對比ncclDevComm_v22902 📎 src/devcomm/devcomm_v22902.cc:38-62和ncclDevComm_v22907 📎 src/devcomm/devcomm_v22907.cc:13-41,可以看到欄位的演化:

欄位v22902v22907v23000
magic/version無無有(偏移 0/4)
ginContextCountuint8_tuint32_tuint32_t
ginNetDeviceTypes[4][NCCL_GIN_MAX_CONNECTIONS][NCCL_GIN_MAX_CONNECTIONS]
ginIsRailed無bool拆分為ginConnectionsRailed + ginContextsRailed
hybridWorldGinBarrier無無有(偏移 112)
結構體大小200224240
〔設計推斷與架構權衡〕

這個演化路徑揭示了 NCCL 的版本策略:只在必要時增加欄位,且盡量利用填充區。v22902 到 v22907 增加了ginSignalBase、ginCounterBase、ginContextBase、ginIsRailed等 GIN 相關欄位;v22907 到 v23000 增加了magic/version校驗欄位和hybridWorldGinBarrier,同時把ginIsRailed拆成兩個更精確的標誌位。

---

二、版本化 ABI 的註冊與匹配:ncclDevCommCompat 結構

直覺模型

把版本化 ABI 想像成一套「翻譯外掛」:當應用程式用 NCCL 2.29.2 編譯,但執行時連結的是 2.31.0 的函式庫,函式庫需要知道「2.29.2 的 kernel 期望什麼樣的ncclDevComm佈局」,然後把當前版本的ncclDevComm翻譯成舊佈局。每個版本區間對應一個翻譯外掛,註冊在一個全域表中。

核心結構:ncclDevCommCompat

每個devcomm_vXXXXX.cc檔案末尾都定義了一個ncclDevCommCompat結構體。以 v23000 為例📎 src/devcomm/devcomm_v23000.cc:192-199:

c
struct ncclDevCommCompat ncclDevCommCompat_v23000 = {
  NCCL_VERSION(2, 30, 0),               // minVersion
  NCCL_VERSION(2, 30, 7),               // maxVersion
  nullptr,                              // commPropertiesFilter
  ncclDevCommRequirementsFilter_v23000, // devCommRequirementsFilter
  ncclDevCommCopyNewToOld_v23000,       // devCommCopyNewToOld
  ncclDevCommCopyOldToNew_v23000,       // devCommCopyOldToNew
};

六個欄位的含義:

1. minVersion / maxVersion:這個外掛負責的版本區間。v23000 覆蓋 2.30.0 到 2.30.7。

2. commPropertiesFilter:可選的過濾器,用於調整ncclCommProperties中暴露給舊版本的能力標誌。v23000 設為nullptr,表示不需要過濾。

3. devCommRequirementsFilter:檢查應用程式請求的裝置側資源是否與舊版本相容。v23000 的實現📎 src/devcomm/devcomm_v23000.cc:95-98只是把ginType從comm->sharedRes複製到reqs。

4. devCommCopyNewToOld:把當前版本的ncclDevComm拷貝到舊版本佈局。

5. devCommCopyOldToNew:把舊版本佈局拷貝回當前版本。

版本區間的劃分

四個檔案的版本區間:

檔案minVersionmaxVersion備註
devcomm_v22902.cc2.29.22.29.3最早的版本化實現
devcomm_v22907.cc2.29.52.29.7增加 GIN 欄位,但不提供 GIN 向後相容
devcomm_v23000.cc2.30.02.30.7增加 magic/version 校驗
devcomm_v23100.cc2.31.0當前版本所有過濾器為 nullptr,表示完全相容

📎 src/devcomm/devcomm_v23100.cc:10-17的 v23100 外掛所有回呼都是nullptr,這意味著從 2.31.0 開始,ncclDevComm的佈局已經穩定,不需要任何轉換。

〔設計推斷與架構權衡〕

注意 v22902 和 v22907 之間的版本區間有「空隙」(2.29.4 和 2.29.6 沒有對應的外掛)。這可能是因為這些版本沒有發布,或者它們的佈局與相鄰版本完全一致,可以複用。

匹配流程

當應用程式呼叫ncclCommGetDeviceHandle或類似 API 時,NCCL 需要:

1. 讀取應用程式編譯時嵌入的 NCCL 版本號(透過reqs->version)。

2. 在全域的ncclDevCommCompat表中查找覆蓋該版本的外掛。

3. 如果找到,呼叫外掛的devCommCopyNewToOld把當前佈局轉換為舊佈局。

4. 如果沒找到,返回錯誤或使用預設行為。

下面的流程圖展示了這個匹配與轉換過程:

mermaid
flowchart TD
    start["应用请求设备侧通信器"] --> read_ver["读取 reqs->version<br/>(应用编译时版本)"]
    read_ver --> find_compat{"在 ncclDevCommCompat 表中<br/>查找覆盖该版本的插件?"}
    find_compat -->|找到| check_filter["调用 devCommRequirementsFilter<br/>检查资源请求兼容性"]
    find_compat -->|未找到| err_unsupported["返回 ncclInvalidUsage<br/>版本不兼容"]
    check_filter --> filter_ok{"过滤器返回<br/>ncclSuccess?"}
    filter_ok -->|是| copy_new_to_old["调用 devCommCopyNewToOld<br/>把当前布局转为旧布局"]
    filter_ok -->|否| err_gin["返回 ncclInvalidUsage<br/>GIN 资源不兼容"]
    copy_new_to_old --> done["返回旧布局 ncclDevComm"]
    err_unsupported --> done_err["应用收到错误"]
    err_gin --> done_err

---

三、欄位級轉換:新舊佈局如何互轉

直覺模型

版本轉換就像「翻譯」:新版本的ncclDevComm是一篇現代漢語文章,舊版本的佈局是一篇文章文。翻譯器需要逐欄位對應——有些欄位直接對應(rank對rank),有些欄位需要「意譯」(ginConnectionStride > 1翻譯成ginConnectionsRailed = true),有些欄位在舊版本中不存在(直接丟棄)。

NewToOld 轉換:從當前版本到舊版本

以ncclDevCommCopyNewToOld_v23000為例📎 src/devcomm/devcomm_v23000.cc:114-152:

c
static ncclResult_t ncclDevCommCopyNewToOld_v23000(ncclComm_t comm, void* oldDevComm,
                                                   struct ncclDevComm const* newDevComm) {
  struct ncclDevComm_v23000* old = (struct ncclDevComm_v23000*)oldDevComm;

  memset(old, '\0', sizeof(*old));  // 先清零,防止未初始化字段泄露
  old->magic = newDevComm->magic;
  old->version = newDevComm->version;
  old->rank = newDevComm->rank;
  ...
  old->ginConnectionsRailed = (newDevComm->ginConnectionStride > 1);
  old->ginStrongLegacySignals = newDevComm->ginStrongLegacySignals;
  old->ginContextsRailed = (newDevComm->ginContextStride > 1);
  ...
}

關鍵步驟:

1. memset清零 📎 src/devcomm/devcomm_v23000.cc:118:這是安全防護——舊結構體中可能有新版本不存在的欄位,清零可以防止未初始化記憶體洩露到裝置側。

2. 直接欄位拷貝:rank、nRanks、lsaRank等直接賦值。

3. 內聯視窗轉換:呼叫ncclDevCommCopyResourceWindowNewToOld_v23000 📎 src/devcomm/devcomm_v23000.cc:100-105,逐欄位拷貝lsaFlatBase、stride4G、mcOffset4K。

4. 語意轉換:ginConnectionsRailed = (newDevComm->ginConnectionStride > 1) 📎 src/devcomm/devcomm_v23000.cc:142。新版本用ginConnectionStride(一個整數步長)表示是否 railed,舊版本用布林值。當步長大於 1 時,說明連接是 railed 的。

5. 陣列拷貝:memcpy拷貝ginNetDeviceTypes和ginHandles陣列📎 src/devcomm/devcomm_v23000.cc:135-136。

OldToNew 轉換:從舊版本到當前版本

反向轉換在📎 src/devcomm/devcomm_v23000.cc:154-190:

c
static ncclResult_t ncclDevCommCopyOldToNew_v23000(ncclComm_t comm, struct ncclDevComm* newDevComm,
                                                   void const* oldDevComm) {
  struct ncclDevComm_v23000 const* old = (struct ncclDevComm_v23000 const*)oldDevComm;

  newDevComm->magic = old->magic;
  ...
  newDevComm->ginConnectionStride = old->ginConnectionsRailed ? old->lsaSize : 1;
  newDevComm->ginContextStride = old->ginContextsRailed ? old->lsaSize : 1;
  ...
}
〔設計推斷與架構權衡〕

注意📎 src/devcomm/devcomm_v23000.cc:180-181的語意轉換:如果舊版本中ginConnectionsRailed為真,則新版本的ginConnectionStride設為lsaSize;否則設為 1。這裡用lsaSize作為步長,是因為 railed 模式下每個 LSA 組內的 rank 共享一個 GIN 連接,步長等於 LSA 組的大小。

v22902 的特殊處理

ncclDevCommCopyOldToNew_v22902 📎 src/devcomm/devcomm_v22902.cc:149-167有一個重要註解:

c
// Note: this callback will be used with v22907 as well because, prior to 2.30.0, ncclDevComm was unversioned,
// so v22902 and v22907 variants are indistinguishable.
〔設計推斷與架構權衡〕

這意味著在 2.30.0 之前,ncclDevComm沒有magic/version欄位,所以函式庫無法區分一個舊結構體到底是 v22902 還是 v22907。因此,v22907 的devCommCopyOldToNew被設為nullptr 📎 src/devcomm/devcomm_v22907.cc:128,實際使用的是 v22902 的版本。由於兩者都不支援 GIN 向後相容,GIN 相關欄位的差異不影響正確性。

資源視窗的版本化

ncclWindow_vidmem_v22902的定義在devcomm_v22902.h中(本章未提供該檔案內容),但從📎 src/devcomm/devcomm_v22902.cc:141和📎 src/devcomm/devcomm_v22902.cc:164可以看到,v22902 使用ncclDevCommCopyResourceWindow_v22902進行視窗轉換。這個函式在devcomm_v22902.h中宣告,具體實作未在本章原始碼中展示。

📎 src/devcomm/devcomm_v23000.cc:11-18的static_assert驗證了 v23000 的視窗佈局與當前版本一致,所以 v23000 的轉換函式可以直接逐欄位拷貝。

---

四、能力過濾與資源檢查:防止舊 kernel 存取不支援的特性

直覺模型

版本轉換不只是「欄位搬家」——還需要檢查舊版本是否支援應用程式請求的特性。比如,一個用 2.29.2 編譯的 kernel 請求 GIN 資源,但 2.29.2 的ncclDevComm佈局中 GIN 欄位不完整,直接轉換會導致 kernel 讀到垃圾資料。所以需要一個「過濾器」在轉換前攔截這種請求。

commPropertiesFilter:能力旗標過濾

ncclCommPropertiesFilter_v22907 📎 src/devcomm/devcomm_v22907.cc:69-77:

c
static ncclResult_t ncclCommPropertiesFilter_v22907(ncclComm_t comm, struct ncclCommProperties* props) {
  // We don't provide backwards compatibility for GIN with 2.29.7.  If a communicator needs it, we indicate that
  // the Device API is not available.
  props->deviceApiSupport = (props->deviceApiSupport && ncclTeamLsa(comm).nRanks == comm->nRanks);
  props->ginType = NCCL_GIN_TYPE_NONE;
  props->railedGinType = NCCL_GIN_TYPE_NONE;
  return ncclSuccess;
}

三個操作:

1. deviceApiSupport降級:如果 LSA 組的 rank 數不等於總 rank 數(即存在跨節點通訊),則停用裝置 API。這是因為 2.29.7 的 GIN 不支援跨節點。

2. ginType設為 NONE:明確告訴應用程式「這個版本不支援 GIN」。

3. railedGinType設為 NONE:同上。

ncclCommPropertiesFilter_v22902 📎 src/devcomm/devcomm_v22902.cc:86-96類似,但多了一個細節:

c
// v22902 ncclCommProperties is _almost_ compatible with newer ones, with the exception of ginType, which in that
// version was based on uint_8, not an int.
((struct ncclCommProperties_v22902*)props)->ginType = NCCL_GIN_TYPE_NONE_v22902;

📎 src/devcomm/devcomm_v22902.cc:13-17定義了 v22902 的 GIN 類型列舉:

c
typedef enum : uint8_t {
  NCCL_GIN_TYPE_NONE_v22902 = 0,
  NCCL_GIN_TYPE_PROXY_v22902 = 2,
  NCCL_GIN_TYPE_GDAKI_v22902 = 3,
} ncclGinType_t_v22902;

注意這是uint8_t類型,而新版本中ginType是int。所以 v22902 的過濾器需要把props強制轉換為ncclCommProperties_v22902*,然後寫入uint8_t類型的ginType。📎 src/devcomm/devcomm_v22902.cc:35-36的static_assert驗證了ginType在偏移 34,結構體大小為 40 位元組。

devCommRequirementsFilter:資源請求檢查

ncclDevCommRequirementsFilter_v22907 📎 src/devcomm/devcomm_v22907.cc:79-98檢查應用程式是否請求了 GIN 資源:

c
static ncclResult_t ncclDevCommRequirementsFilter_v22907(ncclComm_t comm, ncclDevCommRequirements_t* reqs) {
  bool requestedGinResources =
    reqs->ginSignalCount > 0 || reqs->ginCounterCount > 0 || reqs->barrierCount > 0 || reqs->railGinBarrierCount > 0;
  struct ncclDevResourceRequirements* node = reqs->resourceRequirementsList;
  while (!requestedGinResources && node != nullptr) {
    requestedGinResources = node->ginSignalCount > 0 || node->ginCounterCount > 0;
    node = node->next;
  }
  if (requestedGinResources && (reqs->ginConnectionType != NCCL_GIN_CONNECTION_NONE || reqs->ginForceEnable)) {
    // 打印警告并返回错误
    return ncclInvalidUsage;
  }
  return ncclSuccess;
}

邏輯分兩步:

1. 檢查頂層請求:reqs->ginSignalCount、ginCounterCount、barrierCount、railGinBarrierCount任一大於 0,說明請求了 GIN 資源。

2. 遍歷資源需求鏈結串列:如果頂層沒有請求,繼續遍歷resourceRequirementsList鏈結串列,檢查每個節點的ginSignalCount和ginCounterCount。

如果確實請求了 GIN 資源,且ginConnectionType不是NONE或ginForceEnable為真,則傳回ncclInvalidUsage並列印警告,提示應用程式需要重新編譯。

ncclDevCommRequirementsFilter_v22902 📎 src/devcomm/devcomm_v22902.cc:98-126更複雜,除了 GIN 檢查外,還處理了barrierCount的語意變化:

c
// Prior to 2.29.4, a non-zero barrierCount did not imply GIN, but it does since.
if (reqs->barrierCount) {
  reqs->lsaBarrierCount = std::max(reqs->lsaBarrierCount, reqs->barrierCount);
  reqs->barrierCount = 0;
}
// Strangely, neither did railGinBarrierCount.
reqs->railGinBarrierCount = 0;
〔設計推斷與架構權衡〕

在 2.29.4 之前,barrierCount只表示 LSA barrier,不隱含 GIN 需求。從 2.29.4 開始,barrierCount隱含 GIN 需求。為了相容舊版本,過濾器把barrierCount轉換為lsaBarrierCount,並清零barrierCount和railGinBarrierCount。

下面的時序圖展示了從應用程式請求到版本轉換的完整互動:

mermaid
sequenceDiagram
    participant App as 应用层
    participant Host as Host 侧 NCCL 库
    participant Compat as ncclDevCommCompat 插件
    participant Dev as 设备侧 ncclDevComm

    App->>Host: ncclCommGetDeviceHandle(comm, &devComm)
    Host->>Host: 读取 reqs->version(应用编译版本)
    Host->>Compat: 查找覆盖该版本的插件
    Compat-->>Host: 返回 ncclDevCommCompat_vXXXXX
    Host->>Compat: devCommRequirementsFilter(comm, reqs)
    alt 请求了不支持的 GIN 资源
        Compat-->>Host: ncclInvalidUsage
        Host-->>App: 返回错误 + 警告日志
    else 资源兼容
        Compat-->>Host: ncclSuccess
        Host->>Compat: devCommCopyNewToOld(comm, oldDevComm, newDevComm)
        Compat->>Compat: memset(old, 0, sizeof(*old))
        Compat->>Compat: 逐字段拷贝 + 语义转换
        Compat-->>Host: ncclSuccess
        Host->>Dev: 返回旧布局 ncclDevComm
        Dev-->>App: 设备侧可访问的通信器
    end

---

五、生產避坑指南與故障復原鏈

陷阱一:GIN 資源請求與舊版本 kernel 的衝突

場景:應用程式用 NCCL 2.29.2 編譯,但執行時連結了 2.31.0 的函式庫。應用程式在 kernel 中呼叫了 GIN 相關的裝置側 API(如ncclGinPut)。

會發生什麼:ncclDevCommRequirementsFilter_v22902 📎 src/devcomm/devcomm_v22902.cc:98-126偵測到ginForceEnable或ginSignalCount > 0,傳回ncclInvalidUsage,並列印警告:

code
The application was compiled with too old version of NCCL. It was compiled with NCCL version 2.29.2, but is
running with NCCL library version 2.31.0. Because of its use of GIN device kernels, it needs to be recompiled,
preferably with the same NCCL version that it will be running with.

根因:2.29.2 的ncclDevComm_v22902佈局中,GIN 欄位(ginContextCount、ginNetDeviceTypes、ginHandles等)與 2.31.0 的佈局不相容。如果強行轉換,kernel 會讀到錯誤的偏移,導致未定義行為。

正確做法:應用程式必須用與執行時函式庫相同(或相容)的 NCCL 版本重新編譯。如果無法重新編譯,應避免在 kernel 中使用 GIN API。

陷阱二:跨節點通訊時裝置 API 被靜默停用

場景:應用程式用 2.29.7 編譯,通訊域包含跨節點 rank(ncclTeamLsa(comm).nRanks != comm->nRanks)。

會發生什麼:ncclCommPropertiesFilter_v22907 📎 src/devcomm/devcomm_v22907.cc:69-77把props->deviceApiSupport設為false。應用程式如果檢查了這個旗標,會知道裝置 API 不可用;但如果不檢查,直接呼叫裝置側 API,會得到未定義行為。

根因:2.29.7 的 GIN 不支援跨節點。LSA(Local SHARP Aggregation)組內的 rank 才能使用裝置側 API。

正確做法:應用程式應在初始化後檢查ncclCommProperties.deviceApiSupport,如果為false,回退到 host 側 API。

陷阱三:memset 清零與未初始化欄位洩漏

場景:ncclDevCommCopyNewToOld_v23000 📎 src/devcomm/devcomm_v23000.cc:118在拷貝前執行memset(old, '\0', sizeof(*old))。

為什麼需要:舊結構體中可能有新版本不存在的欄位(如 v22902 中的ginSignalBase、ginCounterBase)。如果不清零,這些欄位會保留堆疊上的垃圾值,可能被 kernel 誤讀為有效資料。

踩坑點:如果開發者手動實作版本轉換而忘記清零,可能導致 kernel 讀到隨機值,表現為間歇性錯誤——難以復現和除錯。

正確做法:始終在轉換前清零整個目標結構體。NCCL 的所有CopyNewToOld實作都遵循這個模式📎 src/devcomm/devcomm_v22902.cc:132 📎 src/devcomm/devcomm_v22907.cc:104 📎 src/devcomm/devcomm_v23000.cc:118。

陷阱四:版本區間空隙導致的匹配失敗

場景:應用程式用 NCCL 2.29.4 編譯。查看版本區間表:

檔案minVersionmaxVersion
v229022.29.22.29.3
v229072.29.52.29.7

2.29.4 沒有對應的插件。

〔設計推斷與架構權衡〕

會發生什麼: 如果匹配邏輯嚴格按區間查找,2.29.4 會匹配失敗,返回錯誤。但實際實作中,可能有一個「最近匹配」策略——2.29.4 可能被路由到 v22902 或 v22907 的插件。

正確做法:應用程式應盡量使用與執行時庫相同的主版本號。如果必須跨版本,應測試目標版本區間是否有對應的相容插件。

故障恢復鏈

當版本轉換失敗時,NCCL 的錯誤恢復鏈:

1. 過濾器返回錯誤:devCommRequirementsFilter返回ncclInvalidUsage。

2. 上層 API 捕獲錯誤:ncclCommGetDeviceHandle檢查返回值,如果非ncclSuccess,不填充devComm結構。

3. 應用程式處理:應用程式應檢查返回值,如果失敗,回退到 host 側 API 或終止通訊。

4. 日誌記錄:NCCL 列印WARN級別的日誌,包含編譯版本和執行時版本,幫助定位問題。

〔設計推斷與架構權衡〕

目前 NCCL 沒有提供「自動降級」機制——如果版本轉換失敗,不會自動回退到 host 側 API。應用程式需要自己實作回退邏輯。

---

設計思考

為什麼用版本化結構體而不是「穩定 ABI」?

〔設計推斷與架構權衡〕

一個替代方案是設計一個「永不改變」的ncclDevComm佈局,所有新欄位都通過間接指標訪問。但這會帶來兩個問題:一是間接訪問增加延遲(kernel 需要額外解引用),二是無法利用填充區優化佈局。NCCL 選擇版本化結構體,是在「效能」和「相容性」之間的權衡——每個版本區間內的 kernel 獲得最優佈局,跨版本時通過轉換層保證相容。

為什麼 v22907 的devCommCopyOldToNew設為 nullptr?

📎 src/devcomm/devcomm_v22902.cc:153-155的註釋解釋了原因:2.30.0 之前ncclDevComm沒有版本欄位,所以 v22902 和 v22907 的舊佈局無法區分。由於兩者都不支援 GIN 向後相容,GIN 欄位的差異不影響正確性,所以複用 v22902 的轉換函數。

為什麼nRanks_rcp32用定點數而不是浮點數?

〔設計推斷與架構權衡〕

GPU 的浮點除法精度可能不足以精確表示1/nRanks,特別是當nRanks不是 2 的冪時。定點數(32 位整數表示的小數)可以提供足夠的精度,且整數乘法比浮點乘法更快。

---

本章小結

本章拆解了src/devcomm目錄下的版本化 ABI 實作:

1. ncclDevComm的記憶體佈局:每個版本有精確的欄位偏移,用static_assert在編譯期驗證。關鍵欄位包括rank、nRanks、nRanks_rcp32、lsaRank、lsaSize、windowTable、resourceWindow等。

2. 版本化 ABI 的註冊:每個版本區間對應一個ncclDevCommCompat結構體,包含minVersion、maxVersion、過濾器函數和轉換函數。

3. 欄位級轉換:CopyNewToOld和CopyOldToNew逐欄位拷貝,並處理語義變化(如ginConnectionStride > 1轉換為ginConnectionsRailed = true)。

4. 能力過濾:commPropertiesFilter調整暴露給舊版本的能力標誌,devCommRequirementsFilter檢查資源請求是否與舊版本相容。

5. 生產陷阱:GIN 資源請求與舊版本 kernel 的衝突、跨節點通訊時裝置 API 被禁用、memset 清零的必要性、版本區間空隙導致的匹配失敗。

下一章我們將進入裝置側 API 與內核融合,看nccl_device標頭檔如何組織裝置側函數,以及 kernel fusion 如何把多個集合通訊操作合併到一個 kernel 中執行。

本章思考與自測

Q1: 如果將ncclDevCommCopyNewToOld_v23000中的memset(old, '\0', sizeof(*old))去掉,在什麼場景下會導致 kernel 讀到錯誤資料?請結合 v22902 和 v23000 的欄位差異分析。

參考解析:

ncclDevComm_v22902的結構體大小為 200 位元組📎 src/devcomm/devcomm_v22902.cc:84,而ncclDevComm_v23000為 240 位元組📎 src/devcomm/devcomm_v23000.cc:95-98。v22902 中有ginSignalBase(偏移 176)、ginCounterBase(偏移 184)、ginContextBase(偏移 204)等欄位,這些欄位在 v23000 中不存在或語義不同。

如果去掉memset,當從 v23000 轉換到 v22902 時,old結構體中 v23000 不存在的欄位(如ginSignalBase、ginCounterBase)會保留堆疊上的垃圾值。如果 kernel 恰好讀取了這些欄位(例如舊 kernel 的 GIN 程式碼路徑),會得到隨機值,導致:

  • 訊號基底位址錯誤,GIN 操作寫入錯誤的記憶體位置。
  • 計數器基底位址錯誤,導致計數器溢位或下溢。
  • 在極端情況下,可能觸發非法記憶體存取,導致 kernel 崩潰。

memset清零確保所有未顯式賦值的欄位都是 0,這是一個安全的預設值。NCCL 的所有CopyNewToOld實作都包含這個步驟📎 src/devcomm/devcomm_v22902.cc:132 📎 src/devcomm/devcomm_v22907.cc:104 📎 src/devcomm/devcomm_v23000.cc:118。

Q2: 假設應用程式用 NCCL 2.29.4 編譯,執行時連結 2.31.0 的庫。根據本章的版本區間表,2.29.4 沒有對應的ncclDevCommCompat插件。請分析 NCCL 可能如何處理這種情況,以及應用程式應該如何規避。

參考解析:

版本區間表:

  • v22902:2.29.2 - 2.29.3
  • v22907:2.29.5 - 2.29.7
  • v23000:2.30.0 - 2.30.7
  • v23100:2.31.0 - 當前

2.29.4 落在 v22902 和 v22907 之間的空隙。可能的處理方式:

1. 最近匹配:NCCL 可能選擇小於等於請求版本的最大區間,即 v22902。但 v22902 的maxVersion是 2.29.3,嚴格來說不覆蓋 2.29.4。

2. 返回錯誤:如果匹配邏輯嚴格按區間,2.29.4 會匹配失敗,返回ncclInvalidUsage。

3. 向上匹配:選擇大於等於請求版本的最小區間,即 v22907。但 v22907 的minVersion是 2.29.5,也不覆蓋 2.29.4。

〔設計推斷與架構權衡〕

實際實作中,NCCL 可能有一個「容錯」策略——如果找不到精確匹配,嘗試使用相鄰區間的插件。但這不是可靠的保證。

應用程式的規避方法:

  • 使用與執行時程式庫相同的主版本號(如 2.31.x)。
  • 如果必須跨版本,測試目標版本區間是否有對應的相容插件。
  • 在初始化後檢查ncclCommProperties.deviceApiSupport,如果為false,回退到 host 側 API。
Q3: ncclDevCommRequirementsFilter_v22902中有一段邏輯:if (reqs->barrierCount) { reqs->lsaBarrierCount = std::max(reqs->lsaBarrierCount, reqs->barrierCount); reqs->barrierCount = 0; }。請解釋為什麼需要這個轉換,以及如果不轉換會發生什麼。

參考解析:

📎 src/devcomm/devcomm_v22902.cc:117-121的註解說明:「Prior to 2.29.4, a non-zero barrierCount did not imply GIN, but it does since.」

在 2.29.4 之前,barrierCount只表示 LSA barrier 的數量,不隱含 GIN 需求。從 2.29.4 開始,barrierCount隱含 GIN 需求(即請求 barrier 就意味著需要 GIN 資源)。

當應用程式用 2.29.2 編譯時,它可能設定了barrierCount > 0來表示 LSA barrier 需求,但並不知道這會隱含 GIN 需求。如果 NCCL 程式庫(2.31.0)直接按新語意處理,會認為應用程式請求了 GIN 資源,然後ncclDevCommRequirementsFilter_v22902會偵測到 GIN 請求並返回ncclInvalidUsage——這是一個誤報。

轉換邏輯把barrierCount轉換為lsaBarrierCount(取兩者最大值),並清零barrierCount。這樣:

  • lsaBarrierCount保留了應用程式的 barrier 需求。
  • barrierCount = 0避免了誤報 GIN 需求。
  • railGinBarrierCount = 0同理,因為舊版本中它也不隱含 GIN 需求。

如果不轉換,應用程式用 2.29.2 編譯且設定了barrierCount > 0時,會被錯誤地拒絕,無法使用裝置 API。

至此,我們看清了 devcomm 如何透過版本化 ABI 把 host 側通訊域的關鍵元資料安全地映射到裝置側,讓 kernel 無需 host 指標也能取得 rank、位址和連線狀態。這套機制解決了 kernel 存取通訊域的基本問題,但裝置側的能力遠不止於此。當使用者希望在自己的 kernel 中直接呼叫通訊原語,甚至將通訊與計算融合到同一個 kernel 時,就需要更上層的裝置側 API 和核心融合技術。下一章將深入 nccl_device 目錄與相關範例,探索 ncclBarrier、ncclLsaBarrier、ncclGinBarrier 等裝置側 API 如何讓使用者 kernel 參與通訊,以及核心融合如何減少啟動開銷,從而將 NCCL 從程式庫推向程式設計模型。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 20

第 20 章:裝置端原生 API 與算子融合:nccl_device 與 kernel fusion 實踐

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 20 章 / 共 25 章

上一章我們看清了 devcomm 如何把 host 側 ncclComm 的元資料版本化地映射到裝置側,讓 kernel 能讀到 rank、位址和連線狀態。但「能讀到元資料」和「能發起通訊」是兩回事。如果只有元資料,使用者 kernel 頂多能自己算算位址、自己寫寫標誌位,一旦涉及跨 rank 的同步、跨機的訊號傳遞,還是得回到 host 側呼叫 ncclAllReduce 之類的集合 API——而每一次這樣的呼叫都意味著一次 kernel 啟動、一次 host-device 往返。本章要拆解的 src/nccl_device 目錄,正是 NCCL 從「一個被呼叫的函式庫」走向「一套可被程式設計的模型」的關鍵。它提供的不是新的集合通訊演算法,而是一組裝置側原語:讓使用者自己的 kernel 內部就能呼叫 ncclBarrier、ncclLsaBarrier、ncclGinBarrier 這類同步操作,從而把「通訊」和「計算」塞進同一個 kernel,省掉中間的啟動開銷。本章原始碼材料聚焦於這組原語在 host 側的需求宣告(CreateRequirement)與團隊(Team)抽象,這正是裝置側 API 的入口。理解本章的一個關鍵前提:裝置側 API 的設計哲學是「host 側宣告資源需求,device 側消費資源」。host 側不直接建立 barrier,而是告訴 NCCL「我需要 nBarriers 個 barrier,團隊有 team.nRanks 個成員」,NCCL 據此算出需要多少緩衝區、多少 GIN 訊號,然後在 device 側把這些資源實例化。這種「宣告-消費」分離,是裝置側程式碼能在沒有 host 指標的情況下運作的根本原因。

一、Team 抽象:裝置側 API 的座標系

直覺模型

想像一個跨國公司的組織架構。你要發一封郵件,首先得知道「發給誰」——是發給全公司(World)、發給同一個辦公室的同事(LSA)、還是發給同一條業務線的跨辦公室團隊(Rail)。ncclTeam_t就是這套「收件人範圍」的描述符。若沒有 Team 抽象,每個裝置側 API 都得自己重新計算「我在這個通訊域裡排第幾、一共有幾個人」,程式碼會重複且極易出錯。

資料結構與記憶體佈局

ncclTeam_t是裝置側 API 的座標系,它的三個欄位定義了一個等差數列:

欄位含義類比
nRanks團隊內成員總數群裡有多少人
rank當前 rank 在團隊內的編號我在群裡的序號
stride團隊內相鄰成員在 world 中的步長群裡相鄰兩人學號差多少

stride是最容易被忽略但最關鍵的欄位。World 團隊裡stride = 1,因為所有 rank 連續排列;但 Rail 團隊裡stride = lsaSize,因為同一個 rail 上的 rank 在 world 中每隔lsaSize個才出現一次。

📎 src/nccl_device/core.cc:13-19展示了 World 團隊的建構:直接取comm->nRanks和comm->rank,stride固定為 1。這是唯一不需要ncclDevrInitOnce的團隊,因為它的資訊全在 host 側comm裡。

📎 src/nccl_device/core.cc:22-33是 LSA 團隊。注意 L26 的ncclDevrInitOnce(comm)——這是裝置側資源初始化的冪等入口。L23-25 的註解非常關鍵:這裡故意忽略錯誤,因為如果初始化失敗,返回的 team 是「垃圾值」,但下一個真正需要資源的 API 呼叫會再次觸發ncclDevrInitOnce並報告錯誤。這是一種「延遲報錯」策略,避免在團隊查詢這種輕量操作上拋出重錯誤。

場景驅動 Walkthrough:從 World 到 Rail 的座標變換

假設一個 8 卡機器,lsaSize = 4(每 4 卡一個 LSA 域),nRanks = 8。我們來看ncclTeamRail如何建構:

📎 src/nccl_device/core.cc:70-79中,nRanks = 8 / 4 = 2,rank = comm->rank / 4,stride = 4。如果當前 rank 是 5,那麼它在 Rail 團隊裡的rank = 5 / 4 = 1,stride = 4,意味著 Rail 團隊的成員是 world 中的 rank 1 和 rank 5。

再看ncclTeamRankToWorld的換算公式:

📎 src/nccl_device/core.cc:82-84的comm->rank + (rank - team.rank) * team.stride是一個相對偏移計算:先算出目標 rank 相對於當前 rank 在團隊內的偏移(rank - team.rank),再乘以步長stride,加上當前 rank 的 world 編號。這個公式對所有團隊通用,因為stride已經編碼了團隊的排列規律。

ncclTeamRankToLsa則不同:

📎 src/nccl_device/core.cc:87-92用的是comm->devrState.lsaSelf + (rank - team.rank) * team.stride。注意這裡用的是lsaSelf而不是comm->rank——因為 LSA 編號是裝置側資源初始化後才知道的,可能與 world rank 不同。

mermaid
flowchart TD
    start["用户调用 ncclTeamRail(comm)"] --> init{"ncclDevrInitOnce(comm)<br/>成功?"}
    init -->|"否"| empty["返回 ncclTeam_t{}<br/>空团队"]
    init -->|"是"| calc["计算 nRanks = comm->nRanks / lsaSize<br/>rank = comm->rank / lsaSize<br/>stride = lsaSize"]
    calc --> ret["返回 ncclTeam_t"]
    empty --> caller["调用方继续<br/>下一个 API 会报错"]
    ret --> caller

這張圖揭示了「延遲報錯」策略的執行路徑:初始化失敗時返回空團隊,但不中斷呼叫方;錯誤會在下一個真正需要資源的 API(如ncclLsaBarrierCreateRequirement)處暴露。

設計思考與踩坑

為什麼ncclTeamWorld不呼叫ncclDevrInitOnce?因為 World 團隊的資訊完全來自 host 側comm,不需要任何裝置側資源。如果強行呼叫,會讓一個純 host 查詢操作依賴裝置側初始化,增加不必要的失敗點。

踩坑點:ncclTeamRankToLsa在初始化失敗時返回-1(📎 src/nccl_device/core.cc:87-92),而ncclTeamRankToWorld永遠不會失敗。呼叫方如果混用這兩個函式且不檢查返回值,可能在 LSA 初始化失敗時拿到-1當作合法 rank 使用,導致越界存取。生產程式碼中應當把ncclTeamRankToLsa的返回值當作可能失敗的操作處理。

---

二、Barrier 需求宣告:host 側如何「預訂」裝置資源

直覺模型

裝置側 API 的資源分配像預訂會議室:你不能直接衝進會議室開會,得先向前台(host 側CreateRequirement)提交申請——「我要開 3 場會,每場 8 個人參加」。前台據此算出需要多大的場地(bufferSize)、需要多少把椅子(ginSignalCount),然後把場地編號(outBufferHandle)給你。若沒有這套預訂機制,裝置側 kernel 就不知道自己的 barrier 緩衝區在哪裡、有多大,無法安全地讀寫。

資料結構與記憶體佈局

三個 barrier 的CreateRequirement函式共享同一個模式:清零需求結構體 → 填充緩衝區大小/對齊 → 填充輸出句柄指標。但它們的資源類型不同:

Barrier 類型資源類型大小公式對齊
LSA Barrier緩衝區(3*n + n*team.nRanks) * sizeof(uint32_t)alignof(uint32_t)
CFT Barrier緩衝區(3*n + n*team.nRanks) * NCCL_CFT_BARRIER_GRANNCCL_CFT_BARRIER_ALIGN
GIN BarrierGIN 信號n * team.nRanks個信號不涉及緩衝區

先看 LSA Barrier 的大小公式:

📎 src/nccl_device/lsa_barrier.cc:14-22的(3 * nBarriers + nBarriers * team.nRanks) * sizeof(uint32_t)可以拆解為兩部分:

  • 3 * nBarriers:每個 barrier 需要 3 個uint32_t的控制欄位([INFERENCE] 通常是「到達計數」「輪次」「狀態標誌」)。
  • nBarriers * team.nRanks:每個 barrier 需要為團隊內每個成員預留一個uint32_t的到達槽位。

所以單個 barrier 的總大小是3 + team.nRanks個uint32_t。這個公式在 LSA 和 CFT 中完全一致,只是 CFT 用NCCL_CFT_BARRIER_GRAN作為粒度單位(可能是為了對齊到更大的邊界)。

GIN Barrier 則完全不同:

📎 src/nccl_device/gin_barrier.cc:14-20不分配緩衝區,而是設定ginSignalCount = nBarriers * team.nRanks,並把outGinSignalStart指向句柄裡的signal0。這是因為 GIN barrier 走的是網路信號路徑,不需要共享記憶體緩衝區,而是需要網卡能識別的信號槽位。

場景驅動 Walkthrough:一次 LSA Barrier 的完整預訂

假設使用者要在一個 4 卡 LSA 團隊上建立 2 個 barrier:

1. 呼叫 ncclLsaBarrierCreateRequirement(team, 2, &handle, &req)。

2. 清零:memset(outReq, 0, sizeof(*outReq))(📎 src/nccl_device/lsa_barrier.cc:14-22)——保證未設定的欄位是確定值,避免呼叫方讀到堆疊上的垃圾。

3. 記錄 barrier 數量:outHandle->nBarriers = 2(📎 src/nccl_device/lsa_barrier.cc:14-22)。

4. 計算緩衝區大小:(3*2 + 2*4) * 4 = (6 + 8) * 4 = 56位元組(📎 src/nccl_device/lsa_barrier.cc:14-22)。

5. 設定對齊:alignof(uint32_t) = 4(📎 src/nccl_device/lsa_barrier.cc:14-22)。

6. 回填句柄指標:outReq->outBufferHandle = &outHandle->bufHandle(📎 src/nccl_device/lsa_barrier.cc:14-22)——讓 NCCL 在真正分配緩衝區後,把位址寫回句柄。

mermaid
flowchart LR
    subgraph host["host 侧声明阶段"]
        req["ncclLsaBarrierCreateRequirement<br/>team, nBarriers=2"]
        calc["bufferSize = (3*2 + 2*4)*4 = 56<br/>bufferAlign = 4"]
        handle["outHandle->nBarriers = 2<br/>outReq->outBufferHandle = &handle->bufHandle"]
    end
    subgraph dev["device 侧消费阶段"]
        buf["缓冲区 56 字节<br/>3 控制字段 + 4 到达槽位"]
        bar["ncclLsaBarrier 实例"]
    end
    req --> calc --> handle
    handle -.->|"NCCL 分配后回填"| buf
    buf --> bar

這張資料流圖展示了「宣告」與「消費」的分離:host 側只算出大小和指標,真正的緩衝區分配和實例化發生在 NCCL 內部,device 側 kernel 拿到的是已經填充好的句柄。

設計思考與踩坑

為什麼用memset清零整個outReq?因為ncclDevResourceRequirements_t是一個多欄位結構體,不同 barrier 類型只填充其中一部分欄位。清零保證未使用的欄位(如 LSA barrier 不用的ginSignalCount)是 0,NCCL 內部據此判斷「這個資源不需要」。如果不清零,堆疊上的隨機值可能被誤認為「需要 GIN 資源」,觸發上一章提到的誤報問題。

踩坑點:outReq->outBufferHandle = &outHandle->bufHandle把句柄內部欄位的位址交給了 NCCL。這意味著outHandle必須在 NCCL 完成緩衝區分配之前保持有效(不能被堆疊回收或移動)。如果使用者把outHandle放在一個會被提前釋放的作用域裡,NCCL 回填時就會寫入野指標。

〔設計推斷與架構權衡〕

CFT Barrier 的粒度差異:📎 src/nccl_device/cft_barrier.cc:13-21用NCCL_CFT_BARRIER_GRAN和NCCL_CFT_BARRIER_ALIGN替代了 LSA 的sizeof(uint32_t)和alignof(uint32_t)。這說明 CFT( 可能是 Cross-Fabric Team 或類似的跨域團隊)的 barrier 需要更大的對齊粒度,可能因為要跨多播記憶體區域,硬體對位址對齊有更嚴格的要求。

---

三、三種 Barrier 的語意分工:LSA、CFT、GIN 各管什麼

直覺模型

三種 barrier 像三種不同範圍的「集合哨」:

  • LSA Barrier:同一個辦公室內的同事集合,走共享記憶體,最快。
  • CFT Barrier:跨辦公室但同一棟樓內的集合,走多播記憶體,中等。
  • GIN Barrier:跨城市甚至跨國的集合,走網路信號,最慢但覆蓋最廣。

選錯 barrier 類型不會導致錯誤,但會帶來巨大的效能損失——用 GIN barrier 做同辦公室同步,等於用國際快遞送隔壁工位的文件。

資料結構與記憶體佈局對比

從 host 側需求宣告看,三者的資源需求截然不同:

維度LSA BarrierCFT BarrierGIN Barrier
需要comm參數否否是
緩衝區有有無
GIN 信號無無有
大小單位uint32_tNCCL_CFT_BARRIER_GRAN信號個數
輸出句柄欄位bufHandlebufHandlesignal0

注意 GIN Barrier 是唯一需要comm參數的:

📎 src/nccl_device/gin_barrier.cc:14-20的函式簽名包含ncclComm_t comm,而 LSA 和 CFT 的簽名只有ncclTeam_t team。這是因為 GIN 訊號需要綁定到具體的網路連線,而網路連線資訊在comm裡。

場景驅動 Walkthrough:GIN Barrier 的訊號分配

📎 src/nccl_device/gin_barrier.cc:14-20的邏輯比 LSA 更簡單,但語意更微妙:

1. 清零:memset(outReq, 0, sizeof(*outReq))(L16)。

2. 設定訊號數:outReq->ginSignalCount = nBarriers * team.nRanks(L17)——每個 barrier 需要為團隊內每個成員分配一個訊號槽。

3. 回填訊號起始指標:outReq->outGinSignalStart = &outHandle->signal0(L18)——注意這裡沒有設定bufferSize,因為 GIN barrier 不用共享記憶體緩衝區。

〔設計推斷與架構權衡〕

signal0這個名字暗示句柄裡可能有一組連續的訊號欄位(signal0, signal1, ...),outGinSignalStart指向第一個,NCCL 據此知道從哪裡開始分配nBarriers * team.nRanks個訊號。

並發控制與硬體互動

三種 barrier 的並發控制機制完全不同:

  • LSA Barrier:基於共享記憶體的原子操作。3 + team.nRanks個uint32_t中,到達槽位用原子加或原子寫來標記「我到了」,控制欄位用原子讀來檢查「是否所有人都到了」。這是純 GPU 內的同步,不涉及網路。
  • CFT Barrier:基於多播記憶體(multimem)。[INFERENCE] 多播記憶體允許一次寫操作同時更新多個 rank 的視圖,所以 CFT barrier 可能用更少的控制欄位實現更廣的同步。
  • GIN Barrier:基於網路訊號。ginSignalCount個訊號透過網卡發送,接收方輪詢訊號槽位。這是唯一涉及跨機硬體的 barrier。
mermaid
sequenceDiagram
    participant K as "用户 Kernel"
    participant LSA as "LSA 共享内存"
    participant CFT as "CFT 多播内存"
    participant NIC as "网卡 GIN 信号"
    K->>LSA: "原子写到达槽位"
    LSA-->>K: "轮询所有槽位"
    Note over K,LSA: LSA barrier 完成
    K->>CFT: "多播写控制字段"
    CFT-->>K: "读多播状态"
    Note over K,CFT: CFT barrier 完成
    K->>NIC: "发送 GIN 信号"
    NIC-->>K: "轮询信号槽位"
    Note over K,NIC: GIN barrier 完成

這張時序圖展示了三種 barrier 的硬體互動層次:從純 GPU 內同步,到多播記憶體,再到網卡訊號,延遲依次遞增,覆蓋範圍也依次擴大。

設計思考與踩坑

為什麼 LSA 和 CFT 不需要comm參數?因為它們的資源(共享記憶體、多播記憶體)已經在ncclDevrInitOnce階段綁定到了團隊上,team本身就隱含了資源位置資訊。而 GIN 訊號需要動態分配網路資源,必須透過comm存取網路連線狀態。

踩坑點:GIN Barrier 的ginSignalCount是nBarriers * team.nRanks,如果團隊很大(如 1024 個 rank)且 barrier 很多(如 100 個),訊號總數會達到 102400。網卡的訊號槽位是有限資源,超量申請可能導致ncclDevrInitOnce失敗。生產程式碼應當根據實際需要的最小 barrier 數量申請,而不是一次性申請大量備用。

---

四、從需求宣告到裝置側消費:完整生命週期

直覺模型

CreateRequirement只是「下單」,真正的「發貨」和「收貨」發生在 NCCL 內部和裝置側 kernel 裡。整個生命週期像網購:你下單(CreateRequirement)→ 商家備貨(NCCL 分配資源)→ 快遞送達(資源綁定到 DevComm)→ 你簽收使用(device 側 kernel 呼叫 barrier)。

資料結構與記憶體佈局:句柄的欄位演化

以ncclLsaBarrierHandle_t為例,它在生命週期中經歷三個階段:

階段nBarriersbufHandle其他欄位
CreateRequirement 後已設定位址已回填,但內容未分配未設定
NCCL 分配後已設定指向實際緩衝區已設定
Device 側使用唯讀唯讀唯讀

📎 src/nccl_device/lsa_barrier.cc:14-22設定nBarriers,📎 src/nccl_device/lsa_barrier.cc:14-22回填bufHandle的位址。這兩個操作之間,NCCL 內部會完成緩衝區的實際分配。

場景驅動 Walkthrough:一次完整的 barrier 使用

1. Host 側宣告:使用者呼叫ncclLsaBarrierCreateRequirement(team, 2, &handle, &req),得到req.bufferSize = 56。

2. Host 側提交:使用者把req交給ncclDevCommCreate(上一章的內容),NCCL 分配 56 位元組緩衝區,把位址寫入handle.bufHandle。

3. Device 側初始化:使用者 kernel 啟動時,從 DevComm 裡取出handle,用bufHandle定位緩衝區。

4. Device 側同步:kernel 呼叫ncclLsaBarrier(handle, barrierIndex),在緩衝區的對應槽位寫入到達標記,輪詢其他槽位。

5. Device 側完成:所有 rank 到達後,barrier 返回,kernel 繼續執行。

mermaid
flowchart TD
    a["ncclLsaBarrierCreateRequirement<br/>算出 bufferSize=56"] --> b["ncclDevCommCreate<br/>分配 56 字节缓冲区"]
    b --> c{"分配成功?"}
    c -->|"否"| err["返回 ncclSystemError<br/>句柄无效"]
    c -->|"是"| d["回填 handle.bufHandle<br/>指向实际缓冲区"]
    d --> e["用户 kernel 启动<br/>从 DevComm 取 handle"]
    e --> f["ncclLsaBarrier(handle, idx)<br/>写到达槽位 + 轮询"]
    f --> g{"所有 rank 到达?"}
    g -->|"否"| f
    g -->|"是"| h["barrier 返回<br/>kernel 继续"]
    err --> i["用户需检查返回值<br/>不可使用无效句柄"]

這張決策圖展示了從宣告到使用的完整路徑,以及分配失敗時的錯誤分支。注意ncclLsaBarrierCreateRequirement本身永遠返回ncclSuccess(📎 src/nccl_device/lsa_barrier.cc:14-22),真正的失敗發生在後續的資源分配階段。

並發控制與硬體互動

裝置側 barrier 的並發控制核心是原子操作 + 記憶體屏障。以 LSA barrier 為例:

  • 到達階段:每個 rank 用原子寫(或原子加)更新自己的到達槽位。這一步必須用 release 語意,保證 barrier 之前的所有記憶體操作對其他 rank 可見。
  • 輪詢階段:每個 rank 用原子讀(或 volatile 讀)檢查所有槽位。這一步必須用 acquire 語意,保證看到「所有人都到了」之後,能讀到其他人 barrier 之前寫入的資料。
  • 重置階段:barrier 完成後,需要重置槽位供下次使用。這一步的並發控制最微妙——如果重置太快,可能覆蓋還沒讀到的 rank 的標記。
〔設計推斷與架構權衡〕

3 * nBarriers個控制欄位很可能就是用來處理這種「輪次」問題的:一個欄位記錄當前輪次,一個欄位記錄到達計數,一個欄位作為重置標誌。這樣多個 barrier 可以複用同一組槽位而不會混淆輪次。

生產避坑指南

坑 1:句柄生命週期管理。outReq->outBufferHandle = &outHandle->bufHandle把句柄內部欄位的位址交給了 NCCL。如果使用者在ncclDevCommCreate返回之前就銷毀了outHandle,NCCL 回填時會寫入已釋放的記憶體。正確做法是把outHandle的生命週期綁定到 DevComm,而不是綁定到建立它的函式作用域。

坑 2:barrier 數量與團隊大小的乘積。bufferSize = (3*n + n*team.nRanks) * sizeof(uint32_t)中,n*team.nRanks項在大團隊時會主導大小。1024 個 rank、100 個 barrier 需要100*1024*4 = 409600位元組,約 400KB。如果每個 rank 都申請這麼多,顯存壓力不可忽視。應當按實際並行使用的 barrier 數量申請,而不是按總 barrier 數量。

坑 3:GIN barrier 的訊號耗盡。GIN 訊號是網卡資源,數量有限。如果多個 DevComm 同時申請大量 GIN 訊號,可能耗盡網卡槽位。生產程式碼應當在 DevComm 建立失敗時檢查是否是 GIN 訊號不足,並考慮減少nBarriers或改用 LSA barrier。

坑 4:初始化失敗的延遲暴露。ncclTeamLsa等函式在ncclDevrInitOnce失敗時返回空團隊(📎 src/nccl_device/core.cc:22-33),不報錯。如果使用者程式碼不檢查後續 API 的回傳值,可能在空團隊上繼續操作,導致難以定位的錯誤。建議在第一次使用裝置側 API 時顯式檢查團隊的有效性(如team.nRanks > 0)。

---

五、核心融合:為什麼要把通訊和計算塞進一個 kernel

直覺模型

傳統模式下,一次「AllReduce + 激活函式」需要兩個 kernel:一個做通訊,一個做計算。兩個 kernel 之間有一次隱式的全域同步——通訊 kernel 必須完全結束,計算 kernel 才能開始。這就像接力賽:第一棒跑完必須把棒交給第二棒,交接瞬間兩人都在等。核心融合則是讓同一個 kernel 既跑通訊又跑計算,像一個人邊跑邊換鞋,省掉了交接的等待。

資料結構與記憶體佈局

核心融合的關鍵在於:通訊原語(如 barrier)和計算邏輯共享同一個 kernel 的暫存器和共享記憶體。這意味著:

  • 暫存器壓力:通訊原語的原子操作和輪詢迴圈會佔用暫存器,擠壓計算邏輯的暫存器預算。
  • 共享記憶體競爭:LSA barrier 的緩衝區如果放在共享記憶體裡,會和計算邏輯的共享記憶體需求競爭。
  • Occupancy 影響:融合 kernel 的 occupancy 通常低於純計算 kernel,因為通訊原語需要額外的資源。
〔設計推斷與架構權衡〕

裝置側 API 的設計(host 側宣告資源、device 側消費)正是為了緩解這些壓力:資源在 host 側預先分配好,device 側 kernel 只需要讀寫,不需要動態申請,減少了暫存器佔用。

場景驅動 Walkthrough:融合 kernel 的執行流

假設使用者要寫一個「AllReduce + ReLU」的融合 kernel:

1. Host 側準備:呼叫ncclLsaBarrierCreateRequirement申請 barrier,呼叫ncclDevCommCreate分配資源。

2. Kernel 啟動:使用者 kernel 接收 DevComm 和 barrier 句柄作為參數。

3. 通訊階段:kernel 內呼叫ncclLsaBarrier同步所有 rank,然後各 rank 交換資料(透過對稱記憶體直接讀寫)。

4. 計算階段:同步完成後,kernel 直接對本地資料做 ReLU,不需要額外的 kernel 啟動。

5. 完成:kernel 退出,host 側無需等待額外的通訊 kernel。

mermaid
flowchart LR
    subgraph old["传统模式:两个 kernel"]
        k1["通信 kernel<br/>AllReduce"] --> sync["隐式全局同步<br/>kernel 边界"]
        sync --> k2["计算 kernel<br/>ReLU"]
    end
    subgraph fused["融合模式:一个 kernel"]
        f1["通信阶段<br/>ncclLsaBarrier + 数据交换"]
        f1 --> f2["计算阶段<br/>ReLU"]
    end
    old -.->|"融合后省掉"| fused

這張對比圖展示了融合的核心收益:省掉 kernel 邊界處的隱式全域同步。在傳統模式下,這個同步的代價是兩次 kernel 啟動的延遲加上 GPU 流水線的排空。

設計思考與踩坑

為什麼裝置側 API 不直接提供「融合 AllReduce」?因為融合的具體形式取決於使用者的計算邏輯。NCCL 提供的是原語(barrier、訊號、對稱記憶體存取),而不是成品(融合的 AllReduce+ReLU)。使用者需要自己組合這些原語,才能實現符合自己需求的融合 kernel。這是「程式設計模型」而非「函式庫」的本質區別。

踩坑點:融合 kernel 的除錯難度遠高於分離 kernel。如果 barrier 邏輯有 bug,可能導致 kernel 掛起(死鎖),而 GPU kernel 掛起不像 host 程序掛起那樣容易診斷。建議在融合 kernel 中加超時機制,或者先用小規模團隊驗證 barrier 邏輯。

踩坑點:融合 kernel 的 occupancy 下降可能導致計算性能損失超過通信節省的收益。在決定融合之前,應當測量融合前後的端到端時間,而不是只看通信延遲的降低。

本章思考與自測

Q1:如果把ncclTeamLsa中 L26 的ncclDevrInitOnce調用去掉,直接返回comm->devrState.lsaSize和lsaSelf,在什麼場景下會導致設備側 kernel 讀到錯誤的團隊信息?

參考解析:ncclDevrInitOnce是設備側資源初始化的冪等入口。如果去掉它,comm->devrState.lsaSize和lsaSelf可能還是初始值(通常是 0 或未定義)。在首次使用設備側 API 的場景下,用戶調用ncclTeamLsa會拿到nRanks = 0的空團隊。後續如果用戶不檢查團隊有效性,直接用這個團隊調用ncclLsaBarrierCreateRequirement,會算出bufferSize = (3*n + n*0) * 4 = 12n字節——比實際需要的小,因為n*team.nRanks項變成了 0。這會導致緩衝區溢出:barrier 運行時試圖寫入team.nRanks個到達槽位,但緩衝區只分配了3n個uint32_t的空間。更隱蔽的是,如果lsaSelf也是 0,ncclTeamRankToLsa會返回錯誤的 rank 編號,導致 barrier 的到達槽位寫錯位置,可能永遠等不到所有 rank 到達,造成 kernel 掛起。這正是 L23-25 註釋所說的「返回垃圾值,下一個 API 報錯」策略要防止的情況——但前提是下一個 API 確實會報錯,而不是靜默地使用錯誤的大小。

Q2:ncclLsaBarrierCreateRequirement的大小公式是(3*nBarriers + nBarriers*team.nRanks) * sizeof(uint32_t)。如果團隊有 8 個 rank,用戶申請 1 個 barrier,緩衝區是 44 字節。假設 barrier 實現中「3 個控制字段」分別是「到達計數」「輪次」「重置標誌」,請推演:當 8 個 rank 同時到達時,如果「到達計數」用非原子的++操作,會發生什麼?

參考解析:非原子的++在 GPU 上是「讀-改-寫」三步,不是原子操作。8 個 rank 同時執行count++時,可能出現多個 rank 讀到相同的舊值(如都讀到 0),然後都寫回 1。最終count只增加了 1 而不是 8,導致 barrier 永遠認為「還沒到齊」,所有 rank 在輪詢階段死循環。這就是為什麼 LSA barrier 的到達槽位必須用原子操作(如atomicAdd)或每個 rank 寫自己的獨立槽位(nBarriers * team.nRanks項正是為每個 rank 預留獨立槽位)。如果採用「每個 rank 寫自己的槽位」方案,就不需要原子加,只需要原子寫 + 內存屏障,因為每個槽位只有一個寫入者。這也解釋了為什麼大小公式裡有nBarriers * team.nRanks項——它是用空間換原子性,避免多寫者競爭。

Q3:ncclGinBarrierCreateRequirement需要comm參數而ncclLsaBarrierCreateRequirement不需要。如果強行給 LSA barrier 也加上comm參數(假設為了統一接口),會引入什麼設計問題?反過來,如果給 GIN barrier 去掉comm參數,在什麼場景下會失敗?

參考解析:給 LSA barrier 加comm參數的問題是引入了不必要的依賴。LSA barrier 的資源(共享內存)已經在ncclDevrInitOnce階段綁定到團隊上,team本身就隱含了資源位置。加comm會讓一個純團隊操作依賴通信域狀態,增加失敗點(如comm無效時 LSA barrier 也無法創建),且違反「最小權限」原則。反過來,給 GIN barrier 去掉comm參數會失敗,因為 GIN 信號需要綁定到具體的網絡連接。ncclGinBarrierCreateRequirement的ginSignalCount需要知道往哪個網卡、哪個 QP(Queue Pair)發送信號,這些信息在comm的網絡傳輸層狀態裡。沒有comm,NCCL 無法確定信號應該分配到哪個網卡的槽位,也無法保證信號能正確路由到目標 rank。這體現了設備側 API 的一個設計原則:資源需求聲明只依賴它真正需要的上下文——LSA 只需要團隊拓撲,GIN 需要網絡連接。

---

設備側 API 和內核融合把 NCCL 從「一個你調用的庫」變成了「一套你編程的模型」。ncclTeam_t提供了坐標系,CreateRequirement提供了資源預訂機制,三種 barrier 覆蓋了從共享內存到網絡信號的全部同步範圍。但聲明了資源、寫好了融合 kernel,並不等於性能就好——barrier 的數量、團隊的大小、融合的粒度,每一個選擇都會影響端到端性能。下一章我們將進入性能調優實戰,看看 tuning 參數如何影響算法選擇,以及如何用真實 benchmark 驗證調優效果。

至此,我們已經走完了從 devcomm 元資料映射到 nccl_device 裝置側原語的全過程,看到了 NCCL 如何透過「host 宣告、device 消費」的模型,讓使用者 kernel 直接呼叫 barrier 類同步操作,把通訊與計算融合進同一個 kernel。但掌握了這些機制之後,一個更實際的問題自然浮現:當真實訓練任務效能不達標時,我們該如何判斷是演算法選擇不當、協定不匹配,還是通道數配置不合理?下一章將把前 20 章的機制串成一套可操作的調優方法論,結合效能報告、代價模型與環境變數,給出從現象到根因的排查路徑。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 21

第 21 章:效能調優實戰:tuning 實操、benchmark 工具與調優方法論

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 21 章 / 共 25 章

上一章我們看到,使用者自訂 kernel 如何透過裝置側 API 與 NCCL 通訊原語協作,甚至將通訊與計算融合進同一個 kernel。這打開了 NCCL 作為程式設計模型的可能性,但也帶來一個現實問題:當通訊效能不如預期時,該從哪裡入手?NCCL 暴露了上百個 NCCL_PARAM,但真正決定一次集合通訊走哪條路的,其實只有三個旋鈕:演算法(Algo)、協定(Proto)、通道數(nChannels)。本章把前 20 章的機制串成一條可操作的排查路徑——先看效能報告定位現象,再讀代價模型理解 NCCL 自己怎麼選,最後用環境變數和 benchmark 驗證你的假設。

21.1 效能報告:先建立「正常」的基準線

調優的第一步不是改參數,而是知道「正常」長什麼樣。如果你連當前系統的峰值頻寬是多少都不清楚,任何調參都是盲猜。

NCCL 官方在docs/perf下發布參考效能資料,它的定位非常明確——不是產品級保證,而是對齊預期的參照點。

📎 docs/perf/README.md:3-14

code
NCCL publishes reference performance data to:

1. Provide reference points that help users align performance expectations.
2. Help users validate their system setup.
3. Reduce repeated requests to the NCCL team for basic performance numbers.

These results are references, and NOT product-level guarantees that the same
performance is achievable on every system. Performance depends on a complex
combination of software versions, system configuration, hardware, and operating
conditions, including factors outside NCCL's control. A difference within 5% is
generally considered acceptable variance due to differences in the underlying
systems.

這裡有兩個關鍵資訊,小白容易忽略:

第一,5% 以內的差異屬於正常波動。這意味著你測出比官方低 3% 時,不要急著調參——先確認是不是量測雜訊、GPU 時脈抖動、或者鄰居任務干擾。

第二,官方只發布峰值頻寬,不發布延遲。

📎 docs/perf/README.md:24-24

code
We publish peak bandwidth for a selection of commonly used platforms. We do not
currently publish latency because it is typically more sensitive to factors
outside NCCL's control.
〔設計推斷與架構權衡〕

為什麼延遲不發布?因為延遲對系統狀態極度敏感——CPU 頻率、PCIe 鏈路狀態、網卡韌體版本、甚至 BIOS 的電源策略都會影響它。頻寬在大訊息下趨於飽和,相對穩定;延遲在小訊息下由無數個微小環節疊加而成,任何一環抖動都會放大。所以調優時,大訊息看頻寬,小訊息看延遲,這是兩條不同的排查路徑。

📎 docs/perf/README.md:24-24

code
If your workload differs significantly from the published results, open an
issue in the [NCCL repository](https://github.com/NVIDIA/nccl/issues) or contact
NVIDIA Support. We will try our best to help.

排查順序的第一條:先跑一個標準 benchmark(如nccl-tests的all_reduce_perf),把結果和官方報告對比。如果差距在 5% 以內,說明系統配置沒問題,效能瓶頸在你的應用層(比如通訊頻率、訊息切分方式);如果差距顯著,才進入 NCCL 參數調優。

21.2 代價模型:NCCL 自己怎麼選演算法和協定

要調參,先得理解 NCCL 預設是怎麼選的。它內部有一套「代價模型」(cost model),本質是一張查表 + 公式計算:給定訊息大小、拓撲類型、rank 數,估算每種「演算法 × 協定」組合的耗時,選最小的那個。

直覺模型

把代價模型想像成導航軟體。你輸入起點終點(訊息大小、拓撲),它內部對每條路線(演算法/協定組合)估算時間,然後推薦最快的那條。導航的估算基於歷史資料和道路等級,NCCL 的估算基於一張硬編碼的延遲/頻寬參數表。

如果沒有這個模型,NCCL 就只能對所有場景用同一個固定演算法——小訊息會因啟動開銷過大而變慢,大訊息會因頻寬利用不足而變慢,系統會在兩個極端都表現糟糕。

資料結構:模型表與調優上下文

代價模型的核心是modelMap陣列,每個元素對應一種「演算法/協定/對稱核心」組合。

📎 src/tuning/cost_model.cc:230-277

code
static struct ncclTuningModelEntry_t modelMap[] = {
    /*
Initialize default, static models here
{mod_init, mod_sim, mod_final, enabled}
Enable order: Broadcast, Reduce, AllGather, ReduceScatter, AllReduce
*/
  {ncclTuningTreeModelInit, ncclTuningTreeModelSim, nullptr, {0, 0, 0, 0, 1}},       // Tree/LL
  {ncclTuningTreeModelInit, ncclTuningTreeModelSim, nullptr, {0, 0, 0, 0, 1}},       // Tree/LL128
  {ncclTuningTreeModelInit, ncclTuningTreeModelSim, nullptr, {0, 0, 0, 0, 1}},       // Tree/Simple
  {ncclTuningRingModelInit, ncclTuningRingModelSim, nullptr, {1, 1, 1, 1, 1}},       // Ring/LL
  ...

每個條目有四個欄位:mod_init(初始化函式)、mod_sim(模擬函式)、mod_final(清理函式)、enabled(5 個函式各自的啟用標誌)。enabled陣列的順序是{Broadcast, Reduce, AllGather, ReduceScatter, AllReduce}——注意這個順序,後面讀程式碼時會反覆用到。

〔設計推斷與架構權衡〕

關鍵觀察:Tree 只在 AllReduce 上啟用({0,0,0,0,1}),而 Ring 在所有函式上都啟用({1,1,1,1,1})。這是因為 Tree 演算法的優勢在於 AllReduce 的規約階段可以平行,但對 AllGather/ReduceScatter 這類本質是環形流水的操作,Ring 更自然。

模型的具體參數存在ncclTunerConstants_t裡,包含各拓撲下的基礎延遲和頻寬。

📎 src/tuning/cost_model.cc:142-152

code
static const ncclTunerConstants_t ncclTunerConstantsDefaults = {
    // baseLatencies
  {
    {6.8, 14.0, 8.4},  // Tree
    {6.6, 14.0, 8.4},  // Ring
    {0, 0, 0},         // Collnet Direct
    {0, 0, 0},         // Collnet Chain
    {0, 0, 0},         // NVLS
    {0, 0, 0},         // NVLS Tree
    {8.0, 8.0, 8.0}    // PAT
  },

每個演算法有三個基礎延遲值,對應 LL / LL128 / Simple 三種協定。比如 Ring 的{6.6, 14.0, 8.4}意味著:LL 協定基礎延遲 6.6 微秒,LL128 是 14.0,Simple 是 8.4。這些數字是 NVIDIA 在真實硬體上測出來的經驗值。

硬體延遲則按拓撲類型(NVLink / PCI / NET)分別給出。

📎 src/tuning/cost_model.cc:153-184

code
    // hwLatencies
  {
    /* NVLINK */
    {
      {0.6, 1.25, 4.0}, // Tree (LL/LL128/Simple)
      {0.6, 1.9, 3.4},  // Ring (LL/LL128/Simple)
      ...
    },
    /* PCI */
    {
      {1.0, 1.9, 4.0}, // Tree (LL/LL128/Simple)
      {1.0, 2.5, 5.7}, // Ring (LL/LL128/Simple)
      ...
    },
    /* NET */
    {
      {5.0, 8.5, 14},   // Tree (LL/LL128/Simple)
      {2.7, 4.0, 14.0}, // Ring (LL/LL128/Simple)
      ...
    },
  },

對比一下就能看出拓撲差異:NVLink 上 Ring/Simple 的每跳延遲是 3.4 微秒,PCI 上是 5.7,NET 上是 14.0。這就是為什麼跨機通訊慢——每一跳都要多花 10 微秒。

頻寬參數按 GPU 架構分代給出。

📎 src/tuning/cost_model.cc:183-183

code
    // llMaxBws
  {
    {39.0, 39.0, 20.4}, /* Volta-N1/Intel-N2/Intel-N4) */
    {87.7, 22.5 /*avg of ring & tree*/, 19.0}, /* Ampere-N1/AMD-N2/AMD-N4) */
    {141.0, 45.0 /*avg of ring & tree*/, 35.0}, /* Hopper-N1/AMD-N2/AMD-N4) */
    {2 * 141.2, 2 * 45.0 /*avg of ring & tree*/, 2 * 35.0}, /* Blackwell-N1/AMD-N2/AMD-N4) */
  },

每行對應一代架構,三個值分別是單機(N1)、雙機(N2)、四機(N4)場景下的 LL 協定最大頻寬。Hopper 單機 141 GB/s,Blackwell 翻倍到 282 GB/s——這解釋了為什麼新卡上同樣的演算法表現會好很多。

調優上下文:per-comm 的狀態

每個通訊域(communicator)持有一份ncclTuningContext_t,保存這個 comm 的調優狀態。

📎 src/include/tuning.h:81-95

code
struct ncclTuningContext_t {
  // Persistant tuning parameters tied to a communicator.
  ncclTunerConstants_t tuningConstants;
  // State of the tuning models
  // Forced function is set via env var
  int forced[NCCL_NUM_FUNCTIONS];
  // Disabled tuning models are not execute and excluded from implemetation selection.
  int enabled[NCCL_TUNING_COUNT][NCCL_NUM_FUNCTIONS];
  // Store of model contexts per communicator.
  float generalLatencies[NCCL_NUM_FUNCTIONS][NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
  float generalBandwidths[NCCL_NUM_FUNCTIONS][NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];

  ssize_t threadThresholds[NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
  int maxThreads[NCCL_NUM_ALGORITHMS][NCCL_NUM_PROTOCOLS];
};

四個關鍵欄位:

  • forced[NCCL_NUM_FUNCTIONS]:標記哪些函數被環境變數強制指定了演算法/協定。這是NCCL_ALGO/NCCL_PROTO生效的落點。
  • enabled[NCCL_TUNING_COUNT][NCCL_NUM_FUNCTIONS]:二維布林表,標記某個模型對某個函數是否啟用。被禁用的模型不參與選擇。
  • generalLatencies / generalBandwidths:三維陣列,按「函數 × 演算法 × 協定」儲存估算的延遲和頻寬。這是ncclTuningInit列印那張大表的來源。
  • threadThresholds / maxThreads:執行緒數相關的閾值,決定每個 block 用多少執行緒。

場景驅動 Walkthrough:一次 AllReduce 的演算法選擇

假設你呼叫ncclAllReduce,訊息大小 1MB,8 卡單機 NVLink。NCCL 內部會建構一個ncclTuningInput_t,然後呼叫ncclTuningCompute。

📎 src/tuning/tuning.cc:180-202

code
ncclResult_t ncclTuningCompute(struct ncclTuningInput_t* const input, struct ncclTuningResult_t* const result) {
  ncclResult_t ret = ncclSuccess;
  TRACE(NCCL_TUNING, ...);
  struct ncclTuningResultList_t tunings;
  tunings.head = nullptr;
  struct ncclTuningResult_t bestTuning = NCCL_TUNING_RESULT_INIT;
  // Set tuning to Ring/Simple for single rank case
  if (input->comm->nRanks <= 1) {
    bestTuning.algo = NCCL_ALGO_RING;
    bestTuning.proto = NCCL_PROTO_SIMPLE;
    ...
  } else {
    NCCLCHECKGOTO(ncclTuningComputeAllTunings(input, &tunings), ret, exit);

第一步:單 rank 直接返回 Ring/Simple,不做任何計算。這是短路優化——單卡沒有通訊,選什麼演算法都一樣。

第二步:多 rank 時呼叫ncclTuningComputeAllTunings,遍歷所有候選組合。

📎 src/tuning/tuning.cc:128-149

code
ncclResult_t ncclTuningComputeAllTunings(struct ncclTuningInput_t* const input,
                                         struct ncclTuningResultList_t* const tunings) {
  ncclResult_t ret = ncclSuccess;

  for (int i = 0; i < NCCL_TUNING_COUNT; i++) {
    struct ncclTuningResult_t tuning = NCCL_TUNING_RESULT_INIT;
    tuning.id = i;
    tuning.valid = 1;

    if (!(input->tuningMask & (1ULL << i))) {
      tuning.valid = 0;
      continue;
    }
    NCCLCHECK(ncclTuningExpandId(i, &tuning.algo, &tuning.proto, &tuning.symKernelId, &tuning.ceMethodId));
    NCCLCHECKGOTO(ncclTuningComputeTuning(i, input, &tuning), ret, fail);
    if (tuning.valid) NCCLCHECKGOTO(ncclTuningResultListPushFront(tunings, tuning), ret, fail);
  }

這裡有個精妙的設計:tuningMask是一個 64 位元遮罩,每一位對應一個候選組合。NCCL_TUNING_MASK_GENERAL_KERNELS、NCCL_TUNING_MASK_SYM_KERNELS、NCCL_TUNING_MASK_CE分別圈定不同類別的候選。

📎 src/include/tuning.h:17-25

code
#define NCCL_TUNING_SYM_KERNEL_ID_OFFSET (NCCL_NUM_ALGORITHMS * NCCL_NUM_PROTOCOLS)
#define NCCL_TUNING_CE_METHOD_ID_OFFSET (NCCL_TUNING_SYM_KERNEL_ID_OFFSET + ncclSymkKernelId_Count)
#define NCCL_TUNING_COUNT (NCCL_TUNING_CE_METHOD_ID_OFFSET + ncclCeMethodId_Count)

#define NCCL_TUNING_MASK_GENERAL_KERNELS ((1ULL << NCCL_TUNING_SYM_KERNEL_ID_OFFSET) - 1ULL)
#define NCCL_TUNING_MASK_SYM_KERNELS \
  ((1ULL << NCCL_TUNING_CE_METHOD_ID_OFFSET) - 1ULL - NCCL_TUNING_MASK_GENERAL_KERNELS)
#define NCCL_TUNING_MASK_CE ((1ULL << NCCL_TUNING_COUNT) - (1ULL << NCCL_TUNING_CE_METHOD_ID_OFFSET))
#define NCCL_TUNING_MASK_ALL ((1ULL << NCCL_TUNING_COUNT) - 1ULL)

遮罩的佈局是:低NCCL_NUM_ALGORITHMS × NCCL_NUM_PROTOCOLS位是傳統「演算法×協定」組合,中間ncclSymkKernelId_Count位是對稱核心,高位是 CE(Copy Engine)方法。用位元遮罩而不是陣列,是為了在ncclTuningCompute裡快速判斷「這個候選是否在本次調優範圍內」。

第三步:對每個候選呼叫ncclTuningComputeTuning,它轉調ncclTuningCostModelSimModel。

📎 src/tuning/cost_model.cc:470-497

code
ncclResult_t ncclTuningCostModelSimModel(int id, struct ncclTuningInput_t* const input,
                                         struct ncclTuningResult_t* const result) {
  struct ncclTuningModelEntry_t* model = nullptr;
  ncclResult_t ret = ncclSuccess;
  result->forced = input->comm->tuningContext.forced[input->func];
  NCCLCHECKGOTO(getModelEntry(id, &model), ret, not_valid);
  if (model == nullptr) {
    ret = ncclInternalError;
    goto not_valid;
  }
  if (input->comm->tuningContext.enabled[id][input->func] == 0) {
    goto not_valid;
  }
  if (model->model != nullptr) {
    NCCLCHECKGOTO(model->model(input, result), ret, not_valid);
    if (result->timeUs <= 0.0) {
      goto not_valid;
    }
  } else {
    goto not_valid;
  }
exit:
  return ret;
not_valid:
  result->timeUs = NCCL_TUNING_IGNORE;
  result->valid = 0;
  goto exit;
}

注意not_valid標籤的處理:任何一步失敗(模型不存在、被禁用、模擬返回非正時間),都會把timeUs設為NCCL_TUNING_IGNORE、valid設為 0。這個候選就被排除在後續選擇之外。

第四步:從所有有效候選中選耗時最小的。

📎 src/tuning/tuning.cc:155-173

code
static ncclResult_t ncclTuningSelectBestTuning(struct ncclTuningResultList_t* tunings,
                                               struct ncclTuningResult_t* const bestTuning) {
  bestTuning->timeUs = FLT_MAX;
  float bestSelectionTimeUs = FLT_MAX;
  struct ncclTuningResultListNode* node = tunings->head;
  while (node != nullptr) {
    const struct ncclTuningResult_t& tuning = node->result;
    float selectionTimeUs = tuning.selectionTimeUs > 0.0f ? tuning.selectionTimeUs : tuning.timeUs;
    TRACE(NCCL_TUNING, "A/P/S %s/%s/%s, time: %f, selection time: %f", ...);
    if (selectionTimeUs < bestSelectionTimeUs) {
      *bestTuning = tuning;
      bestSelectionTimeUs = selectionTimeUs;
    }
    node = node->next;
  }
  return ncclSuccess;
}

這裡有個細節:選擇用的是selectionTimeUs,如果它大於 0 就用它,否則回退到timeUs。selectionTimeUs是「選擇時間」,可能包含了額外的懲罰項(比如某些演算法在特定場景下要額外開銷)。這給了代價模型一個「估算時間」和「選擇時間」分離的能力。

流程圖

mermaid
flowchart TD
    start["ncclTuningCompute(input, result)"] --> check_ranks{"comm->nRanks <= 1?"}
    check_ranks -->|是| single["bestTuning = Ring/Simple<br/>nChannels = 0"]
    check_ranks -->|否| all["ncclTuningComputeAllTunings()"]
    all --> loop{"遍历 i in NCCL_TUNING_COUNT"}
    loop -->|mask 未命中| skip["tuning.valid = 0<br/>continue"]
    loop -->|mask 命中| expand["ncclTuningExpandId(i, ...)"]
    expand --> sim["ncclTuningComputeTuning()<br/>→ ncclTuningCostModelSimModel()"]
    sim --> sim_check{"enabled[id][func] != 0<br/>且 model->model != nullptr?"}
    sim_check -->|否| invalid["timeUs = NCCL_TUNING_IGNORE<br/>valid = 0"]
    sim_check -->|是| push["ncclTuningResultListPushFront()"]
    skip --> loop
    invalid --> loop
    push --> loop
    loop -->|遍历结束| tuner_check{"comm->tuner != NULL?"}
    tuner_check -->|是| plugin["tuner->getCollInfo()<br/>覆盖 generalTable"]
    tuner_check -->|否| select["ncclTuningSelectBestTuning()"]
    plugin --> select
    select --> channels["ncclTuningGetChannels()"]
    channels --> eff{"CTAPolicy & EFFICIENCY<br/>且 NCCL_ALGO/NCCL_PROTO 未设置?"}
    eff -->|是| nvls["尝试 NVLS 覆盖<br/>ncclNvlsRegResourcesQuery()"]
    eff -->|否| done["*result = bestTuning"]
    nvls --> done
    single --> done

這張圖完整畫出了從入口到最終結果的決策路徑,包括單 rank 短路、遮罩過濾、模型禁用、tuner 外掛介入、CTAPolicy 覆蓋等所有分支。

21.3 環境變數:真正影響效能的三個旋鈕

理解了代價模型,就知道環境變數是怎麼介入的。NCCL_ALGO、NCCL_PROTO、NCCL_SYM_KERNEL這三個變數透過parseList解析後,直接修改enabled表,把不符合使用者意圖的候選全部禁用。

解析語法

parseList支援的語法比大多數人想像的複雜。

📎 src/tuning/cost_model.cc:14-32

code
// Parse a map of prefixes to a list of elements. The first prefix is
// optional and, if not present, the list of elements will be applied
// to all prefixes. Only the first list of elements can lack a
// prefix. Prefixes (if present) are followed by a colon. Lists of
// elements are comma delimited. Mappings of prefix to the lists of
// elements are semi-colon delimited.
//
// For example:
//
//     NCCL_ALGO="ring,collnetdirect;allreduce:tree,collnetdirect;broadcast:ring"
// Enable ring and collnetdirect for all functions, then select tree
// and collnetdirect for allreduce and ring for broadcast.
//
//     NCCL_PROTO="LL,Simple;allreduce:^LL"
// Enable LL and Simple for all functions, but everything except LL
// for allreduce.
//
//     NCCL_PROTO="^LL128;allreduce:LL128"
// Enable everything but LL128, but only LL128 for allreduce.

三種用法:

1. 全域列表:NCCL_ALGO="ring,tree"—— 所有函數只用 ring 和 tree。

2. 按函數前綴:NCCL_ALGO="ring;allreduce:tree"—— 預設 ring,但 allreduce 用 tree。

3. 排除語法:NCCL_PROTO="^LL128"—— 除了 LL128 其他都啟用。

^前綴是關鍵——它表示「unset」,即從預設全啟用中排除某個選項。

📎 src/tuning/cost_model.cc:59-67

code
    int unset, set;
    if (elemList[0] == '^') {
      unset = 1;
      set = 0;
      elemList++;
    } else {
      unset = 0;
      set = 1;
    }

解析到^時,unset=1、set=0。隨後對匹配的 prefix,先把整個列表填成unset(全排除),再把列出的元素設為set。

📎 src/tuning/cost_model.cc:69-96

code
    bool foundPrefix = false;
    for (int p = 0; p < nprefixes; p++) {
      if (prefix && strcasecmp(prefix, prefixElems[p]) != 0) continue;
      foundPrefix = true;
      for (int e = 0; e < nelems; e++) list[p * nelems + e] = unset;

      tokStr = strdup(elemList);
      char* tmpStr;
      char* elem = strtok_r(tokStr, ",", &tmpStr);
      while (elem) {
        int e;
        for (e = 0; e < nelems; e++) {
          if (strcasecmp(elem, elems[e]) == 0) {
            list[p * nelems + e] = set;
            forced[p] = 1;
            break;
          }
        }
        if (e == nelems) {
          WARN("Unrecognized element token \"%s\" when parsing \"%s\"", elem, str);
          ret = ncclInvalidUsage;
          goto fail;
        }
        elem = strtok_r(NULL, ",", &tmpStr);
      }

注意forced[p] = 1這一行——只要使用者顯式列了某個元素,對應的函數就被標記為「強制」。這個標記後面會用來判斷是否允許代價模型自由選擇。

強制與禁用的互動

ncclTuningCostModelInit裡有一段關鍵邏輯,處理使用者強制與環境變數、平台能力的互動。

📎 src/tuning/cost_model.cc:363-384

code
    for (int f = 0; f < NCCL_NUM_FUNCTIONS; f++) {
      // Disable LL128 when 1) it is not supported on the platform, and 2) user did not explicitly request it.
      // protoEnable[..] == 2 indicates that user did not set NCCL_PROTO=LL128 explicitly.
      if (proto == NCCL_PROTO_LL128 && protoEnable[f * NCCL_NUM_PROTOCOLS + proto] == 2 &&
          !isLL128Enabled(comm->minCompCap, comm->maxCompCap, comm->graphs[algo].typeInter,
                          comm->graphs[algo].typeIntra, comm->nRanks, f, algo, comm->minDriverVersion)) {
        comm->tuningContext.enabled[i][f] = 0;
      }
      //  Check the user env vars only for functions that have a forced configuration and not already disabled.
      if (comm->tuningContext.forced[f] == 0 || comm->tuningContext.enabled[i][f] == 0) continue;
      comm->tuningContext.enabled[i][f] = 0;
      TRACE(NCCL_TUNING, "a/p/s %s/%s/%s enabled %d/%d/%d", ...);
      if (((algo != NCCL_ALGO_UNDEF && algoEnable[f * NCCL_NUM_ALGORITHMS + algo] != 0) &&
           (proto != NCCL_PROTO_UNDEF && protoEnable[f * NCCL_NUM_PROTOCOLS + proto] != 0)) ||
          (symKernelId != ncclSymkKernelId_Count && symKernelIdEnable[f * ncclSymkKernelId_Count + symKernelId] != 0)) {
        comm->tuningContext.enabled[i][f] = 1;
      }
    }

這段邏輯的順序很重要:

1. 先處理 LL128 平台能力:如果平台不支援 LL128(isLL128Enabled返回 0)且使用者沒顯式要求(protoEnable == 2),直接禁用。

2. 再處理使用者強制:如果這個函數被強制了(forced[f] != 0),先把它禁用(enabled[i][f] = 0),然後檢查使用者是否允許這個組合——允許就重新啟用。

protoEnable的值有三種:0(使用者排除)、1(使用者啟用)、2(使用者未提及,預設啟用)。這個三態設計讓「使用者顯式要求」和「平台預設」能區分開。

環境變數讀取的快取機制

所有NCCL_PARAM巨集最終都走ncclLoadParam。

📎 src/misc/param.cc:78-108

code
int64_t ncclLoadParam(char const* env, int64_t deftVal, int64_t uninitialized, int64_t* cache, int8_t* noCache) {
  static std::mutex mutex;
  std::lock_guard<std::mutex> lock(mutex);

  // noCache is only load/stored within the mutex, no need for atomic
  if (*noCache == /*uninitialized*/ -1) ncclGetCachePolicy(env, noCache);

  if (COMPILER_ATOMIC_LOAD(cache, std::memory_order_relaxed) != uninitialized) {
    return COMPILER_ATOMIC_LOAD(cache, std::memory_order_relaxed);
  }

  // Read the environment variable
  const char* str = ncclGetEnv(env);
  int64_t value = deftVal;

  if (str && strlen(str) > 0) {
    errno = 0;
    char* end = nullptr;
    value = strtoll(str, &end, 0);
    // Preserve numeric-prefix parsing while rejecting non-numeric values.
    if (errno || end == str) {
      value = deftVal;
      ATTN("Invalid value %s for %s, using default %lld.", str, env, (long long)deftVal);
    } else {
      INFO(NCCL_ENV, "%s set by environment to %lld.", env, (long long)value);
    }
  }

  if (*noCache == /*cache*/ 0) COMPILER_ATOMIC_STORE(cache, value, std::memory_order_relaxed);
  return value;
}

這段程式碼有幾個值得注意的設計:

全域互斥鎖:static std::mutex mutex保護整個讀取過程。這意味著所有參數的首次讀取是串行的。為什麼用鎖而不是無鎖?因為參數讀取只在初始化階段發生,不在熱路徑上,鎖的開銷可以忽略,而正確性更重要。

雙重檢查:先原子讀cache,如果已初始化就直接返回。這避免了每次讀參數都進鎖——雖然鎖本身在初始化後幾乎不競爭,但原子讀更快。

快取策略:noCache標誌決定是否把讀到的值寫回cache。某些參數(如需要動態回應的)可能停用快取,每次都重新讀環境變數。

錯誤處理:strtoll解析失敗時用預設值,並列印ATTN警告。注意end == str的判斷——如果字串開頭就不是數字,end會等於str,說明完全沒解析出數字。

設定檔支援

環境變數不一定要從 shell 設定,NCCL 支援從設定檔讀取。

📎 src/misc/param.cc:52-67

code
static void initEnvFunc() {
  char confFilePath[1024];
  const char* userFile = std::getenv("NCCL_CONF_FILE");
  if (userFile && strlen(userFile) > 0) {
    snprintf(confFilePath, sizeof(confFilePath), "%s", userFile);
    setEnvFile(confFilePath);
  } else {
    const char* userDir = userHomeDir();
    if (userDir) {
      snprintf(confFilePath, sizeof(confFilePath), "%s/.nccl.conf", userDir);
      setEnvFile(confFilePath);
    }
  }
  snprintf(confFilePath, sizeof(confFilePath), "/etc/nccl.conf");
  setEnvFile(confFilePath);
}

載入順序:NCCL_CONF_FILE指定的檔案(如果設定了)→~/.nccl.conf → /etc/nccl.conf。後載入的會覆蓋先載入的(因為setEnvFile呼叫ncclOsSetEnv)。

📎 src/misc/param.cc:69-72

code
void initEnv() {
  static std::once_flag once;
  std::call_once(once, initEnvFunc);
}

std::call_once保證設定檔只載入一次,即使多個執行緒同時首次呼叫ncclGetEnv。

21.4 通道數:被低估的效能旋鈕

演算法和協定決定「怎麼走」,通道數決定「開幾條路」。很多人調優時只關注前兩個,忽略了通道數——但在大訊息場景下,通道數往往是決定頻寬利用率的關鍵。

通道數從哪來

ncclTuningCompute在選出最佳演算法/協定後,會呼叫ncclTuningGetChannels計算通道數。

📎 src/tuning/tuning.cc:233-235

code
  if (bestTuning.algo != NCCL_ALGO_UNDEF && bestTuning.proto != NCCL_PROTO_UNDEF) {
    NCCLCHECKGOTO(ncclTuningGetChannels(input, &bestTuning), ret, exit);
  }

通道數的計算邏輯不在本章原始碼材料中,但可以從ncclTuningResult_t的欄位看出它的作用。

📎 src/include/tuning.h:42-55

code
struct ncclTuningResult_t {
  int id;
  int valid;
  float timeUs;
  float selectionTimeUs;
  int algo;
  int proto;
  int symKernelId;
  int ceMethodId;
  int nChannels;
  int maxChannels;
  int nWarps;
  int forced;
};

nChannels是最終使用的通道數,maxChannels是上限。nWarps是每個 block 的 warp 數。

CTAPolicy 對通道數的覆蓋

有一段特殊邏輯處理NCCL_CTA_POLICY_EFFICIENCY策略。

📎 src/tuning/tuning.cc:236-257

code
  // NCCL_CTA_POLICY_EFFICIENCY requires user (non-symmetric) buffer registration (currently unsupported with MNNVL).
  // Run after GetChannels so bestTuning.nChannels is valid. Skip when a tuner plugin owns selection
  // (same as pre-rearch). The NVLS-bit guard keeps this bias inside the candidate set: a per-call
  // algSelection may have narrowed tuningMask, so EFFICIENCY must not resurrect NVLS when excluded.
  if (input->comm->tuner == NULL && (input->CTAPolicy & NCCL_CTA_POLICY_EFFICIENCY) &&
      ncclGetEnv("NCCL_ALGO") == NULL && ncclGetEnv("NCCL_PROTO") == NULL && !input->comm->MNNVL &&
      (input->tuningMask & (1ull << (NCCL_ALGO_NVLS * NCCL_NUM_PROTOCOLS + NCCL_PROTO_SIMPLE)))) {
    if (input->regBuff && (input->func == ncclFuncAllGather || input->func == ncclFuncReduceScatter)) {
      if ((input->comm->nNodes > 1 && input->collNetSupport && input->nvlsSupport) ||
          (input->comm->nNodes == 1 && input->nvlsSupport)) {
        int recChannels;
        NCCLCHECKGOTO(ncclNvlsRegResourcesQuery(input->comm, input->func, &recChannels), ret, exit);
        if (recChannels <= bestTuning.nChannels) {
          bestTuning.algo = NCCL_ALGO_NVLS;
          bestTuning.proto = NCCL_PROTO_SIMPLE;
          bestTuning.nChannels = recChannels;
          bestTuning.maxChannels = recChannels;
          bestTuning.nWarps = input->comm->tuningContext.maxThreads[bestTuning.algo][bestTuning.proto] / WARP_SIZE;
        }
      }
    }
  }

這段程式碼的守衛條件非常密集,值得逐條解讀:

1. input->comm->tuner == NULL:沒有 tuner 外掛時才走這段。外掛擁有選擇權時,NCCL 不干預。

2. input->CTAPolicy & NCCL_CTA_POLICY_EFFICIENCY:使用者設定了效率優先策略。

3. ncclGetEnv("NCCL_ALGO") == NULL && ncclGetEnv("NCCL_PROTO") == NULL:使用者沒有強制演算法/協定。如果強制了,尊重使用者選擇。

4. !input->comm->MNNVL:MNNVL 場景不支援。

5. input->tuningMask & (1ull << (NCCL_ALGO_NVLS * NCCL_NUM_PROTOCOLS + NCCL_PROTO_SIMPLE)):NVLS/Simple 在候選集內。這個守衛防止「復活」被排除的選項。

滿足條件後,查詢 NVLS 註冊資源能支援的通道數,如果不超過當前選擇,就切換到 NVLS 演算法。

〔設計推斷與架構權衡〕

為什麼 EFFICIENCY 策略偏向 NVLS?因為 NVLS(NVLink SHARP)利用交換器硬體做規約,能減少 GPU 的計算和通訊開銷,在 AllGather/ReduceScatter 這類操作上效率更高。但它的通道數受限於硬體資源,所以需要ncclNvlsRegResourcesQuery查詢實際可用量。

對稱核心的回退邏輯

對稱核心(symmetric kernel)是較新的特性,當它不可用時需要回退到通用核心。

📎 src/tuning/tuning.cc:258-298

code
  if ((bestTuning.symKernelId != ncclSymkKernelId_Count ||
       (input->tuningMask & NCCL_TUNING_MASK_SYM_KERNELS && bestTuning.symKernelId == ncclSymkKernelId_Count)) &&
      bestTuning.algo == NCCL_ALGO_UNDEF && bestTuning.proto == NCCL_PROTO_UNDEF) {
    bool isLLKernel = (1 << bestTuning.symKernelId) & ncclSymkLLKernelMask();
    bool isOneThreadMultiGpus = input->comm->intraRanks > 1 && !ncclParamSingleProcMemRegEnable();
    bool needFallback = bestTuning.symKernelId != ncclSymkKernelId_Count ? false : true;

    // General kernel tuning structs if fallback is needed
    struct ncclTuningResult_t generalTuning = NCCL_TUNING_RESULT_INIT;
    struct ncclTuningInput_t generalInput = *input;
    generalInput.tuningMask = NCCL_TUNING_MASK_GENERAL_KERNELS;

    // Fallback logic for symmetric LL kernels:
    // - If both src and dst are registered, we don't fall back if a symmetric kernel is available.
    // - Otherwise, we have to fall back to generl kernel if running the selected symmetric LL kernel is
    //   not possible (if the buffers are not registered and we manage multiple GPUs).
    // - If the user forced a symmetric kernel via NCCL_SYM_KERNEL or requested preference for using
    //   symmetric kernels even without symmetric buffers via NCCL_SYM_NOWIN_ENABLE, we respect that.
    // - Otherwise, we query the general cost model and if it selects a non-LL proto, we pick that.
    if (bestTuning.symKernelId != ncclSymkKernelId_Count) {
      if (input->winRegType == ncclSymSendRegRecvReg) {
        needFallback = false;
      } else if (isLLKernel) {
        needFallback = isOneThreadMultiGpus && input->winRegType == ncclSymSendNonregRecvNonreg;
        if (!needFallback && !result->forced) {
          needFallback = !ncclParamSymNoWinEnable() && input->winRegType == ncclSymSendNonregRecvNonreg;
          if (!needFallback) {
            NOWARN(ncclTuningCompute(&generalInput, &generalTuning), NCCL_TUNING);
            needFallback = (generalTuning.proto != NCCL_PROTO_LL);
          }
        }
      }
    }

回退決策樹:

  • 如果傳送和接收緩衝區都註冊了(ncclSymSendRegRecvReg),不回退。
  • 如果是 LL 核心且單執行緒管理多 GPU 且緩衝區未註冊,回退。
  • 如果使用者沒設定NCCL_SYM_NOWIN_ENABLE且緩衝區未註冊,回退。
  • 否則,查詢通用代價模型,如果它選了非 LL 協定,回退。
〔設計推斷與架構權衡〕

這個邏輯的核心是:對稱 LL 核心需要緩衝區註冊才能發揮優勢。未註冊時,LL 核心的優勢(低延遲)可能被額外的位址轉換開銷抵消,所以回退到通用核心更划算。

無可用組合時的錯誤處理

如果所有候選都被排除,NCCL 會報錯並給出診斷資訊。

📎 src/tuning/tuning.cc:308-329

code
  if ((bestTuning.algo == NCCL_ALGO_UNDEF || bestTuning.proto == NCCL_PROTO_UNDEF) &&
      bestTuning.symKernelId == ncclSymkKernelId_Count && bestTuning.ceMethodId == ncclCeMethodId_Count) {
    char ncclAlgoEnvStr[1024] = "";
    char ncclProtoEnvStr[1024] = "";
    char ncclSymKernelIdEnvStr[1024] = "";
    const char* symKernelIdEnv = ncclGetEnv("NCCL_SYM_KERNEL");
    if (symKernelIdEnv) {
      snprintf(ncclSymKernelIdEnvStr, 1023, " NCCL_SYM_KERNEL was set to %s.", symKernelIdEnv);
    }
    const char* algoEnv = ncclGetEnv("NCCL_ALGO");
    if (algoEnv) {
      snprintf(ncclAlgoEnvStr, 1023, " NCCL_ALGO was set to %s.", algoEnv);
    }
    const char* protoEnv = ncclGetEnv("NCCL_PROTO");
    if (protoEnv) {
      snprintf(ncclProtoEnvStr, 1023, " NCCL_PROTO was set to %s.", protoEnv);
    }
    WARN("No algorithm/protocol nor symKernelId available for function %s with datatype %s.%s%s%s",
         ncclFuncToString(input->func), ncclDatatypeToString(input->datatype), ncclAlgoEnvStr, ncclProtoEnvStr,
         ncclSymKernelIdEnvStr);
    ret = (algoEnv || protoEnv || symKernelIdEnv) ? ncclInvalidUsage : ncclInternalError;
  }

錯誤碼的選擇有講究:如果使用者設定了環境變數(algoEnv || protoEnv || symKernelIdEnv),返回ncclInvalidUsage——這是使用者的配置問題;否則返回ncclInternalError——這是 NCCL 內部的問題(所有候選都被意外排除了)。

21.5 生產避坑指南

坑一:環境變數拼寫錯誤導致靜默回退

parseList遇到無法識別的 token 會返回ncclInvalidUsage,但如果你寫的是NCCL_ALGO=RING(大寫),strcasecmp會正確匹配。真正危險的是拼寫錯誤,比如NCCL_ALGO=rnig。

📎 src/tuning/cost_model.cc:87-91

code
        if (e == nelems) {
          WARN("Unrecognized element token \"%s\" when parsing \"%s\"", elem, str);
          ret = ncclInvalidUsage;
          goto fail;
        }

這裡會列印 WARN 並返回錯誤。但如果你沒開NCCL_DEBUG=WARN,可能看不到這條警告。建議:調優時始終設定NCCL_DEBUG=WARN或NCCL_DEBUG=INFO,確保能看到配置解析的結果。

坑二:NCCL_ALGO 和 NCCL_PROTO 的互動

如果你設定NCCL_ALGO=tree但沒設定NCCL_PROTO,NCCL 會在 Tree 演算法下選擇最佳協定。但如果你同時設定NCCL_ALGO=tree和NCCL_PROTO=LL,而 Tree/LL 組合在某些函式上被停用(比如 Tree 只在 AllReduce 啟用),就會觸發「無可用組合」錯誤。

📎 src/tuning/cost_model.cc:379-383

code
      if (((algo != NCCL_ALGO_UNDEF && algoEnable[f * NCCL_NUM_ALGORITHMS + algo] != 0) &&
           (proto != NCCL_PROTO_UNDEF && protoEnable[f * NCCL_NUM_PROTOCOLS + proto] != 0)) ||
          (symKernelId != ncclSymkKernelId_Count && symKernelIdEnable[f * ncclSymkKernelId_Count + symKernelId] != 0)) {
        comm->tuningContext.enabled[i][f] = 1;
      }

只有當演算法和協定同時被允許時,組合才啟用。這是 AND 邏輯,不是 OR。

坑三:LL128 的平台限制

LL128 不是所有平台都支援。isLL128Enabled檢查了計算能力、驅動版本、連接類型。

📎 src/tuning/cost_model.cc:119-139

code
static int isLL128Enabled(int minCompCap, int maxCompCap, int interType, int intraType, int nRanks, int func, int algo,
                          int minDriverVersion) {
  int ret = 1;
  if (ncclParamLl128C2c() && minCompCap >= 90 && (!RUBIN_AND_LATER(minCompCap) || minDriverVersion >= 13030)) {
    // Rubin, Blackwell, and Hopper: Enable LL128 for all P2C and PXN if CUDA supports it.
    ret &= (interType <= PATH_PXN);
  } else {
    // Enable LL128 only up to PXB. Don't enable LL128 over PxN because PxN can encapsulate PxB or P2C links.
    ret &= (interType <= PATH_PXB);
    if (!ncclParamLl128C2c() && minCompCap >= 90)
      INFO(
        NCCL_GRAPH | NCCL_TUNING,
        "Disabling LL128 over all PxN connections (PXB and C2C). This ensures that no C2C link will be used by LL128.");
  }
  ret &= (intraType <= PATH_NVB);
  // Enable LL128 for interoperability between GPUs with different compcap (Hopper and above)
  ret &= (minCompCap == maxCompCap || minCompCap >= 90);
  ret &= !(minCompCap < 70 || (minCompCap == 90 && CUDART_VERSION == 11080 && func == ncclFuncAllReduce &&
                               algo == NCCL_ALGO_RING && nRanks == 2));
  return ret;
}

幾個關鍵限制:

  • minCompCap < 70:Volta 之前的 GPU 不支援 LL128。
  • intraType <= PATH_NVB:機內連接必須是 NVLink 級別。
  • Hopper + CUDA 11.8 + AllReduce + Ring + 2 ranks:這是一個已知的 bug 場景,被顯式排除。

建議:如果你的平台不支援 LL128,不要強行設定NCCL_PROTO=LL128,否則會觸發錯誤。讓 NCCL 自動選擇。

坑四:通道數與顯存

通道數越多,需要的緩衝區越大。在顯存緊張的場景下,過多的通道可能導致 OOM。

📎 src/tuning/tuning.cc:246-253

code
        int recChannels;
        NCCLCHECKGOTO(ncclNvlsRegResourcesQuery(input->comm, input->func, &recChannels), ret, exit);
        if (recChannels <= bestTuning.nChannels) {
          bestTuning.algo = NCCL_ALGO_NVLS;
          bestTuning.proto = NCCL_PROTO_SIMPLE;
          bestTuning.nChannels = recChannels;
          bestTuning.maxChannels = recChannels;
          bestTuning.nWarps = input->comm->tuningContext.maxThreads[bestTuning.algo][bestTuning.proto] / WARP_SIZE;
        }

NVLS 的通道數由ncclNvlsRegResourcesQuery查詢硬體資源決定,不是隨意設定的。如果硬體資源不足,通道數會被限制。

21.6 調優決策流程

把前面的內容串起來,得到一個可操作的排查流程。

mermaid
flowchart TD
    start["性能不达标"] --> baseline["跑 nccl-tests 对比官方报告"]
    baseline --> diff{"差距 > 5%?"}
    diff -->|否| app["检查应用层:<br/>通信频率、消息切分"]
    diff -->|是| debug["设置 NCCL_DEBUG=INFO<br/>查看算法/协议选择"]
    debug --> check_algo{"选择的算法合理?"}
    check_algo -->|否| force_algo["尝试 NCCL_ALGO 强制<br/>对比不同算法"]
    check_algo -->|是| check_proto{"协议合理?"}
    check_proto -->|否| force_proto["尝试 NCCL_PROTO 强制<br/>小消息 LL,大消息 Simple"]
    check_proto -->|是| check_chan{"通道数合理?"}
    check_chan -->|否| tune_chan["调整 NCCL_NCHANNELS<br/>或检查显存限制"]
    check_chan -->|是| check_topo["检查拓扑:<br/>NCCL_TOPO_DUMP 确认链路"]
    force_algo --> verify["重新 benchmark 验证"]
    force_proto --> verify
    tune_chan --> verify
    check_topo --> verify
    verify --> improved{"性能提升?"}
    improved -->|是| done["固化配置"]
    improved -->|否| escalate["提交 issue 或联系支持"]

這個流程的核心思想是:先定位,再調參,最後驗證。不要一上來就亂設環境變數。

本章小結

本章把 NCCL 的調優路徑拆成了四個層次:

1. 基準線:用官方性能報告建立預期,5% 以內是正常波動,大消息看頻寬、小消息看延遲。

2. 代價模型:NCCL 內部用modelMap表 + 延遲/頻寬參數估算每種組合的耗時,選最小的。理解這個模型是調參的前提。

3. 環境變數:NCCL_ALGO、NCCL_PROTO、NCCL_SYM_KERNEL透過parseList解析後修改enabled表,強制或排除特定組合。語法支援全域、按函數、排除三種模式。

4. 通道數:由ncclTuningGetChannels計算,受硬體資源和 CTAPolicy 影響。

本章思考與自測

Q1: 如果把ncclTuningCompute中單 rank 短路邏輯(input->comm->nRanks <= 1分支)去掉,會發生什麼?在什麼場景下會導致問題?

參考解析:

單 rank 短路在📎 src/tuning/tuning.cc:191-200:

cpp
  // Set tuning to Ring/Simple for single rank case
  if (input->comm->nRanks <= 1) {
    bestTuning.algo = NCCL_ALGO_RING;
    bestTuning.proto = NCCL_PROTO_SIMPLE;
    bestTuning.symKernelId = ncclSymkKernelId_Count;
    bestTuning.ceMethodId = ncclCeMethodId_Count;
    bestTuning.nChannels = 0;
    bestTuning.maxChannels = 0;
    bestTuning.nWarps = 0;
    bestTuning.forced = 0;
  } else {
    NCCLCHECKGOTO(ncclTuningComputeAllTunings(input, &tunings), ret, exit);
    ...

如果去掉這個分支,單 rank 場景會進入ncclTuningComputeAllTunings,遍歷所有候選組合。問題在於:

1. 性能浪費:單 rank 沒有通信,所有算法的耗時估算都是純開銷,選哪個都一樣。遍歷所有候選是純粹的浪費。

2. 可能選不出結果:某些算法在單 rank 下可能被模型判定為無效(比如 Ring 需要至少 2 個 rank 才能形成環),導致tunings列表為空,ncclTuningSelectBestTuning返回FLT_MAX的初始值,最終bestTuning.algo仍是NCCL_ALGO_UNDEF。

3. 觸發錯誤路徑:如果bestTuning.algo == NCCL_ALGO_UNDEF,會進入📎 src/tuning/tuning.cc:308-329的錯誤處理,列印 "No algorithm/protocol available" 警告,並返回ncclInternalError。

所以這個短路不只是優化,更是正確性保證——單 rank 場景必須有一個確定的預設值。

Q2: parseList中forced[p] = 1這行程式碼(📎 src/tuning/cost_model.cc:83)的作用是什麼?如果去掉它,NCCL_ALGO=ring的行為會有什麼變化?

參考解析:

forced[p] = 1在📎 src/tuning/cost_model.cc:80-85:

cpp
        for (e = 0; e < nelems; e++) {
          if (strcasecmp(elem, elems[e]) == 0) {
            list[p * nelems + e] = set;
            forced[p] = 1;
            break;
          }
        }

forced陣列在ncclTuningContext_t中定義([

關鍵旋鈕只有三個:算法、協議、通道數。其他參數大多是輔助診斷或特定場景優化。掌握了這條調優路徑,你已經能讓 NCCL 在多數場景下跑出接近硬體的效能。但效能之外,生產環境還有另一類更棘手的問題:那些看似正常的程式碼,可能在特定條件下掛死或出錯。下一章我們將彙總 NCCL 在生產中的典型踩坑案例——死鎖、超時、版本不匹配與常見誤用,並看看 NCCL 內部是如何檢測和報告這些問題的。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 22

第 22 章:生產排障與踩坑:常見死鎖、超時、版本不匹配與排查方案

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 22 章 / 共 25 章

上一章我們梳理了效能調優的排查順序與關鍵旋鈕,但生產環境中的 NCCL 故障往往不是效能不達標,而是程式直接掛起或崩潰。這些故障的根源通常不是某個函數寫錯了,而是呼叫順序、生命週期或版本契約被破壞。本章聚焦四類最典型的踩坑:group 語意誤用導致的死鎖、參數校驗缺失導致的靜默錯誤、ABI 版本不匹配、以及超時與重試的邊界。我們會沿著 src/group.cc、src/misc/argcheck.cc、src/include/checks.h 和 contrib/nccl_ep/nccl_ep.cc 四條線索,看清 NCCL 內部是如何在錯誤發生前就把它擋住的。

Group 語意誤用:為什麼「少寫一個 GroupEnd」會掛死

直覺模型:Group 是「購物車」,不是「加速開關」

把ncclGroupStart() / ncclGroupEnd()想像成網購的購物車:你把多件商品(多次通訊呼叫)放進購物車,最後一次性結算(ncclGroupEnd)。如果只放不結算,購物車永遠懸在半空——NCCL 內部維護的ncclGroupDepth計數器就不會歸零,後續所有通訊呼叫都會以為「還在攢單」,永遠不真正下發 kernel,於是整個行程掛死。

〔設計推斷與架構權衡〕

這是生產中最常見的死鎖形態:程式碼在某個異常分支裡return了,跳過了ncclGroupEnd,而ncclGroupDepth是thread_local的,不會因為函式返回而自動清理。

資料結構:thread_local 的 group 狀態

NCCL 把 group 狀態全部放在執行緒區域儲存裡,這是理解死鎖的關鍵。

📎 src/group.cc:34-34

cpp
thread_local int ncclGroupDepth = 0; // depth of ncclGroupStart nesting
thread_local ncclResult_t ncclGroupError = ncclSuccess;
thread_local struct ncclComm* ncclGroupCommHead[ncclGroupTaskTypeNum] = {nullptr};
thread_local struct ncclComm* ncclGroupCommPreconnectHead = nullptr;
thread_local struct ncclIntruQueue<struct ncclAsyncJob, &ncclAsyncJob::next> ncclAsyncJobs;
thread_local int ncclGroupBlocking = -1; /* default mode */

逐欄位解讀:

  • ncclGroupDepth:巢狀深度。ncclGroupStart遞增,ncclGroupEnd遞減,只有減到 0 才真正觸發下發。支援巢狀是設計上的便利,但也意味著「漏掉一個 End」會讓深度永遠停在 1。
  • ncclGroupError:本執行緒累積的 group 錯誤。一旦某次呼叫失敗,後續ncclGroupEnd會直接走失敗路徑。
  • ncclGroupCommHead[]:按任務類型(collective / rawTask / mgmtTask / symRegister)分組的通訊域鏈結串列頭。
  • ncclAsyncJobs:待執行的非同步任務佇列(比如 preconnect、symmetric register)。
  • ncclGroupBlocking:-1表示「還沒遇到任何通訊域」,0表示非阻塞,1表示阻塞。這個欄位是後面「阻塞與非阻塞混用」檢測的核心。
〔設計推斷與架構權衡〕

用thread_local而非全域變數的動機很直接:NCCL 允許多執行緒各自持有獨立的 group 上下文,互不干擾。代價是——執行緒退出時這些狀態不會自動清理,如果執行緒在 group 中途退出,狀態就洩漏了。

Step-by-Step:一次 GroupEnd 的完整校驗鏈

代入場景:應用呼叫ncclGroupEnd(),此時ncclGroupDepth為 1。

第一步,檢查是否真的在 group 裡:

📎 src/group.cc:1048-1052

cpp
  if (ncclGroupDepth == 0) {
    WARN("ncclGroupEnd: not in a group call.");
    ret = ncclInvalidUsage;
    goto exit;
  }

如果使用者沒呼叫ncclGroupStart就直接ncclGroupEnd,這裡會列印 "not in a group call" 並返回ncclInvalidUsage。這是最友善的錯誤——立刻報錯,不會掛死。

第二步,遞減深度,判斷是否是最外層:

📎 src/group.cc:1061-1063

cpp
  if ((--ncclGroupDepth) > 0) goto exit;

  if ((ret = ncclGroupError) != ncclSuccess) goto fail;

如果巢狀了多層,內層的End只是遞減深度就返回,不觸發下發。只有最外層才繼續。同時檢查累積錯誤。

第三步,校驗阻塞模式一致性。這是「阻塞與非阻塞混用」的檢測點:

📎 src/group.cc:1095-1101

cpp
  if (hasCommHead || !ncclIntruQueueEmpty(&groupJob->asyncJobs) || ncclGroupCommPreconnectHead != nullptr) {
    /* make sure ncclGroupBlocking has been set. */
    if (ncclGroupBlocking != 0 && ncclGroupBlocking != 1) {
      WARN("Invalid group blocking state %d", ncclGroupBlocking);
      ret = ncclInternalError;
      goto fail;
    }

ncclGroupBlocking必須在{0, 1}之間。如果它還是-1,說明 group 裡既沒有通訊域也沒有非同步任務,邏輯上不該走到這裡。

第四步,根據阻塞模式分叉。非阻塞走執行緒非同步下發,阻塞走同步下發:

📎 src/group.cc:1102-1134

cpp
    if (ncclGroupBlocking == 0) {
      /* nonblocking group */
      if (!ncclIntruQueueEmpty(&groupJob->asyncJobs)) {
        ncclAsyncJob* job = ncclIntruQueueHead(&groupJob->asyncJobs);
        do {
          NCCLCHECKGOTO(ncclCommSetAsyncError(job->comm, ncclInProgress), ret, fail);
          if (job->comm->groupJob == NULL) {
            job->comm->groupJob = groupJob;
            groupJob->groupRefCount++;
          }
          job = job->next;
        } while (job);
      }
      ...
      groupJob->base.func = groupLaunchNonBlocking;
      STDTHREADCREATE_GOTO(groupJob->base.thread, ncclAsyncJobMain, ret, fail, &groupJob->base);
      groupJob->nonBlockingInit = true;
      ret = ncclInProgress;
    }

注意groupRefCount++和ret = ncclInProgress:非阻塞模式下,ncclGroupEnd立刻返回ncclInProgress,真正的下發在背景執行緒裡跑。呼叫方必須後續用ncclCommGetAsyncError輪詢,或者用ncclGroupJobComplete等待。

阻塞與非阻塞混用:為什麼被禁止

回到ncclAsyncLaunch,看混用檢測:

📎 src/group.cc:55-64

cpp
    /* check if there are blocking and nonblocking comms at the same time in group. */
    if (comm->destroyFlag) {
      ncclGroupBlocking = 1;
    } else if (ncclGroupBlocking == -1) {
      /* first met communicator */
      ncclGroupBlocking = comm->config.blocking;
    } else if (ncclGroupBlocking != comm->config.blocking) {
      WARN("Blocking and nonblocking communicators are not allowed in the same group.");
      ret = ncclInvalidArgument;
    }
〔設計推斷與架構權衡〕

為什麼禁止混用?因為阻塞通訊域的下發語義是「呼叫返回時 kernel 已提交」,而非阻塞是「呼叫返回時任務已入佇列但未提交」。如果兩者在同一個 group 裡,ncclGroupEnd無法給出統一的返回語義——到底是等還是不等?NCCL 選擇直接拒絕,把問題暴露在 API 邊界。

生產踩坑:三個真實場景

場景一:異常分支漏掉 GroupEnd。程式碼在ncclGroupStart和ncclGroupEnd之間拋異常或提前return,ncclGroupDepth停在 1。後續所有通訊呼叫都進入「攢單」狀態,永遠不下發。排查方法:在ncclGroupEnd前列印ncclGroupDepth,或者用gdb觀察該 thread_local 變數。

場景二:跨執行緒使用同一個 comm。因為 group 狀態是thread_local,執行緒 A 呼叫ncclGroupStart後,執行緒 B 呼叫ncclAllReduce不會進入 A 的 group。如果 A 和 B 操作同一個 comm,會出現「部分呼叫在 group 內、部分在 group 外」的錯亂。NCCL 不檢測這種情況,因為它假設一個 comm 在任一時刻只被一個執行緒操作。

場景三:CUDA graph capture 與 group 的互動。看doLaunches裡的檢測:

📎 src/group.cc:448-455

cpp
    if (capturingYes && capturingNo) {
      // We have entered barriers but are aborting without leaving them. Thus
      // these comms are permanently trashed. We need a good mechanism for
      // tracking and reporting that.
      WARN("Either none or all communicators in a ncclGroup() can be CUDA graph captured.");
      result = ncclInvalidUsage;
      goto failure;
    }

註解說得很直白:一旦進入 barrier 又中途放棄,這些 comm 就被「永久損壞」了。所以規則是——一個 group 裡的所有通訊域,要麼全部在 capture 中,要麼全部不在。混用會導致 comm 狀態不一致,且 NCCL 目前沒有好的恢復機制。

mermaid
flowchart TD
    start["ncclGroupEnd()"] --> depth_check{"ncclGroupDepth == 0?"}
    depth_check -->|是| err_usage["WARN not in a group call<br/>return ncclInvalidUsage"]
    depth_check -->|否| dec["--ncclGroupDepth"]
    dec --> nested{"depth > 0?"}
    nested -->|是| exit_ok["goto exit 返回"]
    nested -->|否| err_check{"ncclGroupError == success?"}
    err_check -->|否| fail_clean["groupCleanup 清理所有 comm 与 asyncJobs"]
    err_check -->|是| blocking_check{"ncclGroupBlocking in {0,1}?"}
    blocking_check -->|否| err_internal["WARN Invalid group blocking state<br/>return ncclInternalError"]
    blocking_check -->|是| mode_split{"ncclGroupBlocking == 0?"}
    mode_split -->|是 非阻塞| async_launch["STDTHREADCREATE groupLaunchNonBlocking<br/>ret = ncclInProgress"]
    mode_split -->|否 阻塞| sync_launch["groupLaunch 同步下发<br/>delete groupJob"]
    async_launch --> reset["groupLocalResetJobState"]
    sync_launch --> reset
    reset --> exit_ok
    fail_clean --> reset

參數校驗與靜默錯誤:ArgCheck 如何擋住「看起來正常」的呼叫

直覺模型:ArgCheck 是「機場安檢」

參數校驗就像機場安檢:它不負責讓你飛得更快,但能擋住那些「看起來是行李、實際是危險品」的東西。沒有它,一個傳錯裝置的指標會讓 GPU kernel 讀到垃圾資料,或者更糟——靜默寫壞別人的顯存。

資料結構:校驗模式與全域檢查佇列

NCCL 的參數校驗不是「每次都全查」,而是分模式。核心是comm->checkMode:

📎 src/misc/argcheck.cc:227-251

cpp
  if (info->comm->checkMode != ncclCheckModeDefault) {
    if ((info->coll == ncclFuncSend || info->coll == ncclFuncRecv)) {
      if (info->count > 0) NCCLCHECK(CudaPtrCheck(info->recvbuff, info->comm, "buff", info->opName));
    } else if (info->coll == ncclFuncPutSignal || info->coll == ncclFuncSignal || info->coll == ncclFuncWaitSignal) {
      // One-sided RMA ops specify the remote destination via peerWin, not sendbuff/recvbuff,
      // so the standard CUDA pointer checks do not apply here.
      INFO(NCCL_COLL, "%s : skipping sendbuff/recvbuff pointer check (one-sided RMA uses peerWin)", info->opName);
    } else {
      // Check CUDA device pointers
      if (info->coll != ncclFuncBroadcast || info->comm->rank == info->root) {
        NCCLCHECK(CudaPtrCheck(info->sendbuff, info->comm, "sendbuff", info->opName));
      }
      if (info->coll != ncclFuncReduce || info->comm->rank == info->root) {
        NCCLCHECK(CudaPtrCheck(info->recvbuff, info->comm, "recvbuff", info->opName));
      }
    }

    if (info->comm->checkMode == ncclCheckModeDebugGlobal) {
      struct ncclArgsInfo* argsInfo;
      NCCLCHECK(ncclCalloc(&argsInfo, 1));
      argsInfo->info = *info;
      argsInfo->next = NULL;
      ncclIntruQueueEnqueue(&info->comm->argsInfoQueue, argsInfo);
    }
  }

三種模式:

  • ncclCheckModeDefault:只做最便宜的檢查(root 範圍、datatype 範圍、op 範圍),不碰 CUDA API。
  • 非預設模式:呼叫CudaPtrCheck,這會真正呼叫cudaPointerGetAttributes,有效能開銷。
  • ncclCheckModeDebugGlobal:除了本地檢查,還把ncclInfo塞進argsInfoQueue,等 group 結束時做跨 rank 的全局一致性檢查。
〔設計推斷與架構權衡〕

這個設計是效能與正確性的權衡:cudaPointerGetAttributes是同步 CUDA 呼叫,在熱路徑上每次通訊都調會顯著拖慢小訊息。所以預設模式只做「零成本」檢查,把昂貴的指標校驗留給除錯模式。

Step-by-Step:CudaPtrCheck 的三層防線

代入場景:使用者傳入一個sendbuff,NCCL 在除錯模式下校驗它。

第一層,指標是否有效:

📎 src/misc/argcheck.cc:12-18

cpp
ncclResult_t CudaPtrCheck(const void* pointer, struct ncclComm* comm, const char* ptrname, const char* opname) {
  cudaPointerAttributes attr;
  cudaError_t err = cudaPointerGetAttributes(&attr, pointer);
  if (err != cudaSuccess || attr.devicePointer == NULL) {
    WARN("%s : %s %p is not a valid pointer", opname, ptrname, pointer);
    return ncclInvalidArgument;
  }

cudaPointerGetAttributes對無效指標會傳回錯誤,或者devicePointer為 NULL。這擋住了「傳了個 host 堆疊位址」或「傳了個已釋放的指標」。

第二層,裝置是否匹配:

📎 src/misc/argcheck.cc:19-26

cpp
#if CUDART_VERSION >= 10000
  if (attr.type == cudaMemoryTypeDevice && attr.device != comm->cudaDev) {
#else
  if (attr.memoryType == cudaMemoryTypeDevice && attr.device != comm->cudaDev) {
#endif
    WARN("%s : %s allocated on device %d mismatchs with NCCL device %d", opname, ptrname, attr.device, comm->cudaDev);
    return ncclInvalidArgument;
  }

這是最隱蔽的坑:指標是有效的 GPU 指標,但屬於另一塊 GPU。在多卡機器上,如果使用者忘了cudaSetDevice,很容易傳錯。NCCL 在這裡明確拒絕。

第三層,通訊域物件完整性:

📎 src/misc/argcheck.cc:38-45

cpp
ncclResult_t CommCheck(struct ncclComm* comm, const char* opname, const char* ptrname) {
  NCCLCHECK(PtrCheck(comm, opname, ptrname));
  if (comm->startMagic != NCCL_MAGIC || comm->endMagic != NCCL_MAGIC) {
    WARN("Error: corrupted comm object detected");
    return ncclInvalidArgument;
  }
  return ncclSuccess;
}

startMagic / endMagic是放在ncclComm結構體首尾的哨兵值。如果使用者傳了個野指標、或者 comm 已被釋放,magic 就對不上。這是「記憶體損壞偵測」的經典手法——用兩個哨兵夾住結構體,任何越界寫都可能破壞其中一個。

全局一致性檢查:registrationCheck 的跨 rank 校驗

這是 NCCL 裡最「重」的校驗,只在ncclCheckModeDebugGlobal下觸發。它檢查的是——所有 rank 的對稱記憶體註冊狀態是否一致。

📎 src/misc/argcheck.cc:95-111

cpp
  NCCLCHECKGOTO(bootstrapAllGather(comm->bootstrap, bufInfo, sizeof(struct symBufInfo) * 2), ret, fail);

  cmpBufInfo[0] = bufInfo[0];
  cmpBufInfo[1] = bufInfo[1];
  for (int r = 1; r < comm->nRanks; r++) {
    int infoIdx = r * 2;
    if (cmpBufInfo[0].isSymRegistered != bufInfo[infoIdx].isSymRegistered ||
        cmpBufInfo[1].isSymRegistered != bufInfo[infoIdx + 1].isSymRegistered) {
      if (comm->rank == 0) {
        WARN("Coll %s size %ld symmetric registration check failed on rank %d: sendReg %d recvReg %d mismatch with "
             "rank 0 sendReg %d recvReg %d",
             info->opName, size, r, bufInfo[infoIdx].isSymRegistered, bufInfo[infoIdx + 1].isSymRegistered,
             cmpBufInfo[0].isSymRegistered, cmpBufInfo[1].isSymRegistered);
      }
      ret = ncclInvalidArgument;
      goto fail;
    }

它透過 bootstrap 的allGather把每個 rank 的(isSymRegistered, bigOffset, userOffset)收集起來,然後逐 rank 比對。如果 rank 0 的 send buffer 註冊了對稱記憶體,而 rank 3 沒註冊,這裡就會報錯。

〔設計推斷與架構權衡〕

為什麼這個檢查重要?對稱記憶體(symmetric memory)要求所有 rank 用同一套虛擬位址存取緩衝區。如果某個 rank 的 buffer 沒註冊,kernel 裡算出來的位址就是錯的,會讀到垃圾或越界。這種錯誤在執行時表現為「結果偶爾不對」,極難排查。NCCL 選擇在 API 邊界用一次 allGather 的代價把它擋住。

生產踩坑

坑一:預設模式下指標錯誤不報。如果使用者沒開除錯模式,傳了個錯誤裝置的指標,NCCL 不會在ArgsCheck階段報錯,而是等到 kernel 執行時才發現——此時可能已經寫壞了別的 rank 的顯存。建議在開發階段用NCCL_DEBUG=WARN加checkMode除錯。

坑二:ncclCheckModeDebugGlobal的 allGather 開銷。每次通訊都做一次 bootstrap allGather,在小訊息高頻場景下會成為瓶頸。這個模式只適合除錯,不能上生產。

坑三:userRedOp 的生命週期。看這段:

📎 src/misc/argcheck.cc:220-225

cpp
  int opIx = int(ncclUserRedOpMangle(info->comm, info->op)) - int(ncclNumOps);
  if (ncclNumOps <= info->op &&
      (info->comm->userRedOpCapacity <= opIx || info->comm->userRedOps[opIx].freeNext != -1)) {
    WARN("%s : reduction operation %d unknown to this communicator", info->opName, info->op);
    return ncclInvalidArgument;
  }

使用者自訂的 reduction op 是註冊在 comm 上的。如果使用者傳了一個「曾經註冊過但已被釋放」的 op,freeNext != -1會偵測到它已被回收。這是防止「懸空 op 句柄」的檢查。

錯誤傳播巨集:NCCLCHECK 家族如何保證「錯誤不丟」

直覺模型:錯誤傳播巨集是「接力棒」

NCCL 的錯誤處理靠一組巨集接力:底層函式傳回ncclResult_t,上層用NCCLCHECK檢查,非成功就立刻傳回。這就像接力賽——棒子(錯誤碼)必須一路傳到底,任何一棒掉了,整個鏈條就斷了。

資料結構:巨集家族全貌

📎 src/include/checks.h:148-166

cpp
#define NCCLCHECK(call) \
  do { \
    ncclResult_t RES = call; \
    if (RES != ncclSuccess && RES != ncclInProgress) { \
      /* Print the back trace*/ \
      if (ncclDebugNoWarn == 0) INFO_LOC(NCCL_ALL, "-> %d", RES); \
      return RES; \
    } \
  } while (0)

#define NCCLCHECKGOTO(call, RES, label) \
  do { \
    RES = call; \
    if (RES != ncclSuccess && RES != ncclInProgress) { \
      /* Print the back trace*/ \
      if (ncclDebugNoWarn == 0) INFO_LOC(NCCL_ALL, "-> %d", RES); \
      goto label; \
    } \
  } while (0)

關鍵細節:ncclInProgress被視為「非錯誤」。這是非阻塞通訊的核心——ncclGroupEnd傳回ncclInProgress表示「任務已提交,還沒完成」,呼叫方應該繼續輪詢而不是當錯誤處理。

NCCLCHECK直接return,NCCLCHECKGOTO跳到label。後者用於需要清理資源的場景。

清理路徑:NCCLCHECKIGNORE 保留首個錯誤

📎 src/include/checks.h:168-177

cpp
// Report failure but continue - useful for cleanup paths where we want to
// attempt all cleanup steps. Preserves the first error in RES.
#define NCCLCHECKIGNORE(call, RES) \
  do { \
    ncclResult_t TMPRES = call; \
    if (TMPRES != ncclSuccess && TMPRES != ncclInProgress) { \
      if (ncclDebugNoWarn == 0) INFO_LOC(NCCL_ALL, "-> %d", TMPRES); \
      if (RES == ncclSuccess) RES = TMPRES; \
    } \
  } while (0)

註解說得很清楚:清理路徑上要「嘗試所有清理步驟」,不能被第一個錯誤打斷。但錯誤碼要保留第一個——因為第一個錯誤通常是最有診斷價值的根因。

等待與中止:NCCLWAIT 的 abortFlag 檢查

📎 src/include/checks.h:196-205

cpp
#define NCCLWAIT(call, cond, abortFlagPtr) \
  do { \
    uint32_t* tmpAbortFlag = (abortFlagPtr); \
    ncclResult_t RES = call; \
    if (RES != ncclSuccess && RES != ncclInProgress) { \
      if (ncclDebugNoWarn == 0) INFO_LOC(NCCL_ALL, "-> %d", RES); \
      return ncclInternalError; \
    } \
    if (COMPILER_ATOMIC_LOAD(tmpAbortFlag, std::memory_order_acquire)) NEQCHECK(*tmpAbortFlag, 0); \
  } while (!(cond))

這是輪詢等待的模板:每次迴圈呼叫call(推進進度),檢查cond(是否滿足),同時檢查abortFlag(是否被中止)。abortFlag用memory_order_acquire載入,保證看到其他執行緒寫入的中止訊號。

〔設計推斷與架構權衡〕

這個設計解決了一個經典問題:當某個 rank 出錯時,其他 rank 可能還在死等它的資料。abortFlag是跨 rank 傳播中止訊號的機制——一旦設定,所有等待迴圈都會退出。

執行緒建立與記憶體分配的安全巨集

📎 src/include/checks.h:237-256

cpp
#define STDTHREADCREATE_IMPL(var, func, error_action, ...) \
  do { \
    try { \
      (var) = std::thread(func, __VA_ARGS__); \
    } catch (const std::exception& e) { \
      WARN("Thread creation failed: %s", e.what()); \
      error_action; \
    } \
  } while (0)

#define STDTHREADCREATE(var, func, ...) STDTHREADCREATE_IMPL(var, func, return ncclSystemError, __VA_ARGS__)

#define STDTHREADCREATE_GOTO(var, func, RES, label, ...) \
  STDTHREADCREATE_IMPL( \
    var, func, \
    do { \
      RES = ncclSystemError; \
      goto label; \
    } while (0), \
    __VA_ARGS__)

std::thread建構失敗會拋異常(比如執行緒數超限)。這個巨集把異常轉成ncclSystemError,避免異常穿透 C API 邊界。

📎 src/include/checks.h:258-275

cpp
#define NEW_NOTHROW(var, x) \
  do { \
    (var) = new (std::nothrow) x{}; \
    if (!(var)) { \
      WARN("Allocation failed"); \
      return ncclSystemError; \
    } \
  } while (0)

new (std::nothrow)在分配失敗時傳回 nullptr 而非拋異常。這是 C++ 程式碼在 C API 邊界上的標準做法。

生產踩坑

坑一:ncclInProgress被誤當成功。有些使用者程式碼寫if (ret == ncclSuccess)判斷成功,但非阻塞模式下傳回的是ncclInProgress。正確做法是if (ret == ncclSuccess || ret == ncclInProgress),或者用ncclCommGetAsyncError查詢。

坑二:NCCLCHECK在解構函式裡用。如果解構函式裡用NCCLCHECK,錯誤會直接return,跳過後續清理。應該用NCCLCHECKIGNORE。

ABI 版本不匹配:nccl_ep 的 size-based 設計

直覺模型:ABI 是「插座標準」

ABI(應用二進位介面)就像電源插座標準:如果函式庫和呼叫方對「結構體長什麼樣」的理解不一致,就會像把美標插頭插進歐標插座——輕則不工作,重則燒毀。contrib/nccl_ep用了一個巧妙的設計:每個跨邊界結構體都以size欄位開頭。

資料結構:size + magic 雙重校驗

📎 contrib/nccl_ep/nccl_ep.cc:70-76

cpp
// Size-based ABI versioning: every cross-boundary struct starts with a `size`
// field set by the caller to sizeof(struct). The library checks that against
// its own known size; any mismatch means caller and library are from different
// releases. Strict equality for now — see nccl_ep.h for the planned future
// relaxation (all-zero-trailing-bytes escape hatch).
// Immediately after `size` there is a `magic` field pre-filled by NCCL_EP_*_INIT
// to catch unininitialized structures.

設計要點:

  • size欄位由呼叫方填sizeof(struct),函式庫檢查它是否等於自己認識的 size。
  • magic欄位由NCCL_EP_*_INIT巨集預填,用來捕獲「未初始化」的結構體。
  • 當前是嚴格相等,未來計劃支援「尾部全零則允許 size 更小」的寬鬆模式。

Step-by-Step:EP_REQUIRE_STRUCT 的校驗流程

📎 contrib/nccl_ep/nccl_ep.cc:77-80

cpp
#define EP_REQUIRE_STRUCT(ptr) \
    do { \
        assert( \
            (ptr) != nullptr && (ptr)->size == sizeof(*(ptr)) && \

這個巨集在ncclEpDispatch、ncclEpCombine等入口處呼叫:

📎 contrib/nccl_ep/nccl_ep.cc:2827-2830

cpp
    EP_REQUIRE_STRUCT(inputs);
    EP_REQUIRE_STRUCT(outputs);
    EP_OPTIONAL_LAYOUT_INFO(layout_info);
    EP_OPTIONAL_STRUCT(config);

inputs和outputs是必需參數,用EP_REQUIRE_STRUCT;layout_info和config是選用參數,用EP_OPTIONAL_*。

版本安全的欄位讀取:layoutInfoRecvTopkIdxKind

這是最精妙的部分——如何在「呼叫方結構體可能更小」的情況下安全讀取欄位。

📎 contrib/nccl_ep/nccl_ep.cc:139-144

cpp
// Safe field reader for ncclEpLayoutInfo_t::recv_topk_idx_kind. Returns AUTO
// when the caller's struct (size) does not cover the field, preserving the
// pre-flag default.
static inline ncclEpExpertIdKind_t layoutInfoRecvTopkIdxKind(const ncclEpLayoutInfo_t* lip) {
    if (lip == nullptr) return NCCL_EP_EXPERT_ID_AUTO;
    constexpr size_t field_end = offsetof(ncclEpLayoutInfo_t, recv_topk_idx_kind) + sizeof(ncclEpExpertIdKind_t);
    if (lip->size < field_end) return NCCL_EP_EXPERT_ID_AUTO;
    return lip->recv_topk_idx_kind;
}

邏輯是:如果呼叫方的size小於「該欄位結束的偏移」,說明呼叫方用的是舊版本結構體,這個欄位不存在,回傳預設值AUTO。否則正常讀取。

〔設計推斷與架構權衡〕

這是 ABI 相容的標準手法:新欄位只能加在結構體末尾,讀取時用size判斷欄位是否存在。這樣舊呼叫方用舊結構體,新函式庫也能正確處理。

版本號檢查:軟警告而非硬拒絕

📎 contrib/nccl_ep/nccl_ep.cc:1393-1400

cpp
    if (in_config->version != NCCL_EP_API_VERSION) {
        fprintf(
            stderr,
            "NCCL EP WARN: ncclEpGroupConfig_t.version=%u, library API_VERSION=%u; "
            "behavior may differ across versions.\n",
            in_config->version,
            (unsigned)NCCL_EP_API_VERSION);
    }

注意這裡是WARN而非return error。版本號不匹配只是警告,因為size檢查已經保證了記憶體佈局安全。版本號更多是「行為可能不同」的提示。

生產踩坑

坑一:忘記用 INIT 巨集初始化。如果使用者手動memset結構體為 0,magic就是 0,EP_REQUIRE_STRUCT會失敗。必須用NCCL_EP_*_INIT巨集。

坑二:跨版本混用動態函式庫。如果應用連結的是新版libnccl_ep.so,但標頭檔是舊版,sizeof(struct)會不一致,EP_REQUIRE_STRUCT會立刻報錯。這是設計意圖——快速失敗優於靜默錯誤。

坑三:EP_OPTIONAL_LAYOUT_INFO的範圍檢查。看這段:

📎 contrib/nccl_ep/nccl_ep.cc:114-123

cpp
            if ((ptr)->size < kNcclEpLayoutInfoMinSize || (ptr)->size > sizeof(*(ptr))) { \
                fprintf( \
                    stderr, \
                    "NCCL EP: ncclEpLayoutInfo_t size out of supported range: " \
                    "got %u, expected [%zu, %zu]\n", \
                    (ptr)->size, \
                    kNcclEpLayoutInfoMinSize, \
                    sizeof(*(ptr))); \
                return ncclInvalidArgument; \
            } \

layout_info允許 size 在[min, sizeof]範圍內,這比EP_REQUIRE_STRUCT的嚴格相等更寬鬆。原因是layout_info是選用參數,且歷史上欄位有增減。

mermaid
flowchart TD
    entry["ncclEpDispatch(inputs, outputs, layout_info, config)"] --> req_inputs{"EP_REQUIRE_STRUCT(inputs)<br/>size == sizeof?"}
    req_inputs -->|否| err_size["assert 失败 / 返回错误"]
    req_inputs -->|是| req_outputs{"EP_REQUIRE_STRUCT(outputs)"}
    req_outputs -->|否| err_size
    req_outputs -->|是| opt_layout{"layout_info != nullptr?"}
    opt_layout -->|否| skip_layout["跳过 layout 校验"]
    opt_layout -->|是| range_check{"size in [min, sizeof]?"}
    range_check -->|否| err_range["fprintf size out of range<br/>return ncclInvalidArgument"]
    range_check -->|是| magic_check{"magic == NCCL_EP_MAGIC?"}
    magic_check -->|否| err_magic["fprintf magic mismatch<br/>return ncclInvalidArgument"]
    magic_check -->|是| read_field["layoutInfoRecvTopkIdxKind<br/>size < field_end ? AUTO : 实际值"]
    skip_layout --> read_field
    read_field --> proceed["继续执行 dispatch 逻辑"]

逾時、重試與中止:從 NCCLWAIT 到 nccl_ep 的 timeout_cycles

直覺模型:逾時是「保險絲」

分散式通訊裡,一個 rank 卡住會導致所有 rank 死等。逾時機制就像保險絲:正常情況下不動作,一旦電流異常就熔斷,避免整個系統燒毀。

資料結構:abortFlag 與 timeout_cycles

NCCL 核心用abortFlag傳播中止訊號。看ncclAsyncLaunch裡的傳遞:

📎 src/group.cc:49-52

cpp
    job->abortFlag = comm->abortFlag;
    job->abortFlagDev = comm->abortFlagDev;
    job->childAbortFlag = comm->childAbortFlag;
    job->childAbortFlagDev = comm->childAbortFlagDev;

每個 job 持有 comm 的 abortFlag 指標。當 group 偵測到錯誤時:

📎 src/group.cc:118-126

cpp
        if (!job->destroyFlag &&
            (COMPILER_ATOMIC_LOAD(groupAbortFlag, std::memory_order_acquire) || errorJobAbortFlag == true)) {
          COMPILER_ATOMIC_STORE(job->abortFlag, uint32_t(1), std::memory_order_release);
          COMPILER_ATOMIC_STORE(job->abortFlagDev, uint32_t(1), std::memory_order_release);
          if (job->childAbortFlag) {
            COMPILER_ATOMIC_STORE(job->childAbortFlag, uint32_t(1), std::memory_order_release);
            COMPILER_ATOMIC_STORE(job->childAbortFlagDev, uint32_t(1), std::memory_order_release);
          }
        }

一旦groupAbortFlag或errorJobAbortFlag為真,所有 job 的 abortFlag 都被設為 1。memory_order_release保證之前的寫操作對其他執行緒可見。

nccl_ep 的逾時設計:GPU 時鐘週期

nccl_ep用了更精細的逾時——以 GPU 時鐘週期為單位。

📎 contrib/nccl_ep/nccl_ep.cc:1558-1591

cpp
    // Resolve timeout_cycles: env var > config field > compile-time default
    {
        int dev;
        int clock_khz_int;
        CUDA_CHECK(cudaGetDevice(&dev));
        CUDA_CHECK(cudaDeviceGetAttribute(&clock_khz_int, cudaDevAttrClockRate, dev));
        uint64_t clock_khz = static_cast<uint64_t>(clock_khz_int);

        uint64_t resolved = NUM_TIMEOUT_CYCLES;
        const char* source = "compile-time default";
        const uint64_t env_ms = static_cast<uint64_t>(ep_group->env.timeout_ms.value.ul);
        // Only a positive timeout overrides the default.
        const bool have_env_ms = ep_group->env.timeout_ms.is_set && env_ms > 0;

        if (have_env_ms) {
            resolved = clock_khz * 1000ULL * env_ms / 1000ULL;
            source = "NCCL_EP_TIMEOUT_MS env var";
            ...
        } else if (ep_group->config.timeout_ns != 0) {
            resolved = clock_khz * 1000ULL * (ep_group->config.timeout_ns / 1000000ULL) / 1000ULL;
            source = "config.timeout_ns";
        }

        ep_group->timeout_cycles = resolved;

優先級是:環境變數NCCL_EP_TIMEOUT_MS> 配置欄位timeout_ns> 編譯期預設值。轉換公式是clock_khz * 1000 * ms / 1000,即把毫秒轉成時鐘週期。

〔設計推斷與架構權衡〕

為什麼用時鐘週期而非毫秒?因為 GPU kernel 裡的等待迴圈無法呼叫系統時間 API,只能讀clock64()暫存器。用時鐘週期做逾時判斷,kernel 裡可以直接比較,無需 host 介入。

非同步錯誤標誌:host-pinned 記憶體

📎 contrib/nccl_ep/nccl_ep.cc:1767-1778

cpp
    // Allocate mask buffer and async error flag for active-mask support
    if (ep_group->config.enable_mask && ep_group->config.algorithm == NCCL_EP_ALGO_LOW_LATENCY) {
        size_t mask_bytes = ep_group->nRanks * sizeof(int);
        CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&ep_group->mask_buffer), mask_bytes));
        // Initialize all ranks as active (1 = active, 0 = masked/failed)
        std::vector<int> all_active(ep_group->nRanks, 1);
        CUDA_CHECK(
            cudaMemcpyAsync(ep_group->mask_buffer, all_active.data(), mask_bytes, cudaMemcpyHostToDevice, stream));
        CUDA_CHECK(
            cudaHostAlloc(reinterpret_cast<void**>(&ep_group->async_error_flag), sizeof(int), cudaHostAllocMapped));
        *ep_group->async_error_flag = 0;
    }

async_error_flag用cudaHostAllocMapped分配,這是 host-pinned 且映射到裝置位址空間的記憶體。GPU kernel 可以寫它,host 可以讀它,無需顯式拷貝。

讀取非同步錯誤:原子載入

📎 contrib/nccl_ep/nccl_ep.cc:4312-4321

cpp
ncclResult_t ncclEpGetAsyncError(ncclEpGroup_t ep_group, int* error_out) {
    EP_HOST_ASSERT(ep_group != nullptr);
    if (!ep_group->config.enable_mask) {
        return ncclInvalidUsage;
    }
    EP_HOST_ASSERT(ep_group->async_error_flag != nullptr && "ncclEpGetAsyncError: enable_mask must be true");
    EP_HOST_ASSERT(error_out != nullptr);
    *error_out = __atomic_load_n(ep_group->async_error_flag, __ATOMIC_ACQUIRE);
    return ncclSuccess;
}

用__atomic_load_n加__ATOMIC_ACQUIRE,保證讀到的是 GPU 寫入的最新值,而不是快取的舊值。

生產踩坑

坑一:逾時設定過短導致誤報。如果NCCL_EP_TIMEOUT_MS設得太小,正常的網路抖動會被誤判為逾時。建議根據實際網路 RTT 設定,一般不小於 10 秒。

坑二:abortFlag 設定後未清理。一旦 abortFlag 被設為 1,comm 就進入「中止」狀態。如果使用者想繼續用這個 comm,必須先清理 abortFlag。NCCL 的ncclCommAbort會做這個清理。

坑三:ncclEpMaskClean的前置條件。看這段:

📎 contrib/nccl_ep/nccl_ep.cc:4262-4266

cpp
    EP_HOST_ASSERT(ep_group->config.algorithm == NCCL_EP_ALGO_LOW_LATENCY);
    EP_HOST_ASSERT(
        ep_group->rdma_buffer != nullptr &&
        "ncclEpMaskClean: rdma_buffer not yet allocated; create at least one LL handle first");
    EP_HOST_ASSERT(ep_group->sync_buffer != nullptr && ep_group->sync_window != nullptr);

ncclEpMaskClean要求rdma_buffer已分配。如果使用者建立了 group 但還沒建立任何 LL handle,rdma_buffer是 nullptr(因為 LL 是懶分配),這裡會 assert 失敗。

本章小結

本章串起了四類生產踩坑:

1. Group 語意誤用:ncclGroupDepth是 thread_local,漏掉ncclGroupEnd會導致永久掛死;阻塞與非阻塞通信域不能混用;CUDA graph capture 必須全有或全無。

2. 參數校驗:ArgsCheck分模式校驗,預設模式只做零成本檢查;CudaPtrCheck三層防線擋住無效指標、錯誤裝置、損壞的 comm;registrationCheck做跨 rank 的對稱記憶體一致性檢查。

3. 錯誤傳播:NCCLCHECK家族保證錯誤不丟;ncclInProgress不是錯誤;NCCLCHECKIGNORE用於清理路徑保留首個錯誤;NCCLWAIT在輪詢中檢查 abortFlag。

4. ABI 版本:nccl_ep用 size-based 設計,每個跨邊界結構體以size開頭,配合magic捕獲未初始化;新欄位只能加在末尾,讀取時用size判斷是否存在。

5. 超時與中止:核心用abortFlag傳播中止;nccl_ep用 GPU 時鐘週期做超時,async_error_flag用 host-pinned 記憶體實現 GPU→host 異步通知。

本章思考與自測

Q1: 如果把ncclGroupEndInternal中if ((--ncclGroupDepth) > 0) goto exit;(📎 src/group.cc:1061)改成if (ncclGroupDepth > 0) goto exit;(不遞減),會發生什麼?在嵌套 group 場景下會有什麼後果?

參考解析:

原代碼--ncclGroupDepth先遞減再判斷。如果改成不遞減:

cpp
if (ncclGroupDepth > 0) goto exit;  // 错误版本

那麼每次ncclGroupEnd都不會減少深度。假設用戶寫了:

cpp
ncclGroupStart();  // depth = 1
ncclGroupStart();  // depth = 2
ncclAllReduce(...);
ncclGroupEnd();    // 原版: depth = 1, 返回; 错误版: depth = 2, 返回
ncclGroupEnd();    // 原版: depth = 0, 触发下发; 错误版: depth = 2, 返回

錯誤版本下,第二次ncclGroupEnd時ncclGroupDepth仍是 2,> 0成立,直接goto exit,永遠不觸發下發。所有通信調用都停留在「攢單」狀態,進程掛死。

更隱蔽的是:ncclGroupDepth是 thread_local,不會因為函數返回而重置。即使後續代碼不再調用 group API,這個線程上的所有通信都會失效。

這個改動還會破壞ncclGroupStart的配對語義——ncclGroupStart遞增、ncclGroupEnd不遞減,深度只增不減,最終溢出(雖然 int 溢出需要 20 億次調用,實際更可能是邏輯掛死)。

Q2: CudaPtrCheck中attr.type == cudaMemoryTypeDevice && attr.device != comm->cudaDev(📎 src/misc/argcheck.cc:20)這個檢查,如果去掉attr.type == cudaMemoryTypeDevice這個條件,會有什麼問題?在什麼場景下會誤報?

參考答案:

cudaPointerAttributes.type有三個可能值:cudaMemoryTypeDevice(裝置記憶體)、cudaMemoryTypeHost(主機記憶體)、cudaMemoryTypeManaged(統一記憶體)。

如果去掉attr.type == cudaMemoryTypeDevice條件,變成:

cpp
if (attr.device != comm->cudaDev) {  // 错误版本

那麼對於 host 記憶體或 managed 記憶體,attr.device可能是 -1 或 0,與comm->cudaDev不匹配,會誤報「裝置不匹配」。

具體場景:用戶傳入一個cudaMallocManaged分配的指標。managed 記憶體的attr.device通常是分配時的裝置,但如果記憶體被遷移到其他裝置,attr.device可能變化。更常見的是 host 記憶體(比如cudaHostAlloc分配的 pinned 記憶體),attr.device為 -1,與任何cudaDev都不等,會誤報。

NCCL 允許 host 記憶體作為通信緩衝區(通過cudaMemcpy中轉),所以必須區分「裝置記憶體但裝置不對」和「非裝置記憶體」。前者是錯誤,後者是合法的。

Q3: layoutInfoRecvTopkIdxKind(📎 contrib/nccl_ep/nccl_ep.cc:139-144)用lip->size < field_end判斷欄位是否存在。如果新版本在結構體中間插入了一個欄位(而非末尾),這個判斷會怎樣失效?為什麼 ABI 設計規定新欄位只能加在末尾?

參考解析:

假設原結構體是:

c
struct ncclEpLayoutInfo_t {
    unsigned int size;
    unsigned int magic;
    ncclEpExpertIdKind_t recv_topk_idx_kind;  // offset = 8
};

field_end = offsetof(recv_topk_idx_kind) + sizeof(...) = 8 + 4 = 12。

如果新版本在magic和recv_topk_idx_kind之間插入一個欄位:

c
struct ncclEpLayoutInfo_t {
    unsigned int size;
    unsigned int magic;
    unsigned int new_field;                    // 新插入
    ncclEpExpertIdKind_t recv_topk_idx_kind;  // offset 变成 12
};

此時field_end = 12 + 4 = 16。舊調用方的size是 12(舊結構體大小),12 < 16成立,函數返回AUTO——但舊調用方其實是有recv_topk_idx_kind欄位的,只是偏移不同。這會導致舊調用方設置的recv_topk_idx_kind被忽略。

更糟的是,如果舊調用方按舊偏移(8)寫入了recv_topk_idx_kind,新庫按新偏移(12)讀取,會讀到new_field的值,完全錯亂。

所以 ABI 設計的鐵律是:新欄位只能加在結構體末尾。這樣舊調用方的size小於新欄位的field_end,函數正確返回預設值;新調用方的size覆蓋新欄位,正常讀取。中間插入欄位會破壞所有基於offsetof的版本判斷。

本章剖析了生產環境中四類典型踩坑及其內部防禦機制,這些邊界條件提醒我們,NCCL 的穩定運行不僅依賴核心實現,也離不開周邊生態的適配與擴展。下一章我們將轉向生態與擴展,看看 nccl4py、nccl4rust、nccl_ep、nccl_ubx 這些周邊項目如何把 NCCL 的能力帶給更廣泛的用戶。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 23

第 23 章:生態擴展:nccl4py、nccl4rust、nccl_ep、nccl_ubx 等周邊項目

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 23 章 / 共 25 章

上一章我們排查了 NCCL 在生產環境中的典型故障——group 語意誤用、rank 數不匹配、stream 互動、ABI 版本衝突以及網路逾時。這些問題大多發生在直接使用 C ABI 的場景中,而現代大模型訓練框架往往不直接呼叫 C ABI,而是透過 Python、Rust 等語言綁定,或借助針對 MoE、超頻寬通訊等場景的擴充專案來複用 NCCL 的能力。這些周邊專案放在 bindings/ 和 contrib/ 目錄下,定位是實驗性、社群維護,不繼承核心函式庫的發布品質保證。本章逐一剖析 nccl4py、nccl4rust、nccl_ep、nccl_ubx 和 nccl_checkpoint,看它們如何透過語言綁定、裝置 API 擴充和符號攔截三條路徑,在核心之外建構起豐富的生態。

nccl4py:Cython 綁定與命名空間套件設計

直覺模型:把 C ABI 翻譯成 Python 能懂的話

想像 NCCL 核心是一個只會說 C 語言的外交官,而 Python 訓練腳本是一個只會說 Python 的實習生。nccl4py 就是那個翻譯官——它不改變外交官說的話(NCCL 的行為),只是把「ncclAllReduce(sendbuff, recvbuff, count, ...)」翻譯成「nccl.all_reduce(tensor)」。如果沒有這層翻譯,每個 Python 框架都得自己寫 ctypes 綁定,重複勞動且容易出錯。

分層結構:Cython 底層 + Python 高層

nccl4py 的設計是兩層:底層是 Cython 綁定(nccl/bindings/cynccl.pxd),高層是 Python API(nccl.core)。README 裡明確說了這個分層📎 bindings/nccl4py/README.md:4-4:

nccl4py provides low-level Cython bindings and a high-level Python API

Cython 綁定以.pxd檔案形式隨 wheel 分發,供其他 Cython 擴充直接cimport 📎 bindings/nccl4py/README.md:39-43:

cython
from nccl.bindings cimport cynccl
〔設計推斷與架構權衡〕

為什麼要暴露 Cython 層而不只是 Python 層?因為有些框架(比如 DeepSpeed、Megatron)的核心迴圈在 Cython 裡,每次呼叫都走 Python 直譯器開銷太大。直接cimport cynccl可以讓 Cython 擴充以接近 C 的零開銷呼叫 NCCL 函式。這是「分層暴露」的典型設計——高層給普通使用者,底層給效能敏感場景。

命名空間套件:多個發行版共享nccl前綴

這是 nccl4py 最巧妙的設計。nccl是一個 PEP 420 隱式命名空間套件📎 bindings/nccl4py/README.md:50-51:

nccl is a PEP 420 implicit namespace package. nccl4py provides nccl.bindings and nccl.core; other NCCL extension distributions can provide additional nccl.* subpackages.
〔設計推斷與架構權衡〕

傳統 Python 套件裡,nccl/__init__.py會「擁有」整個nccl命名空間。如果 nccl4py 和 nccl_ep 的 Python 綁定都想提供nccl.xxx,就會衝突——誰先安裝誰贏。PEP 420 命名空間套件解決了這個問題:沒有__init__.py,多個發行版可以各自往nccl/目錄裡放子套件,Python 匯入系統會把它們合併。所以 nccl4py 提供nccl.bindings和nccl.core,nccl_ep 提供nccl.ep,兩者可以共存📎 contrib/nccl_ep/README.md:80-82。

這個設計對生態擴充至關重要:未來任何第三方想加nccl.monitoring、nccl.profiling,都不需要改 nccl4py 的程式碼。

CUDA 版本選擇:extra 機制

安裝時用nccl4py[cu12]或nccl4py[cu13]選擇 CUDA 大版本📎 bindings/nccl4py/README.md:13-17。README 解釋了原因:extras 會安裝對應的 NCCL runtime 和 CUDA Python 依賴📎 bindings/nccl4py/README.md:19。已發布的 wheel 不需要CUDA_HOME或本地 CUDA Toolkit,但從原始碼編譯需要📎 bindings/nccl4py/README.md:20-21。

〔設計推斷與架構權衡〕

這是 Python 生態處理 CUDA 版本碎片化的標準做法。CUDA 12 和 13 的 ABI 不相容,不能用一個 wheel 通吃。用 extra 讓 pip 根據使用者環境選擇正確的二進位依賴,避免了執行時才發現版本不匹配。

生產避坑

坑一:命名空間套件與__init__.py衝突。如果某個第三方套件在nccl/下放了__init__.py,PEP 420 命名空間套件機制會被破壞,導致nccl.core匯入失敗。排查方法:python -c "import nccl; print(nccl.__path__)",如果報AttributeError說明nccl不是命名空間套件。

坑二:Cython ABI 版本漂移。 cynccl.pxd是實驗性 API📎 bindings/nccl4py/README.md:32-32,NCCL 升級時.pxd可能變。依賴cimport cynccl的 Cython 擴充必須和 nccl4py 版本嚴格匹配,否則編譯期符號解析失敗。

nccl4rust:RAII 所有權與裝置側邊界

直覺模型:讓編譯器幫你管生命週期

C 語言裡,你ncclCommInitRank拿到一個 communicator,用完必須ncclCommDestroy。忘了銷毀就洩漏,提前銷毀就崩潰。Rust 的 RAII(Resource Acquisition Is Initialization)機制讓編譯器在變數離開作用域時自動呼叫解構函式——就像飯店房卡,你退房時系統自動結算,不用手動去櫃檯。

nccl4rust 的核心價值就是把這套所有權語意套在 NCCL 的 C ABI 上。

分層結構:五個 crate 各司其職

README 的 Layout 表格列出了五個 crate📎 contrib/nccl4rust/README.md:20-28:

PathPurpose
crates/nccl-sysbindgen 生成的原始 host ABI
crates/ncclRust 風格 host 包裝 + RAII 所有權
crates/nccl-device-sysno_stdCUDA-Oxide 裝置宣告
crates/nccl-device型別化DevComm、Team、Window包裝
shim/純 C-ABI 墊片,只用公開標頭檔
〔設計推斷與架構權衡〕

這個拆分是刻意的。README 解釋了動機📎 contrib/nccl4rust/README.md:30-32:host 應用可以只用nccl而不需要 Rust GPU 編譯器;CUDA-Oxide 核心用nccl-device;需要原始 ABI 的消費者可以選-syscrate。這種「按需分層」讓不同使用者只付自己需要的編譯成本。

關鍵設計:用指標而非值傳遞裝置通訊器

這是 nccl4rust 最值得學習的設計決策。README 的 Host/device ownership boundary 一節📎 contrib/nccl4rust/README.md:211-219:

ncclDevCommCreate produces a versioned public structure in host memory. The host DeviceCommunicator wrapper owns that structure and destroys it before its parent communicator. CUDA-Oxide remains responsible for allocating device memory, copying those bytes, and keeping the copy alive while kernels execute. Kernels construct nccl_device::DevComm from a pointer to that device copy. Using a pointer rather than a by-value Rust mirror keeps the versioned C struct layout out of the kernel argument ABI.
〔設計推斷與架構權衡〕

為什麼不用 Rust 結構體鏡像 C 結構體?因為ncclDevComm_t是版本化的——不同 NCCL 版本欄位可能不同。如果核心參數按值傳遞 Rust 鏡像,那麼核心 ABI 就綁定了特定 NCCL 版本的結構體佈局。一旦 NCCL 升級結構體,所有已編譯的核心都要重編。用指標傳遞則只傳一個位址,核心透過指標存取,佈局變化不影響 ABI。這和上一章講的ncclEpLayoutInfo_t的 size-based ABI 是同一個思路——把版本差異隔離在指標背後。

安全邊界:哪些是 unsafe 的

README 的 Current API contracts 一節列了六條契約📎 contrib/nccl4rust/README.md:230-249,其中關鍵幾條:

  • 原始-syscrate 只鏡像 C ABI,不加所有權或生命週期校驗📎 contrib/nccl4rust/README.md:232-233
  • 當前集合通訊和點對點包裝接受原始裝置指標,宣告為unsafe 📎 contrib/nccl4rust/README.md:42-45
  • 指標翻譯方法回傳原始裝置指標,無法校驗偏移邊界、對齊、peer 成員關係、別名或視窗生命週期📎 contrib/nccl4rust/README.md:242-244
〔設計推斷與架構權衡〕

這是 Rust 綁定 NCCL 的根本困難:NCCL 的很多 API 契約是「緩衝區必須在 CUDA stream 完成前保持有效」,但 Rust 的型別系統無法表達「stream 完成」這個非同步事件。所以這些方法只能是unsafe,把責任交回呼叫者。README 也指出了改進方向📎 contrib/nccl4rust/README.md:44-45:一個 stream-aware 的緩衝區抽象可以把這些要求編碼進安全 API。這是未來工作。

裝置側:CUDA-Oxide 與 LTOIR 墊片

裝置側的核心挑戰是:NCCL 的裝置 API 是 C++ 模板,而 Rust 裝置程式碼(CUDA-Oxide)需要 C ABI。解決方案是一個 C++ 墊片📎 contrib/nccl4rust/README.md:26:

shim/ — CUDA C++ C-ABI shim built exclusively from public nccl.h and nccl_device.h

墊片編譯成 LTOIR(LLVM 中間表示),和 Rust PTX 一起連結成 cubin📎 contrib/nccl4rust/README.md:165-167。README 說明了建置流程📎 contrib/nccl4rust/README.md:158-163:

bash
make device \
  NCCL_INCLUDE_DIR="$NCCL_INCLUDE_DIR" \
  CUDA_HOME="$CUDA_HOME" \
  ARCH=90
〔設計推斷與架構權衡〕

LTOIR 是 NVIDIA 的連結時優化中間格式。用 LTOIR 而非直接編譯成 cubin,是為了讓墊片和 Rust 核心在連結期做跨語言優化——比如內聯墊片函式到 Rust 核心裡。這是「C++ 模板 + Rust 核心」混合編程的關鍵技術。

生產避坑

坑一:NCCL 版本必須精確匹配。README 明確要求Matching NCCL 2.31 headers and runtime 📎 contrib/nccl4rust/README.md:80-81,因為原型直接初始化了早期 NCCL 裝置 API 版本中不同的欄位。標頭檔和libnccl.so版本不一致會導致裝置通訊器欄位錯位。

坑二:CUDA graph 與裝置通訊器。裝置通訊器是 host 記憶體裡的版本化結構,拷貝到裝置後核心透過指標存取。如果 CUDA graph 捕獲時把裝置指標烘焙進核心參數,之後重新建立通訊器會導致 graph 裡的指標失效。這和 nccl_ep 的 RDMA buffer 重分配問題同源。

坑三:安全初始化不能和原始 group 混用。README 警告📎 contrib/nccl4rust/README.md:238-239:安全初始化和產生輸出的管理呼叫不能和原始nccl-sysgroup 狀態混用,因為包裝層觀察不到原始 group 狀態。混用會導致包裝層的輪詢邏輯和原始 group 語意衝突。

nccl_ep:專家並行的 dispatch/combine 原語

直覺模型:MoE 的「分揀中心」

MoE(Mixture of Experts)模型裡,每個 token 要被路由到 top-k 個專家。專家分佈在不同 GPU 上,所以 token 需要跨 GPU 傳輸——這就是 dispatch。專家算完後,結果要送回原 token 所在的 GPU——這就是 combine。nccl_ep 就是這套「分揀中心」的通信引擎。

如果沒有它,每個 MoE 框架都得自己實現 dispatch/combine 的通信邏輯,重複且難以優化。nccl_ep 把它做成 NCCL 生態裡的標準原語。

兩種算法:LL 與 HT

README 說明了兩種算法📎 contrib/nccl_ep/README.md:36-40:

  • Low-Latency (LL):小 batch、延遲敏感(LLM 推理)。用直接點對點 all-to-all 通信。
  • High-Throughput (HT):大 batch 訓練和推理預填充。用分層通信——節點內 NVLink 聚合,節點間 RDMA。利用 Hopper 的 warp-specialized pipeline 和 TMA。
〔設計推斷與架構權衡〕

這兩種算法的分野反映了 MoE 推理和訓練的不同瓶頸。推理時 batch 小,延遲是主要矛盾,所以 LL 用直接點對點避免聚合開銷。訓練時 batch 大,帶寬是主要矛盾,所以 HT 用分層聚合減少跨節點流量。這是典型的「按工作負載特徵選算法」設計。

核心數據結構:ncclEpGroupConfig_t

這是 EP 的配置結構,字段很多📎 contrib/nccl_ep/README.md:339-362。關鍵字段:

  • size和version:ABI 版本檢查,和上一章講的 size-based ABI 同源📎 contrib/nccl_ep/README.md:340-341
  • algorithm:HT 或 LL📎 contrib/nccl_ep/README.md:342
  • max_dispatch_tokens_per_rank:單 rank 最多 dispatch 的 token 數📎 contrib/nccl_ep/README.md:344
  • rdma_buffer_size:LL 模式的 RDMA 緩衝區大小📎 contrib/nccl_ep/README.md:356-356
  • alloc:自定義設備內存分配器📎 contrib/nccl_ep/README.md:359
〔設計推斷與架構權衡〕

rdma_buffer_size的NCCL_EP_AUTO語義值得深挖。README 解釋📎 contrib/nccl_ep/README.md:396-406:AUTO 模式下緩衝區不在ncclEpCreateGroup時分配,而是第一次ncclEpInitHandle時按實際(layout, num_topk)分配。後續 handle 需要更大緩衝區時會集體重分配。這個「惰性分配」設計避免了用戶猜測緩衝區大小,但引入了三個約束📎 contrib/nccl_ep/README.md:396-406:

1. 所有 rank 必須用相同(layout, num_topk)同步調用ncclEpInitHandle

2. 重分配會丟棄舊緩衝區內容,send_only暫存的數據會丟失

3. CUDA graph 捕獲會烘焙 RDMA 基址指針,重分配後必須重新捕獲

這是本章最重要的生產陷阱之一。惰性分配換來了易用性,但把「何時重分配」的複雜度轉嫁給了用戶。

張量描述符:靜態與動態兩種形態

ncclEpTensor_t是輕量值類型📎 contrib/nccl_ep/README.md:310-332。README 展示了兩種用法:

靜態描述符(棧上,NCCL_EP_TENSOR_INIT_INLINE)📎 contrib/nccl_ep/README.md:806-809:

c
ncclEpTensor_t expert_counters = { NCCL_EP_TENSOR_INIT_INLINE,
                                   .ndim = 1, .datatype = ncclInt32,
                                   .data = expert_counters_data,
                                   .sizes = expert_counters_dims };

動態描述符(堆上,ncclEpTensorAlloc)📎 contrib/nccl_ep/README.md:793-798:

c
ncclEpTensor_t* topk_idx = nullptr;
{
    size_t dims[2] = { num_tokens, top_k };
    ncclEpTensorAlloc(&topk_idx, 2, ncclInt64, dims, /*config=*/NULL);
    cudaMalloc(&topk_idx->data, num_tokens * top_k * sizeof(int64_t));
}
〔設計推斷與架構權衡〕

兩種形態的區別在sizes數組的所有權。靜態描述符的sizes是調用者擁有的棧數組,必須活得比描述符久📎 contrib/nccl_ep/README.md:325-326。動態描述符的sizes是庫擁有的堆拷貝,由ncclEpTensorDestroy釋放📎 contrib/nccl_ep/README.md:514-514。公共結構體持有ncclEpTensor_t*指針,所以兩種形態可以在同一個調用裡混用📎 contrib/nccl_ep/README.md:514-514。這個設計讓簡單場景零堆分配,複雜場景有庫管理便利。

執行模式:同步與分階段

README 的 Execution Modes 一節📎 contrib/nccl_ep/README.md:701-741說明了兩種模式:

同步模式(默認):整個操作期間佔用 GPU 資源,包括等待數據接收的時間📎 contrib/nccl_ep/README.md:705-709。

分階段模式(僅 LL):操作拆成 send 和 receive 兩階段📎 contrib/nccl_ep/README.md:718-726。用send_only = 1發起,數據傳輸啟動後釋放 GPU 資源,應用可以用這些資源做計算,最後用ncclEpComplete完成📎 contrib/nccl_ep/README.md:728-741。

mermaid
sequenceDiagram
    participant App as 应用线程
    participant EP as ncclEpDispatch
    participant GPU as GPU 内核
    participant Net as RDMA 网卡
    App->>EP: ncclEpDispatch(send_only=1)
    EP->>GPU: 启动发送内核
    GPU->>Net: GIN put/signal 发起传输
    EP-->>App: 立即返回,释放 SM
    Note over App: 应用用释放的 SM 做计算
    App->>EP: ncclEpComplete()
    EP->>GPU: 启动接收内核
    GPU->>Net: 等待数据到达
    Net-->>GPU: 数据写入
    GPU-->>EP: 完成
    EP-->>App: 返回,数据就绪

這張時序圖展示了分階段模式的核心價值:send_only發起後立即返回,SM 資源被釋放給計算,等應用做完其他工作再調ncclEpComplete等待接收完成。這是「計算-通信重疊」的經典模式。

生產避坑

坑一:ncclEpInitHandle的條件集體性。AUTO 模式下,ncclEpInitHandle是條件集體調用📎 contrib/nccl_ep/README.md:396-406。如果某個 rank 因為 layout 不同觸發了重分配,其他 rank 必須同步參與。不同步會導致死鎖或數據錯亂。

坑二:CUDA graph 捕獲期間禁止ncclEpInitHandle。README 明確警告📎 contrib/nccl_ep/README.md:396-406:AUTO 模式下不能在cudaStreamBeginCapture和cudaStreamEndCapture之間調用ncclEpInitHandle。因為重分配會改變 RDMA 基址,而 graph 捕獲已經烘焙了舊指針。

坑三:guard 開銷。README 提到📎 contrib/nccl_ep/README.md:299-303:EP 默認給內部通信緩衝區加 guard,防止相鄰 dispatch/combine 調用互相破壞數據。高級用戶如果已保證連續操作不會競爭,可以用NCCL_EP_DISABLE_GUARD=1關閉以回收開銷。但關錯了會導致數據靜默損壞。

nccl_ubx:融合集合通信與對稱分配器

直覺模型:把「搬家前後的打包拆包」也交給搬家公司

普通集合通訊只負責搬資料。但實際模型裡,AllReduce 之前往往要做殘差加法,之後要做 RMSNorm。如果這些操作分開做,資料要在顯存裡多走幾趟。nccl_ubx 的思路是:把殘差加法、RMSNorm、mxfp8 量化都融合進集合通訊內核📎 contrib/nccl_ubx/README.md:6-9。就像搬家公司不僅搬箱子,還幫你打包和拆包,一趟搞定。

硬體前提:必須有 NVLink 多播

README 明確要求 SM 9.0+(Hopper/Blackwell),且 MC 內核路徑需要 NVLink 多播硬體📎 contrib/nccl_ubx/README.md:24-24。SM 8.0(A100)不支援,因為 Ampere 沒有 NVLink 多播硬體,multimem.*內聯 PTX 無法為 arch 8.0 彙編📎 contrib/nccl_ubx/README.md:24-24。

〔設計推斷與架構權衡〕

這解釋了為什麼 ubx 是「實驗性」的——它依賴 Hopper 才引入的 NVLink 多播能力。multimem.*指令允許一個 GPU 用一條指令把資料寫到多個 GPU 的對稱位址,這是硬體加速的集合通訊基礎。沒有這個硬體,ubx 的核心優化就不成立。

對稱分配器:讓 PyTorch 張量變成 NCCL 視窗

ubx 的核心是自訂對稱分配器📎 contrib/nccl_ubx/README.md:11-14:

A central piece of the design is a custom symmetric allocator that provides zero-copy collective input/output buffers while remaining easy to plug into existing PyTorch code: tensors are ordinary torch.Tensor instances backed by an NCCL-managed symmetric window.
〔設計推斷與架構權衡〕

這是 ubx 最巧妙的地方。NCCL 的對稱記憶體要求所有 rank 用同一套虛擬位址存取緩衝區(第 14 章講過)。但 PyTorch 使用者習慣用torch.Tensor。ubx 讓torch.Tensor的底層儲存直接是 NCCL 對稱視窗,這樣使用者程式碼不用改,但集合通訊可以零拷貝——輸入輸出緩衝區就是對稱記憶體本身,不需要額外的拷貝。

集合通訊變體與自動選擇

README 的 Available collectives 表格📎 contrib/nccl_ubx/README.md:90-90:

OpVariantsAuto-select
AllReducemc, uc, lamport, autoLamport ≤ 0.25 MB, else MC
AllToAlluc, lamport, autoLamport ≤ 0.25 MB, else UC
AllGathermc—
〔設計推斷與架構權衡〕

三種變體的區別:mc用 NVLink 多播硬體,uc用普通單播,lamport是低延遲演算法。自動選擇按 0.25 MB 分界——小訊息用 Lamport 低延遲,大訊息用 MC/UC 高頻寬。這個閾值和 NCCL 核心的 tuning 邏輯類似,但 ubx 簡化成了固定閾值。

融合操作:residual + RMSNorm

README 提到📎 contrib/nccl_ubx/README.md:103-103:

SymmAllocator.allreduce_mc() and allreduce_lamport() accept optional gamma/residual_in parameters to fuse residual addition + RMSNorm into the same kernel.
〔設計推斷與架構權衡〕

這是 ubx 的核心賣點。傳統流程是:AllReduce → 殘差加法 → RMSNorm,三次顯存讀寫。融合後一次內核完成,顯存頻寬節省 2/3。對頻寬受限的大模型訓練,這是實打實的加速。

MoE token dispatch + mxfp8 量化

README 描述了a2av_token_bf16_mxfp8 📎 contrib/nccl_ubx/README.md:103-103:

a single GPU kernel that routes bf16 tokens to remote ranks while quantizing them to mxfp8 (E8M0 scale per 32 elements) on the fly.
〔設計推斷與架構權衡〕

這個內核把「路由 + 量化」融合。bf16 是 16 位,mxfp8 是 8 位,量化後資料量減半,跨節點傳輸頻寬需求減半。在傳輸前量化比傳輸後量化更優——省的是網路頻寬而非顯存頻寬。這是 MoE 推理的關鍵優化。

生產避坑

坑一:TORCH_CUDA_ARCH_LIST必須帶a後綴。README 強調📎 contrib/nccl_ubx/README.md:47-56:用a後綴確保存取完整multimem.*指令集。有些加速專用變體在普通9.0/10.0上不可用,未來內核用這些變體會靜默降性能或彙編失敗。

坑二:UBX_BUILD_TIMEOUT的執行時開銷。README 說明📎 contrib/nccl_ubx/README.md:47-56:設為 1 會在內核側編譯進 spinloop 逾時,增加執行時開銷(額外的clock64()檢查和逾時時的printf)。只在排查掛死時開啟。

坑三:NCCL_NVLS_ENABLE=0的降級。README 列出這個環境變數📎 contrib/nccl_ubx/README.md:202:設為 0 可以在沒有 NVLink 多播的情況下運行。但 MC 內核路徑會失效,只剩 UC/Lamport 變體,性能大幅下降。

nccl_checkpoint:LD_PRELOAD 攔截與狀態重放

直覺模型:給通訊域拍快照

訓練任務跑了幾小時,突然要遷移到另一台機器,或者要保存狀態以便恢復。普通檢查點只保存模型權重和優化器狀態,但 NCCL 通訊域的狀態(rank 編號、連接、緩衝區)沒法直接序列化。nccl_checkpoint 的思路是:攔截所有 NCCL 呼叫,記錄初始化步驟,恢復時重放這些步驟📎 contrib/nccl_checkpoint/README.md:3-7。

就像錄下你組裝家具的每一步,搬家後按錄影重新組裝,而不是試圖把組裝好的家具整體搬走。

核心機制:LD_PRELOAD 符號攔截

README 的 Design 一節📎 contrib/nccl_checkpoint/README.md:17-20:

The application is launched with LD_PRELOAD=/path/to/libnccl-checkpoint-shim.so in the environment. This allows the library to intercept all calls to NCCL functions to capture all resource initialization steps.
〔設計推斷與架構權衡〕

LD_PRELOAD是 Linux 動態連結器的機制:在應用正常載入共享庫之前先載入指定的.so。如果這個.so裡定義了和 NCCL 同名的符號(比如ncclCommInitRank),動態連結器會優先用的.so裡的版本。這樣 shim 就能攔截所有 NCCL 呼叫,記錄參數,然後在恢復時重放。

檢查點流程

README 的 Python 範例📎 contrib/nccl_checkpoint/README.md:44-58展示了完整流程:

python
nccl_checkpoint.checkpoint_prepare()
drv.cuCheckpointProcessLock(os.getpid(), None)
drv.cuCheckpointProcessCheckpoint(os.getpid(), None)
# CRIU dump happens here.
drv.cuCheckpointProcessRestore(os.getpid(), None)
drv.cuCheckpointProcessUnlock(os.getpid(), None)
nccl_checkpoint.checkpoint_restore()
〔設計推斷與架構權衡〕

流程分四步:

1. checkpoint_prepare():銷毀所有 communicator,讓 CUDA Checkpoint 和 CRIU 能安全 dump 進程狀態📎 contrib/nccl_checkpoint/README.md:25-27

2. cuCheckpointProcessLock/Checkpoint:CUDA 驅動鎖定進程並做檢查點

3. CRIU dump:外部工具把進程記憶體和檔案描述符 dump 到磁碟

4. cuCheckpointProcessRestore/Unlock + checkpoint_restore():恢復進程,重放 NCCL 配置📎 contrib/nccl_checkpoint/README.md:29-31

Redis KVS:跨機器 rendezvous

README 解釋了為什麼需要 Redis📎 contrib/nccl_checkpoint/README.md:33-38:

Because it is useful to restore on different hardware, IP addresses may have changed. There is no convenient way to directly inform the NCCL Checkpoint library of all peer addresses during the restore process, so the library depends on a temporary Redis Key-Value store to be made available.
〔設計推斷與架構權衡〕

恢復時可能換機器,IP 變了。NCCL 通訊域重建需要知道所有 peer 的新位址。但 shim 沒法直接知道這些位址,所以用一個 Redis KVS 做 rendezvous——所有進程把新位址寫到 KVS,從 KVS 讀其他進程的位址。這就像搬家後大家約定在一個公共留言板上交換新位址。

README 說明 Redis 只在恢復引導階段需要📎 contrib/nccl_checkpoint/README.md:221-221,checkpoint_restore()返回後就可以停掉。

限制:三個不支援

README 的 Limitations 一節📎 contrib/nccl_checkpoint/README.md:119-129列了三個限制:

1. ncclWinGetUserPtr()返回的指標在恢復後無效📎 contrib/nccl_checkpoint/README.md:125-126

2. 不支援 CUDA graph 捕獲📎 contrib/nccl_checkpoint/README.md:136-136

3. 不支援裝置 API——ncclDevComm物件和裝置可見的ncclWindow_t值無法恢復📎 contrib/nccl_checkpoint/README.md:136-136

〔設計推斷與架構權衡〕

第三個限制最嚴重。裝置 API 是 NCCL 新方向(第 19 章講的 DevComm),但 checkpoint 不支援。這意味著用裝置 API 的應用(比如 nccl_ep、nccl_ubx)無法用 checkpoint 恢復。這是生態碎片化的體現——新特性跑得快,但可靠性工具跟不上。

生產避坑

坑一:NCCL_CHECKPOINT_KVS_PATH在檢查點前設置,恢復時不可改。README 警告📎 contrib/nccl_checkpoint/README.md:221-221:這個環境變數在檢查點準備階段不用,但會被捕獲進檢查點,恢復時無法輕易修改。所以必須在檢查點前就設好,且恢復環境裡 Redis 位址要匹配。

坑二:NCCL_CHECKPOINT_KVS_TIMEOUT只覆蓋 shim 的 Redis rendezvous。README 說明📎 contrib/nccl_checkpoint/README.md:221-221:預設 300 秒。一旦 communicator 重放進入 NCCL 傳輸建立階段,底層 NCCL 傳輸呼叫用它們自己的行為,可能需要傳輸特定的診斷。也就是說,逾時只保護 Redis 階段,傳輸建立階段掛死要靠NCCL_DEBUG排查。

坑三:NCCL 版本必須匹配。README 要求 NCCL 2.31.0 或更新📎 contrib/nccl_checkpoint/README.md:158,且建議NCCL_SRC路徑裡的 NCCL 版本精確匹配執行時 NCCL 函式庫版本📎 contrib/nccl_checkpoint/README.md:156-158。版本不匹配會導致重放時結構體佈局錯位。

設計思考:生態擴展的三種模式

回顧這五個專案,可以歸納出 NCCL 生態擴展的三種模式:

模式一:語言綁定(nccl4py、nccl4rust)。核心挑戰是所有權和生命週期。C 的 ABI 沒有所有權語義,綁定層要自己補。nccl4py 用 Cython 分層,nccl4rust 用 RAII +unsafe邊界。共同點是:把版本差異隔離在指標背後——nccl4rust 用指標傳 DevComm,nccl4py 用命名空間套件隔離版本。

模式二:裝置 API 擴展(nccl_ep、nccl_ubx)。核心挑戰是 ABI 版本管理和資源生命週期。nccl_ep 用 size-based ABI(上一章詳述),nccl_ubx 用對稱分配器。共同點是:惰性分配 + 集體重分配——nccl_ep 的 RDMA buffer 和 nccl_ubx 的對稱池都是按需分配,但重分配需要所有 rank 同步。

模式三:符號攔截(nccl_checkpoint)。核心挑戰是狀態捕獲和重放。用LD_PRELOAD攔截所有 NCCL 呼叫,記錄初始化步驟,恢復時重放。這種模式不改 NCCL 核心,但能透明地給現有應用加檢查點能力。

〔設計推斷與架構權衡〕

三種模式的共同約束是NCCL 版本相容性。所有專案都要求精確匹配的 NCCL 版本,因為 NCCL 的 ABI 在演進。這反映了 NCCL 生態的一個根本張力:核心快速迭代,但周邊專案需要穩定性。size-based ABI、指標傳遞、命名空間套件都是緩解這個張力的技術手段。

mermaid
flowchart TD
    start["用户想扩展 NCCL"] --> q1{"扩展什么?"}
    q1 -->|"语言互操作"| lang["语言绑定"]
    q1 -->|"新通信模式"| dev["设备 API 扩展"]
    q1 -->|"可靠性"| ckpt["符号拦截"]
    lang --> q2{"性能敏感?"}
    q2 -->|"是"| cython["Cython 底层 + Python 高层<br/>nccl4py"]
    q2 -->|"否"| raii["RAII 包装<br/>nccl4rust"]
    dev --> q3{"需要 MoE?"}
    q3 -->|"是"| ep["dispatch/combine<br/>nccl_ep"]
    q3 -->|"否"| ubx["融合集合通信<br/>nccl_ubx"]
    ckpt --> preload["LD_PRELOAD 拦截<br/>nccl_checkpoint"]
    cython --> abi{"ABI 版本管理"}
    raii --> abi
    ep --> abi
    ubx --> abi
    preload --> abi
    abi -->|"指针传递"| safe["版本差异隔离"]
    abi -->|"size-based"| safe
    abi -->|"命名空间包"| safe

這張決策圖展示了擴展 NCCL 的選擇路徑。無論走哪條路,最終都要面對 ABI 版本管理這個核心問題,而三種技術手段(指標傳遞、size-based ABI、命名空間套件)都是把版本差異隔離在穩定介面背後。

本章小結

本章剖析了 NCCL 生態的五個周邊專案:

  • nccl4py用 Cython 分層 + PEP 420 命名空間包,讓 Python 生態能零衝突地擴展nccl.*子包。
  • nccl4rust用 RAII 所有權 + 指標傳遞設備通訊器,把版本化 C 結構體佈局隔離在內核 ABI 之外。
  • nccl_ep用 LL/HT 雙演算法 + 惰性 RDMA 緩衝區分配,為 MoE 提供 dispatch/combine 原語,但引入了條件集體調用和 CUDA graph 失效的約束。
  • nccl_ubx用對稱分配器 + 內核融合,把殘差加法、RMSNorm、mxfp8 量化折進集合通信內核,但依賴 Hopper+ 的 NVLink 多播硬體。
  • nccl_checkpoint用LD_PRELOAD符號攔截 + Redis rendezvous,實現跨機器通信域檢查點,但不支持設備 API 和 CUDA graph。

本章思考與自測

Q1: nccl_ep 的rdma_buffer_size = NCCL_EP_AUTO模式下,如果 rank 0 先調用了ncclEpInitHandle且觸發了緩衝區重分配,而 rank 1 因為 layout 不同沒有觸發重分配,會發生什麼?請結合📎 contrib/nccl_ep/README.md:396-406的約束分析。

參考解析:README 明確說明📎 contrib/nccl_ep/README.md:396-406:All ranks must call ncclEpInitHandle in lockstep with the same (layout, num_topk)。AUTO 模式下ncclEpInitHandle是條件集體調用——是否觸發重分配取決於該 handle 的(layout, num_topk)是否需要比當前緩衝區更大的空間。

如果 rank 0 的 layout 需要更大緩衝區觸發重分配,而 rank 1 的 layout 不需要,那麼 rank 0 會執行「deregister window → free → ncclMemAlloc → register」這套集體操作📎 contrib/nccl_ep/README.md:396-406,而 rank 1 不會。這導致兩個問題:

1. 集合操作不匹配:NCCL 的 window deregister/register 是集合操作,需要所有 rank 參與。rank 0 單方面執行會導致 rank 1 在後續通信中引用舊的窗口句柄,而 rank 0 已經換了新窗口,通信失敗或數據錯亂。

2. 基址不一致:重分配後 rank 0 的 RDMA 基址變了,rank 1 沒變。雖然 README 說「recorded layout offsets on every live handle are pure offsets relative to the group's rdma_buffer and resolve correctly against the new base」📎 contrib/nccl_ep/README.md:396-406,但這只在所有 rank 都重分配的前提下成立。rank 1 的基址沒變,rank 0 的變了,跨 rank 的地址解析會錯位。

正確做法是:所有 rank 用相同的(layout, num_topk)同步調用ncclEpInitHandle,確保重分配決策一致。如果無法保證,應該用顯式rdma_buffer_size > 0模式,在ncclEpCreateGroup時一次性分配足夠大的緩衝區,避免運行期重分配📎 contrib/nccl_ep/README.md:396-406。

Q2: nccl4rust 為什麼用指標而非值傳遞ncclDevComm_t給設備內核?如果改成值傳遞,在 NCCL 升級結構體佈局後會發生什麼?請結合📎 contrib/nccl4rust/README.md:211-219分析。

參考解析:README 明確說明📎 contrib/nccl4rust/README.md:217-219:Kernels construct nccl_device::DevComm from a pointer to that device copy. Using a pointer rather than a by-value Rust mirror keeps the versioned C struct layout out of the kernel argument ABI.

ncclDevComm_t是版本化的公共結構體,不同 NCCL 版本字段可能不同。如果用值傳遞:

1. 內核 ABI 綁定結構體佈局:內核參數按值傳遞時,編譯器會把整個結構體的位元組佈局烘焙進內核的調用約定。NCCL 升級結構體(加字段、改字段順序、改對齊)後,已編譯的內核仍然按舊佈局解析參數,導致字段錯位。

2. 必須重編所有內核:每次 NCCL 升級都要重新編譯所有使用設備通訊器的內核。對於部署在大量機器上的訓練任務,這是巨大的運維負擔。

3. 跨版本不兼容:如果 host 側用新 NCCL 創建通訊器,設備側內核用舊 NCCL 編譯,值傳遞會導致內核讀到錯誤字段。

用指標傳遞則只傳一個 8 位元組地址,內核通過指標訪問結構體。NCCL 升級結構體佈局時,只要 host 側用新版本創建通訊器、拷貝到設備,內核通過指標訪問的就是新佈局。內核本身不需要重編,因為它的參數只是一個地址。這把版本差異隔離在了指標背後——指標是穩定的,指標指向的內容可以變。

這和 nccl_ep 的 size-based ABI 是同一個設計哲學:用一層間接把易變的版本細節隔離在穩定接口背後。

Q3: nccl_checkpoint 用LD_PRELOAD攔截 NCCL 調用,但如果應用同時鏈接了 nccl4py 和 nccl_checkpoint,nccl4py 的 Cython 綁定直接調用libnccl.so的符號,LD_PRELOAD能攔截到嗎?請分析符號解析順序。

參考解析:這取決於符號解析順序。LD_PRELOAD的機制是:動態連結器在載入應用正常依賴的共享函式庫之前,先載入LD_PRELOAD指定的.so。當應用(或它依賴的函式庫)引用一個符號時,動態連結器按「先載入先解析」的順序查找——LD_PRELOAD的.so優先於libnccl.so。

所以理論上,nccl4py 的 Cython 綁定呼叫ncclCommInitRank時,動態連結器會先找到libnccl-checkpoint-shim.so裡的同名符號,攔截成功。

但有幾個邊界情況:

1. 直接dlopen + dlsym:如果 nccl4py 用dlopen("libnccl.so")然後dlsym拿函式指標,LD_PRELOAD攔截不到,因為dlsym直接在指定的.so裡找符號,不走全域符號表。README 提到 C 應用用dlsym解析ncclCheckpointPrepare 📎 contrib/nccl_checkpoint/README.md:109-109,但那是解析 checkpoint 自己的符號,不是 NCCL 符號。

2. 符號綁定時機:如果 nccl4py 在LD_PRELOAD生效前就綁定了 NCCL 符號(比如在__attribute__((constructor))裡),攔截可能失效。但正常情況LD_PRELOAD在行程啟動時就生效,早於任何使用者程式碼。

3. RTLD_DEEPBIND:如果 nccl4py 用dlopen時指定RTLD_DEEPBIND,符號查找會優先在libnccl.so內部解析,繞過LD_PRELOAD。這是常見的坑。

4. 靜態連結:如果 nccl4py 靜態連結了 NCCL,LD_PRELOAD完全無效,因為符號已經在編譯期解析。

所以結論是:正常動態連結場景下LD_PRELOAD能攔截 nccl4py 的呼叫,但如果 nccl4py 用了dlopen + RTLD_DEEPBIND或靜態連結,攔截會失效。生產使用時應該用LD_DEBUG=bindings驗證符號綁定,確認 NCCL 呼叫被 shim 攔截。

下一章我們將轉向架構演進與未來方向,看看 NCCL 如何從集合通訊庫演進為可程式化通訊引擎。

這些周邊專案透過語言綁定、裝置 API 擴展和符號攔截,展示了 NCCL 核心能力在不同場景下的復用方式。而貫穿所有專案的核心約束是 NCCL ABI 版本相容性——size-based ABI、指標傳遞、命名空間包都是把版本差異隔離在穩定介面背後的技術手段。理解這些手段,是安全使用這些周邊專案的前提。當這些擴展專案不斷試探核心的邊界,NCCL 自身也在悄然演進:從固定集合操作走向可程式化通訊引擎,從 host proxy 走向 GPU 直發,從註冊緩衝區走向對稱記憶體。下一章我們將基於原始碼中的演進痕跡,探討這些變化將如何重塑上層框架的通訊方式。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 24

第 24 章:架構演進與未來方向:從靜態通訊到可程式化通訊

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 24 章 / 共 25 章

上一章我們看到社群如何圍繞 NCCL 核心構建周邊生態:Python 綁定、Rust 綁定、專家並行通訊、超頻寬原語、通訊檢查點。這些專案都在復用 NCCL 的穩定 API,但它們的訴求已經超出了傳統集合通訊的範疇——專家並行需要細粒度的點對點收發,檢查點需要暫停/恢復通訊狀態,超頻寬原語需要繞過標準集合操作直接操作網路。這些訴求指向同一個問題:NCCL 的固定集合操作模型,正在被更靈活的通訊需求撐破。本章我們不再看某個單一模組,而是從原始碼中已經出現的演進痕跡出發,討論 NCCL 正在走向何方。具體來說,我們將剖析三股交織的演進力量:通訊原語從固定集合走向可程式化——src/rma/rma.cc 中的 RMA 任務調度,讓上層可以組合 Put/Signal/WaitSignal 原語,而不是只能呼叫 AllReduce;網路發起從 host proxy 走向 GPU 直發——src/gin/gin_host.cc 中的 GIN 後端管理,讓 GPU kernel 直接驅動網卡;記憶體模型從註冊緩衝區走向對稱記憶體——src/sym_kernels.cc 中的對稱記憶體 kernel 選擇,讓所有 rank 用同一套虛擬位址存取彼此的緩衝區。這三股力量不是孤立的,它們共享同一個基礎設施:src/nccl_device/core.cc 中的 team 抽象和 src/devcomm/devcomm_v23100.cc 中的版本化 DevComm。理解它們如何咬合,就理解了 NCCL 從「集合通訊庫」到「可程式化通訊引擎」的演進邏輯。

一、可程式化通訊原語:RMA 如何把「固定菜譜」變成「自助餐」

直覺模型

傳統 NCCL 的集合通訊像一份固定套餐:你點 AllReduce,廚房就按 AllReduce 的流程做完。但專家平行(MoE)場景下,每個 token 要發給不同的專家,發送模式在編譯期根本不知道——這就像自助餐,你得自己決定拿什麼、拿多少、什麼時候拿。

RMA 就是 NCCL 給上層提供的「自助餐台」:Put(把資料寫到對端記憶體)、Signal(通知對端)、WaitSignal(等待對端信號)。上層框架可以自由組合這三個原語,實現任意通訊模式。

如果沒有 RMA,MoE 的 all-to-all 只能靠多次小規模集合操作模擬,每次都要走完整的 kernel 啟動和同步流程,延遲高得無法接受。

資料結構與記憶體佈局

RMA 的核心資料結構是ncclTaskRma(任務描述)和ncclRmaArgs(計劃參數)。我們先看ncclRmaArgs的欄位,它在scheduleRmaTasksToPlan中被初始化。

📎 src/rma/rma.cc:166-171

cpp
plan->isRma = true;
plan->rmaArgs = ncclMemoryStackAlloc<struct ncclRmaArgs>(&comm->memScoped);
plan->rmaArgs->func = firstTask->func;
plan->rmaArgs->nRmaTasks = 0;
plan->rmaArgs->nRmaTasksProxy = 0;
plan->rmaArgs->nRmaTasksCe = 0;

這裡的關鍵欄位是nRmaTasksProxy和nRmaTasksCe。它們把 RMA 任務分成兩條執行路徑:

  • CE 路徑(Copy Engine,拷貝引擎):目標 rank 在 LSA(Local Symmetric Access,本地對稱存取)範圍內,可以用 GPU 的拷貝引擎直接完成,不需要網路。
  • Proxy 路徑:目標 rank 不在 LSA 範圍內,必須走 host proxy 執行緒驅動網路。
〔設計推斷與架構權衡〕

這種二分法的設計動機很直接:LSA 範圍內的通訊走 NVLink 或 PCIe,頻寬高、延遲低,用 CE 非同步拷貝最划算;跨機通訊必須走網卡,只能由 proxy 執行緒驅動。把兩類任務分開排程,才能讓 CE 和 proxy 平行執行,而不是串列等待。

ncclTaskRma本身包含peers、nsignals、signalIdxs三個陣列指標,分別記錄對端 rank、信號數量、信號索引。對於 WaitSignal 任務,一個任務可以等待多個 peer;對於 Put/Signal 任務,一個任務只針對一個 peer。

Step-by-Step Walkthrough:一次 WaitSignal 的排程

我們代入一個具體場景:rank 0 呼叫ncclWaitSignal,等待 rank 1 和 rank 3 的信號。假設 rank 1 在 LSA 範圍內,rank 3 不在。

第一步:找到第一個非空上下文佇列。

📎 src/rma/rma.cc:148-158

cpp
int ctx = -1;
for (int i = 0; i < comm->config.numRmaCtx; i++) {
  if (!ncclIntruQueueEmpty(&planner->rmaTaskQueues[i])) {
    ctx = i;
    break;
  }
}
if (ctx == -1) return ncclSuccess;

RMA 任務按 context 分佇列,每個 context 是一個獨立的 RMA 通道。這裡找到第一個有任務的 context,取出它的佇列。

第二步:取出第一個任務,判斷類型。

📎 src/rma/rma.cc:163-168

cpp
struct ncclTaskRma* firstTask = ncclIntruQueueDequeue(ctxQueue);
plan->isRma = true;
plan->rmaArgs = ncclMemoryStackAlloc<struct ncclRmaArgs>(&comm->memScoped);
plan->rmaArgs->func = firstTask->func;

firstTask->func是ncclFuncWaitSignal,進入 WaitSignal 分支。

第三步:按 LSA 可達性拆分 peer。

📎 src/rma/rma.cc:187-204

cpp
for (int i = 0; i < firstTask->npeers; i++) {
  int peerRank = firstTask->peers[i];
  bool lsaAccessible = isLsaAccessible(comm, peerRank);
  if (lsaAccessible) {
    peersCe[npeersCe] = peerRank;
    nsignalsCe[npeersCe] = firstTask->nsignals[i];
    signalIdxsCe[npeersCe] = firstTask->signalIdxs[i];
    npeersCe++;
  } else {
    peersProxy[npeersProxy] = peerRank;
    nsignalsProxy[npeersProxy] = firstTask->nsignals[i];
    signalIdxsProxy[npeersProxy] = firstTask->signalIdxs[i];
    npeersProxy++;
  }
}

isLsaAccessible遍歷comm->devrState.lsaRankList,判斷 peer 是否在 LSA 團隊內。rank 1 在 LSA 內,進 CE 列表;rank 3 不在,進 Proxy 列表。

第四步:為 CE 和 Proxy 各建立一個新任務。

📎 src/rma/rma.cc:206-246

cpp
if (npeersCe > 0) {
  struct ncclTaskRma* waitSignalTaskCe = ...;
  waitSignalTaskCe->peers = peersCe;
  waitSignalTaskCe->npeers = npeersCe;
  ncclIntruQueueEnqueue(&plan->rmaTaskQueueCe, waitSignalTaskCe);
  plan->rmaArgs->nRmaTasksCe = 1;
}
if (npeersProxy > 0) {
  struct ncclTaskRma* waitSignalTaskProxy = ...;
  waitSignalTaskProxy->peers = peersProxy;
  waitSignalTaskProxy->npeers = npeersProxy;
  ncclIntruQueueEnqueue(&plan->rmaTaskQueueProxy, waitSignalTaskProxy);
  plan->rmaArgs->nRmaTasksProxy = 1;
}

原來的一個 WaitSignal 任務被拆成兩個:CE 任務等 rank 1,Proxy 任務等 rank 3。兩個任務可以平行執行——CE 路徑在 GPU 上等,Proxy 路徑在 host 執行緒上等。

第五步:釋放原任務。

📎 src/rma/rma.cc:249-251

cpp
planner->nTasksRma -= 1;
ncclMemoryPoolFree(&comm->memPool_ncclTaskRma, firstTask);

原任務已經拆成兩個新任務,釋放回記憶體池。

並行控制與硬體互動

RMA 的平行執行體現在ncclRmaWaitSignal中。

📎 src/rma/rma.cc:43-74

cpp
if (plan->rmaArgs->nRmaTasksProxy > 0 && plan->rmaArgs->nRmaTasksCe > 0) {
  cudaStream_t ceStream = comm->rmaState.rmaCeState.ceStream;
  cudaEvent_t ceEvent = comm->rmaState.rmaCeState.ceEvent;
  CUDACHECKGOTO(cudaEventRecord(ceEvent, stream), ret, fail);
  CUDACHECKGOTO(cudaStreamWaitEvent(ceStream, ceEvent, 0), ret, fail);
  NCCLCHECKGOTO(ncclRmaProxyWaitLaunch(comm, plan, stream), ret, fail);
  NCCLCHECKGOTO(ncclRmaCeWaitLaunch(comm, plan, ceStream), ret, fail);
  CUDACHECKGOTO(cudaEventRecord(ceEvent, ceStream), ret, fail);
  CUDACHECKGOTO(cudaStreamWaitEvent(stream, ceEvent, 0), ret, fail);
}

這段程式碼用 CUDA event 做流間同步:先在輸入流上記錄 event,讓 CE 流等待這個 event,然後在兩個流上分別啟動 proxy 和 CE 任務,最後讓輸入流等待 CE 流的 event。這樣兩條路徑平行推進,但對外表現為一個同步操作。

〔設計推斷與架構權衡〕

這裡的設計權衡是:平行執行能降低延遲,但引入了額外的 event 記錄和流同步開銷。對於小訊息,這個開銷可能超過平行收益;對於大訊息,平行收益顯著。NCCL 沒有在這裡做自適應判斷,而是統一走平行路徑——因為 RMA 的典型場景就是大訊息的細粒度通訊。

生產避坑指南

坑 1:LSA 可達性判斷錯誤導致任務走錯路徑。 isLsaAccessible遍歷lsaRankList,如果lsaSize為 0(比如單 rank 通訊域),所有 peer 都會被判為不可達,全部走 Proxy 路徑。這在小規模測試時不會暴露,但在大規模部署時會導致效能驟降。排查方法是看scheduleRmaTasksToPlan的 INFO 日誌中nRmaTasksProxy和nRmaTasksCe的比例。

坑 2:WaitSignal 任務拆分後 peer 陣列的生命週期。CE 路徑的peersCe用ncclMemoryStackAlloc分配,生命週期跟隨comm->memScoped;Proxy 路徑的peersProxy用ncclCalloc分配,在任務執行完後需要手動free。如果 Proxy 任務建立失敗,fail分支會釋放這些陣列。

📎 src/rma/rma.cc:302-308

cpp
exit:
  return ret;
fail:
  free(peersProxy);
  free(nsignalsProxy);
  free(signalIdxsProxy);
  goto exit;

坑 3:Put/Signal 任務的跨 context 批次。在 Put/Signal 分支中,NCCL 會把所有 context 的 put/signal 任務拉進同一個 plan,但遇到 WaitSignal 就停止。

📎 src/rma/rma.cc:279-295

cpp
for (int c = 0; c < comm->config.numRmaCtx; c++) {
  struct ncclIntruQueue<struct ncclTaskRma, &ncclTaskRma::next>* q = &planner->rmaTaskQueues[c];
  while (!ncclIntruQueueEmpty(q)) {
    struct ncclTaskRma* task = ncclIntruQueueHead(q);
    if (!isRmaPutOrSignal(task->func)) break;
    ncclIntruQueueDequeue(q);
    ...
  }
}

這個設計的意圖是:一次 kernel 啟動涵蓋所有 context 的 put/signal,減少啟動開銷。但每個 context 的佇列只消費到第一個 WaitSignal 為止,保證 per-context FIFO 順序。如果上層在同一個 context 裡交替呼叫 put 和 waitSignal,批次效果會大打折扣——這是使用 RMA 時需要注意的模式。

---

二、GPU 直發網路:GIN 如何讓 kernel 繞過 host proxy

直覺模型

傳統 NCCL 的網路通訊像寄信:GPU kernel 把資料放到緩衝區,host proxy 執行緒把資料交給網卡,網卡發出去。GIN 則是讓 GPU kernel 直接把信投進對方信箱——kernel 直接寫網卡的發送佇列,網卡直接讀 GPU 顯存。

如果沒有 GIN,每次網路通訊都要經過 host 記憶體中轉,延遲至少多一個 PCIe 往返。對於 MoE 這種細粒度通訊,這個延遲是致命的。

資料結構與記憶體佈局

GIN 的核心狀態是ncclGinState,它管理多個後端(backend)和多個 DevComm。我們先看後端版本相容表。

📎 src/gin/gin_host.cc:27-33

cpp
const int proxyBackendMinVersions[] = {0, NCCL_VERSION(2, 30, 3), NCCL_VERSION(2, 30, 5), NCCL_VERSION(2, 32, 0)};
const int gdakiBackendMinVersions[] = {0, NCCL_VERSION(2, 30, 3), NCCL_VERSION(2, 30, 5)};
const int gpiBackendMinVersions[] = {0, NCCL_VERSION(2, 30, 5)};
constexpr int efaGdaBackendMinVersions[] = {0, NCCL_VERSION(2, 31, 0), NCCL_VERSION(2, 32, 0)};

這些陣列的索引是後端版本號,值是相容的最低 NCCL 版本。比如proxyBackendMinVersions[3]對應後端版本 3,要求 NCCL 至少 2.32.0。這個設計讓 NCCL 可以在執行時根據裝置程式碼版本選擇合適後端版本,而不是編譯期綁定。

〔設計推斷與架構權衡〕

這種版本相容表的設計動機是:GIN 後端(網卡驅動、韌體)和 NCCL 函式庫的版本演進節奏不同。如果硬編碼版本要求,任何一方升級都會導致不相容。用陣列做版本映射,可以在執行時動態選擇,向後相容舊後端。

ncclGinStateDevComm是每個 DevComm 的 GIN 狀態,包含contextCount、backendIndex、ginCtx[]、devHandles[]等欄位。它被串成鏈結串列掛在ginState->devComms上。

Step-by-Step Walkthrough:一次 GIN 連線建立

我們代入一個場景:rank 0 初始化通訊域,需要建立 GIN 連線。

第一步:檢查 GIN 是否啟用和支援。

📎 src/gin/gin_host.cc:96-107

cpp
if (ginState->connected) return ncclSuccess;
if (ncclParamGinEnable() == 0) {
  WARN("GIN is disabled.");
  return ncclInternalError;
}
if (!ginState->supported) {
  WARN("GIN not supported.");
  return ncclInvalidUsage;
}

ncclParamGinEnable()讀取環境變數NCCL_GIN_ENABLE,預設 1。如果使用者顯式停用,直接回傳錯誤。

第二步:檢查對稱記憶體支援。

📎 src/gin/gin_host.cc:111-114

cpp
if (!comm->symmetricSupport) {
  WARN("Communicator does not support symmetric memory!");
  return ncclInternalError;
}

GIN 依賴對稱記憶體——因為 GPU kernel 需要知道對端緩衝區的虛擬位址,只有對稱記憶體才能保證位址一致。

第三步:取得本地 GIN 裝置列表。

📎 src/gin/gin_host.cc:116-122

cpp
int nLocalGinDevs;
int localGinDevs[NCCL_TOPO_MAX_NODES];
NCCLCHECK(ncclTopoGetLocalGinDevs(comm, localGinDevs, &nLocalGinDevs));
if (nLocalGinDevs > NCCL_GIN_MAX_CONNECTIONS) {
  ATTN("Found %d local devices, but GIN supports at most %d connections. Using the first %d connections.",
       nLocalGinDevs, NCCL_GIN_MAX_CONNECTIONS, NCCL_GIN_MAX_CONNECTIONS);
}

ncclTopoGetLocalGinDevs從拓撲圖中找出所有支援 GIN 的網卡。如果超過NCCL_GIN_MAX_CONNECTIONS,只取前幾個並列印警告。

第四步:計算 GIN 團隊。

📎 src/gin/gin_host.cc:138-149

cpp
ginTeam = ncclTeamWorld(comm);
if (ginState->ginConnectionType != NCCL_GIN_CONNECTION_FULL) {
  ginTeam = {
    .nRanks = comm->nRanks / comm->contiguousRanksPerHost,
    .rank = comm->rank / comm->contiguousRanksPerHost,
    .stride = comm->contiguousRanksPerHost,
  };
}
for (int r = 0; r < ginTeam.nRanks; r++) {
  int worldRank = ncclTeamRankToWorld(comm, ginTeam, r);
  handles[r] = allHandles + worldRank * NCCL_NET_HANDLE_MAXSIZE;
}

如果連線類型是 FULL,GIN 團隊就是整個世界團隊;否則只連接每個 host 的第一個 rank(rail 連線)。ncclTeamRankToWorld把團隊內 rank 轉成世界 rank。

第五步:逐後端建立連線。

📎 src/gin/gin_host.cc:151-202

cpp
for (int backendIdx = 0; backendIdx < ginState->numActiveBackends; backendIdx++) {
  backend = &ginState->backends[backendIdx];
  NCCLCHECKGOTO(backend->ncclGin->devices(&ndev), ret, fail);
  ...
  for (int commIdx = 0; commIdx < backend->ginCommCount; commIdx++) {
    NCCLCHECKGOTO(backend->ncclGin->listen(...), ret, fail);
    NCCLCHECKGOTO(backend->ncclGin->getProperties(...), ret, fail);
    NCCLCHECKGOTO(bootstrapAllGather(comm->bootstrap, allHandles, NCCL_NET_HANDLE_MAXSIZE), ret, fail);
    NCCLCHECKGOTO(backend->ncclGin->connect(...), ret, fail);
    NCCLCHECKGOTO(backend->ncclGin->closeListen(...), ret, fail);
  }
}

每個後端先呼叫devices取得裝置數量,然後對每個連線執行 listen→getProperties→allGather→connect→closeListen 的流程。bootstrapAllGather在所有 rank 之間交換 handle,這樣每個 rank 都知道對端的連線資訊。

並發控制與硬體互動

GIN 的進度執行緒是核心並發機制。

📎 src/gin/gin_host.cc:56-87

cpp
void* ncclGinProgress(struct ncclGinState* ginState, int threadIdx) {
  if (ncclOsCpuCount(ginState->cpuAffinity)) {
    ncclOsSetAffinity(ginState->cpuAffinity);
  }
  while (1) {
    if (ginState->proxyThreadStopSignal.load()) return NULL;
    if (ginState->writePending.load()) {
      std::this_thread::yield();
      continue;
    }
    {
      std::shared_lock<std::shared_timed_mutex> rlock(ginState->devCommRwMutex);
      struct ncclGinStateDevComm* dc = ginState->devComms;
      while (dc) {
        struct ncclGinBackendState* backend = &ginState->backends[dc->backendIndex];
        for (int commIdx = threadIdx; commIdx < backend->ginCommCount; commIdx += ginState->proxyNthreads) {
          if (dc->devHandles[commIdx]->needsProxyProgress) {
            ncclResult_t ret = backend->ncclGin->ginProgress(dc->ginCtx[commIdx]);
            if (ret != ncclSuccess) {
              COMPILER_ATOMIC_STORE(&ginState->asyncResult, ret, std::memory_order_release);
              return NULL;
            }
          }
        }
        dc = dc->next;
      }
    }
    std::this_thread::yield();
  }
}

這裡有幾個關鍵設計:

1. CPU 親和性:ncclOsSetAffinity把進度執行緒綁定到指定 CPU 核,避免執行緒遷移帶來的快取失效。

2. 寫鎖退避:writePending是一個原子標誌,主執行緒要修改devComms鏈結串列時先置位,進度執行緒看到後主動 yield,避免鎖競爭。

3. 讀寫鎖:devCommRwMutex是shared_timed_mutex,進度執行緒持讀鎖遍歷鏈結串列,主執行緒持寫鎖修改鏈結串列。

4. 執行緒分工:執行緒 t 負責連接 t, t+proxyNthreads, t+2*proxyNthreads, ...,透過 stride 迴圈實現負載均衡。

📎 src/gin/gin_host.cc:43-47

cpp
static void ginProgressWriteLock(struct ncclGinState* ginState) {
  ginState->writePending.store(true);
  ginState->devCommRwMutex.lock();
}
static void ginProgressWriteUnlock(struct ncclGinState* ginState) {
  ginState->devCommRwMutex.unlock();
  ginState->writePending.store(false);
}

這個寫鎖的實現假設只有一個寫者(主執行緒),所以不需要額外的互斥。writePending先置位再拿鎖,確保進度執行緒在拿鎖前就能看到寫意圖,主動退避。

生產避坑指南

坑 1:GIN 連線數不匹配導致 AllGather 死鎖。每個 rank 的ginCommCount可能不同(取決於本地網卡數量),NCCL 透過bootstrapAllGather取所有 rank 的最小值。

📎 src/gin/gin_host.cc:176-180

cpp
ginCommCountHandles[comm->rank] = backend->ginCommCount;
NCCLCHECKGOTO(bootstrapAllGather(comm->bootstrap, ginCommCountHandles, sizeof(int)), ret, fail);
for (int r = 0; r < comm->nRanks; r++) {
  backend->ginCommCount = std::min(backend->ginCommCount, ginCommCountHandles[r]);
}

如果某個 rank 的網卡數量少於其他 rank,所有 rank 都會降到最小值。這保證了連接對稱,但會浪費網卡資源。

坑 2:proxyNthreads 超過 ginCommCount 導致執行緒空轉。如果使用者設定了NCCL_GIN_PROXY_NTHREADS大於ginCommCount,多餘的執行緒會在 stride 迴圈中空轉。

📎 src/gin/gin_host.cc:181-183

cpp
// After cross-rank min, proxyNthreads may exceed ginCommCount if ranks disagree
// on NCCL_GIN_PROXY_NTHREADS (atypical — env vars are normally uniform across a job).
// Extra threads simply idle in the stride loop; no correctness issue.

這不是正確性問題,但會浪費 CPU 資源。排查方法是看NCCL_GIN_PROXY_NTHREADS是否大於實際網卡數。

坑 3:DevComm 釋放時的競態。 ncclGinDevCommFree先從鏈結串列摘除 DevComm,再銷毀 context。

📎 src/gin/gin_host.cc:464-475

cpp
ginProgressWriteLock(ginState);
if (prevDc) prevDc->next = dc->next;
else ginState->devComms = dc->next;
ginProgressWriteUnlock(ginState);
struct ncclGinBackendState* backend = &ginState->backends[dc->backendIndex];
for (int commIdx = 0; commIdx < backend->ginCommCount; commIdx++) {
  NCCLCHECK(backend->ncclGin->destroyContext(dc->ginCtx[commIdx]));
}

摘除後,進度執行緒再也看不到這個 DevComm,所以銷毀 context 是安全的。但如果銷毀過程中有 in-flight 的網路操作,可能會導致未定義行為——這是使用 GIN 時需要確保的:釋放 DevComm 前必須確保所有操作已完成。

---

三、對稱記憶體 kernel:從「註冊緩衝區」到「統一地址空間」

直覺模型

傳統 NCCL 的緩衝區是「註冊制」:每個 rank 註冊自己的緩衝區,通訊時透過 handle 交換地址。對稱記憶體則是「統一地址空間」:所有 rank 約定同一套虛擬地址,rank 0 的地址 A 和 rank 1 的地址 A 指向各自的實體記憶體,但程式碼裡用同一個地址就能存取。

這就像大家約定「第 3 排第 5 座」在每個人家里都指同一個位置,找東西時不用先問「你家第 3 排第 5 座在哪」。

如果沒有對稱記憶體,每個 kernel 都要先解析對端地址,增加了指令開銷和暫存器壓力。

資料結構與記憶體佈局

對稱記憶體 kernel 的核心是 kernel mask——一個位元圖,標記哪些 kernel 在當前通訊域中可用。

📎 src/sym_kernels.cc:17-63

cpp
constexpr uint32_t kernelMask_STMC =
  1 << ncclSymkKernelId_AllGather_LLMC | 1 << ncclSymkKernelId_AllGather_STMC |
  ...
constexpr uint32_t kernelMask_LDMC = ...;
constexpr uint32_t kernelMask_LL = ...;
constexpr uint32_t kernelMask_AG = ...;
constexpr uint32_t kernelMask_AR = ...;
constexpr uint32_t kernelMask_RS = ...;
constexpr uint32_t kernelMask_LSA = ...;
constexpr uint32_t kernelMask_Gin = ...;
constexpr uint32_t kernelMask_Tma = ...;

每個 mask 是一個 32 位元整數,第 i 位為 1 表示 kernel i 可用。這些 mask 按不同維度分組:

  • 按協定:STMC(Simple TMA Multimem Copy)、LDMC(Low-latency Direct Multimem Copy)、LL(Low Latency)
  • 按操作:AG(AllGather)、AR(AllReduce)、RS(ReduceScatter)
  • 按硬體:LSA(Local Symmetric Access)、Gin(GPU-Initiated Networking)、Tma(Tensor Memory Accelerator)
〔設計推斷與架構權衡〕

這種位元圖設計的好處是:可以用位元運算快速篩選可用 kernel。比如kmask &= ~kernelMask_STMC一行就能停用所有 STMC kernel,不需要遍歷列表。

Step-by-Step Walkthrough:一次 kernel mask 計算

我們代入一個場景:rank 0 要執行 AllReduce,資料型別是 float16,訊息大小 1MB,通訊域有 8 個 rank,全部 NVLink 互聯。

第一步:取得操作對應的基礎 mask。

📎 src/sym_kernels.cc:304-306

cpp
uint32_t kmask = kernelMask_coll(coll);

kernelMask_coll(ncclFuncAllReduce)回傳kernelMask_AR,包含 5 個 AllReduce kernel。

第二步:檢查 STMC 和 LDMC 可用性。

📎 src/sym_kernels.cc:308-334

cpp
bool hasSTMC = comm->symkState.hasLsaMultimem;
bool hasLDMC = false;
if (comm->symkState.hasLsaMultimem) {
  switch (ty) {
  case ncclFloat16:
  case ncclBfloat16:
    hasLDMC = red == ncclDevSum || red == ncclDevMinMax || red == ncclDevSumPostDiv;
    break;
  ...
  }
}
if (!hasSTMC) kmask &= ~kernelMask_STMC;
if (!hasLDMC) kmask &= ~kernelMask_LDMC;

hasLsaMultimem在ncclSymkInitOnce中計算,要求 NVLS 對稱多播可用且 LSA 團隊大於 2 個 rank。float16 支援 LDMC,所以如果hasLsaMultimem為真,LDMC kernel 保留。

第三步:檢查訊息大小限制。

📎 src/sym_kernels.cc:336-342

cpp
size_t nBytes = alignUp(nElts * ncclTypeSize(ty), NCCL_SYM_KERNEL_CELL_SIZE);
size_t nBusBytes = (coll == ncclFuncAllReduce ? 1 : comm->nRanks) * nBytes;
if (nBusBytes >= (size_t(2) << 30)) kmask &= ~kernelMask_LL;
if (nBusBytes >= 32 * (size_t(2) << 30)) kmask = 0;

LL kernel 用 32 位元整數追蹤元素計數,所以匯流排位元組數超過 2GB 時停用。如果超過 64GB,所有 kernel 都停用(32 位元整數溢位)。

第四步:檢查 TMA 可用性。

📎 src/sym_kernels.cc:344-345

cpp
if (!ncclSymkTmaAvailable(comm)) kmask &= ~kernelMask_Tma;
if (!symAligned16B) kmask &= ~kernelMask_Tma;

TMA 需要 SMEM 容量和計算能力 10.0+,且緩衝區 16 位元組對齊。

第五步:檢查 GIN 需求。

📎 src/sym_kernels.cc:347-350

cpp
bool hasGin = ncclParamSymGinKernelsEnable() != 0;
if (!hasGin) kmask &= ~kernelMask_Gin;
bool needGin = ncclTeamLsa(comm).nRanks < comm->nRanks;
kmask &= needGin ? kernelMask_Gin : ~kernelMask_Gin;

如果 LSA 團隊覆蓋所有 rank,不需要 GIN;否則只保留 GIN kernel。

並行控制與硬體互動

對稱記憶體 kernel 的初始化涉及 DevComm 建立和資源分配。

📎 src/sym_kernels.cc:185-264

cpp
ncclResult_t ncclSymkInitOnce(struct ncclComm* comm) {
  NCCLCHECK(ncclDevrInitOnce(comm));
  struct ncclSymkState* symk = &comm->symkState;
  if (!symk->initialized) {
    symk->initialized = true;
    struct ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER;
    symk->hasLsaMultimem = ncclNvlsSymmetricMultimemEnabled(comm) && ncclTeamLsa(comm).nRanks > 2 && !comm->p2pCrossClique;
    reqs.lsaMultimem = symk->hasLsaMultimem;
    reqs.lsaBarrierCount = ncclSymkMaxBlocks;
    ...
    NCCLCHECK(ncclDevrCommCreateInternal(comm, &reqs, &symk->kcomm.devComm, /*isInternal=*/true, /*deviceCodeVersion=*/NCCL_VERSION_CODE));
  }
  return ncclSuccess;
}

這裡的關鍵是ncclDevrCommCreateInternal,它建立一個內部 DevComm,包含 LSA 多播、GIN inbox/outbox、訊號等資源。reqs.ginConnectionType = NCCL_GIN_CONNECTION_RAIL指定 GIN 用 rail 連接模式。

📎 src/sym_kernels.cc:257-261

cpp
symk->kcomm.workStarted = comm->profiler.symWorkStarted;
symk->kcomm.workCompleted = comm->profiler.symWorkCompleted;
symk->kcomm.workPhases = comm->profiler.symWorkPhases;

對稱記憶體 kernel 使用獨立的 profiler 緩衝區,避免與常規 kernel 的 workCounter 交錯。

生產避坑指南

坑 1:TMA kernel 的 SMEM 需求。TMA 需要每個 warp 約 8KB 的 SMEM scratch,16 個 warp 就是 128KB。

📎 src/sym_kernels.cc:135-142

cpp
bool ncclSymkTmaAvailable(struct ncclComm* comm) {
  if (comm->maxSharedMemOptin < ncclTmaShmemScratchWarpSize() * 16) {
    return false;
  }
  return comm->minCompCap >= 100 && ncclParamSymTmaEnable();
}

如果 GPU 的 SMEM 容量不足(比如 MIG 實例),TMA kernel 會被停用。排查方法是看maxSharedMemOptin是否小於ncclTmaShmemScratchWarpSize() * 16。

坑 2:GIN chunk size 的邊界。ReduceScatter GIN kernel 的 chunk size 有上下限。

📎 src/sym_kernels.cc:148-153

cpp
static constexpr size_t ncclSymkRsGinDefaultChunkBytes = 128 << 10;
static constexpr size_t ncclSymkRsGinMinChunkBytes = 128;
static constexpr size_t ncclSymkRsGinMaxChunkBytes = size_t(1) << 30;
size_t ncclSymkRsGinChunkBytes() {
  int64_t param = ncclParamSymRsGinChunkSize();
  size_t chunkBytes = param > 0 ? (size_t)param : ncclSymkRsGinDefaultChunkBytes;
  chunkBytes = std::max(ncclSymkRsGinMinChunkBytes, std::min(chunkBytes, ncclSymkRsGinMaxChunkBytes));
  return pow2Down(chunkBytes);
}

如果使用者設定的NCCL_SYM_RS_GIN_CHUNK_SIZE超過 1GB,會被截斷到 1GB;如果小於 128 位元組,會被提升到 128 位元組。最終值還會被向下取整到 2 的冪。

坑 3:對稱記憶體註冊類型不匹配。 ncclGetSymRegType根據 sendWin 和 recvWin 的NCCL_WIN_COLL_SYMMETRIC標誌判斷註冊類型。

📎 src/sym_kernels.cc:395-412

cpp
if (!isSendSymmReg && !isRecvSymmReg) {
  *winRegType = ncclSymSendNonregRecvNonreg;
} else if (isSendSymmReg && !isRecvSymmReg) {
  *winRegType = ncclSymSendRegRecvNonreg;
} else if (!isSendSymmReg && isRecvSymmReg) {
  *winRegType = ncclSymSendNonregRecvReg;
} else if (isSendSymmReg && isRecvSymmReg) {
  *winRegType = ncclSymSendRegRecvReg;
}

如果 send 和 recv 的註冊類型不一致,kernel 需要走不同的程式碼路徑。這會影響效能,但不會導致錯誤。

---

四、Team 抽象與版本化 DevComm:演進的基礎設施

直覺模型

Team 抽象就像「分組」:世界團隊是全班,LSA 團隊是同桌,Rail 團隊是同一列的座位。不同的通訊模式需要不同的分組視角。

版本化 DevComm 就像「翻譯官」:不同版本的裝置程式碼說不同的「方言」,DevComm 相容層負責翻譯,讓新舊程式碼能互相理解。

如果沒有 Team 抽象,每個 kernel 都要自己計算 rank 映射;如果沒有版本化 DevComm,任何 ABI 變化都會導致所有裝置程式碼重新編譯。

資料結構與記憶體佈局

Team 是一個簡單的三元組:nRanks、rank、stride。

📎 src/nccl_device/core.cc:13-19

cpp
ncclTeam_t ncclTeamWorld(ncclComm_t comm) {
  ncclTeam_t ans;
  ans.nRanks = comm->nRanks;
  ans.rank = comm->rank;
  ans.stride = 1;
  return ans;
}

世界團隊的 stride 是 1,因為所有 rank 連續排列。

📎 src/nccl_device/core.cc:70-79

cpp
ncclTeam_t ncclTeamRail(ncclComm_t comm) {
  if (ncclSuccess != ncclDevrInitOnce(comm)) return ncclTeam_t{};
  ncclTeam_t ans;
  ans.nRanks = comm->nRanks / comm->devrState.lsaSize;
  ans.rank = comm->rank / comm->devrState.lsaSize;
  ans.stride = comm->devrState.lsaSize;
  return ans;
}

Rail 團隊的 stride 是lsaSize,因為每個 rail 上的 rank 間隔一個 LSA 團隊的大小。

版本化 DevComm 的核心是ncclDevCommCompat結構。

📎 src/devcomm/devcomm_v23100.cc:10-17

cpp
struct ncclDevCommCompat ncclDevCommCompat_v23100 = {
  NCCL_VERSION(2, 31, 0), // minVersion
  NCCL_VERSION_CODE, // maxVersion
  nullptr,           // commPropertiesFilter
  nullptr,           // devCommRequirementsFilter
  nullptr,           // devCommCopyNewToOld
  nullptr,           // devCommCopyOldToNew
};

這個結構定義了版本 2.31.0 的相容性規則。minVersion和maxVersion定義了適用版本範圍,後面四個函式指標定義了屬性過濾和結構轉換邏輯。如果都是 nullptr,表示這個版本沒有特殊相容需求。

Step-by-Step Walkthrough:一次 Team 轉換

我們代入一個場景:rank 5 在 8 rank 通訊域中,LSA 團隊大小是 4。要計算 rank 5 在 Rail 團隊中的 rank。

第一步:初始化 DevR 狀態。

📎 src/nccl_device/core.cc:70-79

cpp
if (ncclSuccess != ncclDevrInitOnce(comm)) return ncclTeam_t{};

ncclDevrInitOnce計算 LSA 團隊、CFT 團隊等衍生資訊。如果失敗,返回空團隊。

第二步:計算 Rail 團隊參數。

📎 src/nccl_device/core.cc:70-79

cpp
ncclTeam_t ans;
ans.nRanks = comm->nRanks / comm->devrState.lsaSize;  // 8 / 4 = 2
ans.rank = comm->rank / comm->devrState.lsaSize;       // 5 / 4 = 1
ans.stride = comm->devrState.lsaSize;                  // 4

rank 5 在 Rail 團隊中的 rank 是 1,團隊有 2 個 rank,stride 是 4。

第三步:轉換回世界 rank。

📎 src/nccl_device/core.cc:82-84

cpp
int ncclTeamRankToWorld(ncclComm_t comm, ncclTeam_t team, int rank) {
  return comm->rank + (rank - team.rank) * team.stride;
}

如果要把 Rail rank 0 轉成世界 rank:5 + (0 - 1) * 4 = 1。驗證:rank 1 和 rank 5 在同一個 rail 上(間隔 4)。

並行控制與硬體互動

Team 抽象本身是無狀態的,不需要並行控制。但ncclDevrInitOnce是懶載入的,第一次呼叫時會計算所有衍生資訊。

📎 src/nccl_device/core.cc:22-33

cpp
ncclTeam_t ncclTeamLsa(ncclComm_t comm) {
  if (ncclSuccess != ncclDevrInitOnce(comm)) return ncclTeam_t{};
  ncclTeam_t ans;
  ans.nRanks = comm->devrState.lsaSize;
  ans.rank = comm->devrState.lsaSelf;
  ans.stride = 1;
  return ans;
}

註解說「Ignoring errors since if it fails ncclDevrInitOnce will try again」——如果初始化失敗,返回空團隊,下次呼叫會重試。

生產避坑指南

坑 1:Team 轉換的 stride 假設。 ncclTeamRankToWorld假設團隊內 rank 是等差數列。

📎 src/nccl_device/core.cc:82-84

cpp
int ncclTeamRankToWorld(ncclComm_t comm, ncclTeam_t team, int rank) {
  return comm->rank + (rank - team.rank) * team.stride;
}

如果團隊不是等差數列(比如自訂的任意分組),這個函式會算錯。NCCL 目前只支援規則團隊。

坑 2:版本化 DevComm 的空指標。 ncclDevCommCompat_v23100的所有函式指標都是 nullptr,表示沒有特殊相容邏輯。如果未來版本需要轉換,必須實作這些函式,否則新舊程式碼無法互操作。

坑 3:CFT 團隊的層級模式。 ncclTeamCft支援三種模式:FLAT、HIER_MULTIMEM、HIER_LSA。

📎 src/nccl_device/core.cc:36-55

cpp
if (mode == NCCL_CFT_TEAM_FLAT) return flatTeam;
int innerSize;
if (mode == NCCL_CFT_TEAM_HIER_MULTIMEM) {
  innerSize = comm->devrState.cftMcSize;
} else if (mode == NCCL_CFT_TEAM_HIER_LSA) {
  innerSize = comm->devrState.lsaSize;
} else {
  return ncclTeam_t{};
}
return ncclTeamOuterFactor(flatTeam, innerSize);

如果傳入無效模式,返回空團隊。使用 CFT 團隊時需要確保模式正確。

---

設計思考

為什麼 NCCL 要同時支援 RMA、GIN、對稱記憶體三條演進路徑?

〔設計推斷與架構權衡〕

這三條路徑解決的是不同層次的問題:

  • RMA解決「通訊模式固定」的問題——讓上層可以組合原語,實作任意通訊模式。
  • GIN解決「網路延遲高」的問題——讓 GPU 直接驅動網卡,繞過 host proxy。
  • 對稱記憶體解決「位址解析開銷」的問題——讓 kernel 直接用統一地址存取對端記憶體。

它們不是替代關係,而是互補關係。RMA 可以用 GIN 作為底層傳輸,GIN 依賴對稱記憶體提供位址一致性。三者共同構成了「可程式化通訊引擎」的基礎設施。

版本化 DevComm 的設計哲學是什麼?

〔設計推斷與架構權衡〕

版本化 DevComm 的核心思想是「ABI 穩定,API 演進」。裝置程式碼(kernel)編譯後嵌入二進位,不能隨 NCCL 函式庫升級而重新編譯。所以 NCCL 必須保證舊裝置程式碼能在新函式庫上運行。ncclDevCommCompat結構就是相容層的入口:新函式庫根據裝置程式碼版本選擇合適的相容規則,必要時做結構轉換。

---

本章小結

本章我們從原始碼中的演進痕跡出發,剖析了 NCCL 從集合通訊函式庫走向可程式化通訊引擎的三股力量:

1. RMA(src/rma/rma.cc):透過 Put/Signal/WaitSignal 原語組合,讓上層實現任意通訊模式。核心設計是按 LSA 可達性把任務拆成 CE 和 Proxy 兩條路徑並行執行。

2. GIN(src/gin/gin_host.cc):透過 GPU 直發網路,繞過 host proxy。核心設計是多後端管理、版本相容表、進度執行緒池。

3. 對稱記憶體 kernel(src/sym_kernels.cc):透過統一地址空間,消除地址解析開銷。核心設計是 kernel mask 位圖和 TMA/GIN 硬體加速。

4. Team 抽象與版本化 DevComm(src/nccl_device/core.cc、src/devcomm/devcomm_v23100.cc):為演進提供基礎設施。Team 提供分組視角,版本化 DevComm 提供 ABI 相容。

這些變化對上層框架的影響是深遠的:PyTorch 的 ProcessGroup 可以直接呼叫 RMA 原語實現自訂通訊模式;Megatron 的專家並行可以利用 GIN 降低 all-to-all 延遲;對稱記憶體讓 kernel 程式碼更簡潔。

本章思考與自測

Q1:如果把scheduleRmaTasksToPlan中 WaitSignal 分支的 LSA 可達性判斷去掉,所有 peer 都走 Proxy 路徑,會有什麼後果?在什麼場景下會觸發效能災難?

參考解析:

LSA 可達性判斷在📎 src/rma/rma.cc:187-204,它把 peer 分成 CE 和 Proxy 兩組。如果去掉這個判斷,所有 peer 都走 Proxy 路徑,nRmaTasksCe始終為 0。

後果是:CE 路徑完全不被使用,所有 WaitSignal 都透過 host proxy 執行緒輪詢網路。對於 LSA 範圍內的 peer(同機 NVLink 互聯),本來可以用 GPU 拷貝引擎非同步等待,現在變成 host 執行緒輪詢,延遲從微秒級升到毫秒級。

效能災難場景:MoE 訓練中,每個 token 要等待多個專家的信號。如果所有信號都走 Proxy,host 執行緒成為瓶頸,GPU 大量時間在等 host 輪詢。在 8 卡全 NVLink 的機器上,這個退化尤其明顯——本來所有通訊都可以走 CE,現在全部擠到 host。

排查方法:看scheduleRmaTasksToPlan的 INFO 日誌,如果nRmaTasksCe始終為 0 而nRmaTasksProxy很大,說明 LSA 判斷有問題。

Q2:ncclGinProgress中writePending標誌和devCommRwMutex讀寫鎖的配合,如果去掉writePending檢查,只保留讀寫鎖,會有什麼問題?

參考解析:

writePending檢查在📎 src/gin/gin_host.cc:63-66,它讓進度執行緒在主執行緒要寫時主動 yield。如果去掉這個檢查,進度執行緒會直接嘗試拿讀鎖。

問題在於:std::shared_timed_mutex的讀鎖是共享的,多個進度執行緒可以同時持有。如果主執行緒要拿寫鎖,必須等所有讀鎖釋放。在高負載下,進度執行緒頻繁拿讀鎖,主執行緒可能長時間拿不到寫鎖,導致ncclGinDevCommSetup或ncclGinDevCommFree阻塞。

更嚴重的是:如果主執行緒在ginProgressWriteLock中先置位writePending再拿鎖,而進度執行緒不檢查writePending,那麼進度執行緒可能在主執行緒置位後仍然拿讀鎖,導致主執行緒等待時間不可預測。

writePending的作用是「軟性通知」:告訴進度執行緒「我要寫了,你們先讓讓」。這比單純依賴鎖的公平性更高效,因為進度執行緒可以主動 yield 而不是阻塞在鎖上。

Q3:ncclSymkMask中,如果nBusBytes >= 32 * (size_t(2) << 30)時把所有 kernel 都禁用(kmask = 0),此時ncclSymkAvailable返回 false,NCCL 會回退到什麼路徑?這個回退路徑有什麼效能影響?

參考解析:

kmask = 0在📎 src/sym_kernels.cc:342,此時ncclSymkAvailable返回 false(📎 src/sym_kernels.cc:354-361)。

回退路徑是:NCCL 會使用傳統的集合通訊 kernel(非對稱記憶體 kernel)。這些 kernel 透過註冊緩衝區的方式存取對端記憶體,需要先解析地址,指令開銷更大。

效能影響:對於超大訊息(超過 64GB 總線位元組),傳統 kernel 的地址解析開銷佔比很小,因為資料傳輸本身佔主導。但在邊界情況下(剛好超過 64GB),傳統 kernel 可能比對稱記憶體 kernel 慢 10-20%。

這個限制的根本原因是:對稱記憶體 kernel 用 32 位元整數追蹤 unrolled loop chunk,每個 chunk 至少 32 位元組,所以最大可定址範圍是 32 * 2^31 = 64GB。超過這個範圍會整數溢位。

實際生產中,單次集合通訊超過 64GB 的場景很少(通常是梯度累積後的 all-reduce),但並非不可能。如果遇到這種場景,可以考慮分片通訊或使用傳統 kernel。

---

章末過渡

本章我們看到 NCCL 正在從「固定集合操作」走向「可程式化通訊引擎」:RMA 提供原語組合,GIN 提供 GPU 直發,對稱記憶體提供統一地址空間,Team 和版本化 DevComm 提供基礎設施。

這些演進不是孤立的,它們共同指向一個目標:讓上層框架能夠以更低的延遲、更高的靈活性實現自訂通訊模式。對於 PyTorch、Megatron 這樣的框架,這意味著它們可以直接在 NCCL 之上構建 MoE all-to-all、流水線並行、專家並行等複雜通訊模式,而不需要繞過 NCCL 自己實現網路層。

下一章是全書最後一章。我們將把一次 AllReduce 的完整鏈路重新走一遍——從ncclAllReduce呼叫開始,經過任務入隊、演算法選擇、kernel 啟動、proxy 推進、網路傳輸,直到結果返回。這次回顧會把前面 24 章的知識點串聯起來,形成一個完整的認知地圖。

至此,我們看清了 NCCL 從固定集合操作向可編程通訊引擎演進的三條主線:RMA 原語組合、GPU 直發網路、對稱記憶體模型,以及支撐它們的 team 抽象與版本化 DevComm。這些機制共同指向一個更靈活、更貼近硬體能力的通訊未來。然而,無論架構如何演進,一次 AllReduce 的完整鏈路始終是理解 NCCL 的基石。下一章我們將不引入新程式碼,而是把第 3 章到第 10 章的端到端流程重新串講一遍——從 ncclAllReduce 呼叫,到通訊域建立、拓撲搜尋、演算法選型、任務入隊、kernel 啟動、裝置側原語執行、結果回寫。你將把分散在各章的機制重新組裝成一個完整心智模型,並得到一份「遇到問題該查哪一章」的索引。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

CHAPTER 25

第 25 章:全景回顧與思考:一個 AllReduce 的終極旅程與設計精髓

Upstream: NVIDIA/nccl · Commit @12df1a11 · 閱讀進度:第 25 章 / 共 25 章

上一章我們基於原始碼中的演進痕跡,展望了 NCCL 從固定集合操作走向可編程、從 host proxy 走向 GPU 直發、從註冊緩衝區走向對稱記憶體的架構趨勢。現在,是時候把這些趨勢放回一個具體的執行流中檢驗了。這一章不引入任何新程式碼,而是將第 3 章到第 10 章的端到端鏈路重新串聯起來——從 ncclAllReduce 這一行呼叫開始,一路走到結果寫回顯存。讀完之後,你應該能清晰地回答:一次 AllReduce 究竟經過了哪些函式?每個函式在哪個檔案、哪一行?遇到問題時該翻哪一章?

一、初始化:通訊域是怎麼「長」出來的

直覺模型

把通訊域想像成一個「群聊」。你調ncclCommInitRank就是「申請加入群聊」,NCCL 要在這時候把群成員名單(peerInfo)、誰和誰走哪條線(拓撲圖)、每條線開幾條流水線(channel)全部確定下來。如果這一步錯了,後面所有通訊都是錯的——就像群聊裡有人沒被拉進來,你發的訊息永遠少一個人收到。

資料結構與記憶體佈局

通訊域的核心結構是ncclComm,它的初始化分兩段:commAlloc負責「分配骨架」,initTransportsRank負責「填充血肉」。

commAlloc裡最值得注意的是共享資源引用計數的設計。當子通訊域(split/shrink 產生)複用父通訊域資源時,不是拷貝一份,而是共享同一個ncclSharedResources並遞增引用計數:

📎 src/init.cc:533-555

cpp
if (parent == NULL || !parent->shareResources) {
    struct ncclSharedResources* sharedRes;
    NEW_NOTHROW(sharedRes, ncclSharedResources);
    sharedRes->owner = comm;
    ...
    comm->sharedRes = sharedRes;
    sharedRes->refCount = 1;
    NCCLCHECK(ncclNetInit(comm));
    NCCLCHECK(ncclRmaInit(comm));
    NCCLCHECK(ncclGinInit(comm));
} else {
    comm->sharedRes = parent->sharedRes;
    ncclAtomicRefCountIncrement(&parent->sharedRes->refCount);
    NCCLCHECK(ncclNetInitFromParent(comm, parent));
    NCCLCHECK(ncclRmaInitFromParent(comm, parent));
}

這段程式碼的意圖很清晰:網路外掛、RMA、GIN 這些「重資源」只初始化一次,子通訊域直接借用。refCount用原子操作遞增,保證多執行緒下不會重複釋放。

另一個關鍵點是commAlloc裡對通道的初始化。所有通道先被標記為「未初始化」(id = -1),後續setupChannel才會真正填內容:

📎 src/init.cc:607-608

cpp
// Mark channels as non initialized.
for (int c = 0; c < MAXCHANNELS; c++) comm->channels[c].id = -1;

這個-1是個哨兵值。任何程式碼如果誤用了未初始化的通道,id == -1會立刻暴露問題,而不是讀到一堆隨機記憶體。

Step-by-Step:從 ncclCommInitRank 到 initTransportsRank

使用者呼叫ncclCommInitRank後,實際執行流是這樣的:

1. ncclCommInitRank先調ncclInitEnv載入環境外掛,再調ncclGroupStartInternal進入 group 語義(這是為了支援「一次 group 裡初始化多個通訊域」)。

2. 接著調ncclCommInitRankDev,它做參數校驗、分配comm結構、解析 config,然後把真正的初始化工作丟給一個非同步 job:

📎 src/init.cc:2923-2929

cpp
if (ncclParamEnqueueRearchEnable()) {
    NCCLCHECKGOTO(ncclMgmtTaskEnqueue((struct ncclAsyncJob*)job, ncclCommInitRankFunc, ncclCommInitJobFree, comm), res, fail);
} else {
    NCCLCHECKGOTO(ncclAsyncLaunch((struct ncclAsyncJob*)job, ncclCommInitRankFunc, NULL, ncclCommInitJobFree, comm), res, fail);
}

注意這裡的ncclParamEnqueueRearchEnable()分支——這是 NCCL 正在進行的「enqueue 重構」的痕跡。預設走ncclAsyncLaunch,開啟重構後走ncclMgmtTaskEnqueue。兩條路徑最終都會呼叫ncclCommInitRankFunc。

3. ncclCommInitRankFunc是初始化的主函式。它先設裝置、查 GPU 屬性、初始化 kernel:

📎 src/init.cc:2119-2127

cpp
timers[TIMER_INIT_TOTAL] = clockNano();
CUDACHECKGOTO(cudaSetDevice(cudaDev), res, fail);
CUDACHECKGOTO(cudaDeviceGetAttribute(&maxSharedMem, cudaDevAttrMaxSharedMemoryPerBlockOptin, cudaDev), res, fail);
CUDACHECKGOTO(cudaDeviceGetAttribute(&archMajor, cudaDevAttrComputeCapabilityMajor, cudaDev), res, fail);
CUDACHECKGOTO(cudaDeviceGetAttribute(&archMinor, cudaDevAttrComputeCapabilityMinor, cudaDev), res, fail);
cudaArch = 100 * archMajor + 10 * archMinor;

timers[TIMER_INIT_KERNELS] = clockNano();
NCCLCHECKGOTO(ncclInitKernelsForDevice(cudaArch, maxSharedMem, &maxLocalSizeBytes), res, fail);

cudaArch = 100 * archMajor + 10 * archMinor這個編碼方式很實用:sm90 變成 900,sm100 變成 1000,方便後續用整數比較判斷架構代際。

4. 然後根據是普通初始化還是 split/shrink/grow,走不同的 bootstrap 路徑:

📎 src/init.cc:2136-2191

cpp
if (job->parent && !job->isGrow) {
    // SPLIT/SHRINK: use bootstrapSplit
    ...
    NCCLCHECKGOTO(bootstrapSplit(comm->commHash, comm, job->parent, job->color, job->key, parentRanks), res, fail);
} else {
    // GROW or NORMAL INIT: use bootstrapInit
    ...
    NCCLCHECKGOTO(bootstrapInit(job->nId, (struct ncclBootstrapHandle*)job->commId, comm, job->parent), res, fail);
}

5. 最後調initTransportsRank,這是整個初始化裡最重的函式(約 800 行)。它內部做了兩次 AllGather:

  • AllGather1:交換ncclPeerInfo(每個 rank 的裝置資訊、host hash、pid hash、GPU UUID 等):

📎 src/init.cc:1236-1239

cpp
NCCLCHECKGOTO(ncclCalloc(&comm->peerInfo, nranks + 1), ret, fail); // Extra rank to represent CollNet root
NCCLCHECKGOTO(fillInfo(comm, comm->peerInfo + rank, comm->commHash), ret, fail);
NCCLCHECKGOTO(bootstrapAllGather(comm->bootstrap, comm->peerInfo, sizeof(struct ncclPeerInfo)), ret, fail);
COMPILER_ATOMIC_STORE(&comm->peerInfoValid, true, std::memory_order_release);

注意nranks + 1這個分配——多出來的一個位置是給 CollNet root 用的。peerInfoValid用 release 語意儲存,保證其他執行緒看到這個標誌時,peerInfo 的內容已經可見。

  • AllGather3:交換拓撲計算結果(每個 rank 算出的 ring/tree 結構、頻寬、通道數等),然後取所有 rank 的最小值來對齊:

📎 src/init.cc:1687-1703

cpp
for (int i = 0; i < nranks; i++) {
    allTopoRanks[i] = &allGather3Data[i].topoRanks;
    // Make sure we align all ranks so that the tuning is consistent across ranks
    for (int a = 0; a < NCCL_NUM_ALGORITHMS; a++) {
        graphs[a]->nChannels = std::min(allGather3Data[i].graphInfo[a].nChannels, graphs[a]->nChannels);
        graphs[a]->sameChannels = std::min(allGather3Data[i].graphInfo[a].sameChannels, graphs[a]->sameChannels);
        graphs[a]->bwIntra = std::min(allGather3Data[i].graphInfo[a].bwIntra, graphs[a]->bwIntra);
        graphs[a]->bwInter = std::min(allGather3Data[i].graphInfo[a].bwInter, graphs[a]->bwInter);
        graphs[a]->typeIntra = std::max(allGather3Data[i].graphInfo[a].typeIntra, graphs[a]->typeIntra);
        graphs[a]->typeInter = std::max(allGather3Data[i].graphInfo[a].typeInter, graphs[a]->typeInter);
        graphs[a]->crossNic = std::max(allGather3Data[i].graphInfo[a].crossNic, graphs[a]->crossNic);
    }
    ...
}

頻寬取 min、類型取 max,這是「木桶原理」:整個通訊域的效能由最慢的那個 rank 決定。如果不對齊,不同 rank 可能算出不同的演算法選擇,導致通訊死鎖。

初始化流程圖

mermaid
flowchart TD
    api["ncclCommInitRank()"] --> env["ncclInitEnv()"]
    env --> grp["ncclGroupStartInternal()"]
    grp --> dev["ncclCommInitRankDev()"]
    dev --> alloc["ncclCalloc(comm) + parseCommConfig()"]
    alloc --> launch{"ncclParamEnqueueRearchEnable()?"}
    launch -->|是| mgmt["ncclMgmtTaskEnqueue(ncclCommInitRankFunc)"]
    launch -->|否| async["ncclAsyncLaunch(ncclCommInitRankFunc)"]
    mgmt --> func["ncclCommInitRankFunc()"]
    async --> func
    func --> kernels["ncclInitKernelsForDevice(cudaArch)"]
    kernels --> branch{"job->parent && !job->isGrow?"}
    branch -->|是 split/shrink| split["bootstrapSplit()"]
    branch -->|否 grow/normal| init["bootstrapInit()"]
    split --> transports["initTransportsRank()"]
    init --> transports
    transports --> ag1["bootstrapAllGather(peerInfo)"]
    ag1 --> topo["ncclTopoGetSystem() + ncclTopoComputePaths()"]
    topo --> graphs["ncclTopoCompute(ringGraph/treeGraph/nvlsGraph)"]
    graphs --> ag3["bootstrapAllGather(allGather3Data)"]
    ag3 --> align["min/max 对齐所有 rank 的图参数"]
    align --> connect["setupChannel() + ncclTransportRingConnect()"]
    connect --> devcomm["devCommSetup()"]
    devcomm --> done["initState = ncclSuccess"]

設計思考與踩坑

為什麼初始化要非同步?因為多 rank 初始化需要跨行程同步(bootstrap),如果同步執行會阻塞呼叫執行緒。非同步化後,使用者可以在 group 裡同時初始化多個通訊域,並行推進。

踩坑點:initTransportsRank末尾有一個 intra-node barrier:

📎 src/init.cc:1968-1971

cpp
/* Local intra-node barrier */
NCCLCHECKGOTO(bootstrapIntraNodeBarrier(comm->bootstrap, comm->localRankToRank, comm->localRank, comm->localRanks, comm->localRankToRank[0]), ret, fail);

這個 barrier 保證同機所有 rank 都完成了資源分配才繼續。如果某個 rank 卡在devCommSetup裡(比如顯示記憶體不足),其他 rank 會在這裡等死。生產環境遇到「初始化 hang 住」,第一件事就是看是不是某個 rank 的devCommSetup失敗了。

二、任務入隊:從 API 呼叫到內部任務物件

直覺模型

使用者調ncclAllReduce就像在餐廳點菜。ncclEnqueueCheck是服務生,它把你的訂單翻譯成廚房能看懂的「工單」(ncclTaskColl),放進comm->planner這個「訂單池」裡。如果沒有這一層,NCCL 就沒法把多次呼叫合併成一次 kernel 啟動——每次點菜都單獨開火,效率極低。

資料結構與記憶體佈局

任務入隊的核心是ncclKernelPlanner,它掛在comm->planner上。關鍵欄位包括:

  • collSorter:按流量大小排序的集合通訊任務佇列
  • collTaskQueue:最終排好序的任務佇列
  • peers[]:每個 peer 的 send/recv 佇列(P2P 用)
  • wipPlan:正在構建的 kernel plan

任務物件ncclTaskColl的關鍵欄位在collTaskAppend裡填充:

📎 src/enqueue/enqueue.cc:2800-2847

cpp
struct ncclTaskColl* t = ncclMemoryPoolAlloc<struct ncclTaskColl>(&comm->memPool_ncclTaskColl, &comm->memPermanent);
t->func = info->coll;
t->sendbuff = info->sendbuff;
t->recvbuff = info->recvbuff;
t->count = info->count;
t->root = info->root;
t->datatype = info->datatype;
size_t elementSize = ncclTypeSize(t->datatype);
if (t->func == ncclFuncAllGather || t->func == ncclFuncBroadcast) {
    t->count *= elementSize;
    t->datatype = ncclInt8;
    elementSize = 1;
}
t->trafficBytes = t->count * elementSize * ncclFuncTrafficPerByte(t->func, comm->nRanks);
...
t->aggIsolate = ncclCollConfigNeedAggIsolate(&info->collConfig) || info->collConfig.CTAPolicy != comm->config.CTAPolicy;
NCCL_CONFIG_SET(t, minCTAs, ncclParamMinCTAs(), info->collConfig.minCTAs, comm->config.minCTAs, 1, MAXCHANNELS);
NCCL_CONFIG_SET(t, maxCTAs, ncclParamMaxCTAs(), (std::min(info->collConfig.maxCTAs, comm->config.maxCTAs)), comm->config.maxCTAs, 1, MAXCHANNELS);
...
planner->nTasksColl += 1;
ncclTaskCollSorterInsert(&planner->collSorter, t, t->trafficBytes);

注意幾個細節:

1. AllGather/Broadcast 的特殊處理:把 count 乘以元素大小,datatype 改成ncclInt8。這是因為這兩個操作的語意是「搬運位元組」,不需要關心原始類型。

2. trafficBytes的計算:ncclFuncTrafficPerByte回傳每個位元組需要傳輸幾次。AllReduce 回傳 2(reduce + broadcast),AllGather 回傳 nRanks:

📎 src/enqueue/enqueue.cc:123-134

cpp
static inline int ncclFuncTrafficPerByte(ncclFunc_t func, int nRanks) {
  switch (func) {
  case ncclFuncAllReduce:
    return 2;
  case ncclFuncAllGather:
    return nRanks;
  case ncclFuncReduceScatter:
    return nRanks;
  default:
    return 1;
  }
}

3. NCCL_CONFIG_SET巨集:這是「env > per-call > comm」三級配置解析。環境變數優先級最高,其次是單次呼叫的 config,最後是通訊域級別的預設值。

Step-by-Step:ncclAllReduce 的入隊路徑

1. ncclEnqueueCheck先做通訊域校驗和 group 進入:

📎 src/enqueue/enqueue.cc:3478-3495

cpp
ncclResult_t ncclEnqueueCheck(struct ncclInfo* info) {
  ncclResult_t ret = CommCheck(info->comm, info->opName, "comm");
  if (ret != ncclSuccess) return ncclGroupErrCheck(ret);
  if (info->comm->revokedFlag) {
    WARN("%s: communicator was revoked", info->opName);
    return ncclGroupErrCheck(ncclInvalidUsage);
  }
  ...
  NCCLCHECK(ncclGroupStartInternal());
  ret = ncclSuccess;
  int devOld = -1;
  NCCLCHECKGOTO(ncclCommEnsureReady(info->comm), ret, fail);

2. 然後調taskAppend,它根據操作類型分派:

📎 src/enqueue/enqueue.cc:3337-3348

cpp
static ncclResult_t taskAppend(struct ncclComm* comm, struct ncclInfo* info) {
  ncclFunc_t collAPI = info->coll;
  bool hasLaunchCompletionEvent = ncclInfoHasLaunchCompletionEvent(info);

  if (ncclParamEnqueueRearchEnable()) {
    NCCLCHECK(rawTaskAppend(comm, info));
  } else if (info->coll == ncclFuncSend || info->coll == ncclFuncRecv) {
    NCCLCHECK(p2pTaskAppend(comm, info, info->coll, collAPI, (void*)info->recvbuff, info->count, info->datatype, info->root, true));
  } else if (info->coll == ncclFuncPutSignal || info->coll == ncclFuncSignal || info->coll == ncclFuncWaitSignal) {
    NCCLCHECK(rmaTaskAppend(comm, info));
  } else {
    ...
  }
}

對於 AllReduce,走的是最後的else分支,最終調collTaskAppend。

3. collTaskAppend把任務插入collSorter,按trafficBytes排序。排序的目的是讓排程器優先處理大任務,避免小任務碎片化通道資源。

任務入隊資料流

mermaid
flowchart LR
    api["ncclAllReduce()"] --> info["ncclInfo 填充"]
    info --> enq["ncclEnqueueCheck()"]
    enq --> check["CommCheck + ncclCommEnsureReady()"]
    check --> append["taskAppend()"]
    append --> coll["collTaskAppend()"]
    coll --> task["ncclTaskColl 分配"]
    task --> sorter["ncclTaskCollSorterInsert(collSorter)"]
    sorter --> prepare["ncclPrepareTasks()"]
    prepare --> algo["ncclGetAlgoInfo() 选算法"]
    algo --> schedule["scheduleCollTasksToPlan()"]
    schedule --> plan["ncclKernelPlan"]

設計思考與踩坑

為什麼用ncclMemoryPoolAlloc而不是malloc?因為任務物件生命週期短、分配頻繁。記憶體池避免了每次malloc/free的系統呼叫開銷。注意ncclMemoryPoolAlloc的第二個參數是&comm->memPermanent——這意味著任務物件在通訊域銷毀時才統一釋放,而不是每個任務單獨釋放。

踩坑點:ncclPrepareTasks裡有一個「聚合」邏輯,把大小相近(4 倍以內)的任務合併:

📎 src/enqueue/enqueue.cc:506-512

cpp
// We aggregate operations that are within 4X size of each other.
while (aggEnd != nullptr && aggEnd->trafficBytes < 4 * aggBeg->trafficBytes && !aggBeg->aggIsolate && !aggEnd->aggIsolate) {
    agg.count += aggEnd->count;
    agg.trafficBytes += aggEnd->trafficBytes;
    aggEnd = aggEnd->next;
}

這個聚合是為了讓演算法選擇更穩定——如果每個小任務單獨選演算法,可能選出一堆不同的演算法,導致 kernel 碎片化。但aggIsolate標誌會阻止聚合,用於那些「必須單獨排程」的任務(比如帶 per-call config 的)。

三、演算法選型:代價模型怎麼挑出最佳解

直覺模型

演算法選型就像導航軟體選路線。NCCL 的「代價模型」(tuning 模組)會估算每種演算法/協定組合在給定訊息大小和拓撲下的耗時,然後選最快的那個。如果沒有代價模型,NCCL 只能寫死一套演算法,在小訊息上浪費頻寬、在大訊息上浪費延遲。

資料結構與記憶體佈局

演算法選型的入口是ncclGetAlgoInfo:

📎 src/enqueue/enqueue.cc:2159-2185

cpp
ncclResult_t ncclGetAlgoInfo(struct ncclComm* comm, struct ncclTaskColl* info, int collNetSupport, int nvlsSupport,
                             int numPipeOps, ncclSimInfo_t* simInfo) {
  size_t elementSize = ncclTypeSize(info->datatype);
  size_t nBytes = elementSize * ncclFuncMaxSendRecvCount(info->func, comm->nRanks, info->count);
  info->algorithm = NCCL_ALGO_UNDEF;
  info->protocol = NCCL_PROTO_UNDEF;
  struct ncclTuningInput_t input;
  input.comm = comm;
  input.tuningMask = NCCL_TUNING_MASK_GENERAL_KERNELS;
  uint64_t effAlgMask = comm->tuningContext.forced[info->func] ? 0 : info->algMask;
  if (effAlgMask != 0) {
    input.tuningMask = effAlgMask & NCCL_TUNING_MASK_GENERAL_KERNELS;
  }
  input.CTAPolicy = info->CTAPolicy;
  input.func = info->func;
  input.redOp = info->opHost;
  input.devRedOp = info->opDev.op;
  input.datatype = info->datatype;
  input.nBytes = nBytes;
  input.numPipeOps = numPipeOps;
  input.collNetSupport = collNetSupport;
  input.nvlsSupport = nvlsSupport;
  input.count = info->count;
  NCCLCHECK(ncclGetRegBuff(comm, info, &input.regBuff));
  ...
}

注意effAlgMask的邏輯:如果環境變數強制指定了演算法(comm->tuningContext.forced[info->func]非零),則忽略使用者的algMask,用環境變數的。這是「env > per-call」優先級的體現。

然後調ncclTuningCompute得到最佳結果:

📎 src/enqueue/enqueue.cc:2213-2224

cpp
} else {
    NCCLCHECK(ncclTuningCompute(&input, &bestTuning));
}
INFO(NCCL_TUNING, "Best tuning, algorithm, %s, protocol, %s", ncclAlgoToString(bestTuning.algo), ncclProtoToString(bestTuning.proto));
info->algorithm = bestTuning.algo;
info->protocol = bestTuning.proto;
info->nWarps = bestTuning.nWarps;
if (simInfo) simInfo->estimatedTime = bestTuning.timeUs;
TRACE(NCCL_COLL, "%ld Bytes -> Algo %d proto %d time %f", nBytes, info->algorithm, info->protocol, bestTuning.timeUs);
info->nMaxChannels = bestTuning.maxChannels == 0 ? info->nMaxChannels : bestTuning.maxChannels;

Step-by-Step:一次 AllReduce 的演算法選擇

假設 8 卡單機、訊息大小 1MB、AllReduce:

1. nBytes = 1MB,numPipeOps是當前 plan 裡已有的任務數。

2. collNetSupport和nvlsSupport由ncclGetCollNetSupport和ncclNvlsTransportEnabled決定。

3. ncclTuningCompute遍歷所有可用的 (algo, proto) 組合,用代價模型估算時間。

4. 對於 1MB 單機場景,通常 NVLS 或 Tree+LL128 會勝出。

5. 結果寫回info->algorithm、info->protocol、info->nWarps。

演算法選擇決策圖

mermaid
flowchart TD
    start["ncclGetAlgoInfo()"] --> nbytes["计算 nBytes = elementSize * count"]
    nbytes --> forced{"comm->tuningContext.forced[func]?"}
    forced -->|是| envMask["effAlgMask = 0, 用环境变量强制"]
    forced -->|否| userMask{"info->algMask != 0?"}
    userMask -->|是| useUser["tuningMask = algMask"]
    userMask -->|否| full["tuningMask = GENERAL_KERNELS"]
    envMask --> compute["ncclTuningCompute(input, bestTuning)"]
    useUser --> compute
    full --> compute
    compute --> result{"bestTuning.algo == UNDEF?"}
    result -->|是| fallback["重算全量菜单"]
    fallback --> force{"forceAlgSelection?"}
    force -->|是| err["返回 ncclInvalidArgument"]
    force -->|否| auto["回退到自动选择"]
    result -->|否| assign["info->algorithm = bestTuning.algo"]
    auto --> assign
    assign --> done["返回 ncclSuccess"]

設計思考與踩坑

為什麼演算法選擇要「跨 rank 對齊」?因為不同 rank 如果選了不同演算法,通訊模式就不匹配,會死鎖。所以initTransportsRank裡用 min/max 對齊了所有圖參數,保證每個 rank 的代價模型輸入一致。

踩坑點:ncclGetAlgoInfo裡有一個「重算」邏輯——如果使用者指定了algMask但沒有任何演算法匹配,會先靜默重算全量選單,再判斷是硬錯誤還是軟回退:

📎 src/enqueue/enqueue.cc:2192-2208

cpp
NOWARN(ncclTuningCompute(&input, &bestTuning), NCCL_TUNING);
if (bestTuning.algo == NCCL_ALGO_UNDEF) {
    input.tuningMask = NCCL_TUNING_MASK_GENERAL_KERNELS;
    bestTuning = NCCL_TUNING_RESULT_INIT;
    bestTuning.maxChannels = 0;
    NCCLCHECK(ncclTuningCompute(&input, &bestTuning));
    if (info->forceAlgSelection) {
        WARN("algSelection: no algorithm in the selected set is available for %s", ncclFuncToString(info->func));
        return ncclInvalidArgument;
    }
    INFO(NCCL_TUNING, "algSelection: selected set unavailable for %s; falling back to automatic selection", ncclFuncToString(info->func));
}

NOWARN巨集臨時抑制警告,因為「沒有演算法匹配」可能是正常情況(使用者選的集合確實不可用)。只有forceAlgSelection為真時才報錯。

四、任務排程與 kernel plan 構建

直覺模型

任務排程就像把一堆訂單分配到幾條流水線上。scheduleCollTasksToPlan決定每個任務用幾條通道、每條通道處理多少資料,最終生成一個ncclKernelPlan——這就是要傳給 GPU 的「工單」。

資料結構與記憶體佈局

ncclKernelPlan的核心欄位:

  • channelMask:這個 plan 用到哪些通道(位圖)
  • workBytes:所有 work 結構的總位元組數
  • nWorkBatches:work batch 數量
  • kernelArgs:kernel 啟動參數
  • workStorageType:work 資料存哪裡(args/fifo/persistent)

finishPlan決定 work 資料的儲存位置:

📎 src/enqueue/enqueue.cc:244-255

cpp
// If we can fit everything into the kernel args we do so.
if (sizeof(ncclDevKernelArgs) + batchBytes + workBytes <= comm->workArgsBytes) {
    plan->workStorageType = ncclDevWorkStorageTypeArgs;
}
plan->kernelArgsSize = sizeof(struct ncclDevKernelArgs) + batchBytes;
plan->kernelArgsSize += (plan->workStorageType == ncclDevWorkStorageTypeArgs) ? workBytes : 0;
plan->kernelArgsSize = alignUp(plan->kernelArgsSize, 16);
plan->kernelArgs = (struct ncclDevKernelArgs*)ncclMemoryStackAlloc(&comm->memScoped, plan->kernelArgsSize, /*align=*/16);
plan->kernelArgs->comm = comm->devComm;
plan->kernelArgs->channelMask = plan->channelMask;
plan->kernelArgs->workStorageType = plan->workStorageType;

三種儲存類型的權衡:

  • Args:最快,但 kernel 參數大小有限(通常 4KB)
  • Fifo:環形緩衝區,適合中等大小
  • Persistent:獨立顯存分配,適合 CUDA Graph 場景

Step-by-Step:scheduleCollTasksToPlan 的通道分配

1. 先估算這個 plan 能裝多少任務:

📎 src/enqueue/enqueue.cc:654-687

cpp
do {
    size_t workBytes = 0;
    struct ncclTaskColl* task = ncclIntruQueueHead(&planner->collTaskQueue);
    struct ncclWorkList* workNode = ncclIntruQueueHead(&planner->collWorkQueue);
    while (task != nullptr) {
        int nBatches = divUp(nPlanColls, 4); // Rough guess: 4 colls per batch.
        if (!ncclTestBudget(budget, nBatches, workBytes + workNode->size)) goto plan_full;
        bool taskAggIsolate = task->aggIsolate;
        if (taskAggIsolate && nPlanColls > 0) goto plan_full;
        nPlanColls += 1;
        workBytes += workNode->size;
        int kind = 2 * task->isCollnet + task->isNvls;
        trafficBytes[kind] += std::max(MinTrafficPerChannel, task->trafficBytes);
        ...
    }
plan_full:;
} while (0);

2. 然後按流量把通道分配給任務。對於非 CollNet 任務,用「cell」為單位切分:

📎 src/enqueue/enqueue.cc:742-759

cpp
int trafficPerByte = ncclFuncTrafficPerByte(task->func, comm->nRanks);
if (task->protocol == NCCL_PROTO_LL) trafficPerByte *= 4;
size_t cellSize = divUp(divUp(MinTrafficPerChannel, (size_t)trafficPerByte), 16) * 16;
int elementsPerCell = cellSize / elementSize;
size_t cells = divUp(task->count * elementSize, cellSize);
size_t trafficPerElement = elementSize * trafficPerByte;
size_t trafficPerCell = cellSize * trafficPerByte;
size_t cellsPerChannel = std::min(cells, divUp(trafficPerChannel, trafficPerCell));
size_t cellsLo;
if (channelId + 1 == nMaxChannels[kind]) {
    cellsLo = cells;
} else {
    cellsLo = std::min(cells, divUp((trafficPerChannel - currentTraffic), trafficPerCell));
}
int nMidChannels = (cells - cellsLo) / cellsPerChannel;
size_t cellsHi = (cells - cellsLo) % cellsPerChannel;
int nChannels = (cellsLo != 0 ? 1 : 0) + nMidChannels + (cellsHi != 0 ? 1 : 0);

這段程式碼把資料切成「低/中/高」三段:countLo、countMid、countHi。低段和高段是邊界通道,中段是中間通道。這樣切分是為了讓每條通道處理的資料量盡量均勻。

3. 最後調calcCollChunking計算每條通道的 chunk 大小:

📎 src/enqueue/enqueue.cc:2228-2275

cpp
static ncclResult_t calcCollChunking(struct ncclComm* comm, struct ncclTaskColl* info, int nChannels, size_t nBytes,
                                     uint32_t* outChunkSize, uint32_t* outDirectFlags, struct ncclProxyOp* proxyOp) {
  ncclPattern_t pattern;
  size_t grainSize = ncclProtoGrainSize(info->protocol);
  switch (info->func) {
  case ncclFuncAllReduce:
    pattern = info->algorithm == NCCL_ALGO_NVLS           ? ncclPatternNvls :
              info->algorithm == NCCL_ALGO_NVLS_TREE      ? ncclPatternNvlsTree :
              info->algorithm == NCCL_ALGO_COLLNET_DIRECT ? ncclPatternCollnetDirect :
              info->algorithm == NCCL_ALGO_COLLNET_CHAIN  ? ncclPatternCollnetChain :
              info->algorithm == NCCL_ALGO_TREE           ? ncclPatternTreeUpDown :
                                                            ncclPatternRingTwice;
    break;
  ...
  }
  int stepSize = comm->buffSizes[info->protocol] / NCCL_STEPS;
  int chunkSteps = (info->protocol == NCCL_PROTO_SIMPLE && info->algorithm == NCCL_ALGO_RING) ? info->chunkSteps : 1;
  int sliceSteps = (info->protocol == NCCL_PROTO_SIMPLE && info->algorithm == NCCL_ALGO_RING) ? info->sliceSteps : 1;
  int chunkSize = stepSize * chunkSteps;
  if (info->protocol == NCCL_PROTO_LL) chunkSize /= 2;
  if (info->protocol == NCCL_PROTO_LL128) chunkSize = (chunkSize / NCCL_LL128_LINEELEMS) * NCCL_LL128_DATAELEMS;
  ...
}

排程流程圖

mermaid
flowchart TD
    prep["ncclPrepareTasks()"] --> sort["collSorter 按 trafficBytes 排序"]
    sort --> agg["按 (fn,op,ty) 聚合任务"]
    agg --> algo["ncclGetAlgoInfo() 选算法"]
    algo --> bins["按 isCollnet/isNvls 分箱"]
    bins --> sched["scheduleCollTasksToPlan()"]
    sched --> budget{"ncclTestBudget()?"}
    budget -->|否| full["plan_full: 停止添加"]
    budget -->|是| kind{"task->isCollnet?"}
    kind -->|是| collnet["calcCollChunking + 全通道分配"]
    kind -->|否| cells["cell 切分: countLo/Mid/Hi"]
    collnet --> batch["ncclAddWorkBatchToPlan()"]
    cells --> batch
    batch --> proxy["ncclAddProxyOpIfNeeded()"]
    proxy --> finish["finishPlan()"]
    finish --> storage{"workBytes 能放进 args?"}
    storage -->|是| args["ncclDevWorkStorageTypeArgs"]
    storage -->|否| fifo["ncclDevWorkStorageTypeFifo"]

設計思考與踩坑

為什麼 CollNet 任務單獨處理?因為 CollNet 用的是網路交換器做歸約,通道分配邏輯和普通 ring/tree 完全不同。CollNet 任務直接佔用所有可用通道,而普通任務需要按流量切分。

踩坑點:ncclTestBudget的估算用了一個粗略公式nBatches = divUp(nPlanColls, 4)——假設每 4 個集合操作產生一個 batch。這個估算可能不準,所以後面還有精確檢查:

📎 src/enqueue/enqueue.cc:711-714

cpp
// Ensure room for worst case of one new batch per channel
if (!ncclTestBudget(budget, plan->nWorkBatches + nChannels, plan->workBytes + workNode->size)) {
    return ncclSuccess;
}

如果精確檢查失敗,直接返回(不報錯),讓上層再開一個新 plan。

五、Kernel 啟動與裝置側執行

直覺模型

Kernel 啟動就像把工單交給工廠。ncclLaunchKernel把ncclKernelPlan翻譯成 CUDA kernel 啟動參數,然後調cuLaunchKernelEx。裝置側 kernel 收到工單後,按演算法執行資料搬運。

資料結構與記憶體佈局

ncclLaunchKernel的關鍵步驟:

📎 src/enqueue/enqueue.cc:1886-1909

cpp
ncclResult_t ncclLaunchKernel(struct ncclComm* comm, struct ncclKernelPlan* plan) {
  ncclResult_t ret = ncclSuccess;
  struct ncclKernelPlanner* planner = &comm->planner;
  int nChannels = countOneBits(plan->channelMask);
  void* sym = plan->kernelFn;
  dim3 grid = {(unsigned)nChannels, 1, 1};
  dim3 block = {(unsigned)plan->threadPerBlock, 1, 1};
  int smem = plan->isSymColl ? plan->kernelDynSmem : ncclShmemDynamicSize(comm->cudaArch);
  cudaStream_t launchStream = planner->streams->stream;
  ...
  void* extra[] = {CU_LAUNCH_PARAM_BUFFER_POINTER, plan->kernelArgs, CU_LAUNCH_PARAM_BUFFER_SIZE, &plan->kernelArgsSize, CU_LAUNCH_PARAM_END};
  ...
  CUfunction fn;
  CUDACHECKGOTO(cudaGetFuncBySymbol(&fn, sym), ret, do_return);

注意grid.x = nChannels——每個通道一個 block。block.x = plan->threadPerBlock——每個 block 的執行緒數由任務決定。

Step-by-Step:從 plan 到 kernel 啟動

1. 先調uploadWork把 work 資料寫到目標位置(args/fifo/persistent):

📎 src/enqueue/enqueue.cc:1365-1407

cpp
static ncclResult_t uploadWork(struct ncclComm* comm, struct ncclKernelPlan* plan) {
  if (plan->isSymColl || plan->isCeColl || plan->isRma) return ncclSuccess;
  size_t workBytes = plan->workBytes;
  size_t batchBytes = plan->nWorkBatches * sizeof(struct ncclDevWorkBatch);
  void* fifoBufHost;
  uint32_t fifoCursor, fifoMask;
  switch (plan->workStorageType) {
  case ncclDevWorkStorageTypeArgs:
    plan->kernelArgs->workBuf = nullptr;
    fifoBufHost = (void*)plan->kernelArgs;
    fifoCursor = sizeof(ncclDevKernelArgs) + batchBytes;
    fifoMask = ~0u;
    break;
  case ncclDevWorkStorageTypeFifo:
    fifoBufHost = comm->workFifoBuf;
    fifoCursor = comm->workFifoProduced;
    fifoMask = comm->workFifoBytes - 1;
    NCCLCHECK(waitWorkFifoAvailable(comm, fifoCursor + workBytes));
    plan->kernelArgs->workBuf = comm->workFifoBufDev;
    break;
  ...
  }
}

2. 然後構造 CUDA launch 屬性。對於 sm90+,會設定 cluster 維度:

📎 src/enqueue/enqueue.cc:1929-1936

cpp
if (clusterSize) {
    // Grid dimension must be divisible by clusterSize
    if (grid.x % clusterSize) clusterSize = 1;
    launchAttrs[attrs].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
    launchAttrs[attrs++].value.clusterDim = {clusterSize, 1, 1};
    launchAttrs[attrs].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE;
    launchAttrs[attrs++].value.clusterSchedulingPolicyPreference = CU_CLUSTER_SCHEDULING_POLICY_SPREAD;
}

3. 最後調cuLaunchKernelEx:

📎 src/enqueue/enqueue.cc:1992

cpp
CUCHECKGOTO(cuLaunchKernelEx(&launchConfig, fn, nullptr, extra), ret, do_return);

裝置側:runRing 的執行

裝置側 kernel 收到工單後,根據演算法呼叫對應的RunWorkColl特化。以 Ring AllReduce 為例:

📎 src/device/all_reduce.h:14-83

cpp
template <typename T, typename RedOp, typename Proto>
__device__ __forceinline__ void runRing(int tid, int nthreads, struct ncclDevWorkColl* work) {
  ncclRing* ring = &ncclShmem.channel.ring;
  int ringIx = ring->index;
  const int nranks = ncclShmem.comm.nRanks;
  ssize_t gridOffset;
  ssize_t channelCount;
  ssize_t chunkCount;
  ncclCollCbdPart(work, ncclShmem.channelId, Proto::Id, sizeof(T), (ssize_t*)nullptr, &gridOffset, &channelCount, &chunkCount);
  const ssize_t loopCount = nranks * chunkCount;
  ...
  Primitives<T, RedOp, FanSymmetric<1>, 1, Proto, 0> prims(tid, nthreads, &ring->prev, &ring->next, work->sendbuff, work->recvbuff, work->redOpArg, 0, 0, 0, work);

  for (ssize_t elemOffset = 0; elemOffset < channelCount; elemOffset += loopCount) {
    ssize_t remCount = channelCount - elemOffset;
    ssize_t chunkOffset;
    if (remCount < loopCount) chunkCount = alignUp(divUp(remCount, nranks), 16 / sizeof(T));
    auto modRanks = [&] __device__(int r) -> int { return r - (r >= nranks ? nranks : 0); };

    // step 0: push data to next GPU
    chunk = modRanks(ringIx + nranks - 1);
    chunkOffset = chunk * chunkCount;
    offset = gridOffset + elemOffset + chunkOffset;
    nelem = (int)min(chunkCount, remCount - chunkOffset);
    prims.directSend(offset, offset, nelem);

    // k-2 steps: reduce and copy to next GPU
    for (int j = 2; j < nranks; ++j) {
      chunk = modRanks(ringIx + nranks - j);
      chunkOffset = chunk * chunkCount;
      offset = gridOffset + elemOffset + chunkOffset;
      nelem = (int)min(chunkCount, remCount - chunkOffset);
      prims.directRecvReduceDirectSend(offset, offset, nelem);
    }

    // step k-1: reduce this buffer and data, which will produce the final result
    chunk = ringIx + 0;
    chunkOffset = chunk * chunkCount;
    offset = gridOffset + elemOffset + chunkOffset;
    nelem = (int)min(chunkCount, remCount - chunkOffset);
    prims.directRecvReduceCopyDirectSend(offset, offset, nelem, /*postOp=*/true);

    // k-2 steps: copy to next GPU
    for (int j = 1; j < nranks - 1; ++j) {
      chunk = modRanks(ringIx + nranks - j);
      chunkOffset = chunk * chunkCount;
      offset = gridOffset + elemOffset + chunkOffset;
      nelem = (int)min(chunkCount, remCount - chunkOffset);
      prims.directRecvCopyDirectSend(offset, offset, nelem);
    }

    // Make final copy from buffer to dest.
    chunk = modRanks(ringIx + 1);
    chunkOffset = chunk * chunkCount;
    offset = gridOffset + elemOffset + chunkOffset;
    nelem = (int)min(chunkCount, remCount - chunkOffset);
    prims.directRecv(offset, nelem);
  }
}

Ring AllReduce 的經典兩階段:

  • Reduce-Scatter 階段(前 nranks-1 步):每個 rank 把自己的資料發給下一個,同時接收上一個的資料並歸約。
  • AllGather 階段(後 nranks-1 步):把歸約好的結果沿環傳播。

modRanks這個 lambda 處理環形索引回繞:當r >= nranks時減 nranks。

Kernel 啟動時序圖

mermaid
sequenceDiagram
    participant Host as Host 线程
    participant Plan as ncclKernelPlan
    participant CUDA as CUDA Driver
    participant Kernel as GPU Kernel
    participant Proxy as Proxy 线程

    Host->>Plan: ncclLaunchPrepare()
    Plan->>Plan: scheduleCollTasksToPlan()
    Plan->>Plan: finishPlan() 分配 kernelArgs
    Host->>Plan: ncclLaunchKernelBefore_NoUncapturedCuda()
    Plan->>Plan: uploadWork() 写 work 数据
    Host->>CUDA: cuLaunchKernelEx(fn, grid, block, smem)
    CUDA->>Kernel: 启动 nChannels 个 block
    Kernel->>Kernel: runRing() 执行 Ring AllReduce
    Host->>Plan: ncclLaunchKernelAfter_NoCuda()
    Plan->>Proxy: hostStreamPlanTask() + uploadProxyOps()
    Proxy->>Proxy: ncclProxyStart() 推进网络 I/O
    Kernel-->>Host: kernel 完成
    Host->>Plan: ncclLaunchFinish()
    Plan->>Plan: reclaimPlan() 释放资源

設計思考與踩坑

為什麼用cuLaunchKernelEx而不是cudaLaunchKernel?因為需要設定 launch 屬性(cluster 維度、mem sync domain、launch completion event)。這些屬性在 CUDA 12.0+ 才支援。

踩坑點:uploadWork裡對 persistent 模式的處理很複雜——它需要分配顯存、拷貝數據、記錄事件,還要在 CUDA Graph 捕獲模式下正確工作:

📎 src/enqueue/enqueue.cc:1445-1478

cpp
CUDACHECKGOTO(cudaThreadExchangeStreamCaptureMode(&mode), result, fail);
NCCLCHECKGOTO(ncclStrongStreamAcquire(ncclCudaGraphNone(comm->config.graphUsageMode), &comm->sharedRes->deviceStream, /*concurrent=*/false, &deviceStream), result, fail);
if (comm->memPool) {
    CUDACHECKGOTO(cudaMallocAsync(&fifoBufDev, workBytes, comm->memPool, deviceStream), result, fail);
} else {
    CUDACHECKGOTO(cudaMalloc(&fifoBufDev, workBytes), result, fail);
}
plan->workBufPersistent = fifoBufDev;
plan->kernelArgs->workBuf = fifoBufDev;
CUDACHECKGOTO(cudaMemcpyAsync(fifoBufDev, fifoBufHost, workBytes, cudaMemcpyDefault, deviceStream), result, fail);
cudaEvent_t memcpyDone;
CUDACHECKGOTO(cudaEventCreateWithFlags(&memcpyDone, cudaEventDisableTiming), result, fail);
CUDACHECKGOTO(cudaEventRecord(memcpyDone, deviceStream), result, fail);

cudaThreadExchangeStreamCaptureMode是為了在捕獲模式下臨時切換到 relaxed 模式,允許分配顯存。拷貝完成後記錄事件,後續通過ncclCommPollEventCallbacks回收。

六、生產避坑指南

坑 1:初始化 hang 住

現象:ncclCommInitRank卡住不返回。

排查:看NCCL_DEBUG=INFO日誌,找到最後一個打印的 rank。如果所有 rank 都打印了 "Init START" 但沒有 "Init COMPLETE",說明卡在initTransportsRank裡。

常見原因:

  • 某個 rank 的devCommSetup失敗(顯存不足、CUDA 錯誤)
  • bootstrap 網絡不通(防火牆、端口佔用)
  • 不同 rank 的 NCCL 版本不一致

源碼依據:initTransportsRank末尾的 intra-node barrier 會等待所有本機 rank:

📎 src/init.cc:1968-1971

cpp
/* Local intra-node barrier */
NCCLCHECKGOTO(bootstrapIntraNodeBarrier(comm->bootstrap, comm->localRankToRank, comm->localRank, comm->localRanks, comm->localRankToRank[0]), ret, fail);

坑 2:work FIFO 溢出

現象:kernel 啟動後 hang 住,或者報ncclInternalError。

原因:waitWorkFifoAvailable在等 FIFO 空間,但消費端(kernel)沒有推進。

📎 src/enqueue/enqueue.cc:1333-1349

cpp
static ncclResult_t waitWorkFifoAvailable(struct ncclComm* comm, uint32_t desiredProduced) {
  bool hasRoom = (desiredProduced - comm->workFifoConsumed) <= comm->workFifoBytes;
  if (!hasRoom) {
    while (true) {
      // Check abort flag to break deadlock when abort is signaled
      if (COMPILER_ATOMIC_LOAD(comm->abortFlag, std::memory_order_acquire)) {
        return ncclInternalError;
      }
      NCCLCHECK(ncclCommPollEventCallbacks(comm, /*waitSome=*/true));
      hasRoom = (desiredProduced - comm->workFifoConsumed) <= comm->workFifoBytes;
      if (hasRoom) break;
      std::this_thread::yield();
    }
  }
  return ncclSuccess;
}

注意 abort flag 檢查——這是唯一的逃生通道。如果 abort 也沒設,就會死循環。

避坑:調大NCCL_WORK_FIFO_BYTES,或者減少單次 group 裡的操作數。

坑 3:CUDA Graph 捕獲失敗

現象:在 CUDA Graph 捕獲期間調 NCCL,報 "operation not permitted"。

原因:捕獲模式下不能做某些 CUDA 操作(如cudaMalloc)。NCCL 用cudaThreadExchangeStreamCaptureMode臨時切換模式,但不是所有操作都能繞過。

源碼依據:uploadWork的 persistent 分支:

📎 src/enqueue/enqueue.cc:1445

cpp
CUDACHECKGOTO(cudaThreadExchangeStreamCaptureMode(&mode), result, fail);

避坑:用NCCL_GRAPH_MIXING_SUPPORT=1開啟 graph 混合模式,或者預分配 work buffer。

本章小結

這一章我們把一次 AllReduce 的完整鏈路重新走了一遍:

1. 初始化:ncclCommInitRank → ncclCommInitRankFunc → initTransportsRank,建立通信域、搜索拓撲、對齊圖參數。

2. 任務入隊:ncclEnqueueCheck → taskAppend → collTaskAppend,把 API 調用翻譯成ncclTaskColl。

3. 算法選型:ncclGetAlgoInfo → ncclTuningCompute,用代價模型選出最優 (algo, proto)。

4. 任務調度:ncclPrepareTasks → scheduleCollTasksToPlan → finishPlan,把任務分配到通道,生成ncclKernelPlan。

5. Kernel 啟動:ncclLaunchKernel → cuLaunchKernelEx,把 plan 翻譯成 CUDA 啟動參數。

6. 設備側執行:runRing / runTreeUpDown / runNvls,按算法執行數據搬運。

本章思考與自測

Q1: 如果把initTransportsRank裡 AllGather3 之後的 min/max 對齊邏輯(L1690-L1698)去掉,在什麼場景下會導致通信死鎖?為什麼?

參考解析:這段邏輯保證所有 rank 對每個算法的nChannels、bwIntra、bwInter等參數達成一致。如果去掉,每個 rank 會用自己的本地拓撲計算結果。考慮一個異構集群:rank 0 在 8 卡 NVLink 機器上,rank 8 在 4 卡 PCIe 機器上。rank 0 算出 ring 有 8 條通道,rank 8 算出 4 條。當它們執行 Ring AllReduce 時,rank 0 會等 rank 8 在 8 條通道上發數據,但 rank

至此,我們完成了對一次 AllReduce 完整鏈路的回顧。從初始化、拓撲搜索、算法選擇、任務入隊、kernel 啟動,到設備側執行與網絡傳輸,每個環節都對應著前面章節的深入剖析。這份鏈路圖不僅是理解 NCCL 的骨架,也是排查問題的索引:初始化失敗查第 3、4 章,算法選錯查第 5 章,任務入隊報錯查第 6、7 章,kernel 啟動失敗查第 8 章,設備側 hang 查第 9、10 章,網絡問題查第 12、13 章。隨著 NCCL 向可編程通信、GPU 直發和對稱內存演進,這條鏈路還將繼續延伸——而你已經掌握了追蹤它的方法。

AI 賦能程式碼庫精讀 · 本地優先架構

讀完了本章?為你自己的私有專案生成專屬架構全景書

基於 Tauri 2 + Rust 本地原生引擎,100% 源碼離線隱私安全,零程式碼上傳雲端。像閱讀一本傳世專著一樣拆解你的複雜系統。

⚡ Tauri 2 · Rust 原生引擎 · 100% 離線私密安全 · 適配超百萬行程式碼庫

讀懂任何複雜專案,你真正需要的是一本專著

本書由 AiReadCode 掃描官方開源倉庫全自動編撰,結合真實不可變 Commit 節點與 FACT 藥丸行號溯源,提供純靜態、零服務依賴的極致雙欄互動式線上閱讀體驗。

在 GitHub 上 Star 本專案 ★ 瀏覽更多架構專著 →