jax-triton

When a Triton kernel performs an in-place operation (where an input argument is also the output), use the input_output_aliases parameter in jt.triton_call.

input_output_aliases is a dictionary mapping the index of the input argument to the index of the output argument. For example, {1: 0} indicates that the argument at index 1 is also the first output argument.

To improve efficiency and avoid XLA making copies of non-donated in-out arguments (since JAX arrays are immutable by default), use jax.jit with donate_argnames.

Example: In-place addition

@triton.jit
def add_inplace_y_kernel(x_ptr, y_inout_ptr, length, block_size: tl.constexpr):
  # ... kernel logic ...
  tl.store(y_inout_ptr + offsets, output, mask=mask)

@partial(jax.jit, donate_argnames="y")
def add_inplace_y(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
  block_size = 8
  return jt.triton_call(
      x,
      y,
      x.size,
      kernel=add_inplace_y_kernel,
      input_output_aliases={1: 0}, # arg index 1 (y) is the first output
      out_shape=x,
      grid=(x.size // block_size,),
      block_size=block_size
  )

@partial(jax.jit, donate_argnames="y")
def add_inplace_y(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
  block_size = 8
  return jt.triton_call(
      x,
      y,
      x.size,
      kernel=add_inplace_y_kernel,
      input_output_aliases={1: 0},  # input arg idx 1 (y) is the first output arg
      out_shape=x,
      grid=(x.size // block_size,),
      block_size=block_size)

The primary API for integrating Triton kernels with JAX is jax_triton.triton_call. This function allows you to apply a Triton kernel to JAX arrays and can be used inside jax.jit-compiled functions.

Basic Usage Pattern

1. Define a Triton kernel using @triton.jit.
2. Use jt.triton_call to invoke the kernel.

Arguments for triton_call:

• Input arguments: The first arguments passed to triton_call correspond to the kernel's input arguments.
• kernel: The Triton kernel function.
• out_shape: The shape/template for the output array.
• grid: The grid dimensions for the kernel execution.
• kwargs: Any additional keyword arguments (like block_size) are passed as tl.constexpr parameters to the Triton kernel.

Note: The output argument is passed implicitly after the input arguments in the Triton kernel signature.

import jax
import jax.numpy as jnp
import jax_triton as jt
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)

def add(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
  block_size = 8
  return jt.triton_call(
      x,
      y,
      x.size,
      kernel=add_kernel,
      out_shape=x,
      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))

The provided example demonstrates a fused Flash Attention implementation using Triton. The kernel _fwd_kernel uses tiling to compute attention scores and updates an accumulator in a single pass to minimize memory bandwidth usage.

Key components of the kernel implementation:

• Tiling: Uses BLOCK_M and BLOCK_N to process chunks of the sequence.
• Online Softmax: Implements the online softmax algorithm to update the running maximum (m_i) and the running sum of exponentials (l_i).
• Scratchpad Buffer: Uses a TMP buffer to work around specific compiler bugs when storing/loading intermediate scaling factors.
• Strides: Explicitly handles strides for all dimensions (Batch, Head, Seq, Dim) to allow for flexible memory layouts.

@triton.jit
def _fwd_kernel(
    Q, K, V,
    TMP, L, M,  # Scratchpad, L, and M buffers
    Out,
    stride_qz, stride_qh, stride_qm, stride_qk,
    # ... other strides
    Z, H, N_CTX,
    BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
):
    # ... kernel implementation logic ...
    # 1. Load Q, K, V pointers
    # 2. Loop over blocks of K and V
    # 3. Compute dot product and online softmax updates
    # 4. Update accumulator and store results