Zhenwei Cao 的技术笔记

第7章-高级主题

Properties1
tagsgpu-computing, gpu-programming, concept-note

11分钟阅读

第7章：高级主题

深入 GPU 高级特性，掌握专家级优化技术

本章定位

本章负责建立高级主题的地图：Tensor Core、PTX、低精度、多 GPU、算子融合、CUTLASS。初学阶段不要求直接手写最底层 PTX，重点是知道每个主题解决什么问题、何时该使用库或 DSL。

配套主文档：

学习目标

  • 掌握 Tensor Core 高级用法
  • 理解多 GPU 编程模型
  • 了解量化与低精度计算
  • 掌握算子融合技术
  • 学习 CUTLASS 等高性能库的使用

7.1 Tensor Core 深度使用

7.1.1 Tensor Core 概述

Tensor Core 是专门用于矩阵运算的硬件单元：

  • 单周期完成 D = A × B + C（4×4×4 或更大矩阵）
  • 支持 FP16、BF16、TF32、FP64、INT8 等多种精度
  • 相比 CUDA Core 可获得 3-8 倍加速

支持的精度（A100）：

输入精度累加精度峰值性能
FP16FP32312 TFLOPS
BF16FP32312 TFLOPS
TF32FP32156 TFLOPS
FP64FP6419.5 TFLOPS
INT8INT32624 TOPS

7.1.2 WMMA API 详解

#include <mma.h>
using namespace nvcuda::wmma;
 
// 定义矩阵片段大小
const int WMMA_M = 16;
const int WMMA_N = 16;
const int WMMA_K = 16;
 
__global__ void wmma_gemm(half* a, half* b, float* c, int M, int N, int K) {
    // 声明矩阵片段
    fragment<matrix_a, WMMA_M, WMMA_N, WMMA_K, half, row_major> a_frag;
    fragment<matrix_b, WMMA_M, WMMA_N, WMMA_K, half, col_major> b_frag;
    fragment<accumulator, WMMA_M, WMMA_N, WMMA_K, float> c_frag;
 
    // 获取当前 warp 的位置
    int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;
    int warpN = blockIdx.y * blockDim.y + threadIdx.y;
 
    // 初始化累加器
    fill_fragment(c_frag, 0.0f);
 
    // 遍历 K 维度
    for (int i = 0; i < K; i += WMMA_K) {
        // 加载矩阵 A 的片段
        load_matrix_sync(a_frag, a + warpM * WMMA_M * K + i, K);
 
        // 加载矩阵 B 的片段
        load_matrix_sync(b_frag, b + i * N + warpN * WMMA_N, N);
 
        // 执行矩阵乘加
        mma_sync(c_frag, a_frag, b_frag, c_frag);
    }
 
    // 存储结果
    store_matrix_sync(c + warpM * WMMA_M * N + warpN * WMMA_N,
                      c_frag, N, mem_row_major);
}

7.1.3 PTX 级别的 Tensor Core 操作

// 使用内联 PTX 直接控制 Tensor Core
__device__ void mma_m16n8k16_fp16(
    uint32_t* d0, uint32_t* d1,
    uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
    uint32_t b0, uint32_t b1,
    uint32_t c0, uint32_t c1
) {
    asm volatile(
        "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
        "{%0, %1}, "
        "{%2, %3, %4, %5}, "
        "{%6, %7}, "
        "{%8, %9};"
        : "=r"(*d0), "=r"(*d1)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3),
          "r"(b0), "r"(b1),
          "r"(c0), "r"(c1)
    );
}

7.2 多 GPU 编程

7.2.1 CUDA 多设备管理

// 获取可用 GPU 数量
int device_count;
cudaGetDeviceCount(&device_count);
printf("可用 GPU 数量: %d\n", device_count);
 
// 设置当前设备
cudaSetDevice(0);  // 使用 GPU 0
 
// 获取设备属性
cudaDeviceProp prop;
cudaGetDeviceProperties(&prop, 0);
printf("设备名称: %s\n", prop.name);
printf("计算能力: %d.%d\n", prop.major, prop.minor);

7.2.2 Peer-to-Peer 内存访问

// 检查并启用 P2P 访问
int can_access;
cudaDeviceCanAccessPeer(&can_access, 0, 1);
if (can_access) {
    cudaSetDevice(0);
    cudaDeviceEnablePeerAccess(1, 0);
    cudaSetDevice(1);
    cudaDeviceEnablePeerAccess(0, 0);
}
 
// 现在 GPU 0 可以直接访问 GPU 1 的内存
float* d_ptr_on_gpu1;
cudaSetDevice(1);
cudaMalloc(&d_ptr_on_gpu1, size);
 
cudaSetDevice(0);
kernel<<<blocks, threads>>>(d_ptr_on_gpu1);  // 直接从 GPU 0 访问

7.2.3 NCCL 多 GPU 通信

#include <nccl.h>
 
// 初始化 NCCL
ncclComm_t comm;
ncclUniqueId id;
 
if (rank == 0) {
    ncclGetUniqueId(&id);
}
// 广播 id 给所有进程
MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD);
 
