[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
13 changed files
--- a/extensions/triton_extension.py
+++ b/extensions/triton_extension.py
--- a/colossalai/kernel/extensions/utils.py
+++ b/colossalai/kernel/extensions/utils.py
--- a/op_builder/README.md
+++ b/op_builder/README.md
--- a/op_builder/__init__.py
+++ b/op_builder/__init__.py
--- a/op_builder/builder.py
+++ b/op_builder/builder.py
--- a/op_builder/gptq.py
+++ b/op_builder/gptq.py
--- a/op_builder/multi_head_attn.py
+++ b/op_builder/multi_head_attn.py
--- a/op_builder/smoothquant.py
+++ b/op_builder/smoothquant.py
--- a/op_builder/utils.py
+++ b/op_builder/utils.py
--- a/setup.py
+++ b/setup.py
--- a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py
+++ b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py
--- a/tests/test_optimizer/test_adam_kernel.py
+++ b/tests/test_optimizer/test_adam_kernel.py
--- a/tests/test_utils/test_flash_attention.py
+++ b/tests/test_utils/test_flash_attention.py