Refactor turbomind (low-level abstractions)

.github/workflows/windows-x64-gpu.yml

@@ -50,11 +50,11 @@ jobs:
             INPUT_CUDA_VERSION: ${{ matrix.cudaver }}
       - name: Build wheel
         run: |
-          $env:BUILD_TEST="ON"
+          $env:BUILD_TEST="OFF"
           mkdir build
           cd build
           ..\builder\windows\generate.ps1
-          cmake --build . --config Release -- /m /v:q
+          cmake --build . --config Release -- /m /v:n
           if (-Not $?) {
             echo "build failed"
             exit 1

CMakeLists.txt

@@ -15,22 +15,28 @@
 cmake_minimum_required(VERSION 3.11 FATAL_ERROR) # for PyTorch extensions, version should be greater than 3.13
 project(TurboMind LANGUAGES CXX CUDA)
 
-find_package(CUDA 10.2 REQUIRED)
+if (MSVC)
+    # use standard conformant preprocessor
+    add_compile_options($<$<COMPILE_LANGUAGE:CXX>:/Zc:preprocessor>)
+    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler=/Zc:preprocessor")
+endif ()
 
 find_package(CUDAToolkit REQUIRED)
 
-if(${CUDA_VERSION_MAJOR} VERSION_GREATER_EQUAL "11")
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "11")
   add_definitions("-DENABLE_BF16")
-  message("CUDA_VERSION ${CUDA_VERSION_MAJOR}.${CUDA_VERSION_MINOR} is greater or equal than 11.0, enable -DENABLE_BF16 flag")
 endif()
 
 set(CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/Modules)
 
 option(BUILD_MULTI_GPU "Build multi-gpu support" ON)
 option(BUILD_PY_FFI "Build python ffi" ON)
 option(BUILD_TEST "Build tests" OFF)
+option(SPARSITY_SUPPORT "Build project with Ampere sparsity feature support" OFF)
+option(BUILD_FAST_MATH "Build in fast math mode" ON)
 
 include(FetchContent)
+
 if (BUILD_TEST)
   FetchContent_Declare(
     repo-cutlass
@@ -45,6 +51,14 @@ if (BUILD_TEST)
 
   set(CUTLASS_HEADER_DIR ${PROJECT_SOURCE_DIR}/3rdparty/cutlass/include)
   set(CUTLASS_EXTENSIONS_DIR ${PROJECT_SOURCE_DIR}/src/turbomind/cutlass_extensions/include)
+
+
+  FetchContent_Declare(
+    Catch2
+    GIT_REPOSITORY https://github.com/catchorg/Catch2.git
+    GIT_TAG        v3.8.0
+  )
+  FetchContent_MakeAvailable(Catch2)
 endif()
 
 FetchContent_Declare(
@@ -56,10 +70,6 @@ set(YAML_BUILD_SHARED_LIBS OFF CACHE BOOL "Build static library of yaml-cpp")
 FetchContent_MakeAvailable(yaml-cpp)
 
 
-option(SPARSITY_SUPPORT "Build project with Ampere sparsity feature support" OFF)
-
-option(BUILD_FAST_MATH "Build in fast math mode" ON)
-
 # the environment variable
 #   ASAN_OPTIONS=protect_shadow_gap=0,intercept_tls_get_addr=0
 # must be set at runtime
@@ -112,13 +122,13 @@ set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -Wall -ldl") # -Xptxas -v
 # TODO: build for sm_72 & sm_87 on aarch64 platform (Jetson devices)
 if (NOT CMAKE_CUDA_ARCHITECTURES)
   set(CMAKE_CUDA_ARCHITECTURES 70-real 75-real)
-  if (${CUDA_VERSION} VERSION_GREATER_EQUAL "11")
+  if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "11")
     list(APPEND CMAKE_CUDA_ARCHITECTURES 80-real)
   endif ()
-  if (${CUDA_VERSION} VERSION_GREATER_EQUAL "11.1")
+  if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "11.1")
     list(APPEND CMAKE_CUDA_ARCHITECTURES 86-real)
   endif ()
