ColossalAI/colossalai/kernel/__init__.py
2022-04-06 11:40:59 +08:00

4 lines
150 B
Python

from .cuda_native import LayerNorm, FusedScaleMaskSoftmax, MultiHeadAttention
__all__ = ["LayerNorm", "FusedScaleMaskSoftmax", "MultiHeadAttention"]