Files
ColossalAI/colossalai/kernel/op_builder/__init__.py
2023-01-03 20:29:39 +08:00

8 lines
340 B
Python

from .cpu_adam import CPUAdamBuilder
from .fused_optim import FusedOptimBuilder
from .moe import MOEBuilder
from .multi_head_attn import MultiHeadAttnBuilder
from .scaled_upper_triang_masked_softmax import ScaledSoftmaxBuilder
__all__ = ['CPUAdamBuilder', 'FusedOptimBuilder', 'MultiHeadAttnBuilder', 'ScaledSoftmaxBuilder', 'MOEBuilder']