-  if (${CUDA_VERSION} VERSION_GREATER_EQUAL "11.8")
+  if (${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "11.8")
     list(APPEND CMAKE_CUDA_ARCHITECTURES 89-real 90-real)
   endif ()
   if (MSVC)
@@ -132,19 +142,23 @@ set(CMAKE_CUDA_RUNTIME_LIBRARY Shared)
 set(CMAKE_C_FLAGS_DEBUG    "${CMAKE_C_FLAGS_DEBUG}    -Wall -O0")
 set(CMAKE_CXX_FLAGS_DEBUG  "${CMAKE_CXX_FLAGS_DEBUG}  -Wall -O0")
 # set(CMAKE_CUDA_FLAGS_DEBUG "${CMAKE_CUDA_FLAGS_DEBUG} -O0 -G -Xcompiler -Wall  --ptxas-options=-v --resource-usage")
-set(CMAKE_CUDA_FLAGS_DEBUG "${CMAKE_CUDA_FLAGS_DEBUG} -O0 -G -Xcompiler -Wall -DCUDA_PTX_FP8_F2FP_ENABLED")
+set(CMAKE_CUDA_FLAGS_DEBUG "${CMAKE_CUDA_FLAGS_DEBUG} -O0 -G -Xcompiler -Wall")
 
 set(CMAKE_CXX_STANDARD "${CXX_STD}")
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-extended-lambda")
 set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
-set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --std=c++${CXX_STD} -DCUDA_PTX_FP8_F2FP_ENABLED")
+set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --std=c++${CXX_STD}")
+
+string(REPLACE "-O2" "" CMAKE_CXX_FLAGS_RELEASE         "${CMAKE_CXX_FLAGS_RELEASE}")
+string(REPLACE "-O2" "" CMAKE_CUDA_FLAGS_RELEASE        "${CMAKE_CUDA_FLAGS_RELEASE}")
+string(REPLACE "-O2" "" CMAKE_CXX_FLAGS_RELWITHDEBINFO  "${CMAKE_CXX_FLAGS_RELWITHDEBINFO}")
+string(REPLACE "-O2" "" CMAKE_CUDA_FLAGS_RELWITHDEBINFO "${CMAKE_CUDA_FLAGS_RELWITHDEBINFO}")
 
-set(CMAKE_CXX_FLAGS_RELEASE        "${CMAKE_CXX_FLAGS_RELEASE}        -O3")
-set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} -O3")
-# set(CMAKE_CUDA_FLAGS_RELEASE "${CMAKE_CUDA_FLAGS_RELEASE} -Xcompiler -O3 --ptxas-options=--verbose")
-set(CMAKE_CUDA_FLAGS_RELEASE        "${CMAKE_CUDA_FLAGS_RELEASE}        -Xcompiler -O3 -DCUDA_PTX_FP8_F2FP_ENABLED")
-set(CMAKE_CUDA_FLAGS_RELWITHDEBINFO "${CMAKE_CUDA_FLAGS_RELWITHDEBINFO} -Xcompiler -O3 -DCUDA_PTX_FP8_F2FP_ENABLED")
+set(CMAKE_CXX_FLAGS_RELEASE         "${CMAKE_CXX_FLAGS_RELEASE}         -O3")
+set(CMAKE_CXX_FLAGS_RELWITHDEBINFO  "${CMAKE_CXX_FLAGS_RELWITHDEBINFO}  -O3")
+set(CMAKE_CUDA_FLAGS_RELEASE        "${CMAKE_CUDA_FLAGS_RELEASE}        -O3")
+set(CMAKE_CUDA_FLAGS_RELWITHDEBINFO "${CMAKE_CUDA_FLAGS_RELWITHDEBINFO} -O3")
 
 if(BUILD_FAST_MATH)
     set(CMAKE_CUDA_FLAGS_RELEASE        "${CMAKE_CUDA_FLAGS_RELEASE}        --use_fast_math")
@@ -207,13 +221,11 @@ link_directories(
   ${COMMON_LIB_DIRS}
 )
 
-# add_subdirectory(3rdparty)
 add_subdirectory(src)
-# add_subdirectory(examples)
 
-if(BUILD_TEST)
-    add_subdirectory(tests/csrc)
-endif()
+# if(BUILD_TEST)
+#     add_subdirectory(tests/csrc)
+# endif()
 
 # install python api
 if (BUILD_PY_FFI)

builder/windows/generate.ps1

