jax-triton

jax-triton 仓库包含了 JAX 和 Triton 之间的集成。

文档可以在这里找到。

这不是 Google 官方支持的产品。

快速入门

主要的函数是 jax_triton.triton_call，用于将 Triton 函数应用于 JAX 数组，包括在 jax.jit 编译的函数内部。例如，我们可以定义一个来自 Triton 教程的内核：

import triton
import triton.language as tl


@triton.jit
def add_kernel(
    x_ptr,
    y_ptr,
    length,
    output_ptr,
    block_size: tl.constexpr,
):
  """添加两个向量。"""
  pid = tl.program_id(axis=0)
  block_start = pid * block_size
  offsets = block_start + tl.arange(0, block_size)
  mask = offsets < length
  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)

然后我们可以使用 jax_triton.triton_call 将其应用于 JAX 数组：

import jax
import jax.numpy as jnp
import jax_triton as jt

def add(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
  out_shape = jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype)
  block_size = 8
  return jt.triton_call(
      x,
      y,
      x.size,
      kernel=add_kernel,
      out_shape=out_shape,
      grid=(x.size // block_size,),
      block_size=block_size)

x_val = jnp.arange(8)
y_val = jnp.arange(8, 16)
print(add(x_val, y_val))
print(jax.jit(add)(x_val, y_val))

查看示例目录，特别是 fused_attention.py 和融合注意力 ipynb。

安装

$ pip install jax-triton

确保安装了与 CUDA 兼容的 jaxlib。 例如，你可以运行：

$ pip install "jax[cuda11_cudnn82]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

安装最新版本

JAX-Triton 和 Pallas 在 JAX 和 Jaxlib 的最新版本以及接近 Triton 的最新版本上开发。要获取 JAX-Triton 的最新版本，请运行：

$ pip install 'jax-triton @ git+https://github.com/jax-ml/jax-triton.git'

这应该会安装兼容版本的 JAX 和 Triton。

JAX-Triton 确实依赖于 Jaxlib，但它通常是一个更稳定的依赖项。你可能可以使用最近的 jaxlib 发布版本：

$ pip install jaxlib[cuda11_pip]
$ # 或
$ pip install jaxlib[cuda12_pip]

如果发现最新的 Jaxlib 发布版本有问题，可以尝试使用 Jaxlib 的每日构建版本。 要安装新的 jaxlib，你可以找到 CUDA 11 每日构建 或 CUDA 12 每日构建 的链接。然后通过以下方式安装：

$ pip install 'jaxlib @ <每日构建链接>'

或者自动通过 pip 安装 CUDA，你可以这样做：

$ pip install 'jaxlib[cuda11_pip] @ <每日构建链接>'
$ # 或
$ pip install 'jaxlib[cuda12_pip] @ <每日构建链接>'

开发

要开发 jax-triton，你可以通过以下方式克隆仓库：

$ git clone https://github.com/jax-ml/jax-triton.git

然后进行可编辑安装：

$ cd jax-triton
$ pip install -e .

要运行 jax-triton 测试，你需要安装 pytest：

$ pip install pytest
$ pytest tests/