2. NVLink带宽基准测试

任务说明

本案例目标为：在双GPU（通过NVLink互联）的服务器上，测量点对点NVLink的实际带宽，并与同一对GPU之间的PCIe带宽进行对比。测量结果用于验证NVLink互联是否生效及其性能水平。

实现方案

采用CUDA运行时API提供的cudaMemcpyPeerAsync进行设备间数据拷贝，分别通过NVLink和PCIe路径传输相同大小的数据块，记录耗时并计算带宽。为区分两条路径，利用cudaDeviceEnablePeerAccess启用对等访问时，驱动会自动选择最快的互联路径（NVLink优先）；而禁用对等访问时，拷贝将回退到PCIe（通过主机暂存或驱动管理的PCIe路径）。通过对比两种模式下的带宽，即可量化NVLink相对PCIe的优势。

代码实现

// nvlink_bw_test.cu
// 编译: nvcc -o nvlink_bw_test nvlink_bw_test.cu
// 运行: ./nvlink_bw_test [data_size_mb] [iterations]
#include <cstdio>
#include <cstdlib>
#include <cuda_runtime.h>
#include <chrono>

#define CUDA_CHECK(call) \
    do { \
        cudaError_t err = (call); \
        if (err != cudaSuccess) { \
            fprintf(stderr, "CUDA error at %s:%d: %s\n", __FILE__, __LINE__, \
                    cudaGetErrorString(err)); \
            exit(EXIT_FAILURE); \
        } \
    } while(0)

// 计算传输耗时(ms)
double time_ms(cudaEvent_t start, cudaEvent_t stop) {
    float ms = 0.0f;
    cudaEventElapsedTime(&ms, start, stop);
    return (double)ms;
}

int main(int argc, char** argv) {
    int dev_count = 0;
    CUDA_CHECK(cudaGetDeviceCount(&dev_count));
    if (dev_count < 2) {
        fprintf(stderr, "需要至少2块GPU\n");
        return 1;
    }

    // 默认参数: 256MB数据, 100次迭代
    size_t data_mb = (argc > 1) ? atoi(argv[1]) : 256;
    int iterations = (argc > 2) ? atoi(argv[2]) : 100;
    size_t bytes = data_mb * 1024 * 1024;

    // 检查对等访问支持
    int can_access = 0;
    CUDA_CHECK(cudaDeviceCanAccessPeer(&can_access, 0, 1));
    if (!can_access) {
        fprintf(stderr, "GPU0 和 GPU1 之间不支持对等访问\n");
        return 1;
    }

    // 分配设备内存
    int dev0 = 0, dev1 = 1;
    CUDA_CHECK(cudaSetDevice(dev0));
    char* buf0 = nullptr;
    CUDA_CHECK(cudaMalloc(&buf0, bytes));
    CUDA_CHECK(cudaMemset(buf0, 0xAB, bytes));

    CUDA_CHECK(cudaSetDevice(dev1));
    char* buf1 = nullptr;
    CUDA_CHECK(cudaMalloc(&buf1, bytes));
    CUDA_CHECK(cudaMemset(buf1, 0xCD, bytes));

    // 事件用于计时
    cudaEvent_t start, stop;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop));

    // ========== 测试1: NVLink路径 (启用对等访问) ==========
    CUDA_CHECK(cudaSetDevice(dev0));
    CUDA_CHECK(cudaDeviceEnablePeerAccess(dev1, 0));
    CUDA_CHECK(cudaSetDevice(dev1));
    CUDA_CHECK(cudaDeviceEnablePeerAccess(dev0, 0));

    // 预热
    CUDA_CHECK(cudaMemcpyPeerAsync(buf1, dev1, buf0, dev0, bytes, 0));
    CUDA_CHECK(cudaDeviceSynchronize());

    double nvlink_bw = 0.0;
    for (int i = 0; i < iterations; i++) {
        CUDA_CHECK(cudaEventRecord(start, dev0));
        CUDA_CHECK(cudaMemcpyPeerAsync(buf1, dev1, buf0, dev0, bytes, 0));
        CUDA_CHECK(cudaEventRecord(stop, dev1));
        CUDA_CHECK(cudaDeviceSynchronize());
        double ms = time_ms(start, stop);
        double bw = (double)bytes / (ms * 1e-3) / 1e9; // GB/s
        nvlink_bw += bw;
    }
    nvlink_bw /= iterations;

    // 禁用对等访问
    CUDA_CHECK(cudaSetDevice(dev0));
    CUDA_CHECK(cudaDeviceDisablePeerAccess(dev1));
    CUDA_CHECK(cudaSetDevice(dev1));
    CUDA_CHECK(cudaDeviceDisablePeerAccess(dev0));

    // ========== 测试2: PCIe路径 (禁用对等访问) ==========
    // 预热
    CUDA_CHECK(cudaMemcpyPeerAsync(buf1, dev1, buf0, dev0, bytes, 0));
    CUDA_CHECK(cudaDeviceSynchronize());

    double pcie_bw = 0.0;
    for (int i = 0; i < iterations; i++) {
        CUDA_CHECK(cudaEventRecord(start, dev0));
        CUDA_CHECK(cudaMemcpyPeerAsync(buf1, dev1, buf0, dev0, bytes, 0));
        CUDA_CHECK(cudaEventRecord(stop, dev1));
        CUDA_CHECK(cudaDeviceSynchronize());
        double ms = time_ms(start, stop);
        double bw = (double)bytes / (ms * 1e-3) / 1e9;
        pcie_bw += bw;
    }
    pcie_bw /= iterations;

    // 输出结果
    printf("数据大小: %zu MB, 迭代次数: %d\n", data_mb, iterations);
    printf("NVLink 平均带宽: %.2f GB/s\n", nvlink_bw);
    printf("PCIe   平均带宽: %.2f GB/s\n", pcie_bw);
    printf("NVLink/PCIe 加速比: %.2fx\n", nvlink_bw / pcie_bw);

    // 清理
    CUDA_CHECK(cudaEventDestroy(start));
    CUDA_CHECK(cudaEventDestroy(stop));
    CUDA_CHECK(cudaFree(buf0));
    CUDA_CHECK(cudaFree(buf1));
    return 0;
}

落地与运行

依赖：CUDA Toolkit（≥11.0），NVIDIA驱动，双GPU且支持NVLink（如A100、V100、RTX 3090等）。

编译与运行：

nvcc -o nvlink_bw_test nvlink_bw_test.cu
./nvlink_bw_test 256 100   # 256MB数据块，100次迭代

验证方法与判定标准

运行程序后观察输出：

• 若NVLink带宽显著高于PCIe带宽（通常加速比>3x），且数值接近该GPU型号的NVLink理论峰值（如A100 600GB/s），说明NVLink链路工作正常。
• 若两者带宽接近，可能原因：NVLink未正确连接、驱动未启用、或数据块过小导致延迟主导。可尝试增大数据块（如1024MB）并减少迭代次数。
• 使用nvidia-smi topo -m可查看GPU拓扑，确认NVLink连接状态。

预期输出示例

数据大小: 256 MB, 迭代次数: 100
NVLink 平均带宽: 512.34 GB/s
PCIe   平均带宽: 24.56 GB/s
NVLink/PCIe 加速比: 20.86x

注：以上数值仅为示意，实际带宽取决于GPU型号、NVLink版本及系统配置。