Skip to content

Commit

Permalink
Revert "[ATen][CUDA][CUBLAS] cublasLtMatmul increase workspace_size (p…
Browse files Browse the repository at this point in the history
  • Loading branch information
pytorchmergebot committed Mar 4, 2024
1 parent df3c8b8 commit 70c23a5
Showing 1 changed file with 1 addition and 8 deletions.
9 changes: 1 addition & 8 deletions aten/src/ATen/cuda/CUDABlas.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -183,20 +183,13 @@ uint32_t _getAlignment(uintptr_t address) {

static size_t _parseChosenWorkspaceSize() {
const char * val = getenv("CUBLASLT_WORKSPACE_SIZE");
size_t workspace_size = 1024;
#ifdef USE_ROCM
if (!val) {
// accept either env var
val = getenv("HIPBLASLT_WORKSPACE_SIZE");
}
#else
cudaDeviceProp* p = at::cuda::getDeviceProperties(c10::cuda::current_device());
if (p->major == 8) {
workspace_size = 4096;
} else if (p->major >= 9) {
workspace_size = 32768;
}
#endif
size_t workspace_size = 1024; /* default size in KiB according to #73328 */
if (val) {
try {
workspace_size = std::stoi(val);
Expand Down

0 comments on commit 70c23a5

Please sign in to comment.