### Install JAX-Triton from HEAD

Source: https://github.com/jax-ml/jax-triton/blob/main/docs/index.md

Install the latest development version of JAX-Triton directly from its GitHub repository.

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

--------------------------------

### Install Triton and Chex

Source: https://github.com/jax-ml/jax-triton/blob/main/examples/JAX_+_Triton_Flash_Attention.ipynb

Installs the Triton library and Chex for testing JAX components. Ensure you have cmake installed.

```python
!pip install -U --pre triton
!pip install chex
!pip install cmake
```

--------------------------------

### Install JAX-Triton

Source: https://github.com/jax-ml/jax-triton/blob/main/docs/index.md

Install JAX-Triton using pip. This command also installs compatible versions of JAX and Triton.

```bash
pip install jax-triton
```

--------------------------------

### Install pytest

Source: https://github.com/jax-ml/jax-triton/blob/main/README.md

Installs pytest, a testing framework required for running the jax-triton tests.

```bash
pip install pytest
```

--------------------------------

### Multiple Config Compilation Setup

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/06-autotuning.md

This snippet shows the setup for compiling multiple configurations of a kernel separately. It iterates through different configurations, normalizes grids, and prepares parameters for each.

```python
config_params = []
for config in configs:
    config_metaparams = {**metaparams, **config.kwargs}
    config_grid = normalize_grid(grid, config_metaparams)
    # ... setup for this config
    config_params.append({...})

kernel_calls = []
for params in config_params:
    kernel, specialization_attr = jtfu.get_or_create_triton_kernel(...)
    # ... create kernel call for this config
    kernel_calls.append(...)
```

--------------------------------

### Install JAX with Triton Support

Source: https://github.com/jax-ml/jax-triton/blob/main/examples/JAX_+_Triton_Flash_Attention.ipynb

Installs a specific version of JAX from a Git repository with Triton support, along with pybind11. This process involves cloning the JAX repository, checking out the 'triton' branch, and then building and installing JAX with CUDA and Triton configurations.

```python
#@title Install JAX + Triton { vertical-output: true }
%cd /root
!pip uninstall jax -y
!git clone https://github.com/sharadmv/jax.git
%cd jax
!git checkout triton
!pip install -I -U ".[cuda11_cudnn82]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
!pip install pybind11
%cd triton
!make
!pip install .
%cd /root
```

--------------------------------

### Install JAX with CUDA support

Source: https://github.com/jax-ml/jax-triton/blob/main/README.md

Installs JAX with CUDA support, which is a prerequisite for using jax-triton with GPU acceleration.

```bash
pip install "jax[cuda12]"
```

--------------------------------

### Editable install of jax-triton

Source: https://github.com/jax-ml/jax-triton/blob/main/README.md

Performs an editable installation of jax-triton after cloning the repository. This allows for development and testing of the library.

```bash
cd jax-triton
pip install -e .
```

--------------------------------

### Install latest stable jaxlib with CUDA

Source: https://github.com/jax-ml/jax-triton/blob/main/docs/index.md

Install the latest stable release of jaxlib with CUDA support.

```bash
pip install "jaxlib[cuda]"
```

--------------------------------

### Block Size Optimization Example

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/02-utils.md

Shows the import of necessary utilities for block size optimization, often involving next_power_of_2.

```python
import jax_triton as jt
import math

```

--------------------------------

### Install JAX with CUDA support

Source: https://github.com/jax-ml/jax-triton/blob/main/docs/index.md

Ensure you have a CUDA-compatible jaxlib installed for JAX-Triton to function on GPU.

```bash
pip install "jax[cuda]"
```

--------------------------------

### Example: Autotuner with Fixed BLOCK_SIZE

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/06-autotuning.md

Demonstrates an autotuner-decorated kernel where the user explicitly sets BLOCK_SIZE, resulting in only the matching configuration being compiled.

```python
# Autotuner with multiple BLOCK_SIZE configs
@triton.autotune(
    configs=[
        triton.Config(kwargs={"BLOCK_SIZE": 64}),
        triton.Config(kwargs={"BLOCK_SIZE": 128}),
        triton.Config(kwargs={"BLOCK_SIZE": 256}),
    ],
    key=["n"],
)
@triton.jit
def kernel(input_ptr, output_ptr, n: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    pass

# User explicitly sets BLOCK_SIZE
result = jt.triton_call(
    x,
    kernel=kernel,
    out_shape=x,
    grid=(256,),
    n=x.shape[0],
    BLOCK_SIZE=128,  # Fixed by user
)

# Only the BLOCK_SIZE=128 config will be compiled
```

--------------------------------

### Running Tests

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/10-package-reference.md

Install pytest and run tests from the tests directory.

```bash
pip install pytest
pytest tests/
```

--------------------------------

### Install JAX and Triton Dependencies

Source: https://github.com/jax-ml/jax-triton/blob/main/examples/JAX_+_Triton_Flash_Attention.ipynb

Installs JAX, Triton, and other necessary libraries using pip. Ensure your environment meets the CUDA and cuDNN version requirements specified by jaxlib.

