Kernels:
Trusted publisher
File size: 8,030 Bytes
e19323e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | # Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import contextlib
import functools
import logging
import os
import platform
import sys
import warnings
from enum import Enum
from functools import cache, lru_cache
import torch
import triton
from packaging import version as package_version
logger = logging.getLogger(__name__)
@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.
"""
# Check Operating System
if sys.platform == 'win32':
# Check if triton-windows is installed
try:
from importlib.metadata import PackageNotFoundError, metadata
metadata('triton-windows')
# triton-windows is installed, no warning needed
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.",
)
# Check Python version
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()
def _cpu_device_warning():
warnings.warn(('Triton is not supported on current platform, roll back to CPU.'), stacklevel=2)
@cache
def check_pytorch_version(version_s: str = '2.4') -> bool:
return package_version.parse(torch.__version__) >= package_version.parse(version_s)
@cache
def get_multiprocessor_count(tensor_idx: int = 0) -> int:
try:
return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['multiprocessor_count']
except Exception:
# Maybe we use a NPU device.
try:
if triton.runtime.driver.active.get_current_target().backend == 'npu':
return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['num_vectorcore']
except Exception:
logger.debug('Failed to get NPU multiprocessor count, falling back to 1.', exc_info=True)
return 1
@cache
def get_available_device() -> str:
try:
return triton.runtime.driver.active.get_current_target().backend
except Exception:
_cpu_device_warning()
return 'cpu'
def map_triton_backend_to_torch_device() -> str:
backend = get_available_device() # 'cuda' | 'hip' | 'xpu' | 'cpu' | ...
return {'cuda': 'cuda', 'hip': 'cuda', 'xpu': 'xpu'}.get(backend, backend)
# For AMD GPUs, the triton backend is 'hip', while for Nvidia GPUs, the triton backend is 'cuda'.
# However, the torch backend is 'cuda' for both Nvidia and AMD GPUs.
# Therefore, we need to check the triton backend to determine the actual GPU vendor.
device = get_available_device() if get_available_device() != 'hip' else 'cuda'
device_torch_lib = getattr(torch, device)
device_platform = get_available_device()
device_name = map_triton_backend_to_torch_device()
IS_AMD = (device_platform == 'hip')
IS_ARM = platform.machine().lower() in ('aarch64', 'arm64')
IS_INTEL = (device_platform == 'xpu')
IS_INTEL_ALCHEMIST = (IS_INTEL and 'Intel(R) Arc(TM) A' in torch.xpu.get_device_name(0))
IS_NVIDIA = (device_platform == 'cuda')
IS_NPU = (device_platform == 'npu')
IS_NVIDIA_BLACKWELL = (IS_NVIDIA and torch.cuda.get_device_capability()[0] in (10, 12))
IS_NVIDIA_HOPPER = (IS_NVIDIA and ('NVIDIA H' in torch.cuda.get_device_name(0) or torch.cuda.get_device_capability()[0] == 9))
# Nvidia Ampere or newer, haven't check AMD and intel yet.
IS_TF32_SUPPORTED = (IS_NVIDIA and torch.cuda.get_device_capability(0)[0] >= 8)
IS_GATHER_SUPPORTED = hasattr(triton.language, 'gather')
IS_TMA_SUPPORTED = (
IS_NVIDIA
and torch.cuda.get_device_capability(0)[0] >= 9
and os.environ.get('FLA_USE_TMA', '0') == '1'
and (hasattr(triton.language, '_experimental_make_tensor_descriptor') or hasattr(triton.language, 'make_tensor_descriptor'))
)
if IS_NVIDIA and not IS_TF32_SUPPORTED:
# Make old card happy, since triton will use tf32 by default.
# This is a workaround for old nvidia card.
os.environ['TRITON_F32_DEFAULT'] = 'ieee'
def _default_alloc_fn(size: int, alignment: int, stream: int | None):
return torch.empty(size, device=torch.device(device_name, device_torch_lib.current_device()), dtype=torch.int8)
if IS_TMA_SUPPORTED:
logger.info('TMA is supported, using TMA by default.')
triton.set_allocator(_default_alloc_fn)
elif IS_NVIDIA_BLACKWELL:
# Blackwell (SM100 datacenter / SM120 consumer): Triton compiler may emit global_scratch for
# autotuned kernels even without TMA. Register a default allocator to
# prevent NullAllocator crashes. See triton-lang/triton#10002.
logger.info('Blackwell detected: registering default global_scratch allocator.')
triton.set_allocator(_default_alloc_fn)
def get_all_max_shared_mem():
try:
return [
triton.runtime.driver.active.utils.get_device_properties(i)['max_shared_mem']
for i in range(device_torch_lib.device_count())
]
except Exception:
_cpu_device_warning()
return [-1]
class Backend(Enum):
ADA = 101376 # RTX 4090
AMPERE = 166912 # A100
HOPPER = 232448 # H100
DEFAULT = 102400 # Default
@classmethod
def get_shared_memory(cls, arch: str) -> int:
try:
return cls[arch.upper()].value
except KeyError:
return cls.DEFAULT.value
@cache
def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool:
try:
device_shared_mem_list = get_all_max_shared_mem()
max_shared_memory = device_shared_mem_list[tensor_idx]
return max_shared_memory >= Backend.get_shared_memory(arch)
except Exception:
return False
if check_pytorch_version('2.4'):
if device == 'cpu':
device = 'cuda'
device_torch_lib = getattr(torch, device)
autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device)
autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device)
def custom_device_ctx(index: int):
if index is None:
return contextlib.nullcontext()
try:
return device_torch_lib.device(index)
except (AttributeError, AssertionError, RuntimeError):
return contextlib.nullcontext()
else:
assert device == 'cuda', 'Only cuda device is supported for PyTorch version < 2.4.0.'
autocast_custom_fwd = device_torch_lib.amp.custom_fwd
autocast_custom_bwd = device_torch_lib.amp.custom_bwd
def custom_device_ctx(index: int):
if index is None:
return contextlib.nullcontext()
try:
return torch.cuda.device(index)
except (AttributeError, AssertionError, RuntimeError):
return contextlib.nullcontext()
|