10-CUDA 编程与算子开发

章节概述

本章深入讲解 CUDA 编程的核心知识和高性能算子开发技术，从 GPU 架构原理到 Tensor Core 编程，从手写 Kernel 到使用 Triton 高效开发，为 AI Infra 架构师提供坚实的底层技术基础。

知识体系

┌─────────────────────────────────────────────────────────────────────────┐
│                     CUDA 编程与算子开发知识体系                           │
├─────────────────────────────────────────────────────────────────────────┤
│                                                                         │
│  ┌─────────────────────────────────────────────────────────────────┐   │
│  │                      GPU 硬件架构                                │   │
│  │                                                                  │   │
│  │   SM 架构      │  内存层次     │  Tensor Core  │  执行模型       │   │
│  │   ├─ CUDA Core│  ├─ 寄存器    │  ├─ WMMA      │  ├─ SIMT       │   │
│  │   ├─ Warp调度 │  ├─ 共享内存  │  ├─ MMA PTX   │  ├─ Warp同步   │   │
│  │   └─ 寄存器堆 │  ├─ L1/L2    │  └─ FP8/INT8  │  └─ 占用率     │   │
│  │              │  └─ HBM       │               │                │   │
│  └─────────────────────────────────────────────────────────────────┘   │
│                                │                                        │
│  ┌─────────────────────────────┴───────────────────────────────────┐   │
│  │                      CUDA 编程模型                               │   │
│  │                                                                  │   │
│  │   线程层次          │  内存管理        │  同步机制               │   │
│  │   ├─ Grid/Block    │  ├─ cudaMalloc   │  ├─ __syncthreads     │   │
│  │   ├─ Warp/Thread   │  ├─ cudaMemcpy   │  ├─ __syncwarp        │   │
│  │   └─ 索引计算      │  ├─ 统一内存     │  ├─ 原子操作           │   │
│  │                    │  └─ 异步拷贝     │  └─ Cooperative Groups│   │
│  └─────────────────────────────────────────────────────────────────┘   │
│                                │                                        │
│  ┌─────────────────────────────┴───────────────────────────────────┐   │
│  │                      性能优化技术                                │   │
│  │                                                                  │   │
│  │   内存优化          │  计算优化        │  并行优化               │   │
│  │   ├─ 合并访问      │  ├─ 指令级并行  │  ├─ 占用率优化          │   │
│  │   ├─ Bank Conflict │  ├─ 循环展开    │  ├─ 流水线              │   │
│  │   ├─ 向量化访问    │  ├─ 寄存器分块  │  ├─ 多流并发            │   │
│  │   └─ 预取/缓存     │  └─ 融合计算    │  └─ 异步执行            │   │
│  └─────────────────────────────────────────────────────────────────┘   │
│                                │                                        │
│  ┌─────────────────────────────┴───────────────────────────────────┐   │
│  │                      核心算子实现                                │   │
│  │                                                                  │   │
│  │   GEMM             │  Reduce          │  Attention              │   │
│  │   ├─ Tiled GEMM   │  ├─ Warp Shuffle│  ├─ Flash Attention     │   │
│  │   ├─ Register Tile│  ├─ Block Reduce│  ├─ Online Softmax      │   │
│  │   └─ TC GEMM      │  └─ Grid Reduce │  └─ Fused Attention     │   │
│  │                    │                  │                         │   │
│  │   LayerNorm        │  Softmax         │  算子融合               │   │
│  │   ├─ Welford      │  ├─ Safe Softmax│  ├─ Epilogue Fusion    │   │
│  │   └─ Vectorized   │  └─ Online      │  └─ Kernel Fusion      │   │
│  └─────────────────────────────────────────────────────────────────┘   │
│                                │                                        │
│  ┌─────────────────────────────┴───────────────────────────────────┐   │
│  │                      开发工具与生态                              │   │
│  │                                                                  │   │
│  │   Triton          │  CUTLASS         │  分析工具               │   │
│  │   ├─ Block编程    │  ├─ 模板库       │  ├─ Nsight Compute     │   │
│  │   ├─ 自动调优     │  ├─ CuTe        │  ├─ Nsight Systems     │   │
│  │   └─ JIT编译      │  └─ 高性能GEMM   │  └─ nvprof             │   │
│  └─────────────────────────────────────────────────────────────────┘   │
│                                                                         │
└─────────────────────────────────────────────────────────────────────────┘

章节内容

01-CUDA编程模型与内存层次

GPU 架构基础

SM (Streaming Multiprocessor) 结构详解
CUDA Core vs Tensor Core
Warp 调度与 SIMT 执行模型

CUDA 编程模型

Grid/Block/Warp/Thread 层次
线程索引计算
Kernel 编写基础

内存层次

寄存器、共享内存、L1/L2、全局内存
共享内存 Bank Conflict
全局内存合并访问
统一内存与异步传输

同步机制

__syncthreads、__syncwarp
原子操作
Cooperative Groups