@@ -3,6 +3,5 @@ cmake .. -A x64 -T "v142,cuda=$env:CUDA_PATH" `
     -DCMAKE_INSTALL_PREFIX=install `
     -DBUILD_PY_FFI=ON `
     -DBUILD_MULTI_GPU=OFF `
-    -DCMAKE_CUDA_FLAGS="-lineinfo" `
-    -DUSE_NVTX=ON `
+    -DUSE_NVTX=OFF `
     -DBUILD_TEST="$env:BUILD_TEST"

builder/windows/setup_cuda.ps1

@@ -24,6 +24,8 @@ if ($CUDA_VERSION_FULL -eq "12.1.0") {
     $downloadUrl = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_531.14_windows.exe"
 } elseif ($CUDA_VERSION_FULL -eq "11.8.0") {
     $downloadUrl = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_522.06_windows.exe"
+} elseif ($CUDA_VERSION_FULL -eq "12.5.0") {
+    $downloadUrl = "https://developer.download.nvidia.com/compute/cuda/12.5.0/local_installers/cuda_12.5.0_555.85_windows.exe"
 } else {
     Write-Output "Unsupported CUDA version specified"
     exit 1
@@ -84,6 +86,8 @@ $msBuildExtensions = (Get-ChildItem  "$src\visual_studio_integration\CUDAVisualS
     }
 }
 
+$CUDA_FLAGS="-allow-unsupported-compiler -D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH=1"
+
 # Add to Github env
 Write-Output "Setting environment variables for GitHub Actions..."
 
@@ -97,7 +101,7 @@ Write-Output "CUDA_PATH_V$($CUDA_MAJOR)_$($CUDA_MINOR)=$dst" >> $env:GITHUB_ENV
 Write-Output "CUDA_PATH_VX_Y=CUDA_PATH_V$($CUDA_MAJOR)_$($CUDA_MINOR)" >> $env:GITHUB_ENV
 Write-Output "CudaToolkitDir=$dst" >> $env:GITHUB_ENV
 Write-Output "CMAKE_CUDA_COMPILER=$dst\bin\nvcc.exe" >> $env:GITHUB_ENV
-Write-Output "NVCC_APPEND_FLAGS=-allow-unsupported-compiler" >> $env:GITHUB_ENV
+Write-Output "NVCC_APPEND_FLAGS=$CUDA_FLAGS" >> $env:GITHUB_ENV
 
 Write-Output "CUDA_VERSION=$CUDA_VERSION_FULL" >> $env:GITHUB_ENV
 Write-Output "Setup completed."

lmdeploy/turbomind/deploy/module.py

@@ -319,7 +319,8 @@ def pad_weight(tensor: torch.Tensor, tp: int):
         if output_weight is not None:
             tp = self.model.attn_tp_size
             output_weight = pad_weight(output_weight, tp=tp)
-            self.model.save_split(output_weight, 'output.weight', split_dim=0, split_num=tp)
+            # transpose
+            self.model.save_split(output_weight.t(), 'output.weight', split_dim=1, split_num=tp)
 
 
 class Transformer:

lmdeploy/turbomind/turbomind.py

@@ -241,7 +241,7 @@ def _from_hf(self, model_source: ModelSource, model_path: str, engine_config: Tu
 
         model_comm = _tm.AbstractTransformerModel.create_llama_model(model_dir='',
                                                                      config=yaml.safe_dump(self.config_dict),
-                                                                     data_type=self.config.model_config.weight_type)
+                                                                     weight_type=self.config.model_config.weight_type)
 
         # create empty weight
         self._create_weight(model_comm)
@@ -275,7 +275,7 @@ def _from_workspace(self, model_path: str, engine_config: TurbomindEngineConfig)
         weight_dir = osp.join(model_path, 'triton_models', 'weights')
         model_comm = _tm.AbstractTransformerModel.create_llama_model(model_dir=weight_dir,
                                                                      config=yaml.safe_dump(self.config_dict),
-                                                                     data_type=self.config.weight_type)
+                                                                     weight_type=self.config.weight_type)
 
         # create weight and load params
         self._create_weight(model_comm)

src/turbomind/CMakeLists.txt

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 add_subdirectory(utils)
+add_subdirectory(core)
 add_subdirectory(kernels)
 add_subdirectory(layers)
 add_subdirectory(comm)

src/turbomind/comm/CMakeLists.txt

@@ -3,10 +3,11 @@
 cmake_minimum_required(VERSION 3.8)
 
 add_library(host_comm STATIC host_comm.cc thread_comm.cc)
+target_link_libraries(host_comm PRIVATE core logger)
 set_property(TARGET host_comm PROPERTY POSITION_INDEPENDENT_CODE ON)
 
 add_library(device_comm STATIC device_comm.cc)
-target_link_libraries(device_comm PRIVATE logger)
+target_link_libraries(device_comm PRIVATE core logger)
 set_property(TARGET device_comm PROPERTY POSITION_INDEPENDENT_CODE ON)
 set_property(TARGET device_comm PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
 
@@ -21,7 +22,7 @@ if (BUILD_MULTI_GPU)
 
     if (BUILD_TEST)
         add_executable(test_comm test_comm.cu)
-        target_link_libraries(test_comm PRIVATE device_comm host_comm pthread nvtx_utils)
+        target_link_libraries(test_comm PRIVATE device_comm host_comm core pthread nvtx_utils)
         target_compile_options(test_comm PRIVATE -O3 -march=native -mtune=native)
     endif ()
 endif ()

src/turbomind/comm/cuda_ipc/CMakeLists.txt

@@ -12,6 +12,8 @@ add_library(cuda_ipc_comm STATIC
 target_link_libraries(cuda_ipc_comm PRIVATE
         rms_norm
         host_comm
+        core
+        cuda_utils
         CUDA::cuda_driver
         logger)
 

src/turbomind/comm/cuda_ipc/allgather.cu

@@ -4,7 +4,6 @@
 #include "src/turbomind/comm/cuda_ipc/device_semaphore.h"
 
 #include "src/turbomind/kernels/core/meta.h"
-#include "src/turbomind/utils/Tensor.h"
 #include "src/turbomind/utils/cuda_utils.h"
 
 namespace turbomind::comm {
@@ -51,7 +50,7 @@ __global__ void __launch_bounds__(1024, 1) Allgather_Simple_Pull(T*
 void CudaIpcCommImpl::AllGather(
     const void* sendbuff, void* recvbuff, size_t sendcount, DataType type, int group, cudaStream_t stream)
 {
-    const size_t bytesize = get_elem_size(type) * sendcount;
+    const size_t bytesize = turbomind::byte_size(type) * sendcount;
 
     const int peers = this->n_ranks(group) - 1;
     const int rank  = this->rank(group);
@@ -165,9 +164,9 @@ void CudaIpcCommImpl::AllGather2D(const void*  sendbuff,
                                   int          group,
                                   cudaStream_t stream)
 {
-    const size_t byte_width  = get_elem_size(type) * width;
-    const size_t byte_pitch  = get_elem_size(type) * pitch;
-    const size_t byte_stride = get_elem_size(type) * stride;
+    const size_t byte_width  = byte_size(type, width);
+    const size_t byte_pitch  = byte_size(type, pitch);
+    const size_t byte_stride = byte_size(type, stride);
 
     void*  base{};
     size_t offset{};

src/turbomind/comm/cuda_ipc/allreduce.cu

@@ -6,9 +6,9 @@
 #include "src/turbomind/comm/cuda_ipc/cuda_ipc_comm.h"
 #include "src/turbomind/comm/cuda_ipc/device_semaphore.h"
 
+#include "src/turbomind/core/data_type.h"
 #include "src/turbomind/kernels/core/array_ops.h"
 #include "src/turbomind/kernels/core/meta.h"
-#include "src/turbomind/utils/Tensor.h"
 
 #include "src/turbomind/utils/cuda_utils.h"
 
@@ -423,14 +423,7 @@ void CudaIpcCommImpl::AllReduceSum(
         }
     };
 
-    switch (type) {
-        case DataType::TYPE_FP16:
-            return invoke(half{});
-        case DataType::TYPE_BF16:
-            return invoke(nv_bfloat16{});
-        default:
-            throw std::runtime_error("not implemented");
-    }
+    TM_DISPATCH_PRIMARY_DTYPES(type, invoke);
 }
 
 }  // namespace turbomind::comm

src/turbomind/comm/cuda_ipc/cuda_ipc_comm.cu

@@ -1,8 +1,7 @@
 // Copyright (c) OpenMMLab. All rights reserved.
 
 #include <memory>
-#include <mutex>
-#include <type_traits>
+#include <numeric>
 #include <vector>
 
 #include <cuda.h>

src/turbomind/comm/cuda_ipc/cuda_ipc_comm.h

@@ -10,7 +10,6 @@
 
 #include "src/turbomind/kernels/core/array.h"
 
-#include "src/turbomind/utils/Tensor.h"
 #include "src/turbomind/utils/cuda_utils.h"
 
 namespace turbomind::comm {

src/turbomind/comm/cuda_ipc/fused_allreduce.cu

@@ -8,13 +8,13 @@
 #include "src/turbomind/comm/cuda_ipc/device_semaphore.h"
 #include "src/turbomind/comm/cuda_ipc/group_sum.h"
 
+#include "src/turbomind/core/data_type.h"
 #include "src/turbomind/kernels/core/array_ops.h"
 #include "src/turbomind/kernels/core/common.h"
 #include "src/turbomind/kernels/core/meta.h"
 
 #include "src/turbomind/kernels/norm/rms_norm.h"
 
-#include "src/turbomind/utils/Tensor.h"
 #include "src/turbomind/utils/cuda_utils.h"
 
 namespace turbomind::comm {
@@ -424,7 +424,7 @@ void CudaIpcCommImpl::AllreduceResidualBiasRMSnorm(void*        hidden,
                                                    cudaStream_t stream)
 {
 
-    const size_t elemsize = get_elem_size(dtype);
+    const size_t elemsize = byte_size(dtype);
     const size_t bytesize = elemsize * token_num * dim;
 
     const int n_ranks = this->n_ranks(group);
@@ -504,19 +504,10 @@ void CudaIpcCommImpl::AllreduceResidualBiasRMSnorm(void*        hidden,
         return false;  // > 1024 vdim
     };
 
-    auto dispatch = [&] {
-        switch (dtype) {
-            case DataType::TYPE_FP16:
-                return dispatch_D(half{});
-            case DataType::TYPE_BF16:
-                return dispatch_D(nv_bfloat16{});
-            default:
-                return false;
-        }
-    };
+    auto dispatch = [&]() -> bool { TM_DISPATCH_PRIMARY_DTYPES_RET(dtype, dispatch_D); };
 
     if (bytesize > (1 << 19)) {
-        if (auto success = dispatch()) {
+        if (dispatch()) {
             return;
         }
     }

src/turbomind/comm/cuda_ipc/fused_allreduce_ex.cu

@@ -5,6 +5,7 @@
 #include "src/turbomind/comm/cuda_ipc/group_sum.h"
 
 #include "src/turbomind/comm/cuda_ipc/mscclpp.h"
+#include "src/turbomind/core/data_type.h"
 #include "src/turbomind/kernels/core/array_ops.h"
 #include "src/turbomind/kernels/core/common.h"
 #include "src/turbomind/kernels/core/meta.h"
@@ -279,18 +280,11 @@ void CudaIpcCommImpl::AllreduceResidualBiasRMSnormEx(void*        hidden,
         return false;  // > 1024 vdim
     };
 
-    auto dispatch = [&] {
-        switch (dtype) {
-            case DataType::TYPE_FP16:
-                return dispatch_D(half{});
-            case DataType::TYPE_BF16:
-                return dispatch_D(nv_bfloat16{});
-            default:
-                return false;
-        }
+    auto dispatch = [&]() -> bool {  //
+        TM_DISPATCH_PRIMARY_DTYPES_RET(dtype, dispatch_D);
     };
 
-    FT_CHECK(dispatch());
+    TM_CHECK(dispatch());
 }
 
 }  // namespace turbomind::comm

src/turbomind/comm/device_comm.cc

@@ -25,7 +25,7 @@ DeviceComm CreateDeviceCommunicator(const std::string& backend, int n_ranks, int
     }
 #endif
 
-    FT_CHECK_WITH_INFO(0, fmtstr("Unknown communication backend: %s", backend.c_str()));
+    TM_CHECK(0) << "Unknown communication backend: " << backend;
     return {};
 }