openxla/tokamax
GitHub: openxla/tokamax
Tokamax 是一个同时支持 NVIDIA GPU 和 Google TPU 的自定义加速器 kernel 库，提供高性能算子实现与自动调优工具，解决大模型训练推理中的 kernel 性能优化难题。
Stars: 250 | Forks: 40
# Tokamax
[](https://github.com/openxla/tokamax/actions/workflows/ci-build.yml)
[](https://pypi.org/project/tokamax/)
Tokamax 是一个自定义加速器 kernel 库，同时支持 NVIDIA GPU 和 Google [TPU](https://cloud.google.com/tpu/docs/intro-to-tpu)。Tokamax 提供了基于 [JAX](https://docs.jax.dev/en/latest/index.html) 和 [Pallas](https://docs.jax.dev/en/latest/pallas/index.html) 构建的、最先进的自定义 kernel 实现。
Tokamax 还为用户提供了工具，以便构建和自动调优（autotune）他们自己的自定义加速器 kernel。
## 状态
Tokamax 仍在大力开发中。可能会出现功能不完整和 API 变更的情况。
我们目前支持以下 GPU kernel：
* `tokamax.dot_product_attention`（[FlashAttention](https://arxiv.org/abs/2205.14135)）。
* `tokamax.gated_linear_unit`（[门控线性单元](https://arxiv.org/abs/2002.05202) (SwiGLU 等)）。
* `tokamax.layer_norm`（[层归一化](https://arxiv.org/abs/1607.06450)和[均方根归一化](https://arxiv.org/abs/1910.07467)）。
同时支持以下 GPU 和 TPU 的 kernel：
* `tokamax.ragged_dot`（[混合专家模型](https://arxiv.org/abs/2211.15841)）。
以及以下 TPU kernel：
* `tokamax.linear_softmax_cross_entropy_loss`（[内存高效的线性交叉熵损失 Kernel](https://arxiv.org/abs/2410.10989v2)）
## 安装说明
最新的 Tokamax [PyPI 发布版本](https://pypi.org/project/tokamax/)：
```
pip install -U tokamax
```
来自 Github 的最新前沿版本，不提供稳定性保证：
```
pip install git+https://github.com/openxla/tokamax.git
```
## 使用 Tokamax
考虑一个包含 Tokamax 函数并在 H100 GPU 上运行的过程：
```
import jax
import jax.numpy as jnp
import tokamax
def loss(x, scale):
x = tokamax.layer_norm(
x, scale=scale, offset=None, implementation="triton"
)
x = tokamax.dot_product_attention(x, x, x, implementation="xla_chunked")
x = tokamax.layer_norm(x, scale=scale, offset=None, implementation=None)
x = tokamax.dot_product_attention(x, x, x, implementation="mosaic")
return jnp.sum(x)
f_grad = jax.jit(jax.grad(loss))
```
当 `implementation=None` 时，Tokamax 可以针对每种 kernel 形状选择最佳实现。它甚至可以为前向传播和梯度选择不同的实现。它也将始终受支持，因为它可以回退到 XLA 实现 `implementation='xla'`。
但是，您可能希望选择特定的 kernel 实现，并在不受支持时报错。例如，`implementation="mosaic"` 会尝试尽可能使用 [Pallas:Mosaic GPU](https://docs.jax.dev/en/latest/pallas/gpu/index.html) kernel，如果因任何原因不受支持，则会抛出异常。例如，使用 FP64 输入或较旧的 GPU 是不受支持的。
### 评估梯度
```
channels, seq_len, batch_size, num_heads = (64, 2048, 32, 16)
scale = jax.random.normal(jax.random.key(0), (channels,), dtype=jnp.float32)
x = jax.random.normal(
jax.random.key(1),
(batch_size, seq_len, num_heads, channels),
dtype=jnp.bfloat16,
)
out = f_grad(x, scale)
```
### 自动调优
为了获得最佳性能，可以对 `f_grad` 中的所有 Tokamax kernel 进行自动调优：
```
autotune_result: tokamax.AutotuningResult = tokamax.autotune(f, x, scale)
```
`autotune_result` 可以用作上下文管理器（context-manager），为 `f_grad` 中的所有 Tokamax kernel 使用调优后的配置：
```
with autotune_result:
out_autotuned = f_grad(x, scale)
```
要对可能耗时的 `tokamax.autotuning` 调用结果进行序列化和重用：
```
autotune_result_json: str = autotune_result.dumps()
autotune_result = tokamax.AutotuningResult.loads(autotune_result_json)
```
用户可以通过继承 `tokamax.Op` 类并重写 `tokamax.Op._get_autotuning_configs` 方法来定义自动调优的搜索空间，从而使用 `tokamax.autotune` 对自己的 kernel 进行调优。
请注意，自动调优本质上是非确定性的：测量 kernel 执行时间会产生噪声。由于在自动调优期间选择不同的配置可能会导致不同的数值结果，因此这是产生数值非确定性的一个潜在原因。对固定的自动调优结果进行序列化和重用，是确保在不同会话中获得相同数值结果的一种方法。
### 序列化
kernel 可以被序列化为 [StableHLO](https://openxla.org/stablehlo)。kernel 调用属于 JAX 自定义调用（custom calls），默认情况下在 `jax.export` 中是被禁止的，[需要使用](https://docs.jax.dev/en/latest/export/export.html#compatibility-guarantees-for-custom-calls)`tokamax.DISABLE_JAX_EXPORT_CHECKS` 才能允许导出所有 Tokamax kernel：
```
from jax import export
f_grad_exported = export.export(f_grad, disabled_checks=tokamax.DISABLE_JAX_EXPORT_CHECKS)(
jax.ShapeDtypeStruct(x.shape, x.dtype),
jax.ShapeDtypeStruct(scale.shape, scale.dtype),
)
```
请注意，使用 Tokamax kernel 序列化的函数会失去标准 StableHLO 的设备独立性。Tokamax 提供两项序列化保证：
1. 针对特定设备序列化的反序列化函数，将保证可以在其序列化时所针对的确切设备上运行。
2. Tokamax 提供[与 JAX 相同的兼容性保证](https://docs.jax.dev/en/latest/export/export.html#compatibility-guarantees-for-custom-calls)：6 个月的向后兼容性。
### 基准测试
JAX 的 Python 开销通常比实际的加速器 kernel 执行时间大得多。这意味着通过测量 `jax.block_until_ready(f_grad(x, scale))` 耗时的常规方法将起不到作用。Tokamax 提供了仅测量加速器执行时间的工具：
```
f_std, args = tokamax.standardize_function(f, kwargs={'x': x, 'scale': scale})
bench: tokamax.BenchmarkData = tokamax.benchmark(f_std, args)
```
有不同的测量技术：例如，在 GPU 上，可以通过 `tokamax.benchmark(f_std, args, method='cupti')` 指定 [CUPTI 分析器](https://docs.nvidia.com/cupti)。这会对 kernel 进行插桩（instruments）并增加少量开销。默认的 `method=None` 允许 Tokamax 自行选择方法，并且同时适用于 TPU 和 GPU。可以通过增加迭代次数来减少基准测试的噪声：
```
tokamax.benchmark(f_std, args, iterations=10)
```
## 免责声明
这不是一款官方的 Google 产品。
标签：AI算子库, GPU, JAX, TPU, Vectored Exception Handling, 底层优化, 深度学习, 逆向工具, 高性能计算