mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
15 lines
423 B
15 lines
423 B
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",
|
|
]
|