Repository: ggml-org/llama.cpp
Pull request: Basic Vulkan Multi-GPU implementation

I'm not fully done, need to finish some details and check whether everything allocated gets cleaned properly, but it works already.

Right now it copies all data between them across RAM, which is not particularly fast, but I may be able to fix that in the future using Vulkan DeviceGroups. But maybe that only works between devices of the same vendor, I haven't tried it yet and the information I found about it is sparse.

This change also cleaned up most of the global variables the backend had, so it's a step towards allowing using it multiple times in the same program.

Edit: Forgot to mention, you have to set `GGML_VK_VISIBLE_DEVICES` to the indices of the devices you want, for example with `export GGML_VK_VISIBLE_DEVICES=0,1,2`. The indices correspond to the device order in `vulkaninfo --summary`. By default it will still only use the first device it finds.

<details>
<summary>Benchmarks</summary>

13B q6_k
```
ggml_vulkan: Using NVIDIA GeForce RTX 3090 | uma: 0 | fp16: 1 | warp size: 32
llama_print_timings: prompt eval time =    1962,96 ms /   622 tokens (    3,16 ms per token,   316,87 tokens per second)
llama_print_timings:        eval time =    5036,72 ms /   127 runs   (   39,66 ms per token,    25,21 tokens per second)
```

```
ggml_vulkan: Using NVIDIA GeForce RTX 3090 | uma: 0 | fp16: 1 | warp size: 32
ggml_vulkan: Using AMD Radeon Pro VII (RADV VEGA20) | uma: 0 | fp16: 1 | warp size: 64
llama_print_timings: prompt eval time =    2788,92 ms /   622 tokens (    4,48 ms per token,   223,03 tokens per second)
llama_print_timings:        eval time =    9158,65 ms /   127 runs   (   72,12 ms per token,    13,87 tokens per second)
```

```
ggml_vulkan: Using AMD Radeon Pro VII (RADV VEGA20) | uma: 0 | fp16: 1 | warp size: 64
ggml_vulkan: Using Intel(R) Arc(tm) A770 Graphics (DG2) | uma: 0 | fp16: 1 | warp size: 32
llama_print_timings: prompt eval time =    5597,99 ms /   622 tokens (    9,00 ms per token,   111,11 tokens per second)
llama_print_timings:        eval time =   15391,61 ms /   127 runs   (  121,19 ms per token,     8,25 tokens per second)
```

```
ggml_vulkan: Using NVIDIA GeForce RTX 3090 | uma: 0 | fp16: 1 | warp size: 32
ggml_vulkan: Using Intel(R) Arc(tm) A770 Graphics (DG2) | uma: 0 | fp16: 1 | warp size: 32
llama_print_timings: prompt eval time =    4132,89 ms /   622 tokens (    6,64 ms per token,   150,50 tokens per second)
llama_print_timings:        eval time =   12393,73 ms /   127 runs   (   97,59 ms per token,    10,25 tokens per second)
```

```
ggml_vulkan: Using NVIDIA GeForce RTX 3090 | uma: 0 | fp16: 1 | warp size: 32
ggml_vulkan: Using AMD Radeon Pro VII (RADV VEGA20) | uma: 0 | fp16: 1 | warp size: 64
ggml_vulkan: Using Intel(R) Arc(tm) A770 Graphics (DG2) | uma: 0 | fp16: 1 | warp size: 32
llama_print_timings: prompt eval time =    4338,94 ms /   622 tokens (    6,98 ms per token,   143,35 tokens per second)
llama_print_timings:        eval time =   14024,83 ms /   127 runs   (  110,43 ms per token,     9,06 tokens per second)
```

I can now run 70B q4_k_s across those three GPUs:

```
ggml_vulkan: Using NVIDIA GeForce RTX 3090 | uma: 0 | fp16: 1 | warp size: 32
ggml_vulkan: Using AMD Radeon Pro VII (RADV VEGA20) | uma: 0 | fp16: 1 | warp size: 64
ggml_vulkan: Using Intel(R) Arc(tm) A770 Graphics (DG2) | uma: 0 | fp16: 1 | warp size: 32
[...]
llm_load_tensors: offloading 80 repeating layers to GPU
llm_load_tensors: offloading non-repeating layers to GPU
llm_load_tensors: offloaded 81/81 layers to GPU
llm_load_tensors:        CPU buffer size =   140,62 MiB
llm_load_tensors:     Vulkan buffer size = 17160,81 MiB
llm_load_tensors:     Vulkan buffer size = 11512,00 MiB
llm_load_tensors:     Vulkan buffer size = 10689,80 MiB
....................................................................................................
llama_new_context_with_model: n_ctx      = 2048
llama_new_context_with_model: freq_base  = 10000,0
llama_new_context_with_model: freq_scale = 1
llama_kv_cache_init:     Vulkan KV buffer size =   280,00 MiB
llama_kv_cache_init:     Vulkan KV buffer size =   192,00 MiB
llama_kv_cache_init:     Vulkan KV buffer size =   168,00 MiB
llama_new_context_with_model: KV self size  =  640,00 MiB, K (f16):  320,00 MiB, V (f16):  320,00 MiB
llama_new_context_with_model: Vulkan_Host input buffer size   =    20,01 MiB
llama_new_context_with_model:     Vulkan compute buffer size =   338,80 MiB
llama_new_context_with_model:     Vulkan compute buffer size =   338,80 MiB
llama_new_context_with_model:     Vulkan compute buffer size =   338,80 MiB
llama_new_context_with_model: Vulkan_Host compute buffer size =    17,60 MiB
llama_new_context_with_model: graph splits (measure): 7
[...]
llama_print_timings: prompt eval time =   19742,11 ms /   622 tokens (   31,74 ms per token,    31,51 tokens per second)
llama_print_timings:        eval time =   38055,18 ms /   127 runs   (  299,65 ms per token,     3,34 tokens per second)
```

</details>

Changed files:
filename: common/common.cpp
patch: @@ -46,6 +46,10 @@
 #define GGML_USE_CUBLAS_SYCL
 #endif
 
+#if (defined(GGML_USE_CUBLAS) || defined(GGML_USE_SYCL)) || defined(GGML_USE_VULKAN)
+#define GGML_USE_CUBLAS_SYCL_VULKAN
+#endif
+
 int32_t get_num_physical_cores() {
 #ifdef __linux__
     // enumerate the set of thread siblings, num entries is num cores
@@ -660,8 +664,8 @@ bool gpt_params_parse_ex(int argc, char ** argv, gpt_params & params) {
                     params.tensor_split[i] = 0.0f;
                 }
             }