```bash
pip install --upgrade pip
pip install --upgrade jax jaxlib
pip install triton
```

--------------------------------

### Standard Triton to Gluon Kernel Migration Example

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/09-gpu-dialect-gluon.md

Provides a side-by-side comparison of a kernel written in standard Triton and its equivalent converted to the Gluon dialect. This demonstrates the key syntax changes required for migration.

```python
# Before: Standard Triton
@triton.jit
def kernel(x_ptr, y_ptr, output_ptr, block_size: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * block_size + tl.arange(0, block_size)
    x = tl.load(x_ptr + offsets)
    y = tl.load(y_ptr + offsets)
    tl.store(output_ptr + offsets, x + y)

# After: Gluon
@gl.jit
def kernel(x, y, output, block_size: gl.constexpr):
    pid = gl.program_id(0)
    offsets = pid * block_size + gl.arange(block_size)
    output[offsets] = x[offsets] + y[offsets]
```

--------------------------------

### Install jax-triton Package

Source: https://github.com/jax-ml/jax-triton/blob/main/examples/JAX_+_Triton_Flash_Attention.ipynb

Installs the jax-triton package after its source code has been processed. This step makes the compiled Triton kernels available for use within JAX.

```bash
pip install /root/jax/triton
```

--------------------------------

### Import Examples for jax-triton

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/INDEX.md

Shows different ways to import `triton_call` and other utilities from the `jax_triton` library. You can import specific functions or the entire module.

```python
# Minimal
from jax_triton import triton_call

# With utilities
from jax_triton import triton_call, cdiv, strides_from_shape

# Module import
import jax_triton as jt
jt.triton_call(...)
jt.cdiv(...)
```

--------------------------------

### Standard Triton Kernel Definition

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/04-types.md

Example of defining a standard Triton kernel using the `@triton.jit` decorator.

```python
import triton
import triton.language as tl

# Standard kernel
@triton.jit
def standard_kernel(x_ptr, y_ptr, output_ptr, block_size: tl.constexpr):
    pass
```

--------------------------------

### Example: Heuristics for num_warps

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/06-autotuning.md

Illustrates a kernel using heuristics to dynamically set num_warps based on the input size 'n'. This allows the kernel to adapt its parallelism.

```python
@triton.heuristics({
    "num_warps": lambda args: 4 if args['n'] < 1024 else 8,
})
@triton.jit
def kernel(input_ptr, output_ptr, n: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    pass

# When called with n=100, num_warps is computed as 4
# When called with n=10000, num_warps is computed as 8
result = jt.triton_call(
    x,
    kernel=kernel,
    out_shape=x,
    grid=(256,),
    n=x.shape[0],
    BLOCK_SIZE=128,
)
```

--------------------------------

### Usage Examples for Grid Type

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/04-types.md

Demonstrates how to specify a 1D, 2D, or 3D grid using the Grid type alias in `triton_call`.

```python
jt.triton_call(..., grid=256)           # 1D grid: 256 blocks
jt.triton_call(..., grid=(64, 64))      # 2D grid: 64x64 blocks
jt.triton_call(..., grid=(16, 16, 16))  # 3D grid: 16x16x16 blocks
```

--------------------------------

### Run jax-triton tests

Source: https://github.com/jax-ml/jax-triton/blob/main/README.md

Executes the jax-triton tests using pytest. Ensure pytest is installed before running this command.

```bash
pytest tests/
```

--------------------------------

### Analyze Generated IR Files

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/07-advanced-features.md

Example of enabling debug mode and analyzing the generated IR files for a kernel. Examine TTIR for structure and PTX for GPU instructions.

```python
# Enable debug
result = jt.triton_call(
    x,
    kernel=my_kernel,
    out_shape=out_shape,
    grid=(256,),
    debug=True,  # Print IR
    block_size=256
)

# Examine /tmp/jax_triton_debug/{kernel_hash}/
# Look at TTIR to understand IR structure
# Look at PTX for actual GPU instructions
```

--------------------------------

### Grid Size Calculation for 1D and 2D Kernels

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/02-utils.md

Provides examples for calculating grid sizes for both 1D and 2D Triton kernels using the cdiv utility.

```python
import jax_triton as jt

# For 1D kernels
block_size = 256
# grid_size = jt.cdiv(total_size, block_size)

# For 2D kernels
# block_size_x, block_size_y = 16, 16
# grid_x = jt.cdiv(width, block_size_x)
# grid_y = jt.cdiv(height, block_size_y)
# grid = (grid_x, grid_y)
```

--------------------------------

### Usage Examples for GridOrLambda Type

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/04-types.md

Shows static grid usage and dynamic grid creation using a callable function that computes the grid based on metaparams.

```python
# Static grid
jt.triton_call(..., grid=(64, 64))

# Dynamic grid based on metaparams
def compute_grid(metaparams):
    block_size = metaparams.get('block_size', 256)
    return (jt.cdiv(total_size, block_size),)

jt.triton_call(..., grid=compute_grid, block_size=256)
```

