Calling Triton kernels from JAX

The primary way of using JAX Triton is using jax_triton.triton_call to call handwritten Triton kernels from inside JIT-ted JAX programs.

jax_triton.triton_call

Calls a Triton kernel with jax.Array arguments.

Example usage:

First we define a simple kernel that adds two vectors.

import triton
import triton.language as tl

@triton.jit
def add_kernel(
    x_ptr,
    y_ptr,
    output_ptr,
    block_size: tl.constexpr = 128,
):
  pid = tl.program_id(axis=0)
  block_start = pid * block_size
  offsets = block_start + tl.arange(0, block_size)
  mask = offsets < 8
  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)

Then we use triton_call to call it from JAX.

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

def add(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
  block_size = 8
  return jt.triton_call(
      x,
      y,
      kernel=add_kernel,
      out_type=jax.typeof(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))
print(jax.jit(add)(x_val, y_val))

Parameters:

Name	Type	Description	Default
*args	Array | Ref | StaticScalar	Positional operands for the Triton kernel. Array and Ref arguments are passed as runtime buffers in their positional order before any out_type pointers. Ref arguments, created via jax.new_ref, are read-write buffers that the kernel mutates in-place; they should not be included in out_type. Non-array scalars are baked in as static (specialization) values, and any argument bound to a constexpr parameter becomes a metaparam.	()
kernel	Autotuner | Heuristics | JITFunction	A Triton kernel (e.g. a function decorated with triton.jit). All static values should be annotated with triton.language.constexpr or triton.experimental.gluon.language.constexpr.	required
out_type	Any	An object with shape and dtype attributes, or an arbitrary pytree of such objects. A pointer for each leaf of out_type will be passed into kernel following the inputs, and the result is returned with the same pytree structure.	required
grid	ValueOrFn[Grid]	An integer, tuple of up to 3 integers, or a function that returns a tuple of up to 3 integers. When grid is an integer, kernel is invoked in grid-many parallel executions. When grid is a sequence of integers, kernel is launched in a prod(grid)-many parallel executions. When grid is a function, it is passed **metaparams and should return a tuple of up to 3 integers.	required
name	str	A name for the kernel call.	''
compute_capability	int | None	The GPU compute capability to compile for.	None
num_warps	int | None	The number of warps used to execute the Triton kernel.	None
num_stages	int | None	The number of stages emitted by the Triton compiler.	None
num_ctas	int | None	The size of thread blocks per cluster to be used on GPUs with compute capabilities >= 9.0. It must be less or equal to 8.	None
debug	bool	Prints out intermediate IRs if True for debugging purposes.	False
backend_options	Mapping[str, Any] | None	A mapping of backend-specific compiler options. The available options depend on the Triton backend. The num_warps, num_stages, num_ctas and debug are merged into this mapping. It is an error to specify the same option in both.	None
serialized_metadata	bytes	Arbitrary metadata that will be added into the serialized kernel call.	b''
cost_estimate	CostEstimate | None	An estimate of the number of floating point operations ("flops") and memory bytes accessed ("bytes_accessed") by this kernel invocation. This is used by profiling tools to compute the performance metrics of this custom call (e.g. FLOPs/s and bandwidth).	None
has_side_effect	bool	Whether the Triton kernel has side effects.	False
**kwargs	Any	Keyword arguments for the Triton kernel. A keyword that names a non-constexpr kernel parameter is treated as an operand and is subject to the same scalar-static/runtime buffer separation as positional *args. All other keywords -- constexpr parameters and names that are not kernel parameters -- are treated as metaparams. Missing constexpr arguments are filled from the kernel's declared defaults. Metaparams are also provided to grid when it is a function. A misspelled operand name silently becomes a metaparam rather than raising an error.	{}

Returns:

Type	Description
Any	Outputs from the Triton kernel.