ncclCommInitRank(&comm, nRanks, id, rank);
 
// AllReduce 操作（常用于分布式训练）
ncclAllReduce(sendbuff, recvbuff, count, ncclFloat, ncclSum, comm, stream);
 
// 销毁 NCCL 通信器
ncclCommDestroy(comm);

7.3 量化与低精度计算

7.3.1 INT8 量化

// 量化：FP32 -> INT8
__device__ int8_t quantize(float x, float scale) {
    float scaled = x / scale;
    scaled = fmaxf(-128.0f, fminf(127.0f, scaled));
    return (int8_t)roundf(scaled);
}
 
// 反量化：INT8 -> FP32
__device__ float dequantize(int8_t x, float scale) {
    return x * scale;
}
 
// INT8 矩阵乘法（使用 DP4A 指令）
__global__ void int8_gemm(
    const int8_t* A, const int8_t* B, int32_t* C,
    int M, int N, int K,
    float scale_a, float scale_b
) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
 
    if (row < M && col < N) {
        int32_t sum = 0;
 
        // 使用 __dp4a 进行 4 元素点积
        for (int k = 0; k < K; k += 4) {
            int32_t a4 = *(int32_t*)(A + row * K + k);
            int32_t b4 = *(int32_t*)(B + col * K + k);
            sum = __dp4a(a4, b4, sum);
        }
 
        // 反量化结果
        C[row * N + col] = sum;
    }
}

7.3.2 FP8 计算（Hopper 架构）

#include <cuda_fp8.h>
 
// FP8 矩阵乘法（H100）
__global__ void fp8_gemm(
    __nv_fp8_e4m3* A, __nv_fp8_e4m3* B, float* C,
    int M, int N, int K
) {
    // 使用 Tensor Core 进行 FP8 计算
    // 需要 CUDA 11.8+ 和 Hopper 架构
 
    using FragA = nvcuda::wmma::fragment<
        nvcuda::wmma::matrix_a, 16, 16, 16,
        __nv_fp8_e4m3, nvcuda::wmma::row_major>;
 
    using FragB = nvcuda::wmma::fragment<
        nvcuda::wmma::matrix_b, 16, 16, 16,
        __nv_fp8_e4m3, nvcuda::wmma::col_major>;
 
    using FragC = nvcuda::wmma::fragment<
        nvcuda::wmma::accumulator, 16, 16, 16, float>;
 
    FragA a_frag;
    FragB b_frag;
    FragC c_frag;
 
    nvcuda::wmma::fill_fragment(c_frag, 0.0f);
 
    // 加载并计算
    nvcuda::wmma::load_matrix_sync(a_frag, A, K);
    nvcuda::wmma::load_matrix_sync(b_frag, B, N);
    nvcuda::wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
    nvcuda::wmma::store_matrix_sync(C, c_frag, N, nvcuda::wmma::mem_row_major);
}

7.4 算子融合

7.4.1 为什么需要算子融合？

独立 Kernel（慢）：
  数据 -> [Kernel A] -> 内存 -> [Kernel B] -> 内存 -> [Kernel C]

融合 Kernel（快）：
  数据 -> [Kernel A + B + C] -> 结果

优势：
- 减少全局内存访问
- 减少 Kernel 启动开销
- 更好的数据局部性

7.4.2 典型的融合模式

// 融合 1：BiasAdd + ReLU
__global__ void bias_add_relu(float* input, float* bias, float* output,
                               int batch, int features) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch * features;
 
    if (idx < total) {
        int f = idx % features;
        float val = input[idx] + bias[f];
        output[idx] = fmaxf(0.0f, val);  // ReLU
    }
}
 
// 融合 2：LayerNorm + Residual
__global__ void layernorm_residual(float* input, float* residual,
                                    float* gamma, float* beta,
                                    float* output, int hidden_size) {
    // 计算 LayerNorm 的同时加上 residual
    // ... 实现略
}
 
// 融合 3：GEMM + BiasAdd + Activation
__global__ void gemm_bias_act(float* A, float* B, float* bias,
                               float* C, int M, int N, int K) {
    // 在写回结果时同时加 bias 和激活函数
    // ... 实现略
}

7.4.3 使用 CUDA Graph 减少开销

// 创建 CUDA Graph
cudaGraph_t graph;
cudaGraphExec_t instance;
 
// 开始捕获
cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
 
// 录制一系列 Kernel 调用
kernel1<<<blocks1, threads1, 0, stream>>>(...);
kernel2<<<blocks2, threads2, 0, stream>>>(...);
kernel3<<<blocks3, threads3, 0, stream>>>(...);
 
// 结束捕获
cudaStreamEndCapture(stream, &graph);
 
// 实例化 Graph
cudaGraphInstantiate(&instance, graph, NULL, NULL, 0);
 
// 之后可以重复执行（低开销）
for (int i = 0; i < iterations; i++) {
    cudaGraphLaunch(instance, stream);
}
 
// 清理
cudaGraphExecDestroy(instance);
cudaGraphDestroy(graph);

7.5 CUTLASS 高性能库