--------------------------------

### Flash Attention Benchmarking Setup

Source: https://github.com/jax-ml/jax-triton/blob/main/examples/JAX_+_Triton_Flash_Attention.ipynb

Sets up a Triton benchmark for fused attention, varying sequence length. It configures benchmark parameters like x-axis values, line arguments, and plot details. This is useful for performance analysis of attention mechanisms.

```python
seq_bench = triton.testing.Benchmark(
    x_names=['N_CTX'],
    x_vals=[2**i for i in range(10, 16)],
    line_arg='provider',
    line_vals=['triton'] + (['flash'] if HAS_FLASH else []),
    line_names=['Triton'] + (['Flash'] if HAS_FLASH else []),
    styles=[('red', '-'), ('blue', '-')],
    ylabel='ms',
    plot_name=f'fused-attention-batch{BATCH}-head{N_HEADS}-d{D_HEAD}',
    args={'H': D_HEAD, 'BATCH': BATCH, 'D_MODEL': D_HEAD, 'dtype': torch.float16}
)


@triton.testing.perf_report([batch_bench, seq_bench])
def bench_flash_attention(BATCH, H, N_CTX, D_MODEL, provider, dtype=torch.float16, device="cuda"):
    warmup = 25
    rep = 500
    if provider == "triton":
        q = torch.randn((BATCH, H, N_CTX, D_MODEL), dtype=dtype, device="cuda", requires_grad=True)
        k = torch.randn((BATCH, H, D_MODEL, N_CTX), dtype=dtype, device="cuda", requires_grad=True)
        v = torch.randn((BATCH, H, N_CTX, D_MODEL), dtype=dtype, device="cuda", requires_grad=True)
        fn = lambda: attention(q, k, v)
        ms = triton.testing.do_bench(fn, percentiles=None, warmup=warmup, rep=rep)
        return ms
    if provider == "flash":
        lengths = torch.full((BATCH,), fill_value=N_CTX, device=device)
        cu_seqlens = torch.zeros((BATCH + 1,), device=device, dtype=torch.int32)
        cu_seqlens[1:] = lengths.cumsum(0)
        qkv = torch.randn((BATCH * N_CTX, 3, H, D_MODEL), dtype=dtype, device=device, requires_grad=True)
        fn = lambda: flash_attn_func(qkv, cu_seqlens, 0., N_CTX, causal=True)
        ms = triton.testing.do_bench(fn, percentiles=None, warmup=warmup, rep=rep)
        return ms


#bench_flash_attention.run(save_path='.', print_data=True)
```

--------------------------------

### Calculate Grid Size with jt.cdiv

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/INDEX.md

Example of using `jt.cdiv` for ceiling division to calculate grid size. This is a common utility for determining the number of blocks needed for kernel execution.

```python
result = jt.cdiv(x.size, block_size)  # Grid size calculation
```

--------------------------------

### Enable Floating-Point Fusion (Default)

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/07-advanced-features.md

Example of enabling floating-point fusion optimizations, which is the default behavior for potentially faster computation. Set to False for strict IEEE 754 compliance or debugging.

```python
# Fast computation (may not be IEEE 754 exact)
result = jt.triton_call(
    x, y,
    kernel=add_kernel,
    out_shape=out_shape,
    ...,
    enable_fp_fusion=True,  # Default
)
```

--------------------------------

### Registering Custom Batching Rule

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/08-transformation-support.md

Provides an example of how to register a custom batching rule for Triton kernel calls, which is necessary when `jax.vmap` is not directly supported. This involves using `jax.experimental.custom_batching`.

```python
from jax.experimental import custom_batching
import jax_triton as jt

@custom_batching.custom_vmap
def batched_triton_kernel(x):
    return jt.triton_call(x, kernel=kernel, ...)

@batched_triton_kernel.def_vmap
def _(batch_x):
    # Vmap implementation
    batch_size = batch_x.shape[0]
    results = jax.lax.map(
        lambda x: jt.triton_call(x, kernel=kernel, ...),
        batch_x,
    )
    return results
```

--------------------------------

### Custom Differentiation Rule Implementation

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/08-transformation-support.md

Provides an example of implementing custom differentiation rules for a Triton kernel using `jax.custom_vjp`. This involves defining forward and backward passes.

```python
import jax

def kernel_and_grad(x, y):
    # Forward pass
    z = jt.triton_call(
        x, y,
        kernel=my_kernel,
        out_shape=x,
        grid=(256,),
    )
    return z

# Implement custom backward pass
@jax.custom_vjp
def kernel_with_grad(x, y):
    return kernel_and_grad(x, y)

def kernel_with_grad_fwd(x, y):
    # Forward: compute output and cache inputs
    z = kernel_and_grad(x, y)
    return z, (x, y)

def kernel_with_grad_bwd(res, g):
    # Backward: compute gradients
    x, y = res
    # Implement backward pass using another Triton kernel or JAX operations
    grad_x = g  # Simplified; actual implementation depends on kernel
    grad_y = g
    return grad_x, grad_y

kernel_with_grad.defvjp(kernel_with_grad_fwd, kernel_with_grad_bwd)
```

