Repository: pytorch
Issue: test_workspace_allocation_error fails on my local devgpu
Labels: triaged,module: cuda graphs

### 🐛 Describe the bug

I get this error:

```
======================================================================
ERROR: test_workspace_allocation_error (__main__.CudaGraphTreeTests)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/data/users/ezyang/d/pytorch/test/inductor/test_cudagraph_trees.py", line 778, in test_workspace_allocation_error
    foo(*inps)
  File "/data/users/ezyang/d/pytorch/torch/_dynamo/eval_frame.py", line 292, in _fn
    return fn(*args, **kwargs)
  File "/data/users/ezyang/d/pytorch/test/inductor/test_cudagraph_trees.py", line 770, in foo
    @torch.compile()
  File "/data/users/ezyang/d/pytorch/torch/_dynamo/eval_frame.py", line 292, in _fn
    return fn(*args, **kwargs)
  File "/data/users/ezyang/d/pytorch/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
  File "/data/users/ezyang/d/pytorch/torch/_functorch/aot_autograd.py", line 3721, in forward
    return compiled_fn(full_args)
  File "/data/users/ezyang/d/pytorch/torch/_functorch/aot_autograd.py", line 1439, in g
    return f(*args)
  File "/data/users/ezyang/d/pytorch/torch/_functorch/aot_autograd.py", line 2394, in runtime_wrapper
    all_outs = call_func_with_args(
  File "/data/users/ezyang/d/pytorch/torch/_functorch/aot_autograd.py", line 1463, in call_func_with_args
    out = normalize_as_list(f(args))
  File "/data/users/ezyang/d/pytorch/torch/_functorch/aot_autograd.py", line 1548, in rng_functionalization_wrapper
    return compiled_fw(args)
  File "/data/users/ezyang/d/pytorch/torch/_inductor/compile_fx.py", line 454, in run
    return model(new_inputs)
  File "/data/users/ezyang/d/pytorch/torch/_inductor/compile_fx.py", line 496, in run
    return compiled_fn(new_inputs)
  File "/data/users/ezyang/d/pytorch/torch/_inductor/cudagraph_trees.py", line 360, in deferred_cudagraphify
    fn, out = cudagraphify(model, inputs, static_input_idxs, *args, **kwargs)
  File "/data/users/ezyang/d/pytorch/torch/_inductor/cudagraph_trees.py", line 384, in cudagraphify
    return manager.add_function(
  File "/data/users/ezyang/d/pytorch/torch/_inductor/cudagraph_trees.py", line 1856, in add_function
    return fn, fn(inputs)
  File "/data/users/ezyang/d/pytorch/torch/_inductor/cudagraph_trees.py", line 1676, in run
    out = self._run(new_inputs, function_id)
  File "/data/users/ezyang/d/pytorch/torch/_inductor/cudagraph_trees.py", line 1717, in _run
    return self.run_eager(new_inputs, function_id)
  File "/data/users/ezyang/d/pytorch/torch/_inductor/cudagraph_trees.py", line 1832, in run_eager
    return node.run(new_inputs)
  File "/data/users/ezyang/d/pytorch/torch/_inductor/cudagraph_trees.py", line 600, in run
    check_memory_pool(self.device_index, self.cuda_graphs_pool, new_storages)
  File "/data/users/ezyang/d/pytorch/torch/_inductor/cudagraph_trees.py", line 1543, in check_memory_pool
    raise RuntimeError(msg)
RuntimeError: These live storage data ptrs are in the cudagraph pool but not accounted for as an output of cudagraph trees:

Data Pointer: 140691082575872, history:
  File "??", line 0, in torch::unwind::unwind()
  File "??", line 0, in torch::CapturedTraceback::gather(bool, bool, bool)
  File "Module.cpp", line 0, in gather_with_cpp()
  File "CUDACachingAllocator.cpp", line 0, in c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::malloc(int, unsigned long, CUstream_st*)
  File "crtstuff.c", line 0, in c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::malloc(void**, int, unsigned long, CUstream_st*)
  File "crtstuff.c", line 0, in c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::allocate(unsigned long) const
  File "??", line 0, in at::cuda::getCurrentCUDABlasHandle()
  File "offloadstuff.c", line 0, in void at::cuda::blas::gemm<float>(char, char, long, long, long, at::OpMathType<float>::type, float const*, long, float const*, long, at::OpMathType<float>::type, float*, long)
  File "Blas.cpp", line 0, in at::native::(anonymous namespace)::addmm_out_cuda_impl(at::Tensor&, at::Tensor const&, at::Tensor const&, at::Tensor const&, c10::Scalar const&, c10::Scalar const&, at::native::(anonymous namespace)::Activation) [clone .isra.0]
  File "??", line 0, in at::native::structured_mm_out_cuda::impl(at::Tensor const&, at::Tensor const&, at::Tensor const&)
  File "RegisterCUDA.cpp", line 0, in at::(anonymous namespace)::wrapper_CUDA_mm_out_out(at::Tensor const&, at::Tensor const&, at::Tensor&)
  File "ADInplaceOrViewType_0.cpp", line 0, in c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor& (c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor&), &torch::ADInplaceOrView::(anonymous namespace)::mm_out_out>, at::Tensor&, c10::guts::typelist::typelist<c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor&> >, at::Tensor& (c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor&)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor&)
  File "VariableType_1.cpp", line 0, in torch::autograd::VariableType::(anonymous namespace)::mm_out_out(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, at::Tensor&)
  File "??", line 0, in at::_ops::mm_out::call(at::Tensor const&, at::Tensor const&, at::Tensor&)
  File "python_torch_functions_1.cpp", line 0, in torch::autograd::THPVariable_mm(_object*, _object*, _object*)
  File "/usr/local/src/conda/python-3.10.11/Objects/methodobject.c", line 543, in cfunction_call
  File "/usr/local/src/conda/python-3.10.11/Objects/call.c", line 305, in _PyObject_Call
  File "/usr/local/src/conda/python-3.10.11/Python/ceval.c", line 5917, in do_call_core
  File "/data/users/ezyang/d/pytorch/torch/utils/_device.py", line 76, in __torch_function__
    return func(*args, **kwargs)
  File "/usr/local/src/conda/python-3.10.11/Include/internal/pycore_ceval.h", line 46, in _PyEval_EvalFrame


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/ezyang/local/d/pytorch-env/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/data/users/ezyang/d/pytorch/test/inductor/test_cudagraph_trees.py", line 783, in test_workspace_allocation_error
    ).run(str(e))
RuntimeError: Expected to find "at::cuda::getNewWorkspace" but did not find it
Searched string:
These live storage data ptrs are in the cudagraph pool but not accounted for as an output of cudagraph trees:
~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

Data Pointer: 140691082575872, history:
From CHECK: at::cuda::getNewWorkspace


======================================================================
FAIL: test_workspace_allocation_error (__main__.CudaGraphTreeTests)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/data/users/ezyang/d/pytorch/test/inductor/test_cudagraph_trees.py", line 132, in tearDown
    self.assertEqual(all_live_block_count(), 0)
  File "/data/users/ezyang/d/pytorch/torch/testing/_internal/common_utils.py", line 3096, in assertEqual
    raise error_metas[0].to_error(
AssertionError: Scalars are not equal!

Expected 0 but got 1.
Absolute difference: 1
Relative difference: inf

----------------------------------------------------------------------
Ran 39 tests in 77.199s
```

