🚀 GPU 架构与 Triton 算子开发速查手册#

1. 核心抽象对应关系#

2. 算子基本骨架 (pid, offsets, mask)#

1
@triton.jit
2
def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
3
    # 1. 拿到当前 Block 的编号 (相当于发给当前 SM 的工单号)
4
    pid = tl.program_id(axis=0)
5

6
    # 2. 计算当前 Block 需要处理的数据全局索引 (Offsets)
7
    block_start = pid * BLOCK_SIZE
8
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
9

10
    # 3. 内存越界保护 (Mask) - 极其重要！防止最后一个 Block 读写溢出
11
    mask = offsets < n_elements
12

13
    # 4. 加载、计算、存储
14
    x = tl.load(x_ptr + offsets, mask=mask)
15
    y = tl.load(y_ptr + offsets, mask=mask)
16
    tl.store(out_ptr + offsets, x + y, mask=mask)

1. 边界对齐：triton.cdiv#

• 用途：计算需要启动多少个 Block (Grid 的大小)。
• 逻辑：triton.cdiv(n_elements, BLOCK_SIZE) 等价于向上取整除法 $\lceil N / B \rceil$。必须配合 Kernel 内的 mask 使用，确保尾部不足一个 BLOCK_SIZE 的数据也能被安全处理。

2. 编译期常量：tl.constexpr#

• 用途：标记像 BLOCK_SIZE 这样决定了底层 SRAM (共享内存) 和寄存器分配大小的变量。
• 机制：Triton JIT 编译器必须在运行前知道这些变量的确切值，从而为特定的 BLOCK_SIZE 编译出极致优化的机器码。如果传入不同的值，Triton 会重新触发编译。

3. CPU/GPU 异步执行 (Asynchronous Execution)#

• 机制：add_kernel[grid](...) 只是 CPU 向 GPU 的 Command Queue 里扔了一个任务工单，CPU 会瞬间返回 (耗时几微秒)。此时返回的 output Tensor 往往还没有被填入计算结果。
• 同步约束：如果需要评估算子真实耗时，必须在前后加上 torch.cuda.synchronize() 强行让 CPU 等待 GPU 清空计算队列。

1. 内存层级速查表#

2. SRAM 的黄金法则 (Shared Memory)#

• 作用域铁律：SRAM 是Block 级的资产。tl.load 的本质就是将 DRAM 中的数据协同搬运到了当前 Block 专属的 SRAM 工作台上 (即 Tiling 切块机制)。
• 容量红线：一个 SM 的 SRAM 通常不超过 200KB。如果你定义的 BLOCK_SIZE 过大，导致单 Block 需要的 SRAM 超过了 SM 总容量，程序会直接崩溃 (OOM) 或发生寄存器溢出 (Spilling)，导致性能暴跌。

1. Block 调度器 (宏观)：提供高 Occupancy (占用率)#

• 目标：尽可能多地向同一个 SM 里塞入 Block。
• 机制：只要 SM 的 SRAM 和寄存器没被填满，全局调度器就会把多个 Block 同时派发给同一个 SM (驻留/Co-residency)。这为微观的 Warp 调度提供了充足的“弹药”。

2. Warp 调度器 (微观)：在【整个 SM 尺度】下执行切换#

• 目标：让 ALU (计算核心) 永远不闲着。
• 机制 (极度重要)：Warp 调度器无视 Block 的边界。它在一个拥有所有驻留 Block 的“Warp 大池子”里巡逻。
  • 如果 Block A 的所有 Warp 都在等显存读取 (卡死)。
  • 调度器会瞬间 (0 周期) 把 ALU 交给 Block B 或 Block C 里数据已经就绪的 Warp 去做数学计算。
• 推论：为了实现完美的延迟隐藏，我们不需要把单个 Block 搞得无限大；我们只需要让 Block 大小适中，确保一个 SM 能同时吞下多个 Block，由 Warp 调度器在跨 Block 的尺度上缝合等待时间。