triton.Config

class triton.Config(kwargs, num_warps=4, num_stages=3, num_ctas=1, maxnreg=None, pre_hook=None, ir_override=None)

An object that represents a possible kernel configuration for the auto-tuner to try.

Variables:

• kwargs – a dictionary of meta-parameters to pass to the kernel as keyword arguments.

• num_warps – the number of warps to use for the kernel when compiled for GPUs. For example, if num_warps=8, then each kernel instance will be automatically parallelized to cooperatively execute using 8 * 32 = 256 threads.

• num_stages – the number of stages that the compiler should use when software-pipelining loops. Mostly useful for matrix multiplication workloads on SM80+ GPUs.

• num_ctas – number of blocks in a block cluster. SM90+ only.

• maxnreg – maximum number of registers one thread can use. Corresponds to ptx .maxnreg directive. Not supported on all platforms.

• pre_hook – a function that will be called before the kernel is called. Parameters of this function are args.

• ir_override – filename of a user-defined IR (*.{ttgir|llir|ptx|amdgcn}).

Example

import triton
import triton.language as tl
import torch


@triton.autotune(
    configs=[triton.Config(kwargs={'BLOCK_SIZE': 64}),
             triton.Config(kwargs={'BLOCK_SIZE': 128})],
    key=['n_elements'],
)
@triton.jit
def vec_add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)


def test_config():
    N = 128
    x = torch.randn(N, dtype=torch.float32, device='npu')
    y = torch.randn(N, dtype=torch.float32, device='npu')
    out = torch.empty(N, dtype=torch.float32, device='npu')
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']), )
    vec_add_kernel[grid](x, y, out, N)
    torch.testing.assert_close(out.cpu(), (x + y).cpu())


if __name__ == "__main__":
    test_config()

Special Restrictions

• num_stages: Not applicable on Ascend

• num_ctas: Not applicable on Ascend

• maxnreg: No effect on Ascend

__init__(kwargs, num_warps=4, num_stages=3, num_ctas=1, maxnreg=None, pre_hook=None, ir_override=None)

Methods

__init__(name)	Initialize self.
all_kwargs