Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
__init__.py 350 Bytes
Newer Older
shenggan's avatar
add colossalai kernel module (#55)  
shenggan committed
1 
from .layer_norm import MixedFusedLayerNorm as LayerNorm
ver217's avatar
refactor kernel (#142)  
ver217 committed
2 
from .multihead_attention import MultiHeadAttention
flybird1111's avatar
[shardformer] update shardformer to use flash attention 2 (#4392)  
flybird1111 committed
3 
from .scaled_softmax import AttnMaskType, FusedScaleMaskSoftmax, ScaledUpperTriangMaskedSoftmax
Frank Lee's avatar
[kernel] fixed repeated loading of kernels (#2549)  
Frank Lee committed
4
flybird1111's avatar
[fix] coloattention support flash attention 2 (#4347)  
flybird1111 committed
5 
__all__ = [
Hongxin Liu's avatar
[misc] update pre-commit and run all files (#4752)  
Hongxin Liu committed
6
7
8
9
10 
    "LayerNorm",
    "MultiHeadAttention",
    "FusedScaleMaskSoftmax",
    "ScaledUpperTriangMaskedSoftmax",
    "AttnMaskType",
flybird1111's avatar
[fix] coloattention support flash attention 2 (#4347)  
flybird1111 committed
11 
]