mirror of https://github.com/hpcaitech/ColossalAI
15 lines
423 B
Python
15 lines
423 B
Python
from .cpu_adam_loader import CPUAdamLoader
|
|
from .cuda_native import FusedScaleMaskSoftmax, LayerNorm, MultiHeadAttention
|
|
from .extensions.flash_attention import AttnMaskType
|
|
from .flash_attention_loader import ColoAttention, FlashAttentionLoader
|
|
|
|
__all__ = [
|
|
"LayerNorm",
|
|
"FusedScaleMaskSoftmax",
|
|
"MultiHeadAttention",
|
|
"CPUAdamLoader",
|
|
"FlashAttentionLoader",
|
|
"ColoAttention",
|
|
"AttnMaskType",
|
|
]
|