Merge b3187 + Cuda Unified Memory support from matteoserva/master

b3187 + Cuda Unified Memory support
Nexesenex · Jun 20, 2024 · 1c307c6 · 1c307c6
2 parents 425c5a3 + fce75cc
commit 1c307c6
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 2 deletions.
diff --git a/examples/server/server.cpp b/examples/server/server.cpp
@@ -1594,7 +1594,7 @@ struct server_context {
                     } else {
                         std::string prompt;
                         if (task.data.contains("prompt") && task.data.at("prompt").is_string()) {
-                            json_value(task.data, "prompt", std::string());
+                            prompt = json_value(task.data, "prompt", std::string());
                         }
 
                         slot = get_available_slot(prompt);

diff --git a/ggml-cuda.cu b/ggml-cuda.cu
@@ -557,7 +557,14 @@ GGML_CALL static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_alloc_buffe
     size = std::max(size, (size_t)1); // cudaMalloc returns null for size 0
 
     void * dev_ptr;
-    cudaError_t err = ggml_cuda_device_malloc(&dev_ptr, size, buft_ctx->device);
+    cudaError_t err;
+    if (getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr)
+    {
+        err = cudaMallocManaged(&dev_ptr, size);
+    }
+    else {
+        err = ggml_cuda_device_malloc(&dev_ptr, size, buft_ctx->device);
+    }
     if (err != cudaSuccess) {
         // clear the error
         cudaGetLastError();

diff --git a/ggml-metal.m b/ggml-metal.m
@@ -735,6 +735,12 @@ static void ggml_metal_free(struct ggml_metal_context * ctx) {
 }
 
 static bool ggml_metal_supports_op(const struct ggml_metal_context * ctx, const struct ggml_tensor * op) {
+    for (size_t i = 0, n = 3; i < n; ++i) {
+        if (op->src[i] != NULL && op->src[i]->type == GGML_TYPE_BF16) {
+            return false;
+        }
+    }
+
     switch (op->op) {
         case GGML_OP_UNARY:
             switch (ggml_get_unary_op(op)) {