mirror of https://github.com/hpcaitech/ColossalAI
15 lines
406 B
Python
15 lines
406 B
Python
from .cuda_native import FusedScaleMaskSoftmax, LayerNorm, MultiHeadAttention
|
|
from .triton import llama_context_attn_fwd, bloom_context_attn_fwd
|
|
from .triton import softmax
|
|
from .triton import copy_kv_cache_to_dest
|
|
|
|
__all__ = [
|
|
"LayerNorm",
|
|
"FusedScaleMaskSoftmax",
|
|
"MultiHeadAttention",
|
|
"llama_context_attn_fwd",
|
|
"bloom_context_attn_fwd",
|
|
"softmax",
|
|
"copy_kv_cache_to_dest",
|
|
]
|