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.
17 lines
457 B
17 lines
457 B
3 years ago
|
from .builder import _build_cuda_native_kernel
|
||
|
|
|
||
|
|
CUDA_NATIVE_KERNEL_BUILD = False
|
||
|
|
|
||
|
|
|
||
|
|
def build_cuda_native_kernel():
|
||
|
|
global CUDA_NATIVE_KERNEL_BUILD
|
||
|
|
if CUDA_NATIVE_KERNEL_BUILD == False:
|
||
|
|
_build_cuda_native_kernel()
|
||
|
|
CUDA_NATIVE_KERNEL_BUILD = True
|
||
|
|
|
||
|
|
|
||
|
|
build_cuda_native_kernel()
|
||
|
|
|
||
|
|
from .layer_norm import MixedFusedLayerNorm as LayerNorm
|
||
|
|
from .scaled_softmax import FusedScaleMaskSoftmax
|
||
|
|
from .multihead_attention import MultiHeadAttention
|