Skip to content

Commit 296c945

Browse files
committed
cuda : fix mul_mat_id with multi gpu
1 parent 33e50f1 commit 296c945

File tree

1 file changed

+9
-3
lines changed

1 file changed

+9
-3
lines changed

ggml-cuda.cu

+9-3
Original file line numberDiff line numberDiff line change
@@ -8361,11 +8361,16 @@ static void ggml_cuda_mul_mat_id(const ggml_tensor * src0, const ggml_tensor * s
83618361
src1_row.ne[1] = 1;
83628362
dst_row.ne[1] = 1;
83638363

8364-
if (src1->backend == GGML_BACKEND_GPU) {
8365-
src1_row.extra = &src1_row_extra;
8366-
}
8364+
src1_row.nb[2] = src1_row.nb[1];
8365+
dst_row.nb[2] = dst_row.nb[1];
8366+
8367+
src1_row.nb[3] = src1_row.nb[1];
8368+
dst_row.nb[3] = dst_row.nb[1];
8369+
8370+
src1_row.extra = &src1_row_extra;
83678371
dst_row.extra = &dst_row_extra;
83688372

8373+
83698374
for (int64_t i01 = 0; i01 < ids->ne[1]; i01++) {
83708375
//int32_t row_id;
83718376
//CUDA_CHECK(cudaMemcpyAsync(&row_id, ids_dev + i01*ids->nb[1] + id*ids->nb[0], sizeof(int32_t), cudaMemcpyDeviceToHost, g_cudaStreams[g_main_device][0]));
@@ -8381,6 +8386,7 @@ static void ggml_cuda_mul_mat_id(const ggml_tensor * src0, const ggml_tensor * s
83818386
src1_row.data = (char *) src1->data + i01*src1->nb[1];
83828387

83838388
dst_row_extra.data_device[g_main_device] = (char *) dst_extra->data_device[g_main_device] + i01*dst->nb[1];
8389+
dst_row.data = (char *) dst->data + i01*dst->nb[1];
83848390

83858391
ggml_cuda_mul_mat(src0_row, &src1_row, &dst_row);
83868392
}

0 commit comments

Comments
 (0)