mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-05-16 04:10:14 +00:00
40d376c566
* [setup] support pre-build and jit-build of cuda kernels * polish code * polish code * polish code * polish code * polish code * polish code
8 lines
165 B
Python
8 lines
165 B
Python
from .cuda_native import FusedScaleMaskSoftmax, LayerNorm, MultiHeadAttention
|
|
|
|
__all__ = [
|
|
"LayerNorm",
|
|
"FusedScaleMaskSoftmax",
|
|
"MultiHeadAttention",
|
|
]
|