GPU 高性能编程与内核优化 (High Performance GPU Programming & Kernels)#

1.1 硬件层级 (Hardware Hierarchy)#

• Streaming Multiprocessors (SMs): GPU 的核心计算单元。A100/H100 包含大量的 SM。
• Memory Hierarchy:
  • Global Memory (DRAM): 容量大但速度慢。
  • L2/L1 Caches: 速度较快。
  • Register File: 每个 Thread 独有的极速存储，编写高性能代码时需大量使用。

1.2 执行模型 (Execution Model)#

• Grid & Blocks:
  • 计算任务被划分为多个 Thread Blocks。
  • 一个 Block 被调度到一个 SM 上执行。
  • Thread Blocks 之间通信昂贵（通过 Global Memory），但在同一 Block 内的 Threads 可以通过 Shared Memory 快速通信。
• Threads & Warps:
  • Threads 是执行计算的最小单位。
  • Warp: 32 个 Threads 组成一组，同时执行（SIMT架构）。为了性能，应尽量保证 Waves（即 Warps 的执行波次）负载均衡,。

1.3 核心性能指标 (Key Metric)#

• Arithmetic Intensity: 计算量与内存访问量的比值 (Flops / Bytes)。
• Compute Bound vs. Memory Bound:
  • Matrix Multiplication (矩阵乘法) 通常是 Compute Bound。
  • 其他操作（如 Activation functions, Normalization）通常是 Memory Bound。
  • 目标：提高 Arithmetic Intensity，减少内存搬运,。

2.1 目的#

• 测量代码的端到端 Wall clock time。
• 比较不同实现（PyTorch vs. Triton vs. CUDA）的性能差异。

2.2 关键步骤与陷阱 (Pitfalls)#

• Warm-up (预热):
  • 首次运行 PyTorch 代码时会涉及编译和初始化，需运行几次预热以测量稳定状态下的速度,。
• Synchronization (同步):
  • CPU-GPU Asynchrony: CPU 通常会跑在 GPU 前面，将内核发射（Kernel Launch）到队列中。
  • 解决方案: 必须调用 torch.cuda.synchronize() 确保 GPU 完成所有任务后再停止计时，否则测量的只是 CPU 将任务推入队列的时间,。

2.3 案例分析 (Example)#

• MLP Scaling: 增加层数 (Layers) 或步数 (Steps) 时，运行时间呈线性增长 (Linear scaling),。
• Matrix Multiplication: 只有矩阵足够大时才能观察到预期的性能扩展，小矩阵受限于启动开销 (Overhead)。

3.1 为什么要进行 Profiling?#

• Benchmarking 只能告诉你代码慢，但 Profiling 能告诉你时间花在哪里。
• 揭示 PyTorch 抽象层之下的底层 CUDA Kernels 调用情况。

3.2 工具 (Tools)#

• PyTorch Profiler: 易于使用，显示高级别的 Kernel 耗时。
• NVIDIA Nsight Systems: 专业的系统级分析工具，能展示 CPU 与 GPU 的时间线交互。

3.3 关键发现 (Key Insights)#

• CPU Queueing: CPU 往往比 GPU 跑得快，它会提前将多个 Kernel 放入 GPU 的执行队列中。
• Synchronization Overhead: 某些操作（如打印 loss.item()）会强制 CPU 等待 GPU 计算结果，破坏流水线并行，导致性能下降,。
• Kernel Detail: 对于简单的操作（如 Add），可能会看到大部分时间花在 Kernel Launch 和数据传输上，而非计算本身。对于复杂的 Matrix Multiply，底层可能会调用 cuBLAS 或 CUTLASS 的特定 Kernel。

4.1 动机：Kernel Fusion#

• 问题: 逐个执行操作（如 $0.5 \times x \times (1 + \tanh(\dots))$）会导致多次读写 Global Memory。
• 解决方案: 将多个算子融合为一个 Kernel（Kernel Fusion），数据只从内存读取一次，在寄存器中完成所有计算，再写回。

4.2 案例：GELU (Gaussian Error Linear Unit)#

我们需要计算：$\text{GELU}(x) \approx 0.5x(1 + \tanh(\sqrt{2/\pi}(x + 0.044715x^3)))$

5.1 挑战#

• 不同于 Element-wise 操作，Softmax 需要 Reduction (归约) 操作（求 Max 和 Sum）。

5.2 Triton 实现策略#

• Row-wise Handling: 让一个 Thread Block 处理矩阵的一整行。
• Block Size: 设置为大于列数的最小 2 的幂次 (next_power_of_2)。
• 流程:
  1. 加载整行数据到 SRAM。
  2. 计算该行的 Max。
  3. 计算 Exponentials 并求 Sum。
  4. 归一化并写回-。

5.3 性能对比#

• Manual Softmax: 灾难性的慢，大量中间内存操作。
• Triton/Compiled Softmax: 单个 Fused Kernel，性能大幅提升。