680742bc33
Signed-off-by: mxuax <mxuax@connect.ust.hk> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
60 lines
1.5 KiB
Python
60 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
# Copyright (c) 2024, Jiarui Fang.
|
|
# Adapted from https://github.com/feifeibear/long-context-attention
|
|
|
|
|
|
# test if flash_attn is available
|
|
try:
|
|
import flash_attn # noqa: F401
|
|
from flash_attn.flash_attn_interface import _flash_attn_backward, _flash_attn_forward # noqa: F401
|
|
|
|
HAS_FLASH_ATTN = True
|
|
except ImportError:
|
|
HAS_FLASH_ATTN = False
|
|
|
|
try:
|
|
from flash_attn_interface import _flash_attn_backward as flash_attn_func_hopper_backward # noqa: F401
|
|
from flash_attn_interface import _flash_attn_forward as flash_attn_forward_hopper # noqa: F401
|
|
from flash_attn_interface import flash_attn_func as flash3_attn_func # noqa: F401
|
|
|
|
HAS_FLASH_ATTN_HOPPER = True
|
|
except ImportError:
|
|
HAS_FLASH_ATTN_HOPPER = False
|
|
|
|
try:
|
|
from flashinfer.prefill import single_prefill_with_kv_cache # noqa: F401
|
|
|
|
HAS_FLASHINFER = True
|
|
except ImportError:
|
|
HAS_FLASHINFER = False
|
|
|
|
try:
|
|
import aiter # noqa: F401
|
|
from aiter import flash_attn_func as flash_attn_func_aiter # noqa: F401
|
|
|
|
HAS_AITER = True
|
|
except ImportError:
|
|
HAS_AITER = False
|
|
|
|
try:
|
|
import sageattention # noqa: F401
|
|
|
|
HAS_SAGE_ATTENTION = True
|
|
except ImportError:
|
|
HAS_SAGE_ATTENTION = False
|
|
|
|
try:
|
|
import spas_sage_attn # noqa: F401
|
|
|
|
HAS_SPARSE_SAGE_ATTENTION = True
|
|
except ImportError:
|
|
HAS_SPARSE_SAGE_ATTENTION = False
|
|
|
|
try:
|
|
import torch_npu # noqa: F401
|
|
|
|
HAS_NPU = True
|
|
except ImportError:
|
|
HAS_NPU = False
|