Triton 与 TPU¶

一句话总结: Triton 让你用 Python 写 GPU 核函数，拿到 CUDA 八成的性能却只用两成的代码；TPU 则用脉动阵列走了另一条路，靠 JAX 与 Pallas 编程。

CUDA C 强大但啰嗦。 Triton 让你用 Python 写 GPU 核函数。 TPU 提供了 GPU 之外的另一种选择，权衡取舍各有不同。本节讲 Triton 核函数编程、以 Flash Attention 为案例，讨论 TPU 架构与 JAX/Pallas，以及如何选对工具。 Vulkan 与跨平台 GPU 计算见第 07 节.

上一节用 CUDA C 教了 GPU 编程。本节沿抽象阶梯往上爬一层：Triton 用两成的力气拿到 CUDA 八成的性能，而且是在 Python 里。 TPU 和 Vulkan 则为特定场景提供了另类的硬件目标。

Triton：用 Python 写 GPU 核函数¶

Triton (OpenAI 出品) 是一门基于 Python 的 GPU 核函数编写语言。你不再需要围绕单个线程去推理 (CUDA 那样)，而是围绕数据块 (block) 来思考。 Triton 编译器会自动帮你处理线程映射、内存合并、共享内存管理以及许多其他优化。
Triton 为什么重要: CUDA C 要求你深入了解 warp 调度、共享内存 bank 冲突、寄存器压力、合并访问模式。 Triton 把大部分这些细节都抽象掉了，让懂 Python 但不懂系统编程的 ML 研究者也能上手 GPU 核函数开发。

你的第一个 Triton 核函数¶

import triton
import triton.language as tl
import torch

@triton.jit
def add_kernel(
    x_ptr, y_ptr, output_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,  # 编译期常量
):
    # 每个程序实例处理 BLOCK_SIZE 个元素组成的一块
    pid = tl.program_id(axis=0)  # 我是哪一块?
    block_start = pid * BLOCK_SIZE

    # 本块内的偏移
    offsets = block_start + tl.arange(0, BLOCK_SIZE)

    # 用掩码处理 n_elements 不是 BLOCK_SIZE 整数倍的情况
    mask = offsets < n_elements

    # 载入数据 (带掩码: 越界读取返回 0)
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)

    # 计算
    output = x + y

    # 存回结果
    tl.store(output_ptr + offsets, output, mask=mask)


def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    output = torch.empty_like(x)
    n_elements = output.numel()

    # 启动: 每块一个程序
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),)
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)

    return output


# 使用
x = torch.randn(1000000, device='cuda')
y = torch.randn(1000000, device='cuda')
z = add(x, y)

与 CUDA 的关键区别:
- 没有显式的线程管理。你思考的是块 (程序)，不是线程。
- tl.arange(0, BLOCK_SIZE) 为整个块生成一个偏移向量。该向量上的所有运算都是隐式向量化的。
- mask 处理边界条件 (类似高级向量扩展 (Advanced Vector Extensions, AVX)-512 的掩码寄存器，见第 03 节). 不需要额外的标量清理循环。
- tl.load 与 tl.store 自动处理合并访问。
- @triton.jit 会在首次调用时把函数编译为 PTX (GPU 汇编)，然后把编译好的核函数缓存起来。

Triton Softmax 核函数¶

Softmax 是 Triton 的一个绝佳示例，因为它需要对数据做多次遍历 (求最大、减去、取指数、求和、除法)，又能从"两次遍历之间把数据留在 SRAM (共享内存) 里"中获益:

@triton.jit
def softmax_kernel(
    output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols,
    BLOCK_SIZE: tl.constexpr,
):
    # 每个程序处理一行
    row_idx = tl.program_id(0)
    row_start = input_ptr + row_idx * input_row_stride

    # 载入该行
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols
    row = tl.load(row_start + col_offsets, mask=mask, other=-float('inf'))

    # Softmax: 求最大值保证数值稳定, 再取指数, 再归一化
    row_max = tl.max(row, axis=0)
    numerator = tl.exp(row - row_max)
    denominator = tl.sum(numerator, axis=0)
    softmax_output = numerator / denominator

    # 存回结果
    output_start = output_ptr + row_idx * output_row_stride
    tl.store(output_start + col_offsets, softmax_output, mask=mask)

在 PyTorch 里，F.softmax(x, dim=-1) 会启动 3 个独立的核函数 (max、exp 加 sum、除法)，每个都要从全局内存读、再写回全局内存。 Triton 版本把这一切塞进一个核函数，数据始终留在寄存器/SRAM 里。这就是所谓的算子融合 (kernel fusion)，也是自定义 Triton 核函数能比 PyTorch 内置算子快 2 到 4 倍的原因。

