refactor kernel (#142)

colossalai/kernel/jit/__init__.py

@@ -1,3 +1,8 @@
 from .option import _set_jit_fusion_options
+from .bias_dropout_add import bias_dropout_add_fused_train, bias_dropout_add_fused_inference
+from .bias_gelu import bias_gelu_impl
+_set_jit_fusion_options()
 
-_set_jit_fusion_options()
+__all__ = [
+    "bias_dropout_add_fused_train", "bias_dropout_add_fused_inference", "bias_gelu_impl",
+]

colossalai/kernel/jit/option.py

@@ -2,6 +2,7 @@ import torch
 
 JIT_OPTIONS_SET = False
 
+
 def _set_jit_fusion_options():
     """Set PyTorch JIT layer fusion options."""
     global JIT_OPTIONS_SET