evox.triton_kernels.backend

Module Contents

Functions

_check_triton	Lazily check if Triton is importable. Caches the result.
has_triton	Check whether the Triton package is importable.
triton_device_types	Return the set of device type strings for which Triton kernels are registered.
register_triton_device_type	Register an additional device type for Triton kernel support.
triton_supports_device	Check whether Triton can execute kernels on the given device.

Data

_triton_checked
_triton_available
_triton_device_types

API

evox.triton_kernels.backend._triton_checked

False

evox.triton_kernels.backend._triton_available

False

evox.triton_kernels.backend._triton_device_types: set[str]

None

evox.triton_kernels.backend._check_triton()

Lazily check if Triton is importable. Caches the result.

evox.triton_kernels.backend.has_triton()

Check whether the Triton package is importable.

Returns:

True if Triton can be imported, False otherwise.

evox.triton_kernels.backend.triton_device_types() → frozenset[str]

Return the set of device type strings for which Triton kernels are registered.

Defaults to {"cuda"}; may be extended via

Func:

register_triton_device_type.

Returns:

A frozenset of registered device type strings.

evox.triton_kernels.backend.register_triton_device_type(device_type: str) → None

Register an additional device type for Triton kernel support.

This allows extending Triton coverage beyond the default "cuda" backend, e.g. to support Ascend NPU ("npu"). The input is normalized by stripping surrounding whitespace and lowercasing. Empty strings are ignored.

Parameters:

device_type – The device type string to register (e.g. "npu").

evox.triton_kernels.backend.triton_supports_device(device: torch.device | str) → bool

Check whether Triton can execute kernels on the given device.

Triton kernels run on CUDA devices (including AMD ROCm, which is aliased under the "cuda" device type) and, when registered, additional backends such as Ascend NPU ("npu"). This function returns True only when Triton is importable and the device type is among the registered Triton device types. For "cuda" devices, CUDA must additionally be available on the host; for other backends membership in the registered set is sufficient.

Parameters:

device – A torch.device or device string (e.g. “cuda”, “cpu”, “npu”).

Returns:

True if a Triton kernel can run on this device.