Triton 自动调优¶

Triton 支持自动调优 (auto-tuning)：尝试多种配置，挑最快的:

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32}),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32}),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64}),
    ],
    key=['M', 'N', 'K'],  # 这些参数变化时重新调优
)
@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, ...):
    ...

Triton 会在真实硬件上对每种配置跑基准测试，选出最快的那个。最佳分块尺寸取决于 GPU 架构、矩阵维度和内存布局 —— 自动调优不需要你手动试，就能把它们找出来。

Triton vs CUDA：什么时候用哪个¶

	Triton	CUDA C
语言	Python	C/C++
抽象层级	块级	线程级
开发速度	快 (每个核 10-50 行)	慢 (每个核 100-500 行)
性能上限	手调 CUDA 的 80-95%	100% (完全硬件控制)
共享内存	自动	手动
合并访问	自动	手动
Warp 级原语	有限	完整 (shuffle、vote 等)
硬件支持	仅 NVIDIA (AMD 实验性)	仅 NVIDIA

用 Triton：融合核、自定义注意力模式、激活函数，以及大多数 ML 研究里的核函数需求。
用 CUDA C：榨干最后的 5-20% 性能、warp 级原语、复杂的数据依赖并行，以及 Triton 表达不出来的模式。

案例：Flash Attention¶

Flash Attention (Dao 等，2022) 是近年 ML 领域影响力最大的自定义核函数。它把注意力的显存占用从 \(O(n^2)\) 降到 \(O(n)\)，直接把序列长度做长了一个量级。
问题所在：标准注意力算的是 \(\text{softmax}(QK^T / \sqrt{d}) \cdot V\). 其中 \(QK^T\) 是一个 \(n \times n\) 矩阵，\(n\) 是序列长度。当 \(n = 128K\) 时，这个矩阵是 \(128K \times 128K \times 4\) 字节 = 64 GB. 塞不进 GPU 显存。
关键洞察：你根本不需要把整个 \(n \times n\) 矩阵显式落地。分块 (tile) 算注意力：载入一块 \(Q\)、一块 \(K\)，算它们的部分注意力得分，累加进去，再挪到下一块。那个 \(n \times n\) 矩阵从来没有完整存在过 —— SRAM 里同时只留一块。
在线 softmax：麻烦在 softmax，它需要知道整行的最大值 (为了数值稳定). Flash Attention 用了在线 softmax技巧：维护一个"当前最大值"，一旦遇到更大的值，就对之前算好的结果重新缩放。这样 softmax 就能一块一块地增量算出来。
算法如下:

对 Q 的每一行块:
    对 K 的每一列块:
        1. 把 Q_block 从 HBM 载入 SRAM
        2. 把 K_block 从 HBM 载入 SRAM
        3. 在 SRAM 里算 S_block = Q_block @ K_block.T
        4. 更新当前最大值, 对之前的结果重新缩放
        5. 计算 exp(S_block - running_max)
        6. 更新当前和与输出累加器
    载入 V_block, 算最终输出
    把输出块写回 HBM

为什么快：内层循环全部在 SRAM (共享内存) 里跑。只有为了载入 Q、K、V 的块和写回最终输出时，才碰全局内存 (HBM). 数据的复用倍数正比于 SRAM 大小，而 SRAM 的访问速度比 HBM 快约 100 倍。
Flash Attention 有 Triton 和 CUDA C 两个版本。 CUDA 版更快 (效率高约 10%)，但 Triton 版可读性好得多，也更容易改，这对研究新的注意力变体很关键。

TPU 架构¶

TPU (Tensor Processing Unit，张量处理单元) 是谷歌自研的 ML 加速器。它走的路和 GPU 截然不同:
脉动阵列 (systolic array): TPU 的核心计算单元是 矩阵乘法单元 (Matrix Multiply Unit, MXU)，一个 128×128 或 256×256 的脉动阵列，通过让数据在乘加单元组成的网格中流动来完成矩阵乘法。数据从阵列边缘进入，逐格向内传播，每个单元做一次乘加，再把结果传给下一个。
和 GPU (调度成千上万个独立线程) 不同，脉动阵列是一条单一的、确定性的数据流。没有线程调度，没有 warp 分化，也没有分支预测。这份简单让 MXU 在矩阵乘法上的能效极其突出。
HBM: TPU 和 GPU 一样用高带宽显存 (High Bandwidth Memory, HBM). TPU v5e 每片有 16 GB HBM2e; TPU v5p 每片有 95 GB HBM2e.
ICI (Inter-Chip Interconnect，片间互连): TPU pod 用一张自研的高速网络把成百上千片 TPU 连起来。 JAX 天然支持 TPU pod 上的数据并行和模型并行 (见第 6 章).
BFloat16: TPU 是最早用 bfloat16 (见第 13 章第 02 节) 的。 BF16 的指数范围和 float32 一样宽 (避免训练中溢出)，尾数精度更低。这种取舍非常适合 ML —— 梯度值跨度很大，但不需要 23 位精度。

