[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/csrc/scaled_masked_softmax_cuda.cu
+++ b/colossalai/kernel/cuda_native/csrc/scaled_masked_softmax_cuda.cu
--- a/colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax.cpp
+++ b/colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax.cpp
--- a/colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax.h
+++ b/colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax.h
--- a/colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax_cuda.cu
+++ b/colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax_cuda.cu
--- a/colossalai/kernel/cuda_native/csrc/type_shim.h
+++ b/colossalai/kernel/cuda_native/csrc/type_shim.h
--- a/colossalai/kernel/cuda_native/scaled_softmax.py
+++ b/colossalai/kernel/cuda_native/scaled_softmax.py
--- a/extensions/cuda_extension.py
+++ b/extensions/cuda_extension.py
--- a/extensions/flash_attention/__init__.py
+++ b/extensions/flash_attention/__init__.py
--- a/extensions/flash_attention/flash_attention_dao_cuda.py
+++ b/extensions/flash_attention/flash_attention_dao_cuda.py
--- a/extensions/flash_attention/flash_attention_npu.py
+++ b/extensions/flash_attention/flash_attention_npu.py
--- a/extensions/flash_attention/flash_attention_xformers_cuda.py
+++ b/extensions/flash_attention/flash_attention_xformers_cuda.py
--- a/extensions/layernorm/__init__.py
+++ b/extensions/layernorm/__init__.py
--- a/op_builder/layernorm.py
+++ b/op_builder/layernorm.py
--- a/extensions/moe/__init__.py
+++ b/extensions/moe/__init__.py
--- a/op_builder/moe.py
+++ b/op_builder/moe.py
--- a/extensions/optimizer/__init__.py
+++ b/extensions/optimizer/__init__.py
--- a/op_builder/fused_optim.py
+++ b/op_builder/fused_optim.py
--- a/extensions/softmax/__init__.py
+++ b/extensions/softmax/__init__.py
--- a/op_builder/scaled_masked_softmax.py
+++ b/op_builder/scaled_masked_softmax.py
--- a/op_builder/scaled_upper_triangle_masked_softmax.py
+++ b/op_builder/scaled_upper_triangle_masked_softmax.py