tilelang.utils.device

Attributes

IS_CUDA
IS_MPS
IS_MPS
Event

Functions

get_current_device()
device_synchronize([device])	Synchronize CUDA/HIP or MPS, preferring CUDA when no device is given.
get_available_cpu_count()	CPU cores available to this process (cpuset affinity), at least 1.

Module Contents

tilelang.utils.device.IS_CUDA

tilelang.utils.device.IS_MPS = False

tilelang.utils.device.IS_MPS

tilelang.utils.device.Event

tilelang.utils.device.get_current_device()

tilelang.utils.device.device_synchronize(device=None)

Synchronize CUDA/HIP or MPS, preferring CUDA when no device is given.

Parameters:

device (int | torch.device | None)

Return type:

None

tilelang.utils.device.get_available_cpu_count()

CPU cores available to this process (cpuset affinity), at least 1.

Falls back to os.cpu_count where affinity is unavailable (macOS/Windows).

Return type:

int