torch/utils/_triton.py

import functools

from torch._dynamo.device_interface import get_interface_for_device


@functools.lru_cache(None)
def has_triton_package() -> bool:
    try:
        import triton

        return triton is not None
    except ImportError:
        return False


@functools.lru_cache(None)
def has_triton() -> bool:
    def is_cuda_compatible_with_triton():
        device_interface = get_interface_for_device("cuda")
        return (
            device_interface.is_available()
            and device_interface.Worker.get_device_properties().major >= 7
        )

    return is_cuda_compatible_with_triton() and has_triton_package()