cuda : implement load_tensor with cuFile

Author: rgerganov

ggml/src/ggml-cuda/CMakeLists.txt

@@ -104,10 +104,10 @@ if (CUDAToolkit_FOUND)
             # As of 12.3.1 CUDA Toolkit for Windows does not offer a static cublas library
             target_link_libraries(ggml-cuda PRIVATE CUDA::cudart_static CUDA::cublas CUDA::cublasLt)
         else ()
-            target_link_libraries(ggml-cuda PRIVATE  CUDA::cudart_static CUDA::cublas_static CUDA::cublasLt_static)
+            target_link_libraries(ggml-cuda PRIVATE  CUDA::cudart_static CUDA::cublas_static CUDA::cublasLt_static CUDA::cuFile_static)
         endif()
     else()
-        target_link_libraries(ggml-cuda PRIVATE CUDA::cudart CUDA::cublas CUDA::cublasLt)
+        target_link_libraries(ggml-cuda PRIVATE CUDA::cudart CUDA::cublas CUDA::cublasLt CUDA::cuFile)
     endif()
 
     if (GGML_CUDA_NO_VMM)

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -60,6 +60,12 @@
 #include <stdlib.h>
 #include <string>
 #include <vector>
+#include <cufile.h>
+#ifdef _WIN32
+#else
+    #include <fcntl.h>
+    #include <unistd.h>
+#endif
 
 static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size");
 
@@ -3410,6 +3416,68 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t
     GGML_UNUSED(reg);
 }
 
+static bool ggml_backend_cuda_buffer_load_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const char * path, size_t file_offset, size_t tensor_offset, size_t size) {
+#ifdef _WIN32
+    GGML_UNUSED(buffer);
+    GGML_UNUSED(tensor);
+    GGML_UNUSED(path);
+    GGML_UNUSED(file_offset);
+    GGML_UNUSED(tensor_offset);
+    GGML_UNUSED(size);
+    return false;
+#else
+    static bool initialized = false;
+    static bool use_cufile = false;
+    if (!initialized) {
+        CUfileError_t err = cuFileDriverOpen();
+        initialized = true;
+        if (err.err != CU_FILE_SUCCESS) {
+            use_cufile = false;
+            return false;
+        }
+        CUfileDrvProps_t props;
+        err = cuFileDriverGetProperties(&props);
+        if (err.err != CU_FILE_SUCCESS) {
+            use_cufile = false;
+            return false;
+        }
+        if (props.nvfs.dcontrolflags & (1 << CU_FILE_ALLOW_COMPAT_MODE)) {
+            // do not use CUfile if the driver is in compatibility mode
+            // as we have faster mechanisms in llama-model-loader
+            use_cufile = false;
+            return false;
+        }
+        use_cufile = true;
+    }
+    if (!use_cufile) {
+        return false;
+    }
+    ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context;
+    ggml_cuda_set_device(ctx->device);
+
+    int fd = open(path, O_RDONLY | O_DIRECT);
+    if (fd < 0) {
+        return false;
+    }
+    CUfileDescr_t cf_descr;
+    CUfileHandle_t cf_handle;
+    memset((void *)&cf_descr, 0, sizeof(CUfileDescr_t));
+    cf_descr.handle.fd = fd;
+    cf_descr.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD;
+    CUfileError_t status = cuFileHandleRegister(&cf_handle, &cf_descr);
+    if (status.err != CU_FILE_SUCCESS) {
+        return false;
+    }
+    ssize_t ret = cuFileRead(cf_handle, (char *)tensor->data, size, file_offset, tensor_offset);
+    if (ret < 0) {
+        return false;
+    }
+    cuFileHandleDeregister(cf_handle);
+    close(fd);
+    return true;
+#endif
+}
+
 static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) {
     GGML_UNUSED(reg);
     if (strcmp(name, "ggml_backend_split_buffer_type") == 0) {
@@ -3424,6 +3492,9 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con
     if (strcmp(name, "ggml_backend_get_features") == 0) {
         return (void *)ggml_backend_cuda_get_features;
     }
+    if (strcmp(name, "ggml_backend_tensor_load") == 0) {
+        return (void *)ggml_backend_cuda_buffer_load_tensor;
+    }
     return nullptr;
 }
 

src/llama-model-loader.cpp

@@ -688,9 +688,13 @@ llama_model_loader::llama_model_loader(
 
     this->use_mmap = use_mmap;
     this->check_tensors = check_tensors;
-    ggml_backend_reg_t rpc_reg = ggml_backend_reg_by_name("RPC");
-    if (rpc_reg) {
-        rpc_load_tensor_fn = (ggml_backend_tensor_load_t) ggml_backend_reg_get_proc_address(rpc_reg, "ggml_backend_tensor_load");
+    ggml_backend_reg_t reg = ggml_backend_reg_by_name("RPC");
+    if (reg) {
+        rpc_load_tensor_fn = (ggml_backend_tensor_load_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_tensor_load");
+    }
+    reg = ggml_backend_reg_by_name("CUDA");
+    if (reg) {
+        cuda_load_tensor_fn = (ggml_backend_tensor_load_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_tensor_load");
     }
 }
 
@@ -895,15 +899,21 @@ void llama_model_loader::load_data_for(struct ggml_tensor * cur) const {
 }
 
 bool llama_model_loader::load_tensor(ggml_tensor * cur, const char * path, size_t file_offset, size_t tensor_offset, size_t size) {
-    if (!rpc_load_tensor_fn) {
-        return false;
-    }
     ggml_backend_buffer_t buf = cur->view_src ? cur->view_src->buffer : cur->buffer;
     const char * buf_name = ggml_backend_buffer_name(buf);
-    if (strncmp(buf_name, "RPC", 3) != 0) {
-        return false;
+    if (strncmp(buf_name, "RPC", 3) == 0) {
+        if (!rpc_load_tensor_fn) {
+            return false;
+        }
+        return rpc_load_tensor_fn(buf, cur, path, file_offset, tensor_offset, size);
+    }
+    if (strncmp(buf_name, "CUDA", 4) == 0) {
+        if (!cuda_load_tensor_fn) {
+            return false;
+        }
+        return cuda_load_tensor_fn(buf, cur, path, file_offset, tensor_offset, size);
     }
-    return rpc_load_tensor_fn(buf, cur, path, file_offset, tensor_offset, size);
+    return false;
 }
 
 bool llama_model_loader::load_all_data(

src/llama-model-loader.h

@@ -156,6 +156,7 @@ struct llama_model_loader {
     void load_data_for(struct ggml_tensor * cur) const;
 
     ggml_backend_tensor_load_t rpc_load_tensor_fn = nullptr;
+    ggml_backend_tensor_load_t cuda_load_tensor_fn = nullptr;
     bool load_tensor(ggml_tensor * cur, const char * path, size_t file_offset, size_t tensor_offset, size_t size);
 
     // Returns false if cancelled by progress_callback