colossalai

1
from .cpu_adam import CpuAdamArmExtension, CpuAdamX86Extension
2
from .flash_attention import (
3
    FlashAttentionDaoCudaExtension,
4
    FlashAttentionNpuExtension,
5
    FlashAttentionXformersCudaExtension,
6
)
7
from .layernorm import LayerNormCudaExtension
8
from .moe import MoeCudaExtension
9
from .optimizer import FusedOptimizerCudaExtension
10
from .softmax import ScaledMaskedSoftmaxCudaExtension, ScaledUpperTriangleMaskedSoftmaxCudaExtension
11

12
ALL_EXTENSIONS = [
13
    CpuAdamArmExtension,
14
    CpuAdamX86Extension,
15
    LayerNormCudaExtension,
16
    MoeCudaExtension,
17
    FusedOptimizerCudaExtension,
18
    ScaledMaskedSoftmaxCudaExtension,
19
    ScaledUpperTriangleMaskedSoftmaxCudaExtension,
20
    FlashAttentionDaoCudaExtension,
21
    FlashAttentionXformersCudaExtension,
22
    FlashAttentionNpuExtension,
23
]
24

25
__all__ = [
26
    "CpuAdamArmExtension",
27
    "CpuAdamX86Extension",
28
    "LayerNormCudaExtension",
29
    "MoeCudaExtension",
30
    "FusedOptimizerCudaExtension",
31
    "ScaledMaskedSoftmaxCudaExtension",
32
    "ScaledUpperTriangleMaskedSoftmaxCudaExtension",
33
    "FlashAttentionDaoCudaExtension",
34
    "FlashAttentionXformersCudaExtension",
35
    "FlashAttentionNpuExtension",
36
]
37