diff --git a/mlx/backend/cuda/cublas_utils.cpp b/mlx/backend/cuda/cublas_utils.cpp index 9a2717fd..108f56c8 100644 --- a/mlx/backend/cuda/cublas_utils.cpp +++ b/mlx/backend/cuda/cublas_utils.cpp @@ -40,21 +40,6 @@ cublasLtMatmulPreference_t get_preference(cu::Device& device) { return pref.pref_; } -void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size) { - if (workspace_size == 0) { - return nullptr; - } - - // Ensure workspace is 256-byte aligned - int nbytes = cuda::ceil_div(workspace_size, 256) * 256; - array workspace( - cu::malloc_async(nbytes, encoder), - {static_cast(workspace_size)}, - int8); - encoder.add_temporary(workspace); - return gpu_ptr(workspace); -} - cublasLtMatrixLayout_t create_matrix_layout( cudaDataType_t type, uint64_t rows, @@ -193,8 +178,7 @@ void CublasMatmulBase::execute_matmul( } } - void* workspace_ptr = - cublas_utils::allocate_workspace(encoder, heuristic_.workspaceSize); + void* workspace_ptr = allocate_workspace(encoder, heuristic_.workspaceSize); // Execute matmul auto capture = encoder.capture_context(); diff --git a/mlx/backend/cuda/cublas_utils.h b/mlx/backend/cuda/cublas_utils.h index 11f32283..371ceeb7 100644 --- a/mlx/backend/cuda/cublas_utils.h +++ b/mlx/backend/cuda/cublas_utils.h @@ -12,8 +12,6 @@ namespace cublas_utils { // Get the shared cublas preference for a device cublasLtMatmulPreference_t get_preference(cu::Device& device); -void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size); - cublasLtMatrixLayout_t create_matrix_layout( cudaDataType_t type, uint64_t rows, diff --git a/mlx/backend/cuda/cudnn_utils.cpp b/mlx/backend/cuda/cudnn_utils.cpp index 499e34a4..af2647f8 100644 --- a/mlx/backend/cuda/cudnn_utils.cpp +++ b/mlx/backend/cuda/cudnn_utils.cpp @@ -94,15 +94,7 @@ fe::error_t DnnGraph::encode_capturing( void* DnnGraph::prepare_workspace(cu::CommandEncoder& encoder) { int64_t workspace_size = 0; CHECK_CUDNN_FE_ERROR(get_workspace_size(workspace_size)); - if (workspace_size > 0) { - array workspace( - cu::malloc_async(workspace_size, encoder), - {static_cast(workspace_size)}, - uint8); - encoder.add_temporary(workspace); - return gpu_ptr(workspace); - } - return nullptr; + return allocate_workspace(encoder, workspace_size); } void DnnGraph::set_tensor_attrs( diff --git a/mlx/backend/cuda/utils.cpp b/mlx/backend/cuda/utils.cpp index 934f68ac..7640f903 100644 --- a/mlx/backend/cuda/utils.cpp +++ b/mlx/backend/cuda/utils.cpp @@ -93,4 +93,23 @@ CudaStream::CudaStream(cu::Device& device) { CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&handle_, cudaStreamNonBlocking)); } +void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size) { + if (workspace_size == 0) { + return nullptr; + } + + // Workspace allocation should not be captured. +#ifndef NDEBUG + cudaStreamCaptureStatus status; + CHECK_CUDA_ERROR(cudaStreamIsCapturing(encoder.stream(), &status)); + assert(status == cudaStreamCaptureStatusNone); +#endif + + // Ensure workspace is 256-byte aligned. + int nbytes = cuda::ceil_div(workspace_size, 256) * 256; + array workspace(cu::malloc_async(nbytes, encoder), {nbytes}, int8); + encoder.add_temporary(workspace); + return gpu_ptr(workspace); +} + } // namespace mlx::core diff --git a/mlx/backend/cuda/utils.h b/mlx/backend/cuda/utils.h index b060880b..a20b652d 100644 --- a/mlx/backend/cuda/utils.h +++ b/mlx/backend/cuda/utils.h @@ -43,4 +43,7 @@ struct Dtype; // Convert Dtype to CUDA C++ types. const char* dtype_to_cuda_type(const Dtype& dtype); +// Allocate an empty array and add it as temporary. +void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size); + } // namespace mlx::core