编程 TPU: JAX 和 Pallas¶

TPU 通过 JAX 与 XLA 来编程。你写 Python/JAX 代码，jax.jit 把它编译成 XLA HLO, XLA 再把 HLO 编译成 TPU 专用指令。没有 CUDA，没有 C++.

import jax
import jax.numpy as jnp

@jax.jit
def matmul(a, b):
    return jnp.dot(a, b)

# 这段代码会跑在 CPU、GPU 或 TPU 上, 取决于当前设备
a = jnp.ones((1024, 1024))
b = jnp.ones((1024, 1024))
c = matmul(a, b)

Pallas 是 JAX 的核函数编写 API —— 相当于 JAX 版的 Triton. 它让你写底层核函数，XLA 会把它编译到 GPU 或 TPU:

from jax.experimental import pallas as pl
import jax.numpy as jnp

def add_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] + y_ref[...]

def add_pallas(x, y):
    return pl.pallas_call(
        add_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid=(x.shape[0] // 128,),
        in_specs=[pl.BlockSpec((128,), lambda i: (i,)),
                  pl.BlockSpec((128,), lambda i: (i,))],
        out_specs=pl.BlockSpec((128,), lambda i: (i,)),
    )(x, y)

Pallas 比 Triton 更新、更不成熟，但它是给 TPU 写自定义核函数唯一的路 (TPU 不支持 CUDA).

GPU vs TPU¶

	GPU (NVIDIA)	TPU (谷歌)
可获得性	任何云、本地都可	仅谷歌云
编程	CUDA C、Triton、PyTorch	JAX/XLA、Pallas
灵活性	通用计算	为矩阵密集的 ML 优化
峰值矩阵乘 FLOPS	极高 (Tensor Core)	极高 (MXU)
非矩阵乘算子	好	慢 (走向量单元，不走 MXU)
多片扩展	NVLink (8 GPU)、InfiniBand	ICI (数千片 TPU，集成度更高)
性价比	有竞争力	大规模训练通常更便宜
生态	最大 (PyTorch、TensorFlow、JAX)	以 JAX 为主

用 GPU：大多数 ML 工作负载、基于 PyTorch 的研究、推理服务，以及非矩阵乘计算占比不小的场景。
用 TPU：大规模 JAX 训练 (几千片)、对成本敏感的谷歌云训练，以及计算主要就是矩阵乘的场景。

选对工具¶

工作负载	首选工具	原因
ML 训练 (PyTorch)	NVIDIA GPU + CUDA/Triton	生态最大，工具最全
ML 训练 (JAX，大规模)	TPU 或 NVIDIA GPU	谷歌规模下 TPU 划算，GPU 更灵活
自定义融合核	Triton (Python) 或 CUDA C	Triton 开发快，CUDA 榨性能
JAX 自定义核	Pallas	给 TPU 写核唯一的路，GPU 也能用
跨平台推理	Vulkan (第 07 节) 或 ONNX Runtime	兼容任何 GPU 厂商
移动端/边缘推理	Metal (Apple)、Vulkan (Android)、NNAPI	平台专属加速器
浏览器推理	WebGPU (第 07 节)	浏览器里唯一的路
仅 CPU 推理	ONNX Runtime + AVX/NEON	不用 GPU，靠 SIMD (第 02-03 节)
新硬件	厂商专属 SDK	每种加速器都有自己的工具链

编程作业 (在带 GPU 运行时的 CoLab 里做)¶

写一个 Triton 向量加法核函数并跑起来。和 PyTorch 内置的加法比性能。

import triton
import triton.language as tl
import torch
import time

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask)

n = 10_000_000
x = torch.randn(n, device='cuda')
y = torch.randn(n, device='cuda')

# Triton
out_triton = torch.empty_like(x)
grid = lambda meta: (triton.cdiv(n, meta['BLOCK']),)
add_kernel[grid](x, y, out_triton, n, BLOCK=1024)

# PyTorch
out_torch = x + y

# 校验正确性
assert torch.allclose(out_triton, out_torch, atol=1e-5)

# 基准测试
torch.cuda.synchronize()
start = time.time()
for _ in range(1000):
    add_kernel[grid](x, y, out_triton, n, BLOCK=1024)
torch.cuda.synchronize()
triton_time = (time.time() - start) / 1000

start = time.time()
for _ in range(1000):
    out_torch = x + y
