From 095c070a6eefe1a76fe3483b21986826114d6d17 Mon Sep 17 00:00:00 2001 From: xs_courtesy Date: Mon, 11 Mar 2024 17:06:57 +0800 Subject: [PATCH] refactor code --- extensions/cpu_adam/cpu_adam_x86.py | 2 +- extensions/csrc/cuda/compat.h | 0 .../{layer_norm_cuda_kernel.cu => layer_norm_kernel.cu} | 0 extensions/csrc/cuda/{moe_cuda_kernel.cu => moe_kernel.cu} | 0 .../{multi_tensor_adam.cu => multi_tensor_adam_kernel.cu} | 0 .../{multi_tensor_lamb.cu => multi_tensor_lamb_kernel.cu} | 0 .../inference.cpp} | 0 .../cuda/{layer_norm_cuda.cpp => pybind/layer_norm.cpp} | 0 extensions/csrc/cuda/{moe_cuda.cpp => pybind/moe.cpp} | 0 .../cuda/{colossal_C_frontend.cpp => pybind/optimizer.cpp} | 0 extensions/csrc/cuda/{ => pybind}/scaled_masked_softmax.cpp | 0 .../{ => pybind}/scaled_upper_triang_masked_softmax.cpp | 0 extensions/csrc/cuda/rms_layernorm_kernel.cu | 2 +- ...sked_softmax_cuda.cu => scaled_masked_softmax_kernel.cu} | 0 ...cuda.cu => scaled_upper_triang_masked_softmax_kernel.cu} | 0 extensions/csrc/{cuda => x86}/cpu_adam.cpp | 0 extensions/csrc/{cuda => x86}/cpu_adam.h | 0 extensions/inference/inference_ops_cuda.py | 2 +- extensions/layernorm/layernorm_cuda.py | 2 +- extensions/moe/moe_cuda.py | 2 +- extensions/optimizer/fused_optimizer_cuda.py | 6 +++--- extensions/softmax/scaled_masked_softmax_cuda.py | 2 +- .../softmax/scaled_upper_triangle_masked_softmax_cuda.py | 4 ++-- 23 files changed, 11 insertions(+), 11 deletions(-) delete mode 100644 extensions/csrc/cuda/compat.h rename extensions/csrc/cuda/{layer_norm_cuda_kernel.cu => layer_norm_kernel.cu} (100%) rename extensions/csrc/cuda/{moe_cuda_kernel.cu => moe_kernel.cu} (100%) rename extensions/csrc/cuda/{multi_tensor_adam.cu => multi_tensor_adam_kernel.cu} (100%) rename extensions/csrc/cuda/{multi_tensor_lamb.cu => multi_tensor_lamb_kernel.cu} (100%) rename extensions/csrc/cuda/{colossal_inference_C_frontend.cpp => pybind/inference.cpp} (100%) rename extensions/csrc/cuda/{layer_norm_cuda.cpp => pybind/layer_norm.cpp} (100%) rename extensions/csrc/cuda/{moe_cuda.cpp => pybind/moe.cpp} (100%) rename extensions/csrc/cuda/{colossal_C_frontend.cpp => pybind/optimizer.cpp} (100%) rename extensions/csrc/cuda/{ => pybind}/scaled_masked_softmax.cpp (100%) rename extensions/csrc/cuda/{ => pybind}/scaled_upper_triang_masked_softmax.cpp (100%) rename extensions/csrc/cuda/{scaled_masked_softmax_cuda.cu => scaled_masked_softmax_kernel.cu} (100%) rename extensions/csrc/cuda/{scaled_upper_triang_masked_softmax_cuda.cu => scaled_upper_triang_masked_softmax_kernel.cu} (100%) rename extensions/csrc/{cuda => x86}/cpu_adam.cpp (100%) rename extensions/csrc/{cuda => x86}/cpu_adam.h (100%) diff --git a/extensions/cpu_adam/cpu_adam_x86.py b/extensions/cpu_adam/cpu_adam_x86.py index a38194167..27b06bb65 100644 --- a/extensions/cpu_adam/cpu_adam_x86.py +++ b/extensions/cpu_adam/cpu_adam_x86.py @@ -21,7 +21,7 @@ class CpuAdamX86Extension(_CudaExtension): # necessary 4 functions def sources_files(self): ret = [ - self.csrc_abs_path("cuda/cpu_adam.cpp"), + self.csrc_abs_path("x86/cpu_adam.cpp"), ] return ret diff --git a/extensions/csrc/cuda/compat.h b/extensions/csrc/cuda/compat.h deleted file mode 100644 index e69de29bb..000000000 diff --git a/extensions/csrc/cuda/layer_norm_cuda_kernel.cu b/extensions/csrc/cuda/layer_norm_kernel.cu similarity index 100% rename from extensions/csrc/cuda/layer_norm_cuda_kernel.cu rename to extensions/csrc/cuda/layer_norm_kernel.cu diff --git a/extensions/csrc/cuda/moe_cuda_kernel.cu b/extensions/csrc/cuda/moe_kernel.cu similarity index 100% rename from extensions/csrc/cuda/moe_cuda_kernel.cu rename to extensions/csrc/cuda/moe_kernel.cu diff --git a/extensions/csrc/cuda/multi_tensor_adam.cu b/extensions/csrc/cuda/multi_tensor_adam_kernel.cu similarity index 100% rename from extensions/csrc/cuda/multi_tensor_adam.cu rename to extensions/csrc/cuda/multi_tensor_adam_kernel.cu diff --git a/extensions/csrc/cuda/multi_tensor_lamb.cu b/extensions/csrc/cuda/multi_tensor_lamb_kernel.cu similarity index 100% rename from extensions/csrc/cuda/multi_tensor_lamb.cu rename to extensions/csrc/cuda/multi_tensor_lamb_kernel.cu diff --git a/extensions/csrc/cuda/colossal_inference_C_frontend.cpp b/extensions/csrc/cuda/pybind/inference.cpp similarity index 100% rename from extensions/csrc/cuda/colossal_inference_C_frontend.cpp rename to extensions/csrc/cuda/pybind/inference.cpp diff --git a/extensions/csrc/cuda/layer_norm_cuda.cpp b/extensions/csrc/cuda/pybind/layer_norm.cpp similarity index 100% rename from extensions/csrc/cuda/layer_norm_cuda.cpp rename to extensions/csrc/cuda/pybind/layer_norm.cpp diff --git a/extensions/csrc/cuda/moe_cuda.cpp b/extensions/csrc/cuda/pybind/moe.cpp similarity index 100% rename from extensions/csrc/cuda/moe_cuda.cpp rename to extensions/csrc/cuda/pybind/moe.cpp diff --git a/extensions/csrc/cuda/colossal_C_frontend.cpp b/extensions/csrc/cuda/pybind/optimizer.cpp similarity index 100% rename from extensions/csrc/cuda/colossal_C_frontend.cpp rename to extensions/csrc/cuda/pybind/optimizer.cpp diff --git a/extensions/csrc/cuda/scaled_masked_softmax.cpp b/extensions/csrc/cuda/pybind/scaled_masked_softmax.cpp similarity index 100% rename from extensions/csrc/cuda/scaled_masked_softmax.cpp rename to extensions/csrc/cuda/pybind/scaled_masked_softmax.cpp diff --git a/extensions/csrc/cuda/scaled_upper_triang_masked_softmax.cpp b/extensions/csrc/cuda/pybind/scaled_upper_triang_masked_softmax.cpp similarity index 100% rename from extensions/csrc/cuda/scaled_upper_triang_masked_softmax.cpp rename to extensions/csrc/cuda/pybind/scaled_upper_triang_masked_softmax.cpp diff --git a/extensions/csrc/cuda/rms_layernorm_kernel.cu b/extensions/csrc/cuda/rms_layernorm_kernel.cu index 99d36575d..0ab40f9f7 100644 --- a/extensions/csrc/cuda/rms_layernorm_kernel.cu +++ b/extensions/csrc/cuda/rms_layernorm_kernel.cu @@ -9,7 +9,7 @@ #include "block_reduce.h" -#include "type_shim.h" +#include "../common/micros.h" template __global__ void rms_layernorm_kernel( diff --git a/extensions/csrc/cuda/scaled_masked_softmax_cuda.cu b/extensions/csrc/cuda/scaled_masked_softmax_kernel.cu similarity index 100% rename from extensions/csrc/cuda/scaled_masked_softmax_cuda.cu rename to extensions/csrc/cuda/scaled_masked_softmax_kernel.cu diff --git a/extensions/csrc/cuda/scaled_upper_triang_masked_softmax_cuda.cu b/extensions/csrc/cuda/scaled_upper_triang_masked_softmax_kernel.cu similarity index 100% rename from extensions/csrc/cuda/scaled_upper_triang_masked_softmax_cuda.cu rename to extensions/csrc/cuda/scaled_upper_triang_masked_softmax_kernel.cu diff --git a/extensions/csrc/cuda/cpu_adam.cpp b/extensions/csrc/x86/cpu_adam.cpp similarity index 100% rename from extensions/csrc/cuda/cpu_adam.cpp rename to extensions/csrc/x86/cpu_adam.cpp diff --git a/extensions/csrc/cuda/cpu_adam.h b/extensions/csrc/x86/cpu_adam.h similarity index 100% rename from extensions/csrc/cuda/cpu_adam.h rename to extensions/csrc/x86/cpu_adam.h diff --git a/extensions/inference/inference_ops_cuda.py b/extensions/inference/inference_ops_cuda.py index 042c598fb..f465fe600 100644 --- a/extensions/inference/inference_ops_cuda.py +++ b/extensions/inference/inference_ops_cuda.py @@ -10,7 +10,7 @@ class InferenceOpsCudaExtension(_CudaExtension): ret = [ self.csrc_abs_path(fname) for fname in [ - "cuda/colossal_inference_C_frontend.cpp", + "cuda/pybind/inference.cpp", "cuda/decode_kv_cache_memcpy_kernel.cu", "cuda/activation_kernel.cu", "cuda/rms_layernorm_kernel.cu", diff --git a/extensions/layernorm/layernorm_cuda.py b/extensions/layernorm/layernorm_cuda.py index db5f2fce1..36cf73590 100644 --- a/extensions/layernorm/layernorm_cuda.py +++ b/extensions/layernorm/layernorm_cuda.py @@ -7,7 +7,7 @@ class LayerNormCudaExtension(_CudaExtension): super().__init__(name="layernorm_cuda") def sources_files(self): - ret = [self.csrc_abs_path(fname) for fname in ["cuda/layer_norm_cuda.cpp", "cuda/layer_norm_cuda_kernel.cu"]] + ret = [self.csrc_abs_path(fname) for fname in ["cuda/pybind/layer_norm.cpp", "cuda/layer_norm_kernel.cu"]] return ret def include_dirs(self): diff --git a/extensions/moe/moe_cuda.py b/extensions/moe/moe_cuda.py index 52883e97f..722daae33 100644 --- a/extensions/moe/moe_cuda.py +++ b/extensions/moe/moe_cuda.py @@ -11,7 +11,7 @@ class MoeCudaExtension(_CudaExtension): return ret def sources_files(self): - ret = [self.csrc_abs_path(fname) for fname in ["cuda/moe_cuda.cpp", "cuda/moe_cuda_kernel.cu"]] + ret = [self.csrc_abs_path(fname) for fname in ["cuda/moe.cpp", "cuda/moe_kernel.cu"]] return ret def cxx_flags(self): diff --git a/extensions/optimizer/fused_optimizer_cuda.py b/extensions/optimizer/fused_optimizer_cuda.py index e065cf34a..41c6260aa 100644 --- a/extensions/optimizer/fused_optimizer_cuda.py +++ b/extensions/optimizer/fused_optimizer_cuda.py @@ -10,12 +10,12 @@ class FusedOptimizerCudaExtension(_CudaExtension): ret = [ self.csrc_abs_path(fname) for fname in [ - "cuda/colossal_C_frontend.cpp", + "cuda/pybind/optimizer.cpp", "cuda/multi_tensor_sgd_kernel.cu", "cuda/multi_tensor_scale_kernel.cu", - "cuda/multi_tensor_adam.cu", + "cuda/multi_tensor_adam_kernel.cu", "cuda/multi_tensor_l2norm_kernel.cu", - "cuda/multi_tensor_lamb.cu", + "cuda/multi_tensor_lamb_kernel.cu", ] ] return ret diff --git a/extensions/softmax/scaled_masked_softmax_cuda.py b/extensions/softmax/scaled_masked_softmax_cuda.py index 5b4208dba..797638c3b 100644 --- a/extensions/softmax/scaled_masked_softmax_cuda.py +++ b/extensions/softmax/scaled_masked_softmax_cuda.py @@ -9,7 +9,7 @@ class ScaledMaskedSoftmaxCudaExtension(_CudaExtension): def sources_files(self): ret = [ self.csrc_abs_path(fname) - for fname in ["cuda/scaled_masked_softmax.cpp", "cuda/scaled_masked_softmax_cuda.cu"] + for fname in ["cuda/pybind/scaled_masked_softmax.cpp", "cuda/scaled_masked_softmax_kernel.cu"] ] return ret diff --git a/extensions/softmax/scaled_upper_triangle_masked_softmax_cuda.py b/extensions/softmax/scaled_upper_triangle_masked_softmax_cuda.py index d4f27a921..d48d542ad 100644 --- a/extensions/softmax/scaled_upper_triangle_masked_softmax_cuda.py +++ b/extensions/softmax/scaled_upper_triangle_masked_softmax_cuda.py @@ -13,8 +13,8 @@ class ScaledUpperTriangleMaskedSoftmaxCudaExtension(_CudaExtension): ret = [ self.csrc_abs_path(fname) for fname in [ - "cuda/scaled_upper_triang_masked_softmax.cpp", - "cuda/scaled_upper_triang_masked_softmax_cuda.cu", + "cuda/pybind/scaled_upper_triang_masked_softmax.cpp", + "cuda/scaled_upper_triang_masked_softmax_kernel.cu", ] ] return ret