ggml-cuda: fix ROCm multi-GPU illegal memory access in recurrent state restore

Cherry-pick of ggml-org#21170.

Remove early-return optimization in ggml_cuda_set_device() that caused
hipErrorIllegalAddress on ROCm multi-GPU setups when hipGetDevice()
returns an unexpected value on threads with uninitialized device context.
Also add missing ggml_cuda_set_device() in ggml_backend_cuda_set_tensor_async()
for consistency with all other cudaMemcpyAsync call sites.

Fixes GPU memory access fault during KV cache checkpoint restore on
multi-turn agentic workflows with multi-GPU row split.

Author: kmbandy

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

@@ -99,16 +99,9 @@ void ggml_cuda_error(const char * stmt, const char * func, const char * file, in
     GGML_ABORT(GGML_CUDA_NAME " error");
 }
 
-// this is faster on Windows
-// probably because the Windows CUDA libraries forget to make this check before invoking the drivers
+// always set device explicitly — early-return optimization is unsafe on ROCm multi-GPU
+// with uninitialized thread contexts (see https://github.com/ggml-org/llama.cpp/issues/21140)
 void ggml_cuda_set_device(int device) {
-    int current_device;
-    CUDA_CHECK(cudaGetDevice(&current_device));
-
-    if (device == current_device) {
-        return;
-    }
-
     CUDA_CHECK(cudaSetDevice(device));
 }
 
@@ -3022,6 +3015,7 @@ static void ggml_backend_cuda_set_tensor_async(ggml_backend_t backend, ggml_tens
 
     GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type");
 
+    ggml_cuda_set_device(cuda_ctx->device);
     CUDA_CHECK(cudaMemcpyAsync((char *) tensor->data + offset, data, size, cudaMemcpyHostToDevice, cuda_ctx->stream()));
 }