Fix building with CUDA < 12.8 (#2782)

mlx/backend/cuda/allocator.cpp

@@ -119,7 +119,8 @@ void copy_to_managed(CudaBuffer& buf) {
   buf.data = new_data;
 }
 
-Buffer CudaAllocator::malloc_impl(size_t size, cudaStream_t stream) {
+Buffer
+CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
   if (size == 0) {
     return Buffer{new CudaBuffer{nullptr, 0, -1}};
   }
@@ -134,9 +135,8 @@ Buffer CudaAllocator::malloc_impl(size_t size, cudaStream_t stream) {
     size = page_size * ((size + page_size - 1) / page_size);
   }
 
-  int device = -1;
-  if (size > small_block_size && stream != nullptr) {
-    CHECK_CUDA_ERROR(cudaStreamGetDevice(stream, &device));
+  if (size <= small_block_size || stream == nullptr) {
+    device = -1;
   }
 
   CudaBuffer* buf = buffer_cache_.reuse_from_cache(size);
@@ -182,12 +182,8 @@ Buffer CudaAllocator::malloc_impl(size_t size, cudaStream_t stream) {
   return Buffer{buf};
 }
 
-Buffer CudaAllocator::malloc_async(size_t size, cudaStream_t stream) {
-  return malloc_impl(size, stream);
-}
-
 Buffer CudaAllocator::malloc(size_t size) {
-  return malloc_impl(size, nullptr);
+  return malloc_async(size, -1, nullptr);
 }
 
 void CudaAllocator::free(Buffer buffer) {
@@ -277,8 +273,9 @@ CudaAllocator& allocator() {
   return *allocator_;
 }
 
-Buffer malloc_async(size_t size, cudaStream_t stream) {
-  auto buffer = allocator().malloc_async(size, stream);
+Buffer malloc_async(size_t size, CommandEncoder& encoder) {
+  auto buffer = allocator().malloc_async(
+      size, encoder.device().cuda_device(), encoder.stream());
   if (size && !buffer.ptr()) {
     std::ostringstream msg;
     msg << "[malloc_async] Unable to allocate " << size << " bytes.";