Lecture 01 - How to profile CUDA kernels in PyTorch

这节课的核心并不是系统讲 CUDA 编程模型，而是建立一套从 PyTorch 程序出发进行 GPU 性能分析的工作流：先观察 PyTorch 实际执行了哪些 Operator 和 CUDA Kernel，再逐步深入到 Kernel 内部，分析计算、访存、Occupancy、Launch Configuration 等性能问题；当现有 PyTorch 实现无法满足性能需求时，再考虑 torch.compile、Triton 甚至自定义 CUDA Kernel。

课程强调的是一种偏工程实践的学习方式：不要等掌握完整 CUDA 理论后才开始优化，而是先通过 Profiling 看见问题，再针对具体问题补充对应的 CUDA 知识。 这也是整节课贯穿始终的思路。

1. 从 PyTorch Operator 到 CUDA Kernel

课程首先从一个非常简单的 Element-wise Square 开始：

torch.square(x)
x * x
x ** 2

虽然三种写法在数学上完全等价，但它们进入 PyTorch 后可能经过不同的 Operator Dispatch，最终调用不同的 CUDA Kernel。

例如通过 profiler 可以观察到，torch.square() 最终会落到类似 aten::pow 的实现，而 x * x 则直接调用 multiplication operator。也就是说，我们平时看到的 Python API 并不是 GPU 真正执行的单位，真正决定性能的是更下层的 ATen Operator 与 CUDA Kernel。

整个执行链可以简单理解为：

Python API
    │
    ▼
PyTorch Operator
    │
    ▼
ATen / C++ Implementation
    │
    ▼
CUDA Kernel Launch
    │
    ▼
GPU Execution

因此进行 PyTorch 性能优化时，第一个问题通常不应该是：

torch.square() 理论上应该多快？

而应该先问：

这段 PyTorch 代码最终执行了什么 Kernel？

课程通过这个简单例子展示了 Profiling 的价值：即使不了解 PyTorch 内部实现，也可以通过 profiler 快速观察其底层行为。

2. CUDA Profiling 的前提：CUDA 是异步执行的

CUDA 性能测量最重要的基础之一是理解 Asynchronous Execution。

当 CPU 发起一个 CUDA Kernel 时：

CPU
 │
 ├── Launch Kernel ───────────────┐
 │                               │
 ├── Continue executing           │
 │                               ▼
 │                             GPU Kernel

Kernel Launch 通常只负责把任务提交给 GPU，并不会等待 GPU 完成。因此如果直接使用：

start = time.time()
y = torch.square(x)
end = time.time()

测到的往往主要是 CPU 侧 Kernel Launch 的开销，而不是真正的 GPU Kernel Execution Time。

正确 benchmark 通常需要三个关键步骤：

Warmup：避免 CUDA Context 初始化、JIT 编译、Cache 等首次运行开销干扰结果。
CUDA Event：使用 GPU 时间戳测量 Kernel 执行时间。
Synchronization：在读取结果之前等待 GPU 完成。

典型结构为：

# warmup
for _ in range(10):
    fn(x)

torch.cuda.synchronize()

start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)

start.record()
fn(x)
end.record()

torch.cuda.synchronize()

elapsed_ms = start.elapsed_time(end)

所以 CUDA Benchmark 的基本链路应该是：

Warmup
   ↓
Record Start Event
   ↓
Launch CUDA Kernel
   ↓
Record End Event
   ↓
cudaSynchronize()
   ↓
Calculate GPU Time

这一点是后面所有性能分析的基础。

3. PyTorch 中的三层 Profiling 工具

这节课实际上介绍了一套逐层深入的 Profiling 路径。

工具	主要回答的问题	分析粒度
Autograd Profiler	哪些 PyTorch Operator 最耗时？	Operator
PyTorch Profiler	CPU、Memcpy、Kernel Launch 和 GPU Kernel 如何执行？	Timeline / Kernel
Nsight Compute	某个 CUDA Kernel 为什么慢？	Kernel Microarchitecture

可以把它理解为：

PyTorch Program
      │
      ▼
Autograd Profiler
“哪个 Operator 慢？”
      │
      ▼
PyTorch Profiler
“它调用了哪个 CUDA Kernel？”
      │
      ▼
Nsight Compute
“这个 Kernel 为什么慢？”

Autograd Profiler

Autograd Profiler 更适合作为第一层分析工具。它会给出 Operator 的 CPU Time、CUDA Time、调用次数等信息，因此可以快速判断计算主要集中在哪些 PyTorch Operator 上。

例如：

torch.square
     ↓
aten::square
     ↓
aten::pow
     ↓
CUDA elementwise kernel

这种分析方式非常适合第一次接触一段陌生 PyTorch 代码时使用，因为它能够快速建立“Python API → ATen Operator”的映射。

PyTorch Profiler