--------------------------------

### Gluon Language Control Flow Examples

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/09-gpu-dialect-gluon.md

Illustrates control flow constructs within the Gluon language, specifically conditional statements (if-else) and loops (for). These are essential for implementing more complex kernel logic.

```python
# Conditionals
if condition:
    do_something()
else:
    do_other()

# Loops
for i in range(10):
    result += x[i]
```

--------------------------------

### Gluon Kernel with Heuristics

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/09-gpu-dialect-gluon.md

An example of a Gluon kernel that utilizes Triton's heuristics for dynamic block size determination. This allows the kernel to adapt its performance based on input size.

```python
@triton.heuristics({
    "BLOCK_SIZE": lambda args: 128 if args['n'] < 1024 else 256,
})
@gl.jit
def heuristic_gluon_kernel(x, y, output, BLOCK_SIZE: gl.constexpr):
    idx = gl.program_id(0) * BLOCK_SIZE + gl.arange(BLOCK_SIZE)
    output[idx] = x[idx] + y[idx]
```

--------------------------------

### Example Gluon Kernel Definition

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/09-gpu-dialect-gluon.md

Defines a simple addition kernel using Triton's experimental Gluon dialect. This kernel uses a higher-level, Pythonic syntax for GPU operations.

```python
import triton.experimental.gluon.language as gl

@gl.jit
def gluon_add_kernel(x, y, output):
    idx = gl.program_id(0) * 256 + gl.arange(256)
    output[idx] = x[idx] + y[idx]
```

--------------------------------

### Find Next Power of 2 with jt.next_power_of_2

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/INDEX.md

Example of using `jt.next_power_of_2` to find the smallest power of 2 greater than or equal to a given number. This is often used for block size calculations.

```python
block_size = jt.next_power_of_2(33)  # Returns 64
```

--------------------------------

### Accessing Utilities via Submodule

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/10-package-reference.md

Demonstrates how to import the entire `utils` submodule and access its functions.

```python
import jax_triton.utils
```

```python
jax_triton.utils.cdiv(x, y)
jax_triton.utils.next_power_of_2(x)
jax_triton.utils.strides_from_shape(shape)
```

--------------------------------

### Basic Vector Addition with Triton Kernel

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/01-triton-call.md

Demonstrates a basic vector addition kernel using Triton and how to call it from JAX using `jt.triton_call`. Ensure necessary imports are present.

```python
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, 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 < 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)

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,
        kernel=add_kernel,
        out_shape=out_shape,
        grid=(jt.cdiv(x.size, block_size),),
        block_size=block_size)

x_val = jnp.arange(8, dtype=jnp.float32)
y_val = jnp.arange(8, 16, dtype=jnp.float32)
result = add(x_val, y_val)
jitted_result = jax.jit(add)(x_val, y_val)
```

--------------------------------

### Get Kernel Signature

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/03-jitjitfunction.md

Retrieves the inspect.Signature object for the kernel function. Provides detailed information about the function's signature.

```python
@property
def signature(self) -> inspect.Signature:
    return self.fn.signature
```

--------------------------------

### Get Kernel Argument Names

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/03-jitjitfunction.md

Retrieves the names of kernel parameters in their declaration order. Useful for understanding kernel signature.

```python
import triton
import triton.language as tl

@triton.jit
def my_kernel(x_ptr, y_ptr, output_ptr, block_size: tl.constexpr):
    pass

from jax_triton.triton_lib import JTJITFunction
wrapper = JTJITFunction(my_kernel)
names = wrapper.arg_names  # ['x_ptr', 'y_ptr', 'output_ptr', 'block_size']
```

--------------------------------

### compile_ttir_to_ptx_inplace Compilation Steps

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/05-compilation.md

Details the sequential steps involved in compiling TTIR to PTX for CUDA, including TTIR optimization, TTGIR generation, LLIR generation, and PTX generation.

```text
1. TTIR Optimization: cuda_backend.make_ttir(ttir, metadata, options, cc)
2. TTGIR Generation: cuda_backend.make_ttgir(opt_ttir, ...)
3. LLIR Generation: cuda_backend.make_llir(ttgir, ...)
4. PTX Generation: cuda_backend.make_ptx(llir, ...)
```

--------------------------------

### JIT Compilation with Triton Kernel

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/01-triton-call.md

Demonstrates how to use `jax.jit` to compile a function that includes a Triton kernel call. Kernel compilation occurs during the lowering phase of JAX compilation.

```python
@jax.jit
def jitted_add(x, y):
    return add(x, y)

```

--------------------------------

### Get Triton Type ID

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/04-types.md

A utility function to determine the Triton type ID string for various Python and JAX objects.

```python
def get_type_id(obj: Any) -> str:
    # ... implementation details ...
    pass
```

--------------------------------

### Import from jax_triton.utils

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/02-utils.md

