38 lines
980 B
Python
38 lines
980 B
Python
import platform
|
|
from typing import Literal
|
|
|
|
# flash-attn is a CUDA-only optional dependency. Keep the upstream CUDA
|
|
# default, while making SDPA the source-native dense-attention path on macOS.
|
|
BACKEND = 'sdpa' if platform.system() == 'Darwin' else 'flash_attn'
|
|
DEBUG = False
|
|
|
|
def __from_env():
|
|
import os
|
|
|
|
global BACKEND
|
|
global DEBUG
|
|
|
|
env_attn_backend = os.environ.get('ATTN_BACKEND')
|
|
env_attn_debug = os.environ.get('ATTN_DEBUG')
|
|
|
|
if env_attn_backend is not None and env_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3', 'sdpa', 'naive']:
|
|
BACKEND = env_attn_backend
|
|
if env_attn_debug is not None:
|
|
DEBUG = env_attn_debug == '1'
|
|
|
|
print(f"[ATTENTION] Using backend: {BACKEND}")
|
|
|
|
|
|
__from_env()
|
|
|
|
|
|
def set_backend(
|
|
backend: Literal['xformers', 'flash_attn', 'flash_attn_3', 'sdpa', 'naive'],
|
|
):
|
|
global BACKEND
|
|
BACKEND = backend
|
|
|
|
def set_debug(debug: bool):
|
|
global DEBUG
|
|
DEBUG = debug
|