ColossalAI/colossalai/kernel/op_builder/__init__.py

6 lines
205 B
Python

from .cpu_adam import CPUAdamBuilder
from .fused_optim import FusedOptimBuilder
from .multi_head_attn import MultiHeadAttnBuilder
__all__ = ['CPUAdamBuilder', 'FusedOptimBuilder', 'MultiHeadAttnBuilder']