PyTorch Profiler 则进一步提供 Timeline。它可以通过 Chrome Trace 展示 CPU Operator、CUDA Runtime、Memory Copy 和 CUDA Kernel 之间的执行关系。

例如 Tensor 从 CPU 复制到 GPU：

aten::copy_
     │
     ▼
cudaMemcpyAsync
     │
     ▼
Memcpy HtoD

随后执行计算：

aten::pow
     │
     ▼
CUDA Launch
     │
     ▼
vectorized_elementwise_kernel

其中：

Host：CPU / System Memory
Device：GPU
HtoD：Host to Device
DtoH：Device to Host

PyTorch Profiler 中的 Flow Event 可以把 CPU 侧 Operator 与 GPU 上实际执行的 Kernel 联系起来，这对于理解 PyTorch 的异步执行尤其重要。

4. 将自定义 CUDA Kernel 接入 PyTorch

观察清楚 PyTorch 默认 Kernel 后，课程开始进入第二部分：如何把自己的 CUDA Kernel 接入 PyTorch。

通常调用链为：

Python
   │
   ▼
C++ Binding
   │
   ▼
CUDA Wrapper
   │
   ▼
CUDA Kernel

如果手动完成这一套流程，需要处理 PyBind、编译参数、Build System 等大量 Boilerplate。为了降低入门门槛，课程推荐使用：

torch.utils.cpp_extension.load_inline

它允许直接把 C++ 和 CUDA 源代码编译成 Python Extension：

module = load_inline(
    name="square_extension",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["square_matrix"],
)

随后可以像普通 Python Module 一样调用：

y = module.square_matrix(x)

load_inline 背后会自动生成和管理：

CUDA / C++ Source
       │
       ▼
   load_inline
       │
       ├── main.cpp
       ├── *.cu
       ├── PyBind Binding
       ├── build.ninja
       │
       ▼
Shared Library
       │
       ▼
Python Module

课程特别强调，这不仅是一个方便的 Extension 工具，也是一个很好的学习入口。因为它生成的 main.cpp、PyBind 和 build.ninja 都可以直接查看，从而反过来理解一个完整 PyTorch CUDA Extension 是怎样被构建出来的。

5. 从 CUDA 到 Triton

除了直接编写 CUDA C++，课程还介绍了 Triton。

Triton 是一个面向 GPU Kernel 编程的 Python DSL。它与 CUDA 的主要区别不在于“有没有 GPU 编程”，而是抽象层级不同。

可以粗略理解为：

方式	编程抽象	控制能力	开发难度
PyTorch	Tensor / Operator	低	最低
`torch.compile`	Graph / Compiler	较低	低
Triton	Block / Tile	较高	中
CUDA C++	Thread / Warp / Block	最高	高

Triton Kernel 大致经历：

Triton Python
      │
      ▼
Triton Compiler
      │
      ▼
Intermediate Representation
      │
      ▼
PTX
      │
      ▼
GPU Machine Code

它最大的优势之一是可以直接嵌入 Python / PyTorch 程序，同时又保留对 Memory Access、Block Size、Tiling 等 Kernel 行为的控制能力。

6. 自定义 Kernel 并不天然更快

课程中一个非常重要的实验是自己实现 Triton Square Kernel。

最开始得到的结果却是：

torch.square
    >
custom Triton square

也就是说，自定义 Triton Kernel 反而更慢。

这里揭示了 GPU Kernel Optimization 中非常重要的一点：

写出 Kernel ≠ 写出高性能 Kernel。

GPU 性能高度依赖 Kernel Configuration，例如：

Grid Size
Block Size
Warp 数量
Memory Access Pattern
Register Usage
Occupancy

课程最后发现，主要问题之一是 Kernel 的 Launch Parameter 不合理。修改 block_size 后，性能趋势发生明显反转。

所以一个更合理的 Kernel 开发闭环应该是：

Implement
   ↓
Benchmark
   ↓
Profile
   ↓
Find Bottleneck
   ↓
Modify Kernel
   ↓
Benchmark Again

而不是：

Write CUDA
   ↓
Assume Faster

7. Triton Debugging 与 PTX

Triton 还提供 Interpreter Mode，可以使用普通 Python Debugger 查看 Kernel 中间状态，例如：

breakpoint()

从而观察：

program_id
pointer offset
mask
intermediate tensor

这对于调试索引错误非常有帮助，因为真正运行在 GPU 上的 Kernel 通常比较接近黑盒，很难直接进行逐行调试。

另一方面，Triton 编译产生的 PTX 也是课程推荐观察的内容。

对于最简单的 Square：

Global Load
    ↓
Register
    ↓
MUL
    ↓
Register
    ↓
Global Store

在 PTX 中可以直接看到：

ld.global
mul.f32
st.global

以及 Register 的使用情况。

