第16章 §05 · Triton / TPU / Pallas · 讲解稿

篇幅: ~4000 字 / 阅读 10 分钟

---

一 · 一句话总览

GPU 编程三大新势力: Triton (Pythonic) / TPU (Google 专用) / Pallas (JAX 核函数)。让 GPU/TPU 编程平民化, 不再是 CUDA C++ 一家独大。

---

二 · Triton (OpenAI)

2.1 核心设计

Triton = Pythonic GPU 编程语言 (Philippe Tillet 2019, OpenAI 收购): - 类 Python 语法 - 自动优化 (memory coalescing / shared mem) - 编译为 PTX (NVIDIA GPU)

2.2 编程示例: Softmax

import triton
import triton.language as tl
import torch

@triton.jit
def softmax_kernel(out_ptr, in_ptr, n_cols, BLOCK: tl.constexpr):
    row_idx = tl.program_id(0)
    in_ptr += row_idx * n_cols
    out_ptr += row_idx * n_cols

    # 1. 加载一行
    offsets = tl.arange(0, BLOCK)
    mask = offsets < n_cols
    x = tl.load(in_ptr + offsets, mask=mask, other=-float('inf'))

    # 2. 减 max (数值稳定)
    x_max = tl.max(x, axis=0)
    x = x - x_max

    # 3. exp & sum
    num = tl.exp(x)
    denom = tl.sum(num, axis=0)

    # 4. 写回
    tl.store(out_ptr + offsets, num / denom, mask=mask)

def softmax_triton(x):
    n_rows, n_cols = x.shape
    out = torch.empty_like(x)
    BLOCK = triton.next_power_of_2(n_cols)
    softmax_kernel[(n_rows,)](out, x, n_cols, BLOCK=BLOCK)
    return out

2.3 Triton vs CUDA

维度	Triton	CUDA
语言	Python	C++
学习曲线	低	高
性能	80-95% SOTA	100% SOTA
控制粒度	中	极细
适用	算子 / Kernel	极致优化

2.4 Triton 生态

• Triton-Lang: 官方
• torch.compile (PyTorch 2.0): 默认用 Triton 后端
• vLLM / SGLang: 核心算子用 Triton 实现
• JAX: 通过 jax-triton 桥接

---

三 · Google TPU

3.1 TPU 演进

版本	年份	HBM	性能	特点
v1	2016	8 GB	23 TFLOPS	训练
v2	2017	16 GB	45 TFLOPS	Pod
v3	2018	32 GB	123 TFLOPS	液冷
v4	2021	32 GB	275 TFLOPS	OCS 互联
v5e	2023	16 GB	197 TFLOPS	推理优化
v5p	2023	95 GB	459 TFLOPS	训练
v6 (Trillium)	2024	32 GB	925 TFLOPS	Sparse Core

3.2 TPU 架构特点

• MXU (Matrix Multiply Unit): 128×128 systolic array
• HBM: 高带宽内存 (与 GPU 类似)
• OCS (Optical Circuit Switch): v4+ 互联 (光开关)
• Pod: 1024 chips (v4) / 9216 chips (v5p)

3.3 JAX + TPU

import jax
import jax.numpy as jnp

@jax.jit
def matmul(x, w):
    return jnp.dot(x, w)

# 自动编译到 TPU
x = jnp.ones((1024, 1024))
w = jnp.ones((1024, 1024))
y = matmul(x, w)  # 在 TPU 上运行

3.4 TPU vs GPU

维度	TPU	GPU
性能/瓦	极佳	中
生态	Google Cloud + JAX	通用
可用性	云为主	本地 + 云
软件	XLA / JAX	CUDA / PyTorch
适用	大规模训练	通用

---

四 · Pallas (JAX Kernel)

4.1 Pallas 是什么

Pallas = JAX 扩展 (Google 2024), 让用户写 GPU/TPU kernel, 像 Numba 但更强大: - BlockSpec: 描述内存加载 - Grid: tile 大小 - 自动 + 手动优化

4.2 Pallas 编程示例

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

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

# GPU TPU 通用
def add_pallas(x, y):
    return pl.pallas_call(
        add_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid=(x.shape[0] // 128,),  # 128 元素一个 block
        in_specs=[pl.BlockSpec((128,), lambda i: (i,))] * 2,
        out_specs=pl.BlockSpec((128,), lambda i: (i,)),
    )(x, y)

4.3 Pallas 适用场景

• 自定义算子 (FlashAttention / LinearAttention)
• TPU 优化 (比 XLA 自动调优更细)
• 研究新算法

---

五 · RISC-V 与嵌入式 (Ch16 §06 简述)

5.1 RISC-V

RISC-V = 开源 ISA: - 伯克利 2010 年发起 - 灵活 (可扩展) - 中国大力发展 (阿里玄铁 / 平头哥)

5.2 RISC-V Vector Extension

• RVV 1.0: 可变长度 SIMD (128-65536 bit)
• 类似 ARM SVE
• 适合 ML 推理

5.3 嵌入式 GPU

• Mali (ARM): 移动 GPU
• Adreno (高通): 手机 GPU
• IMG (Imagination): PowerVR

---

六 · Vulkan Compute

6.1 跨平台 GPU

Vulkan = Khronos 标准, 跨厂商: - Windows / Linux / Android / iOS (MoltenVK) - AMD / NVIDIA / Intel / Apple / ARM Mali

6.2 WebGPU

WebGPU = 浏览器 GPU API (Chrome / Firefox 2024+): - WebLLM: 浏览器跑 LLM - Transformers.js: 浏览器 HuggingFace - 未来: 浏览器原生 LLM 推理

---

七 · 一句话总结

Triton + TPU + Pallas + WebGPU = GPU 编程的未来。Triton 让 GPU 平民化, TPU 让大模型训练便宜, Pallas 让 JAX 用户写 kernel, WebGPU 让浏览器跑 LLM。异构计算 + 软件民主化 = 未来 5 年。