torch.cuda.synchronize()
torch_time = (time.time() - start) / 1000

print(f"Triton:  {triton_time*1000:.3f} ms")
print(f"PyTorch: {torch_time*1000:.3f} ms")
print(f"Ratio:   {torch_time/triton_time:.2f}x")

写一个 Triton 融合核，一趟做完 "乘 + 加 + ReLU". 和 PyTorch 三个独立算子对比。

import triton
import triton.language as tl
import torch
import time

@triton.jit
def fused_mul_add_relu_kernel(x_ptr, w_ptr, b_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    w = tl.load(w_ptr + offs, mask=mask)
    b = tl.load(b_ptr + offs, mask=mask)
    result = tl.maximum(x * w + b, 0.0)  # 融合: mul + add + relu
    tl.store(out_ptr + offs, result, mask=mask)

n = 10_000_000
x = torch.randn(n, device='cuda')
w = torch.randn(n, device='cuda')
b = torch.randn(n, device='cuda')

# 融合 (Triton)
out_fused = torch.empty_like(x)
grid = lambda meta: (triton.cdiv(n, meta['BLOCK']),)
fused_mul_add_relu_kernel[grid](x, w, b, out_fused, n, BLOCK=1024)

# 未融合 (PyTorch)
out_unfused = torch.relu(x * w + b)

assert torch.allclose(out_fused, out_unfused, atol=1e-5)

# 基准测试
torch.cuda.synchronize()
start = time.time()
for _ in range(1000):
    fused_mul_add_relu_kernel[grid](x, w, b, out_fused, n, BLOCK=1024)
torch.cuda.synchronize()
fused_time = (time.time() - start) / 1000

start = time.time()
for _ in range(1000):
    out_unfused = torch.relu(x * w + b)
torch.cuda.synchronize()
unfused_time = (time.time() - start) / 1000

print(f"Fused (Triton):    {fused_time*1000:.3f} ms")
print(f"Unfused (PyTorch): {unfused_time*1000:.3f} ms")
print(f"Speedup:           {unfused_time/fused_time:.2f}x")

测量 JAX 的 XLA 编译器如何自动融合算子。对比一串算子加不加 jit 的差异。

import jax
import jax.numpy as jnp
import time

def chain_ops(x):
    x = x * 2.0
    x = x + 1.0
    x = jnp.maximum(x, 0.0)  # ReLU
    x = x / jnp.sum(x)
    return x

chain_jit = jax.jit(chain_ops)
x = jax.random.normal(jax.random.PRNGKey(0), (10000, 1000))

# 预热
_ = chain_jit(x)
jax.block_until_ready(_)

# 立即执行 (每个算子是一次独立的核启动)
start = time.time()
for _ in range(100):
    y = chain_ops(x)
jax.block_until_ready(y)
eager_time = (time.time() - start) / 100

# JIT (XLA 融合算子)
start = time.time()
for _ in range(100):
    y = chain_jit(x)
jax.block_until_ready(y)
jit_time = (time.time() - start) / 100

print(f"Eager: {eager_time*1000:.2f} ms")
print(f"JIT:   {jit_time*1000:.2f} ms")
print(f"Speedup: {eager_time/jit_time:.1f}x (XLA 把 4 个算子融进了 1 个核)")

📌 本节要点¶

Triton: OpenAI 出品的 Python GPU 核函数语言，以"块"为思考单位，编译器自动搞定线程映射与合并访问。
算子融合: Triton softmax 把 max/exp/sum/div 挤进一个核，数据留在 SRAM，比 PyTorch 快 2-4 倍。
自动调优: @triton.autotune 在真机上试多组分块尺寸，自动挑最快，免去手工实验。
Triton vs CUDA: Triton 10-50 行拿到 80-95% 性能，CUDA 100-500 行拿到 100% 但门槛高。
Flash Attention：分块 + 在线 softmax，把注意力显存从 \(O(n^2)\) 压到 \(O(n)\), 128K 长序列成为可能。
TPU 脉动阵列: MXU 是 128×128 或 256×256 的乘加网格，无线程调度、无分支预测，能效极高。
BFloat16: TPU 首创，指数范围同 float32、尾数更短，训练梯度友好且省显存。
JAX 与 Pallas: TPU 唯一的编程入口是 JAX/XLA; Pallas 是 JAX 版 Triton，给 GPU/TPU 写自定义核。
GPU vs TPU: GPU 生态更大更灵活，TPU 在谷歌云大规模训练上更省钱，但仅限矩阵密集场景。
工具选择: PyTorch 训练配 CUDA/Triton, JAX 大规模配 TPU，跨平台推理走 Vulkan/ONNX，浏览器走 WebGPU.