Illustrates importing utility functions directly from the specific 'utils' submodule.

```python
from jax_triton.utils import cdiv, strides_from_shape, next_power_of_2
```

--------------------------------

### Usage Examples for ShapeDtype Protocol

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/04-types.md

Demonstrates using `ShapeDtypeStruct`, JAX arrays, and NumPy arrays to specify output shapes and dtypes for `triton_call`.

```python
import jax
import jax.numpy as jnp
import numpy as np

# Using ShapeDtypeStruct
out_shape = jax.ShapeDtypeStruct(shape=(100, 256), dtype=jnp.float32)

# Using array as shape specification
out_shape = jnp.zeros((100, 256), dtype=jnp.float32)

# Using NumPy array
out_shape = np.zeros((100, 256), dtype=np.float32)

jt.triton_call(..., out_shape=out_shape)
```

--------------------------------

### Optional Dependencies

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/10-package-reference.md

Lists optional dependencies, such as `pytest` for running tests.

```python
pytest  # For running tests
```

--------------------------------

### make_backend

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/05-compilation.md

Creates a Triton backend, GPU target, and determines the compute capability. It's used to initialize the compilation environment for a specific GPU architecture.

```APIDOC
## make_backend

### Description
Initializes the backend and GPU target for compilation, determining the compute capability.

### Method
```python
def make_backend(
  make_gpu_target_func,
  compute_capability: int | None,
  num_ctas: int
) -> tuple[tc.BaseBackend, tc.GPUTarget, int]: 
```

### Parameters
#### Path Parameters
- **make_gpu_target_func** (callable) - Factory to create GPU target (CUDA or HIP)
- **compute_capability** (int or None) - GPU compute capability; auto-detected if None
- **num_ctas** (int) - Cluster thread blocks count

### Returns
Tuple of (backend, gpu_target, compute_capability)

### Raises
- `ValueError`: If num_ctas > 1 and compute_capability < 90
```

--------------------------------

### JAX Wrapper for Triton Attention Kernel

Source: https://github.com/jax-ml/jax-triton/blob/main/examples/JAX_+_Triton_Flash_Attention.ipynb

A JAX-jit-compiled function that wraps a Triton kernel for attention. It requires careful setup of metaparameters and output shapes.

```python
@jax.jit
def jax_attention_triton(q: jnp.ndarray, k: jnp.ndarray, v: jnp.ndarray) -> jnp.ndarray:
  BLOCK = 128
  Lq, Lk = q.shape[-1], k.shape[-2]
  assert Lq == Lk
  grid = lambda _: (triton.cdiv(q.shape[2], BLOCK), q.shape[0] * q.shape[1])
  out_shape = [
      SimpleNamespace(shape=(q.shape[0] * q.shape[1], q.shape[2]), dtype=q.dtype),
      SimpleNamespace(shape=(q.shape[0] * q.shape[1], q.shape[2]), dtype=q.dtype),
      SimpleNamespace(shape=(q.shape[0] * q.shape[1], q.shape[2]), dtype=q.dtype),
      SimpleNamespace(shape=q.shape, dtype=q.dtype)]
  stride_qz, stride_qh, stride_qm, stride_qk = _strides(q.shape)
  stride_kz, stride_kh, stride_kk, stride_kn = _strides(k.shape)
  stride_vz, stride_vh, stride_vk, stride_vn = _strides(v.shape)
  stride_oz, stride_oh, stride_om, stride_on = _strides(out_shape[-1].shape)
  
  metaparams = dict(
    BLOCK_M=BLOCK,
    BLOCK_N=BLOCK,
    BLOCK_DMODEL=64,
    stride_qz=stride_qz, stride_qh=stride_qh, stride_qm=stride_qm, stride_qk=stride_qk,
    stride_kz=stride_kz, stride_kh=stride_kh, stride_kk=stride_kk, stride_kn=stride_kn,
    stride_vz=stride_vz, stride_vh=stride_vh, stride_vk=stride_vk, stride_vn=stride_vn,
    stride_oz=stride_oz, stride_oh=stride_oh, stride_om=stride_om, stride_on=stride_on,
    Z=q.shape[0], H=q.shape[1], N_CTX=q.shape[2],
    num_warps=4, num_stages=1
  )
  _, _, _, output = jt.triton_call(q, k, v, kernel=fused_attention_kernel,
      out_shape=out_shape, grid=grid, **metaparams)
  return output
```

--------------------------------

### Pure Autotuner Usage Pattern

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/06-autotuning.md

Demonstrates a pure autotuner pattern where a kernel is decorated with @triton.autotune and @triton.jit. The kernel is then called using jt.triton_call, allowing the autotuner to select the best configuration based on the provided key.

