CUDA: fix scratch malloced on non-main device

JohannesGaessler · JohannesGaessler · commit 70ffed6b11b5 · 2023-09-16T21:48:29.000+02:00
diff --git a/ggml-cuda.cu b/ggml-cuda.cu
@@ -6969,6 +6969,9 @@ void ggml_cuda_assign_scratch_offset(struct ggml_tensor * tensor, size_t offset)
     if (g_scratch_size == 0) {
         return;
     }
+
+    ggml_cuda_set_device(g_main_device);
+
     if (g_scratch_buffer == nullptr) {
         CUDA_CHECK(cudaMalloc(&g_scratch_buffer, g_scratch_size));
     }

Original file line number	Diff line number	Diff line change
`@@ -6969,6 +6969,9 @@ void ggml_cuda_assign_scratch_offset(struct ggml_tensor * tensor, size_t offset)`
`6969`	`6969`	`if (g_scratch_size == 0) {`
`6970`	`6970`	`return;`
`6971`	`6971`	`}`
	`6972`	`+`
	`6973`	`+ ggml_cuda_set_device(g_main_device);`
	`6974`	`+`
`6972`	`6975`	`if (g_scratch_buffer == nullptr) {`
`6973`	`6976`	`CUDA_CHECK(cudaMalloc(&g_scratch_buffer, g_scratch_size));`
`6974`	`6977`	`}`