Files
ColossalAI/colossalai/kernel/__init__.py
2022-01-13 16:47:17 +08:00

6 lines
156 B
Python

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