| |
| |
| |
| |
| |
| |
|
|
| import functools |
| import importlib.metadata |
| import inspect |
| import logging |
| import os |
| import shutil |
| import sys |
| from functools import cache, lru_cache |
| from importlib.util import find_spec |
| from pathlib import Path |
|
|
| import torch |
| import triton |
| from packaging import version as package_version |
|
|
| logger = logging.getLogger(__name__) |
|
|
| FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1" |
| FLA_CACHE_RESULTS = os.getenv('FLA_CACHE_RESULTS', '1') == '1' |
|
|
| FLA_DISABLE_TENSOR_CACHE = os.getenv('FLA_DISABLE_TENSOR_CACHE', '0') == '1' |
| try: |
| FLA_TENSOR_CACHE_SIZE = int(os.getenv('FLA_TENSOR_CACHE_SIZE', "4")) |
| except ValueError: |
| FLA_TENSOR_CACHE_SIZE = 4 |
|
|
|
|
| @lru_cache(maxsize=1) |
| def check_environments(): |
| """ |
| Checks the current operating system, Triton version, and Python version, |
| issuing warnings if they don't meet recommendations. |
| This function's body only runs once due to lru_cache. |
| """ |
| |
| if sys.platform == 'win32': |
| |
| try: |
| from importlib.metadata import PackageNotFoundError, metadata |
| metadata('triton-windows') |
| |
| except PackageNotFoundError: |
| logger.warning( |
| "Detected Windows operating system. Consider installing triton-windows " |
| "(https://github.com/triton-lang/triton-windows) for better compatibility. " |
| "Without it, some features may not work correctly.", |
| ) |
|
|
| triton_version = package_version.parse(triton.__version__) |
| required_triton_version = package_version.parse("3.3.0") |
|
|
| if triton_version < required_triton_version: |
| logger.warning( |
| f"Current Triton version {triton_version} is below the recommended 3.3.0 version. " |
| "Errors may occur and these issues will not be fixed. " |
| "Please consider upgrading Triton.", |
| ) |
|
|
| |
| py_version = package_version.parse(f"{sys.version_info.major}.{sys.version_info.minor}") |
| required_py_version = package_version.parse("3.11") |
|
|
| if py_version < required_py_version: |
| logger.warning( |
| f"Current Python version {py_version} is below the recommended 3.11 version. " |
| "It is recommended to upgrade to Python 3.11 or higher for the best experience.", |
| ) |
|
|
| return None |
|
|
|
|
| check_environments() |
|
|
|
|
| @cache |
| def check_pytorch_version(version_s: str = '2.4') -> bool: |
| return package_version.parse(torch.__version__) >= package_version.parse(version_s) |
|
|
|
|
| @cache |
| def check_transformers_version(version_s: str) -> bool: |
| try: |
| installed_version = importlib.metadata.version('transformers') |
| except importlib.metadata.PackageNotFoundError: |
| return False |
| return package_version.parse(installed_version) >= package_version.parse(version_s) |
|
|
|
|
| @cache |
| def supports_transformers_next_sequence_length() -> bool: |
| if find_spec_cached('transformers') is None: |
| return False |
| from transformers.generation import GenerationMixin |
|
|
| return 'next_sequence_length' in inspect.signature(GenerationMixin.prepare_inputs_for_generation).parameters |
|
|
|
|
| TRITON_ABOVE_3_4_0 = package_version.parse(triton.__version__) >= package_version.parse("3.4.0") |
| TRITON_ABOVE_3_5_1 = package_version.parse(triton.__version__) >= package_version.parse("3.5.1") |
| TRITON_ABOVE_3_6_0 = package_version.parse(triton.__version__) >= package_version.parse("3.6.0") |
| TRITON_ABOVE_3_7_1 = package_version.parse(triton.__version__) >= package_version.parse("3.7.1") |
| TRITON_ABOVE_3_8_0 = package_version.parse(triton.__version__) >= package_version.parse("3.8.0") |
|
|
| SUPPORTS_AUTOTUNE_CACHE = "cache_results" in inspect.signature(triton.autotune).parameters |
| autotune_cache_kwargs = {"cache_results": FLA_CACHE_RESULTS} if SUPPORTS_AUTOTUNE_CACHE else {} |
|
|
|
|
| @functools.cache |
| def find_spec_cached(name): |
| try: |
| return find_spec(name) |
| except ModuleNotFoundError as error: |
| if error.name is not None and (name == error.name or name.startswith(f'{error.name}.')): |
| return None |
| raise |
|
|
|
|
| @functools.cache |
| def has_usable_nvcc() -> bool: |
| """Whether a usable nvcc compiler is available for TileLang's JIT. |
| |
| Mirrors the guesses in ``tilelang.env._find_cuda_home`` (env |
| CUDA_HOME/CUDA_PATH, nvcc on PATH, the ``nvidia-cuda-nvcc`` wheel, |
| /usr/local/cuda), but verifies the nvcc binary actually exists — |
| only ``nvidia-cuda-nvcc`` >= 13.0 ships it, the ``-cu12`` variant |
| installs just ptxas. |
| """ |
| cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") |
| if cuda_home is not None and (Path(cuda_home) / "bin" / "nvcc").exists(): |
| return True |
| if shutil.which("nvcc") is not None: |
| return True |
| try: |
| files = importlib.metadata.files("nvidia-cuda-nvcc") or [] |
| except importlib.metadata.PackageNotFoundError: |
| files = [] |
| if any(f.name in ("nvcc", "nvcc.exe") for f in files): |
| return True |
| if (Path("/usr/local/cuda") / "bin" / "nvcc").exists(): |
| return True |
|
|
| logger.info( |
| "[FLA Backend] TileLang is installed but no usable nvcc compiler was found; falling back to Triton. " |
| "Install a CUDA toolkit or nvidia-cuda-nvcc, or set FLA_TILELANG=0 to disable TileLang explicitly." |
| ) |
| return False |
|
|