02-高性能Kernel开发实战

GEMM 实现进阶

朴素实现与分析
Tiled GEMM (共享内存优化)
寄存器分块与向量化
双缓冲隐藏延迟

Reduce 操作

Warp Shuffle Reduce
Block Reduce
多维 Reduce

Softmax 实现

基础 Softmax
Online Softmax (单次遍历)
Fused Softmax

LayerNorm 实现

基础 LayerNorm
Welford 算法 (数值稳定)
向量化 LayerNorm

03-TensorCore与矩阵运算

Tensor Core 架构

硬件演进 (Volta → Turing → Ampere → Hopper)
MMA 操作原理
支持的数据类型与形状

WMMA API 编程

Fragment 与 load/store
mma_sync 操作
共享内存优化的 WMMA

MMA PTX 指令

直接使用 PTX
数据布局详解
FP16/TF32/INT8 MMA

CUTLASS 介绍

层次化架构
GEMM 配置与使用
CuTe (CUTLASS 3.x)

FP8 与 Transformer Engine

FP8 数据格式 (E4M3/E5M2)
Transformer Engine 使用
动态量化实现

04-算子融合与优化技术

融合原理

为什么需要融合
融合类型分类
融合决策指南

Element-wise 融合

基础融合实现
向量化融合
LayerNorm + Dropout + Add

GEMM Epilogue 融合

GEMM + Bias + Activation
完整的融合 GEMM Kernel

Attention 融合 (Flash Attention)

Flash Attention 原理
Online Softmax
Flash Attention 2 优化

动态融合

运行时 Kernel 生成
TorchScript Fusion
torch.compile

05-Triton编程入门

Triton 架构

设计理念
编译流程
与 CUDA 对比

基础语法

program_id、arange
load/store 与 mask
原子操作与 reduce

矩阵乘法实现

基础 GEMM
自动调优 GEMM
L2 Cache 优化

核心算子

Fused Softmax
LayerNorm
Flash Attention

性能调优

@triton.autotune
Benchmark
最佳实践

核心代码示例

共享内存 Tiled GEMM

__global__ void gemm_tiled(
    const float* A, const float* B, float* C,
    int M, int N, int K
) {
    __shared__ float As[TILE][TILE], Bs[TILE][TILE];

    int row = by * TILE + ty, col = bx * TILE + tx;
    float sum = 0.0f;

    for (int t = 0; t < K; t += TILE) {
        As[ty][tx] = A[row * K + t + tx];
        Bs[ty][tx] = B[(t + ty) * N + col];
        __syncthreads();

        for (int k = 0; k < TILE; k++)
            sum += As[ty][k] * Bs[k][tx];
        __syncthreads();
    }

    C[row * N + col] = sum;
}

Warp Shuffle Reduce

__device__ float warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset /= 2)
        val += __shfl_down_sync(0xffffffff, val, offset);
    return val;
}

Triton Softmax

@triton.jit
def softmax_kernel(input_ptr, output_ptr, n_cols, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    x = tl.load(input_ptr + row * n_cols + cols, mask=cols < n_cols)

    x_max = tl.max(x)
    x_exp = tl.exp(x - x_max)
    x_sum = tl.sum(x_exp)

    tl.store(output_ptr + row * n_cols + cols, x_exp / x_sum, mask=cols < n_cols)

性能指标参考

A100 GPU 性能天花板

指标	数值
FP32 CUDA Core	19.5 TFLOPS
FP16 Tensor Core	312 TFLOPS
INT8 Tensor Core	624 TOPS
HBM 带宽	2.0 TB/s
L2 Cache	40 MB
共享内存/SM	164 KB

GEMM 性能基准

实现	性能 (TFLOPS)	效率
Naive CUDA	0.5	2.5%
Tiled Shared	5.0	25%
Register Tile	15.0	75%
WMMA FP16	180	58%
CUTLASS FP16	275	88%
cuBLAS FP16	290	93%

学习路径

Week 1-2: CUDA 基础
├── 完成 01-CUDA编程模型与内存层次
├── 编写向量加法、矩阵转置 Kernel
└── 练习内存优化技术

Week 3-4: 高性能 Kernel
├── 完成 02-高性能Kernel开发实战
├── 实现 GEMM、Softmax、LayerNorm
└── 使用 Nsight 分析性能

Week 5-6: Tensor Core
├── 完成 03-TensorCore与矩阵运算
├── 使用 WMMA 实现 GEMM
└── 了解 CUTLASS 架构

Week 7-8: 算子融合与 Triton
├── 完成 04-算子融合与优化技术
├── 完成 05-Triton编程入门
├── 实现 Fused Attention
└── 对比 CUDA 与 Triton 实现

下一步学习

完成本章后，建议继续学习：

11-通信与网络底层 - NCCL、AllReduce、RDMA
12-框架源码解析 - PyTorch、DeepSpeed 源码
13-编译优化与图优化 - XLA、TorchScript