/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
extensions/pybind/flash_attention/__init__.py
14 строк
470 B
傅剑寒
[Inference/Refactor] Refactor compilation mechanism and unified multi hw (#5613)
24 апр 2024, 09:17
Не верифицирован
24 апр 2024, 09:17
279300d
Код
Авторство
О чём код?
from .flash_attention_dao_cuda import FlashAttentionDaoCudaExtension from .flash_attention_npu import FlashAttentionNpuExtension from .flash_attention_sdpa_cuda import FlashAttentionSdpaCudaExtension try: # TODO: remove this after updating openmoe example import flash_attention # noqa HAS_FLASH_ATTN = True except: HAS_FLASH_ATTN = False __all__ = ["FlashAttentionDaoCudaExtension", "FlashAttentionSdpaCudaExtension", "FlashAttentionNpuExtension"]