tilelang.utils.device

Attributes

IS_CUDA
IS_NPU
IS_NPU
IS_MPS
IS_MPS
Event

Functions

get_current_device()
get_max_dynamic_shared_memory([device])	Largest dynamic shared-memory allocation a thread block may request, in bytes.
device_synchronize([device])	Synchronize CUDA/HIP, MPS, or NPU, 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_NPU = False

tilelang.utils.device.IS_NPU

tilelang.utils.device.IS_MPS = False

tilelang.utils.device.IS_MPS

tilelang.utils.device.Event

tilelang.utils.device.get_current_device()

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

Largest dynamic shared-memory allocation a thread block may request, in bytes.

Kernels that request more than this fail at launch with Failed to set the allowed dynamic shared memory size to N.

Parameters:

device (int | torch.device | None)

Return type:

int

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

Synchronize CUDA/HIP, MPS, or NPU, 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