```python
import triton
import triton.language as tl
import jax_triton as jt

@triton.autotune(
    configs=[
        triton.Config(kwargs={"BLOCK_SIZE": 64}),
        triton.Config(kwargs={"BLOCK_SIZE": 128}),
        triton.Config(kwargs={"BLOCK_SIZE": 256}),
    ],
    key=["n"],
)
@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, n: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    # Implementation
    pass

def matmul(a, b, n):
    result = jt.triton_call(
        a, b,
        kernel=matmul_kernel,
        out_shape=jax.ShapeDtypeStruct((n, n), a.dtype),
        grid=(jt.cdiv(n, 32), jt.cdiv(n, 32)),
        n=n,
    )
    return result
```

--------------------------------

### Get Kernel Argument Name to Index Mapping

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/03-jitjitfunction.md

Obtains a dictionary mapping kernel parameter names to their corresponding indices. Useful for parameter lookup.

```python
mapping = wrapper.arg_name_to_index
# {'x_ptr': 0, 'y_ptr': 1, 'output_ptr': 2, 'block_size': 3}
```

--------------------------------

### Module-Level Import

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/10-package-reference.md

Use this pattern for concise code that clearly defines the namespace, making it evident where functions originate.

```python
import jax_triton as jt

jt.triton_call(...)
jt.cdiv(...)
jt.strides_from_shape(...)
```

--------------------------------

### Get Compiled Kernels Cache Size

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/03-jitjitfunction.md

Returns the number of compiled kernel variants currently stored in the cache. Helps monitor compilation activity and cache usage.

```python
@property
def compiled_kernels_cache_size(self) -> int:
    return len(self.fn._jT_kernel_cache) if hasattr(self.fn, "_jT_kernel_cache") else 0
```

--------------------------------

### JAX to Triton Type Mapping Example

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/03-jitjitfunction.md

Illustrates the mapping from JAX NumPy dtypes to Triton type identifiers. This is crucial for correctly specifying argument types during kernel compilation.

```python
_JAX_TO_TRITON_TYPE_MAP = {
    jnp.dtype("bfloat16"): "bf16",
    jnp.dtype("float64"): "fp64",
    # ... (complete mapping in triton_lib.py:76-94)
}
```

--------------------------------

### Individual Imports from jax_triton

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/02-utils.md

Demonstrates importing specific utility functions directly from the jax_triton library.

```python
from jax_triton import cdiv, strides_from_shape, next_power_of_2

# grid = (cdiv(x.size, block_size),)
```

--------------------------------

### Loop-Based Batching Workaround

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/08-transformation-support.md

Demonstrates a simple workaround for batching Triton kernels by using a Python loop. This approach is straightforward but does not leverage vectorization benefits, though GPU parallelism is still utilized.

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

def process_batch_loop(batch_x):
    results = []
    for x in batch_x:
        result = jt.triton_call(x, kernel=kernel, ...)
        results.append(result)
    return jnp.stack(results)

batch_x = jnp.ones((32, 1024))
result = process_batch_loop(batch_x)
```

--------------------------------

### Runtime Selection of Autotuned Kernel

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/06-autotuning.md

This code demonstrates how to create an autotuned kernel call when multiple kernel calls exist. If there's more than one, it initializes a TritonAutotunedKernelCall; otherwise, it uses the first kernel call.

```python
if len(kernel_calls) > 1:
    kernel_call = triton_kernel_call_lib.TritonAutotunedKernelCall(
        f"{kernel_call_name} ({fn.fn.__name__}) {named_scalar_args}",
        [(call, str(config)) for call, config in zip(kernel_calls, configs)],
        input_output_aliases_with_sizes,
    )
else:
    kernel_call = kernel_calls[0]
```

--------------------------------

### Compilation Flow Diagram

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/05-compilation.md

Illustrates the high-level overview of the Triton compilation pipeline, showing the transformation from Triton JITFunction to platform-specific intermediate representations and finally to GPU assembly.

```text
Triton JITFunction (with constexprs)
    ↓
Type Specialization (based on arg dtypes)
    ↓
AST Source (create IR from AST)
    ↓
TTIR (Triton Intermediate Representation)
    ↓
Platform-Specific Compilation:
    CUDA Path:
        TTIR → TTGIR (GPU-specific IR)
           → LLIR (LLVM IR)
           → PTX (GPU assembly) 
    
    ROCm Path:
        TTIR → TTGIR
           → LLIR
           → AMDGCN (AMD assembly)
           → HSACO (AMD compiled binary) 
    ↓
TritonKernel (cached)
```

--------------------------------

### In-place Operations with Array Donations

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/08-transformation-support.md

Shows how to use `donate_argnames` with `jax.jit` for in-place operations using `triton_call`. Requires the kernel to write to an aliased input buffer and `input_output_aliases` to be specified.

```python
from functools import partial

@partial(jax.jit, donate_argnames="y")
def add_inplace(x, y):
    return jt.triton_call(
        x, y,
        kernel=add_inplace_kernel,
        out_shape=x,
        input_output_aliases={1: 0},
        grid=(jt.cdiv(x.size, 256),),
        block_size=256
    )