7.5.1 CUTLASS 简介

CUTLASS 是 NVIDIA 开源的高性能线性代数库：

  • 提供接近 cuBLAS 的性能
  • 高度可配置的模板设计
  • 支持各种精度和布局

7.5.2 使用 CUTLASS 进行 GEMM

#include <cutlass/gemm/device/gemm.h>
 
// 定义 GEMM 配置
using Gemm = cutlass::gemm::device::Gemm<
    cutlass::half_t,               // ElementA
    cutlass::layout::RowMajor,     // LayoutA
    cutlass::half_t,               // ElementB
    cutlass::layout::ColumnMajor,  // LayoutB
    float,                         // ElementC
    cutlass::layout::RowMajor,     // LayoutC
    float,                         // ElementAccumulator
    cutlass::arch::OpClassTensorCore,  // 使用 Tensor Core
    cutlass::arch::Sm80            // 架构（A100 = Sm80）
>;
 
void run_gemm(half* A, half* B, float* C, int M, int N, int K) {
    // 配置 GEMM 参数
    typename Gemm::Arguments args(
        {M, N, K},                          // 问题大小
        {A, K},                             // A 指针和 leading dimension
        {B, K},                             // B 指针和 leading dimension
        {C, N},                             // C 指针（输入）
        {C, N},                             // D 指针（输出）
        {1.0f, 0.0f}                        // alpha, beta
    );
 
    // 创建并运行 GEMM
    Gemm gemm_op;
    cutlass::Status status = gemm_op(args);
 
    if (status != cutlass::Status::kSuccess) {
        printf("GEMM 失败\n");
    }
}

7.6 高级调试和性能分析

7.6.1 使用 Nsight Compute 深入分析

# 详细的内存分析
ncu --metrics \
    sm__sass_average_data_bytes_per_sector_mem_global_op_ld.pct,\
    l1tex__t_sectors_pipe_lsu_mem_global_op_ld.sum,\
    l1tex__t_sectors_pipe_lsu_mem_global_op_st.sum,\
    dram__sectors_read.sum,\
    dram__sectors_write.sum \
    ./my_kernel
 
# 计算利用率分析
ncu --metrics \
    sm__cycles_active.avg,\
    sm__warps_active.avg.pct_of_peak_sustained_active,\
    sm__inst_executed_pipe_tensor.avg.pct_of_peak_sustained_elapsed \
    ./my_kernel

7.6.2 使用 CUDA Sanitizer

# 检测内存错误
cuda-memcheck ./my_program
 
# 检测数据竞争
cuda-memcheck --tool racecheck ./my_program
 
# 检测越界访问
cuda-memcheck --tool initcheck ./my_program

7.7 Blackwell：tcgen05.mma、Tensor Memory 与 FP4 / Block Scaling

Blackwell（B200 / B300）把 Tensor Core 主路径推进到第五代 tcgen05.mma，对应几个关键概念，初学阶段不要求手写 PTX，但应能听懂术语：

概念含义在哪里能用到
tcgen05.mmaBlackwell 第五代 Tensor Core 的 PTX 指令族，异步、围绕 descriptor 组织CUTLASS 4.x Blackwell GEMM 模板
Tensor Memory (TMEM)SM 内新增的片上存储，约 256KB/SM，专用于存放 MMA accumulator / operandtcgen05.mma 的累加目的地、tcgen05.ld/tcgen05.st 与寄存器搬运
cta_group::2让同 cluster 内两个 CTA 协作完成一次更大的 MMA提升单次 MMA 的 N 维等效 tile
FP4 (E2M1) / NVFP44-bit 浮点，主要用于推理；FP6 (E3M2/E2M3) 6-bitLLM 权重量化 + Tensor Core 路径
Block Scaling (MXFP / NVFP4)多个元素共享一个缩放因子，缩放因子格式为 E8M0（8-bit）或 FP8(E4M3)，不是 FP16；MXFP4 块=32，NVFP4 块=16Blackwell low-precision GEMM 主路径

严格定义、代际算力（dense/sparse）、详细 PTX 形态：见 CUDA GEMM 矩阵乘法优化指南 §13 与 GPU 硬件架构背景与编程范式。Hopper 的 FP8 实战通常走 wgmma + Transformer Engine 路径，WMMA FP8 接口仅作概念示意，不推荐复制到生产代码。

💡 本章要点

  1. Tensor Core 可以带来 3-8 倍加速，支持多种精度
  2. 多 GPU 编程 使用 NCCL 进行高效通信
  3. 量化技术（INT8/FP8）可以显著提升推理速度
  4. 算子融合 减少内存访问，提升端到端性能
  5. CUDA Graph 减少 CPU 开销，适合小 Kernel 场景
  6. CUTLASS 提供灵活的高性能线性代数实现
  7. Blackwell 第五代 Tensor Core：tcgen05.mma + TMEM + FP4 / NVFP4 + block scaling 是新的主路径
  8. 使用 Nsight Compute 深入分析性能瓶颈

推荐阅读

上一章：大模型 Kernel →

关系图谱

反向链接