evox.triton_kernels.kernels.fused_add

Example fused element-wise addition kernel.

This module demonstrates the full Triton kernel integration pattern:

1. Define the @triton.jit kernel function.

2. Define a launcher function that sets up the grid and calls the kernel.

3. Define a PyTorch fallback that runs on all devices (CPU, MPS, etc.).

4. Define a fake (abstract evaluation) function for torch.compile tracing.

5. Register everything via :func:~evox.triton_kernels.register_triton_op.

At call time, PyTorch’s dispatcher routes CUDA tensors to the Triton kernel and all other tensors to the PyTorch fallback automatically.

To extend coverage to additional backends such as Ascend NPU, register the device type with :func:~evox.triton_kernels.register_triton_device_type (e.g. register_triton_device_type("npu")) and pass an explicit triton_device_types=["cuda", "npu"] to

func:

~evox.triton_kernels.register_triton_op.

Module Contents

Functions

_triton_fused_add	Launch the Triton addition kernel.
_fused_add_fake	Fake (abstract evaluation) function for torch.compile tracing.
fused_add	Element-wise addition of two tensors with the same shape.

API

evox.triton_kernels.kernels.fused_add._triton_fused_add(x: torch.Tensor, y: torch.Tensor) → torch.Tensor

Launch the Triton addition kernel.

evox.triton_kernels.kernels.fused_add._fused_add_fake(x: torch.Tensor, y: torch.Tensor) → torch.Tensor

Fake (abstract evaluation) function for torch.compile tracing.

evox.triton_kernels.kernels.fused_add.fused_add(x: torch.Tensor, y: torch.Tensor) → torch.Tensor

Element-wise addition of two tensors with the same shape.

Uses a fused Triton kernel on CUDA devices and falls back to PyTorch’s native + operator on all other devices (CPU, MPS, etc.).

Parameters:

• x – First input tensor.

• y – Second input tensor, must have the same shape as x.

Returns:

Element-wise sum x + y.