x = jnp.arange(1024)
y = jnp.arange(1024, 2048)
result = add_inplace(x, y)
```

--------------------------------

### Basic Usage of triton_call

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/10-package-reference.md

Demonstrates the basic structure for calling a Triton kernel with necessary arguments.

```python
result = triton_call(
    *args,
    kernel=my_triton_kernel,
    out_shape=output_shape,
    grid=(grid_size,),
    **metaparams,
)
```

--------------------------------

### Handling Multiple Outputs from Triton Kernel

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/01-triton-call.md

Demonstrates how to define a Triton kernel that produces multiple outputs and how to specify multiple `out_shape` entries in `jt.triton_call` to receive them.

```python
@triton.jit
def split_kernel(input_ptr, out1_ptr, out2_ptr, block_size: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * block_size + tl.arange(0, block_size)
    data = tl.load(input_ptr + offsets)
    tl.store(out1_ptr + offsets, data)
    tl.store(out2_ptr + offsets, data * 2.0)

def split_data(x: jnp.ndarray):
    out_shapes = [
        jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype),
        jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype),
    ]
    return jt.triton_call(
        x,
        kernel=split_kernel,
        out_shape=out_shapes,
        grid=(jt.cdiv(x.size, 8),),
        block_size=8)

x = jnp.arange(16, dtype=jnp.float32)
out1, out2 = split_data(x)
```

--------------------------------

### Clone jax-triton repository

Source: https://github.com/jax-ml/jax-triton/blob/main/README.md

Clones the jax-triton repository from GitHub to set up the development environment.

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

--------------------------------

### Gluon Language Basic Operations

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/09-gpu-dialect-gluon.md

Demonstrates fundamental operations available in the Gluon language, including program ID retrieval, array indexing for loading and storing, reduction operations, and mathematical functions. Ensure Triton's experimental Gluon language is imported.

```python
import triton.experimental.gluon.language as gl

# Program ID and indexing
pid = gl.program_id(0)  # Get program ID in dimension 0
idx = pid * block_size + gl.arange(block_size)

# Array operations
x = array[idx]  # Load
result = x + y  # Arithmetic
result[idx] = value  # Store (via indexing)

# Reductions
sum_val = gl.sum(array)
max_val = gl.max(array)
min_val = gl.min(array)

# Math functions
exp_val = gl.exp(x)
sqrt_val = gl.sqrt(x)
```

--------------------------------

### Autotuner with User Override

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/06-autotuning.md

This pattern shows how to use @triton.autotune for automatic configuration selection while also allowing specific configurations to be overridden manually during the jt.triton_call. This provides flexibility for both automatic optimization and explicit control.

```python
@triton.autotune(
    configs=[
        triton.Config(kwargs={"BLOCK_SIZE": 64}),
        triton.Config(kwargs={"BLOCK_SIZE": 128}),
        triton.Config(kwargs={"BLOCK_SIZE": 256}),
    ],
    key=["n"],
)
@triton.jit
def kernel(input_ptr, output_ptr, n: tl.constexpr, BLOCK_SIZE: tl.constexpr):
    # Implementation
    pass

def process_generic(x):
    # Let autotuner choose
    return jt.triton_call(x, kernel=kernel, ...)

def process_optimized(x):
    # Force specific config
    return jt.triton_call(
        x,
        kernel=kernel,
        ...,
        BLOCK_SIZE=256,  # Override autotuner
    )
```

--------------------------------

### Gluon Language Constexpr Parameters

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/09-gpu-dialect-gluon.md

Shows how to define compile-time constant parameters in Gluon kernels using `gl.constexpr`. These parameters are fixed during compilation and can optimize kernel performance.

```python
@gl.jit
def kernel(x, y, output, block_size: gl.constexpr, use_fast: gl.constexpr):
    # These are compile-time constants
    pass
```

--------------------------------

### Backend Specialization Attribute Creation

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/05-compilation.md

Demonstrates how the backend generates specialization attributes based on argument data types and calculated alignments. This process is part of type and attribute specialization.

```python
specialization = [
  specialize_impl(
    backend,
    types.SimpleNamespace(data_ptr=lambda: alignment, ...),
    is_const=False,
    do_specialize=True,
    alignment > 0,
  )
  for arg_dtype, alignment in zip(arg_dtypes, alignments)
]
```

--------------------------------

### JAX-Triton Heuristics Application

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/06-autotuning.md

Applies heuristic functions to Triton configurations to dynamically compute final metaparameters. Use this to adapt kernel behavior based on runtime arguments.

```python
def apply_heuristics(
  fn: autotuner.Heuristics,
  configs: list[triton.Config],
  orig_kwargs: dict[str, Any],
  named_args: dict[str, Any],
) -> list[triton.Config]:
  # ... implementation details ...
  pass
```

--------------------------------

### Triggering Compilation via JAX Lowering

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/05-compilation.md

Shows how the compilation pipeline is initiated during the JAX lowering process. This function is responsible for generating the GPU target function.

```python
triton_kernel_call_lowering(
    make_gpu_target_func,
    ctx,
    ...,
)
```

--------------------------------

### Utilities Submodule Import

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/10-package-reference.md

Use this pattern when you want to separate utility functions from the main API, promoting a cleaner structure.

```python
from jax_triton import utils

