use a correct device when creating OptionalCUDAGuard (vllm-project#2583)

Huarong · Jan 26, 2024 · 5265631 · 5265631
1 parent 2832e7b
commit 5265631
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu
@@ -34,7 +34,7 @@ void swap_blocks(
   char *dst_ptr = static_cast<char*>(dst.data_ptr());
 
   const int64_t block_size_in_bytes = src.element_size() * src[0].numel();
-  const at::cuda::OptionalCUDAGuard device_guard(src_device);
+  const at::cuda::OptionalCUDAGuard device_guard(src_device.is_cuda() ? src_device : dst_device);
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   // NOTE(woosuk): This can be slow if the number of blocks is large.
   for (const auto& pair : block_mapping) {