[feat] refactored extension module (#5298)

* [feat] refactored extension module * polish * polish * polish * polish * polish * polish * polish * polish * polish * polish

[feat] refactored extension module (#5298)
* [feat] refactored extension module * polish * polish * polish * polish * polish * polish * polish * polish * polish * polish
7cfed5f0 · Frank Lee · GitHub · d7f8db8e · 7cfed5f0 · 7cfed5f0
Unverified Commit 7cfed5f0 authored Jan 25, 2024 by Frank Lee Committed by GitHub Jan 25, 2024
20 changed files
--- a/colossalai/kernel/cuda_native/__init__.py
+++ b/colossalai/kernel/cuda_native/__init__.py
--- a/colossalai/kernel/cuda_native/csrc/cpu_adam_arm.cpp
+++ b/colossalai/kernel/cuda_native/csrc/cpu_adam_arm.cpp
--- a/colossalai/kernel/cuda_native/csrc/cpu_adam_arm.h
+++ b/colossalai/kernel/cuda_native/csrc/cpu_adam_arm.h
--- a/colossalai/kernel/cuda_native/csrc/colossal_C_frontend.cpp
+++ b/colossalai/kernel/cuda_native/csrc/colossal_C_frontend.cpp
--- a/colossalai/kernel/cuda_native/csrc/compat.h
+++ b/colossalai/kernel/cuda_native/csrc/compat.h
--- a/colossalai/kernel/cuda_native/csrc/cpu_adam.cpp
+++ b/colossalai/kernel/cuda_native/csrc/cpu_adam.cpp
--- a/colossalai/kernel/cuda_native/csrc/cpu_adam.h
+++ b/colossalai/kernel/cuda_native/csrc/cpu_adam.h
--- a/colossalai/kernel/cuda_native/csrc/kernels/include/block_reduce.h
+++ b/colossalai/kernel/cuda_native/csrc/kernels/include/block_reduce.h
--- a/colossalai/kernel/cuda_native/csrc/layer_norm_cuda.cpp
+++ b/colossalai/kernel/cuda_native/csrc/layer_norm_cuda.cpp
--- a/colossalai/kernel/cuda_native/csrc/layer_norm_cuda_kernel.cu
+++ b/colossalai/kernel/cuda_native/csrc/layer_norm_cuda_kernel.cu
--- a/colossalai/kernel/cuda_native/csrc/moe_cuda.cpp
+++ b/colossalai/kernel/cuda_native/csrc/moe_cuda.cpp
--- a/colossalai/kernel/cuda_native/csrc/moe_cuda_kernel.cu
+++ b/colossalai/kernel/cuda_native/csrc/moe_cuda_kernel.cu
--- a/colossalai/kernel/cuda_native/csrc/multi_tensor_adam.cu
+++ b/colossalai/kernel/cuda_native/csrc/multi_tensor_adam.cu
--- a/colossalai/kernel/cuda_native/csrc/multi_tensor_apply.cuh
+++ b/colossalai/kernel/cuda_native/csrc/multi_tensor_apply.cuh
--- a/colossalai/kernel/cuda_native/csrc/multi_tensor_l2norm_kernel.cu
+++ b/colossalai/kernel/cuda_native/csrc/multi_tensor_l2norm_kernel.cu
--- a/colossalai/kernel/cuda_native/csrc/multi_tensor_lamb.cu
+++ b/colossalai/kernel/cuda_native/csrc/multi_tensor_lamb.cu
--- a/colossalai/kernel/cuda_native/csrc/multi_tensor_scale_kernel.cu
+++ b/colossalai/kernel/cuda_native/csrc/multi_tensor_scale_kernel.cu
--- a/colossalai/kernel/cuda_native/csrc/multi_tensor_sgd_kernel.cu
+++ b/colossalai/kernel/cuda_native/csrc/multi_tensor_sgd_kernel.cu
--- a/colossalai/kernel/cuda_native/csrc/scaled_masked_softmax.cpp
+++ b/colossalai/kernel/cuda_native/csrc/scaled_masked_softmax.cpp
--- a/colossalai/kernel/cuda_native/csrc/scaled_masked_softmax.h
+++ b/colossalai/kernel/cuda_native/csrc/scaled_masked_softmax.h