utils.cdiv(...)
utils.next_power_of_2(...)
```

--------------------------------

### Differentiate Triton Kernel with JAX Grad

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/08-transformation-support.md

Demonstrates how to compute gradients for a Triton kernel using JAX's `grad` function. Ensure the kernel is defined to be differentiable.

```python
import jax
import jax.numpy as jnp

# Assume kernel_with_grad is a Triton kernel defined elsewhere
def compute_loss(x, y):
    z = kernel_with_grad(x, y)
    return jnp.sum(z)

grad_fn = jax.grad(compute_loss)
grad_x, grad_y = grad_fn(x, y)
```

--------------------------------

### Memory Layout Handling with Strides

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/02-utils.md

Demonstrates using strides_from_shape to handle memory layout for 2D operations, passing stride information to a Triton kernel.

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

# For 2D operations with non-trivial strides
data = jnp.ones((1000, 512), dtype=jnp.float32)
strides = jt.strides_from_shape(data.shape)

# Pass strides to kernel for efficient memory access
# result = jt.triton_call(
#     data,
#     kernel=process_matrix,
#     out_shape=data,
#     grid=(data.shape[0],),
#     input_stride=strides[0],
#     n_cols=data.shape[1]
# )
```

--------------------------------

### Define Constants for Flash Attention

Source: https://github.com/jax-ml/jax-triton/blob/main/examples/JAX_+_Triton_Flash_Attention.ipynb

Sets up constants for batch size, number of heads, context length, and head dimension, which are commonly used parameters for configuring and testing Flash Attention models.

```python
BATCH, N_HEADS, N_CTX, D_HEAD = 4, 64, 2048, 64
```

--------------------------------

### AST Source Selection for Compilation

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/09-gpu-dialect-gluon.md

Shows how JAX-Triton selects the appropriate AST source (Gluon or standard Triton) during compilation. This ensures the correct compilation path is taken based on the kernel type.

```python
real_ASTSource = gl_runtime.GluonASTSource if self.is_gluon else tc.ASTSource
module = real_ASTSource(
    fn, constexprs=constants, signature=signature, attrs=attrs
).make_ir(gpu_target, options, codegen_fns, backend.get_module_map(), context)
```

--------------------------------

### Core API and Utilities

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/COMPLETION_REPORT.txt

Documentation for the core API, including the `triton_call` function and associated utilities, enabling users to leverage Triton kernels within JAX.

```APIDOC
## Core API: triton_call

### Description
This is the primary function for invoking Triton kernels from JAX. It handles the compilation and execution of Triton kernels, integrating seamlessly with JAX's transformation system.

### Method
Not applicable (JAX function call)

### Endpoint
Not applicable (JAX function call)

### Parameters
(Details on parameters for `triton_call` would be found in `01-triton-call.md`)

### Request Example
(Example usage of `triton_call` would be found in `01-triton-call.md`)

### Response
(Details on return types and values would be found in `01-triton-call.md`)

## Utilities

### Description
This section covers various utility functions that support the core API, such as helper functions for kernel definition, data manipulation, and debugging.

### Method
Not applicable (JAX function calls)

### Endpoint
Not applicable (JAX function calls)

### Parameters
(Details on utility function parameters would be found in `02-utils.md`)

### Request Example
(Example usage of utility functions would be found in `02-utils.md`)

### Response
(Details on return types and values of utility functions would be found in `02-utils.md`)
```

--------------------------------

### In-place Operation with JAX JIT and Argument Donation

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/01-triton-call.md

Shows how to optimize memory usage for in-place operations within a `jax.jit` compiled function by using `donate_argnames`. This allows JAX to reuse the memory of the donated argument.

```python
@jax.jit(donate_argnames="y")
def add_inplace(x, y):
    return jt.triton_call(..., input_output_aliases={1: 0}, ...)

```

--------------------------------

### JAX-Triton Lowering: Autotuner and Heuristics Flow

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/06-autotuning.md

Shows the configuration selection logic within jax-triton's lowering process. It handles Autotuner, Heuristics, and explicit configurations to determine which kernels to compile.

```python
if isinstance(fn, autotuner.Autotuner):
    configs = make_autotuner_configs(fn, metaparams, named_args)
    fn = fn.fn
else:
    config = triton.Config(
      {},
      num_warps=num_warps,
      num_stages=num_stages,
      num_ctas=num_ctas,
    )
    configs = [config]

if isinstance(fn, autotuner.Heuristics):
    configs = apply_heuristics(fn, configs, metaparams, named_args)
    fn = fn.fn
```

--------------------------------

### Setting Environment Variable for Debugging

Source: https://github.com/jax-ml/jax-triton/blob/main/_autodocs/10-package-reference.md

Shows how to set the `JAX_TRITON_DUMP_DIR` environment variable to save intermediate IR files for debugging.

```bash
export JAX_TRITON_DUMP_DIR=/tmp/jax_triton_dumps
python my_script.py
```