/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
extensions/__init__.py
39 строк
1 KB
傅剑寒
[Inference/Refactor] Refactor compilation mechanism and unified multi hw (#5613)
24 апр 2024, 09:17
Не верифицирован
24 апр 2024, 09:17
279300d
Код
Авторство
О чём код?
from .pybind.cpu_adam import CpuAdamArmExtension, CpuAdamX86Extension from .pybind.flash_attention import ( FlashAttentionDaoCudaExtension, FlashAttentionNpuExtension, FlashAttentionSdpaCudaExtension, ) from .pybind.inference import InferenceOpsCudaExtension from .pybind.layernorm import LayerNormCudaExtension from .pybind.moe import MoeCudaExtension from .pybind.optimizer import FusedOptimizerCudaExtension from .pybind.softmax import ScaledMaskedSoftmaxCudaExtension, ScaledUpperTriangleMaskedSoftmaxCudaExtension ALL_EXTENSIONS = [ CpuAdamArmExtension, CpuAdamX86Extension, LayerNormCudaExtension, MoeCudaExtension, FusedOptimizerCudaExtension, InferenceOpsCudaExtension, ScaledMaskedSoftmaxCudaExtension, ScaledUpperTriangleMaskedSoftmaxCudaExtension, FlashAttentionDaoCudaExtension, FlashAttentionSdpaCudaExtension, FlashAttentionNpuExtension, ] __all__ = [ "CpuAdamArmExtension", "CpuAdamX86Extension", "LayerNormCudaExtension", "MoeCudaExtension", "FusedOptimizerCudaExtension", "InferenceOpsCudaExtension", "ScaledMaskedSoftmaxCudaExtension", "ScaledUpperTriangleMaskedSoftmaxCudaExtension", "FlashAttentionDaoCudaExtension", "FlashAttentionSdpaCudaExtension", "FlashAttentionNpuExtension", ]