因此从 Kernel 到硬件可以形成一条比较完整的观察链：

PyTorch
   ↓
Triton
   ↓
PTX
   ↓
GPU Instruction

课程的观点并不是要求初学者立即掌握 PTX，而是认为 PTX 是理解“Compiler 最终为 GPU 生成了什么代码”的一个非常有价值的窗口。

8. `torch.compile` 也可以作为 Kernel 学习工具

课程展示了一个非常实用的技巧：

TORCH_LOGS="output_code"

配合：

compiled_fn = torch.compile(fn)

可以直接观察 torch.compile 为 PyTorch Program 生成的 Triton Kernel。

因此一个很实用的学习方法是：

Write PyTorch
      ↓
torch.compile
      ↓
Generated Triton
      ↓
Read Generated Kernel
      ↓
Modify / Optimize

对于刚开始学习 Triton 的人，这比完全从空白开始写 Kernel 更容易，因为 Compiler 已经给出了一个可工作的实现。

这也可以帮助理解另一个非常重要的优化：Kernel Fusion。

假设：

x = x * x
x = x * x

Eager PyTorch 可能执行：

Kernel 1
Load → Compute → Store

Kernel 2
Load → Compute → Store

经过 Fusion 后则可能变成：

        One Kernel
            │
Load → Compute → Compute → Store

这样可以同时减少：

Kernel Launch Overhead
Global Memory Load
Global Memory Store

这也是 torch.compile 提升许多 PyTorch Workload 性能的重要来源之一。

9. Nsight Compute：分析 Kernel 为什么慢

如果说 PyTorch Profiler 负责找到慢 Kernel，那么 Nsight Compute（NCU）负责回答：

为什么这个 Kernel 慢？

基本用法：

ncu python script.py

NCU 会给出大量 GPU Hardware Counter，例如：

指标	关注的问题
Compute Throughput	计算单元是否被充分利用
Memory Throughput	显存带宽是否充分利用
Occupancy	SM 上是否有足够 Active Warps
Warp Stall	Warp 为什么无法继续执行
L1 / L2	Cache 行为
Grid / Block	Launch Configuration 是否合理

除此之外，NCU 还会自动给出 Optimization Hint。

例如课程中出现的问题：

Grid Too Small
      ↓
无法产生足够多的 Thread Blocks
      ↓
部分 SM 空闲
      ↓
GPU Utilization 下降

这时可以尝试调整：

Grid Size
Block Size
Padding
Parallelism

而如果 NCU 显示大量 Long Scoreboard Stall，则可能意味着 Warp 正在等待 Global Memory 数据，此时需要进一步考虑 Memory Coalescing、Shared Memory 等优化。

10. 整节课真正想建立的工作流

把整节课串起来，其实就是下面这条链：

                    PyTorch Program
                           │
                           ▼
                  Autograd Profiler
                           │
                找到昂贵 Operator
                           │
                           ▼
                  PyTorch Profiler
                           │
             找到具体 CUDA Kernel
                           │
                           ▼
                    Benchmark
                           │
                           ▼
                  Nsight Compute
                           │
          ┌────────────────┴────────────────┐
          ▼                                 ▼
      Compute Problem                  Memory Problem
          │                                 │
     Occupancy / Grid               Load / Store / Cache
          │                                 │
          └────────────────┬────────────────┘
                           ▼
                     Optimization
                           │
                           ▼
                      Benchmark

与此同时，Kernel 实现层级也应该逐渐下沉：

PyTorch
   │
   ▼
torch.compile
   │
   ▼
Triton
   │
   ▼
CUDA C++

并不是任何性能问题都值得直接手写 CUDA。更合理的策略是：

PyTorch 已经足够快 → 保持 PyTorch。
torch.compile 能通过 Fusion 优化 → 使用 Compiler。
需要控制 Tile、Block、Memory Access → Triton。
Triton 无法提供足够底层的控制 → CUDA C++。

课程最后将推荐的 Profiling 顺序总结为：

Autograd Profiler
        ↓
PyTorch Profiler
        ↓
Nsight Compute

本质上对应三个逐渐深入的问题：

执行了什么？

它是怎样执行的？

为什么没有跑得更快？

Summary

这节 Lecture 最值得记住的并不是某几个 API，而是一套 GPU Performance Engineering 的基本方法：

不要猜性能，而要测量性能；不要看到 Kernel 慢就直接重写，而要先定位瓶颈。

最终可以把整套思路压缩成：

Profile
   ↓
Locate Bottleneck
   ↓
Understand Why
   ↓
Form Hypothesis
   ↓
Optimize
   ↓
Benchmark
   ↓
Verify

这实际上也是以后学习 CUDA Kernel Optimization、Triton、FlashAttention、GEMM 以及整个 AI Infra 性能优化时都会反复使用的一套方法论。