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. Unlike input_output_aliases, 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	ShapeDtype | Sequence[ShapeDtype] | None	An object with shape and dtype attributes or a sequence of such objects. Pointers for each of the elements of out_type will be passed into kernel following the inputs.	None
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
input_output_aliases	dict[int, int] | None	Deprecated. A dictionary mapping input argument indices to output indices. Providing a mapping will alias the corresponding buffers. Input indices refer to the flattened non-constexpr operands in kernel parameter declaration order (whether passed positionally or as keyword arguments). If operands contain nested tuples, the indices correspond to the flattened leaves. Output indices correspond to the flattened out_type.	None
zeroed_outputs	ValueOrFn[Sequence[int]]	Deprecated. A sequence of indices into the flattened out_type, or a function returning such a sequence, for outputs that should be zeroed before the kernel is launched. Note that this also supports zeroing input-output (i.e. aliased through input_output_aliases) arguments that should be treated as outputs in this argument.	()
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 and zeroed_outputs when either 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.