大模型算法实战教程

Appearance

08. Programming Models CUDA Triton | 编程模型演进 ​

难度： Medium | 环境： GPU optional

🚀 云端运行环境

本章节的实战代码可以点击以下链接在免费 GPU 算力平台上直接运行：

Open In ColabOpen In Studio (国内推荐：魔搭社区免费实例)

标签： CUDA, Triton, Programming Model | 目标人群： kernel 入门者

这一页把从 PyTorch 到 CUDA / Triton 的编程模型演进讲清楚，重点是知道为什么自定义算子会比原生操作更接近硬件。

关键词： grid, block, kernel

前置阅读 ​

导语： 这一页先把从 PyTorch 到 CUDA / Triton 的编程模型演进接上，再看为什么自定义算子会更接近硬件。

相关阅读 ​

导语： 如果想继续把编程模型和 kernel 组织方式补完整，可以接着看这些页。

Q1：为什么我们需要用 CUDA 或 Triton 编写自定义算子？ ​

点击展开查看解析

PyTorch 原生操作方便，但它默认把很多执行细节交给框架和后端。

当一个计算链里出现大量小算子、重复的 HBM 往返、或者对布局和 tile 有强约束时，原生写法往往会留下更多中间张量和 launch 开销。自定义 CUDA 或 Triton 的价值，不是“更底层”本身，而是把算子组织方式重新对齐到硬件执行粒度。

所以这类问题的关键，不是“要不要手写 kernel”，而是“哪些数据流和执行流必须由我们显式控制”。

### Q1小验证：为什么原生操作可能慢

先想是不是中间张量太多。

python
def memory_traffic(num_ops, fused=False):
    # 自定义算子的价值，不在于把名字写得更底层，而在于减少中间张量落地。
    input_reads = num_ops
    output_writes = 1
    intermediate_writes = max(num_ops - 1, 0) if not fused else max(num_ops - 2, 0)
    return input_reads + output_writes + intermediate_writes

cases = [
    ('separate_ops', memory_traffic(3, fused=False)),
    ('fused_ops', memory_traffic(3, fused=True)),
    ('longer_chain', memory_traffic(5, fused=False)),
]
for name, traffic in cases:
    print(name, '-> traffic units:', traffic)
print('fusion wins when it removes intermediate writes, not when it only shortens code')

Q2：CUDA 的线程层级结构与 GPU 硬件执行单元如何对应？ ​

点击展开查看解析

CUDA 的层级不是语法装饰，而是把“任务怎么切”映射到“硬件怎么调度”。

  • grid 决定一次 kernel launch 覆盖多少独立工作；
  • block 决定哪一组线程能共享数据和同步；
  • warp 是硬件实际调度的基本单位；
  • thread 是最细的计算粒度。

理解这层映射后，就能看出为什么同样的代码，线程数、block 数和 warp 利用率不同，性能会差很多。

### Q2小验证：层级怎么对应

把执行层级记清楚，后面读 kernel 会更顺。

python
def block_launch_stats(num_threads, warp_size=32):
    # grid/block/warp/thread 的对应关系，最终会落到 warp 数和空转 lane 上。
    warps = (num_threads + warp_size - 1) // warp_size
    active_lanes = num_threads
    wasted_lanes = warps * warp_size - num_threads
    occupancy = active_lanes / (warps * warp_size)
    return {'warps': warps, 'wasted_lanes': wasted_lanes, 'occupancy': round(occupancy, 2)}

for threads in [48, 64, 96, 128]:
    print(threads, 'threads ->', block_launch_stats(threads))
print('线程层级的关键不是层级名，而是 warp 是否被填满')

Q3：跨线程共享和同步机制如何限制和优化性能？ ​

点击展开查看解析

跨线程共享真正解决的是“同一份数据能不能在片上被复用”。

如果数据只被读一次，shared memory 和同步的开销可能不值；如果同一份数据会被 block 内多个线程反复使用，先放到 shared memory 再同步，通常能换回更多 HBM 访问的节省。

所以这里不是“共享越多越好”，而是“复用收益是否足以覆盖同步代价”。

### Q3小验证：共享为什么会和同步绑定

复用越多，协作就越重要。

python
def shared_tradeoff(reuse_times, sync_points, hbm_cost=10, smem_cost=2, sync_cost=3):
    # 共享只有在复用收益大于同步代价时才值得。
    saved_hbm = max(reuse_times - 1, 0) * (hbm_cost - smem_cost)
    penalty = sync_points * sync_cost
    return {'net_gain': saved_hbm - penalty, 'worth_it': saved_hbm > penalty}

for case in [(1, 0), (2, 1), (4, 1), (4, 3)]:
    print(case, '->', shared_tradeoff(*case))
print('shared memory is good only when reuse amortizes synchronization')

Q4：Triton 为什么能降低算子开发门槛？ ​

点击展开查看解析

Triton 的抽象重点，不是“帮你自动变快”，而是把 tile、block 和布局这些样板性细节收起来，让你先把张量切分方式写清楚。

和 CUDA 相比，它并不是取消底层约束，而是把底层约束收束成更少的显式概念：布局、mask、tile 和 program 关系。

所以 Triton 的门槛更低，是因为它减少了模板代码；但它的性能边界，仍然取决于你对执行粒度和数据布局的理解。

### Q4小验证：Triton 抽象了什么

先看块组织，再看算子细节。

python
def triton_benefit(tile, layout_contiguous=True, boilerplate_lines=80):
    # Triton 的门槛低，靠的是把 tile / layout / launch 的样板代码抽掉。
    if not layout_contiguous:
        return {'boilerplate_saved': 0, 'ready': False}
    saved = max(boilerplate_lines - tile // 2, 0)
    return {'boilerplate_saved': saved, 'ready': True}

for case in [(64, True, 80), (128, True, 80), (128, False, 80)]:
    print(case, '->', triton_benefit(*case))
print('Triton abstracts launch and tiling details, but still depends on layout discipline')

⚠️ 常见误区 ​

  • 自定义算子不是为了炫技，而是为了更好地控制执行路径。
  • CUDA 和 Triton 是不同范式，但目标都在靠近硬件。
  • 共享和同步是一起看的，不能只谈一个。
  • 编程模型的核心是映射，不是语法本身。

Released under the MIT License.

Copyright © 2024-present