mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-11-13 14:43:32 +00:00
17 lines
457 B
Python
17 lines
457 B
Python
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 |