Triton 与 TPU

Triton 与 TPU

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

• 上一节讲了用 CUDA C 做 GPU 编程。本节顺着抽象阶梯往上爬：Triton 让你用 20% 的力气拿到 CUDA 80% 的性能，而且全程在 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,  # 编译期常量
):
    # 每个 program 实例处理一个含 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()

    # 启动：每个块一个 program
    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 的关键差异：
  • 没有显式的线程管理。你思考的是块（program），而不是线程。
  • tl.arange(0, BLOCK_SIZE) 为整个块创建一个偏移向量。对这个向量的所有操作都隐式向量化。
  • mask 处理边界条件（类似第 03 节里 AVX-512 的掩码寄存器）。无需标量收尾循环。
  • tl.load 和 tl.store 自动处理合并访问。
  • @triton.jit 在第一次调用时把函数编译成 PTX（GPU 汇编），然后缓存编译好的内核。

Triton 的 Softmax 内核

• Softmax 是一个绝佳的 Triton 示例，因为它需要对数据做多趟扫描（最大值、减法、exp、求和、相除），并能从多趟之间把数据留在 SRAM（共享内存）中获益：

@triton.jit
def softmax_kernel(
    output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols,
    BLOCK_SIZE: tl.constexpr,
):
    # 每个 program 处理一行
    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：先取最大值保证数值稳定，再 exp，再归一化
    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 个独立的内核（最大值、exp 与求和、相除），每个都要读写全局内存。Triton 版本在一个内核里做完，数据全程待在寄存器/SRAM 里。这种内核融合就是自定义 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：融合内核、自定义注意力模式、激活函数、大多数 ML 研究的内核需求。
• 用 CUDA C：追求极致性能（最后那 5-20%）、warp 级原语、复杂的数据相关并行、Triton 表达不了的模式。

案例：Flash Attention

• Flash Attention（Dao 等，2022）是近年来 ML 中影响最大的自定义内核。它用 O(n) 的内存计算注意力（而不是 O(n^2)），从而支持长得多的序列。

• 问题：标准注意力计算 \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 矩阵具体化。按瓦片计算注意力：加载一块 Q、一块 K，算出它们的部分注意力分数，累加，再移到下一块。n \times n 矩阵永远不会被完整生成——SRAM 里同一时刻只存在一个瓦片。

• 在线 softmax（online 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（共享内存）里进行。全局内存（HBM）只在加载 Q、K、V 的块和写回最终输出时才会访问。数据复用因子正比于 SRAM 的大小，而 SRAM 的访问速度约是 HBM 的 100 倍。

• Flash Attention 同时有 Triton 和 CUDA C 两种实现。CUDA 版本更快（效率约高 10%），但 Triton 版本可读性和可改性都好得多——这对研究新的注意力变体很重要。

TPU 架构

• TPU（Tensor Processing Unit，张量处理单元）是 Google 的定制 ML 加速器。它走了一条与 GPU 截然不同的路：

• 脉动阵列（systolic array）：TPU 的核心计算单元是 MXU（Matrix Multiply Unit，矩阵乘单元），一个 128×128 或 256×256 的脉动阵列，通过让数据流过一片乘加单元网格来计算矩阵乘法。数据从边缘进入，在阵列中传播，每个单元做一次乘加并把结果传给下一个。

• 与 GPU（调度数千个独立线程）不同，脉动阵列是单一、确定的数据流。没有线程调度、没有 warp 分歧、没有分支预测。这种简洁让 MXU 在矩阵乘法上极其节能。

• HBM：TPU 用和 GPU 一样的 HBM。TPU v5e 每芯片有 16 GB HBM2e；TPU v5p 有 95 GB HBM2e。

• ICI（Inter-Chip Interconnect，片间互联）：TPU Pod 用定制的高速网络把数百个 TPU 连起来。跨 TPU Pod 的数据并行和模型并行（第 6 章）由 JAX 原生支持。

• 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：大多数 ML 工作负载、基于 PyTorch 的研究、推理服务、有大量非 matmul 计算的工作负载。
• 用 TPU：大规模 JAX 训练（数千芯片）、Google Cloud 上对成本敏感的训练、以矩阵乘为主的工作负载。

如何挑对工具

编程练习（使用带 GPU 运行时的 CoLab）

1. 写一个向量加法的 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")
​

2. 写一个 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")
​

3. 测量 JAX 的 XLA 编译器如何自动融合操作。比较一串操作在加 jit 与不加 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(_)

# Eager（每个操作都是一次独立的内核启动）
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 fuses the 4 operations into 1 kernel)")
​