-#ifndef GGML_USE_CUBLAS_SYCL
-            fprintf(stderr, "warning: llama.cpp was compiled without cuBLAS/SYCL. Setting a tensor split has no effect.\n");
+#ifndef GGML_USE_CUBLAS_SYCL_VULKAN
+            fprintf(stderr, "warning: llama.cpp was compiled without cuBLAS/SYCL/Vulkan. Setting a tensor split has no effect.\n");
 #endif // GGML_USE_CUBLAS_SYCL
         } else if (arg == "--no-mmap") {
             params.use_mmap = false;

filename: ggml-vulkan.h
patch: @@ -8,24 +8,29 @@ extern "C" {
 #endif
 
 #define GGML_VK_NAME "Vulkan"
+#define GGML_VK_MAX_DEVICES 16
 
-GGML_API void ggml_vk_init(void);
+GGML_API void ggml_vk_init_cpu_assist(void);
 
-GGML_API void ggml_vk_preallocate_buffers_graph(struct ggml_tensor * node);
-GGML_API void ggml_vk_preallocate_buffers(void);
-GGML_API void ggml_vk_build_graph(struct ggml_tensor * node, bool last_node);
-GGML_API bool ggml_vk_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * tensor);
+GGML_API void ggml_vk_preallocate_buffers_graph_cpu_assist(struct ggml_tensor * node);
+GGML_API void ggml_vk_preallocate_buffers_cpu_assist(void);
+GGML_API void ggml_vk_build_graph_cpu_assist(struct ggml_tensor * node, bool last_node);
+GGML_API bool ggml_vk_compute_forward_cpu_assist(struct ggml_compute_params * params, struct ggml_tensor * tensor);
 #ifdef GGML_VULKAN_CHECK_RESULTS
-void ggml_vk_check_results_1(struct ggml_compute_params * params, struct ggml_tensor * tensor);
+void ggml_vk_check_results_1_cpu_assist(struct ggml_compute_params * params, struct ggml_tensor * tensor);
 #endif
-GGML_API void ggml_vk_graph_cleanup(void);
+GGML_API void ggml_vk_graph_cleanup_cpu_assist(void);
+GGML_API void ggml_vk_free_cpu_assist(void);
 
 // backend API
-GGML_API GGML_CALL ggml_backend_t ggml_backend_vk_init(void);
+GGML_API GGML_CALL ggml_backend_t ggml_backend_vk_init(size_t dev_num);
 
 GGML_API GGML_CALL bool ggml_backend_is_vk(ggml_backend_t backend);
+GGML_API GGML_CALL int  ggml_backend_vk_get_device_count(void);
+GGML_API GGML_CALL void ggml_backend_vk_get_device_description(int device, char * description, size_t description_size);
+GGML_API GGML_CALL void ggml_backend_vk_get_device_memory(int device, size_t * free, size_t * total);
 
-GGML_API GGML_CALL ggml_backend_buffer_type_t ggml_backend_vk_buffer_type(void);
+GGML_API GGML_CALL ggml_backend_buffer_type_t ggml_backend_vk_buffer_type(size_t dev_num);
 // pinned host buffer for use with the CPU backend for faster copies between CPU and GPU
 GGML_API GGML_CALL ggml_backend_buffer_type_t ggml_backend_vk_host_buffer_type(void);
 

filename: ggml.c
patch: @@ -2343,7 +2343,7 @@ struct ggml_context * ggml_init(struct ggml_init_params params) {
 #elif defined(GGML_USE_CLBLAST)
         ggml_cl_init();
 #elif defined(GGML_USE_VULKAN)
-        ggml_vk_init();
+        ggml_vk_init_cpu_assist();
 #elif defined(GGML_USE_SYCL)
         ggml_init_sycl();
 #endif
@@ -14850,10 +14850,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
     GGML_ASSERT(tensor->src[0] == NULL || tensor->src[0]->backend == GGML_BACKEND_CPU);
     GGML_ASSERT(tensor->src[1] == NULL || tensor->src[1]->backend == GGML_BACKEND_CPU);
 #elif defined(GGML_USE_VULKAN)
-    const bool skip_cpu = ggml_vk_compute_forward(params, tensor);
+    const bool skip_cpu = ggml_vk_compute_forward_cpu_assist(params, tensor);
 #ifdef GGML_VULKAN_CHECK_RESULTS
     if (skip_cpu) {
-        ggml_vk_check_results_1(params, tensor);
+        ggml_vk_check_results_1_cpu_assist(params, tensor);
     }
 #endif
     if (skip_cpu) {
@@ -17269,12 +17269,12 @@ int ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cplan * cplan) {
 
 #ifdef GGML_USE_VULKAN
     for (int i = 0; i < cgraph->n_nodes; i++) {
-        ggml_vk_preallocate_buffers_graph(cgraph->nodes[i]);
+        ggml_vk_preallocate_buffers_graph_cpu_assist(cgraph->nodes[i]);
     }
-    ggml_vk_preallocate_buffers();
+    ggml_vk_preallocate_buffers_cpu_assist();
 
     for (int i = 0; i < cgraph->n_nodes; i++) {
-        ggml_vk_build_graph(cgraph->nodes[i], i == cgraph->n_nodes - 1);
+        ggml_vk_build_graph_cpu_assist(cgraph->nodes[i], i == cgraph->n_nodes - 1);
     }
 #endif
 
@@ -17330,7 +17330,7 @@ int ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cplan * cplan) {
     }
 
 #ifdef GGML_USE_VULKAN
-    ggml_vk_graph_cleanup();
+    ggml_vk_graph_cleanup_cpu_assist();
 #endif
 
     // performance stats (graph)

filename: llama.cpp
patch: @@ -1355,7 +1355,7 @@ static ggml_backend_buffer_type_t llama_default_buffer_type_offload(int gpu) {
 #elif defined(GGML_USE_CUBLAS)
     buft = ggml_backend_cuda_buffer_type(gpu);
 #elif defined(GGML_USE_VULKAN)
-    buft = ggml_backend_vk_buffer_type();
+    buft = ggml_backend_vk_buffer_type(gpu);
 #elif defined(GGML_USE_SYCL)
     buft = ggml_backend_sycl_buffer_type(gpu);
 #elif defined(GGML_USE_CLBLAST)
@@ -1392,6 +1392,33 @@ static ggml_backend_buffer_type_t llama_default_buffer_type_split(int fallback_g
     GGML_UNUSED(tensor_split);
 }
 
+static size_t llama_get_device_count() {
+#if defined(GGML_USE_CUBLAS)
+    return ggml_backend_cuda_get_device_count();
+#elif defined(GGML_USE_VULKAN)
+    return ggml_backend_vk_get_device_count();
+#else
+    return 1;
+#endif
+}
+
+static size_t llama_get_device_memory(int device) {
+#if defined(GGML_USE_CUBLAS)
+    size_t total;
+    size_t free;
+    ggml_backend_cuda_get_device_memory(device, &total, &free);
+    return free;
+#elif defined(GGML_USE_VULKAN)
+    size_t total;
+    size_t free;
+    ggml_backend_vk_get_device_memory(device, &total, &free);
+    return free;
+#else
+    return 1;
+    GGML_UNUSED(device);
+#endif
+}
+
 //
 // globals
 //
@@ -1763,6 +1790,10 @@ struct llama_context {
             ggml_backend_free(backend);
         }
 
+#ifdef GGML_USE_VULKAN
+        ggml_vk_free_cpu_assist();
+#endif
+
         ggml_backend_buffer_free(buf_input);
         ggml_free(ctx_input);
     }
@@ -3436,22 +3467,18 @@ static bool llm_load_tensors(
         model.buft_layer[i] = llama_default_buffer_type_cpu(true);
     }
 
-#ifdef GGML_USE_CUBLAS
     if (split_mode == LLAMA_SPLIT_LAYER) {
         // calculate the split points
-        int device_count = ggml_backend_cuda_get_device_count();
+        int device_count = llama_get_device_count();
         bool all_zero = tensor_split == nullptr || std::all_of(tensor_split, tensor_split + device_count, [](float x) { return x == 0.0f; });
-        float splits[GGML_CUDA_MAX_DEVICES];
+        std::vector<float> splits(device_count);
         if (all_zero) {
             // default split, by free memory
             for (int i = 0; i < device_count; ++i) {
-                size_t total;
-                size_t free;
-                ggml_backend_cuda_get_device_memory(i, &total, &free);
-                splits[i] = free;
+                splits[i] = llama_get_device_memory(i);
             }
         } else {
-            std::copy(tensor_split, tensor_split + device_count, splits);
+            std::copy(tensor_split, tensor_split + device_count, splits.begin());
         }
 
         // sum and normalize the splits to get the split points
@@ -3467,19 +3494,17 @@ static bool llm_load_tensors(
         // assign the repeating layers to the devices according to the splits
         int act_gpu_layers = std::min(n_gpu_layers, (int)n_layer + 1);
         for (int64_t i = i_gpu_start; i < n_layer; ++i) {
-            int layer_gpu = std::upper_bound(splits, splits + device_count, float(i - i_gpu_start)/act_gpu_layers) - splits;
+            int layer_gpu = std::upper_bound(splits.begin(), splits.begin() + device_count, float(i - i_gpu_start)/act_gpu_layers) - splits.begin();
             model.buft_layer[i] = llama_default_buffer_type_offload(layer_gpu);
         }
         // assign the output layer
         if (n_gpu_layers > n_layer) {
-            int layer_gpu = std::upper_bound(splits, splits + device_count, float(act_gpu_layers - 1)/act_gpu_layers) - splits;
+            int layer_gpu = std::upper_bound(splits.begin(), splits.begin() + device_count, float(act_gpu_layers - 1)/act_gpu_layers) - splits.begin();
             model.buft_output = llama_default_buffer_type_offload(layer_gpu);
         } else {
             model.buft_output = llama_default_buffer_type_cpu(true);
         }
-    } else
-#endif
-    {
+    } else {
         ggml_backend_buffer_type_t split_buft;
         if (split_mode == LLAMA_SPLIT_ROW) {
             split_buft = llama_default_buffer_type_split(main_gpu, tensor_split);
@@ -10483,6 +10508,8 @@ size_t llama_max_devices(void) {
     return GGML_CUDA_MAX_DEVICES;
 #elif defined(GGML_USE_SYCL)
     return GGML_SYCL_MAX_DEVICES;
+#elif defined(GGML_USE_VULKAN)
+    return GGML_VK_MAX_DEVICES;
 #else
     return 1;
 #endif
@@ -10690,13 +10717,15 @@ struct llama_context * llama_new_context_with_model(
         }
 #elif defined(GGML_USE_VULKAN)
         if (model->n_gpu_layers > 0) {
-            ggml_backend_t backend = ggml_backend_vk_init();
-            if (backend == nullptr) {
-                LLAMA_LOG_ERROR("%s: failed to initialize Vulkan backend\n", __func__);
-                llama_free(ctx);
-                return nullptr;
+            for (int device = 0; device < ggml_backend_vk_get_device_count(); ++device) {
+                ggml_backend_t backend = ggml_backend_vk_init(device);
+                if (backend == nullptr) {
+                    LLAMA_LOG_ERROR("%s: failed to initialize Vulkan%d backend\n", __func__, device);
+                    llama_free(ctx);
+                    return nullptr;
+                }
+                ctx->backends.push_back(backend);
             }
-            ctx->backends.push_back(backend);
         }
 #elif defined(GGML_USE_SYCL)
         if (model->n_gpu_layers > 0) {