Build config

```
$ python -c "import torch.__config__; print(torch.__config__.show())"
PyTorch built with:
  - GCC 11.3
  - C++ Version: 201703
  - Intel(R) oneAPI Math Kernel Library Version 2023.1-Product Build 20230303 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v2.7.3 (Git Hash 6dbeffbae1f23cbbeae17adb7b5b13f1f37c080e)
  - OpenMP 201511 (a.k.a. OpenMP 4.5)
  - LAPACK is enabled (usually provided by MKL)
  - NNPACK is enabled
  - CPU capability usage: AVX2
  - CUDA Runtime 12.0
  - NVCC architecture flags: -gencode;arch=compute_80,code=sm_80
  - Magma 2.6.1
  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=12.0, CXX_COMPILER=/usr/bin/c++, CXX_FLAGS= -D_GLIBCXX_USE_CXX11_ABI=1 -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOROCTRACER -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wall -Wextra -Werror=return-type -Werror=non-virtual-dtor -Werror=range-loop-construct -Werror=bool-operation -Wnarrowing -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-unused-parameter -Wno-unused-function -Wno-unused-result -Wno-strict-overflow -Wno-strict-aliasing -Wno-stringop-overflow -Wno-psabi -Wno-error=pedantic -Wno-error=old-style-cast -Wno-invalid-partial-specialization -Wno-unused-private-field -Wno-aligned-allocation-unavailable -Wno-missing-braces -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Werror=cast-function-type -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_DISABLE_GPU_ASSERTS=ON, TORCH_VERSION=2.1.0, USE_CUDA=ON, USE_CUDNN=OFF, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=ON, USE_NNPACK=ON, USE_OPENMP=ON, USE_ROCM=OFF,
```

cc @mcarilli @eellison 

### Versions

main