Revert "[ATen][CUDA][CUBLAS] cublasLtMatmul increase workspace_size (#120925)"

This reverts commit 0a38a6ac8046e4d3f9cfaba86b7ec6517038646f.

Reverted https://github.com/pytorch/pytorch/pull/120925 on behalf of https://github.com/clee2000 due to broke inductor models and caused accuracy regression on nightly dashboard 0a38a6ac80 https://github.com/pytorch/pytorch/actions/runs/8118465367/job/22193590228 ([comment](https://github.com/pytorch/pytorch/pull/120925#issuecomment-1977556485))
This commit is contained in:
PyTorch MergeBot
2024-03-04 22:13:23 +00:00
parent df3c8b8390
commit 70c23a51ac

View File

@ -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);