How to get the `shared memory capacity` for sm_89?

vgoklani: <p>Hey there,</p>
<p>Is there a pytorch method for fetching the <code>shared memory capacity</code> for the sm_89 series of chips (specifically the NVIDIA A6000 Ada cards)?</p>
<p>It’s my understanding that the sm_86 has a <code>shared memory capacity</code> size of ~100k, which caps the head-dimension to 64 for flash attention. The A100s (sm_80) have 164kB, which allows a head-dimension of 128. I’d like to check this value for sm_89, so I could figure out the max head-dimension.</p>
<p>Is there another name for the <code>shared memory capacity</code>, I don’t think it’s the same as the L1 cache. What should I look for in the Ampere/AdaLovelace/Hopper white papers?</p>
<p>reference: <a href="https://docs.nvidia.com/cuda/ampere-tuning-guide/index.html" class="inline-onebox" rel="noopener nofollow ugc">NVIDIA Ampere GPU Architecture Tuning Guide</a></p>
<p>One more thing: Is there a torch equivalent to this:</p>
<pre><code class="lang-python">    from numba.cuda.cudadrv import enums
    from numba import cuda

    device = cuda.get_current_device()
    attribs = [
        name.replace("CU_DEVICE_ATTRIBUTE_", "")
        for name in dir(enums)
        if name.startswith("CU_DEVICE_ATTRIBUTE_")
    ]
    for attr in attribs:
        print(attr, "=", getattr(device, attr))
</code></pre>
<p>I tried:</p>
<pre><code class="lang-python">torch.cuda.get_device_properties()
</code></pre>
<p>but it doesn’t return all the attributes. In the numba implementation, the device attributes are appended <code>lazily</code>, hence the iteration through the <code>enum</code>.</p>
<p>Thank you!</p>

eqy: <p>sm89 should have the same restriction as sm86, barring future improvements to flash attention implementations. You can find the documented shared memory capacity in the Ada tuning guide similar to the one you linked for Ampere here: <a href="https://docs.nvidia.com/cuda/ada-tuning-guide/index.html#streaming-multiprocessor" class="inline-onebox" rel="noopener nofollow ugc">NVIDIA Ada GPU Architecture Tuning Guide</a></p>
<p>This limitation is enforced in upstream PyTorch’s flash attention here: <a href="https://github.com/pytorch/pytorch/blob/8994d9e6109c541a1d581c383e4de9ed68205d91/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp#L504" class="inline-onebox" rel="noopener nofollow ugc">pytorch/sdp_utils.cpp at 8994d9e6109c541a1d581c383e4de9ed68205d91 · pytorch/pytorch · GitHub</a></p>

vgoklani: <p>Thanks <a class="mention" href="/u/eqy">@eqy</a> - but then how are they able to support a head_dim of 96 here:</p>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/98140#issuecomment-1531977277">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/98140#issuecomment-1531977277" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/98140#issuecomment-1531977277" target="_blank" rel="noopener nofollow ugc">Expected is_sm80 to be true, but got false on 2.0.0+cu118 and Nvidia 4090</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2023-04-01" data-time="13:10:26" data-timezone="UTC">01:10PM - 01 Apr 23 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2023-04-24" data-time="22:13:37" data-timezone="UTC">10:13PM - 24 Apr 23 UTC</span>
        </div>

      <div class="user">
        <a href="https://github.com/Nurffe" target="_blank" rel="noopener nofollow ugc">
          <img alt="Nurffe" src="https://avatars.githubusercontent.com/u/34690104?v=4" class="onebox-avatar-inline" width="20" height="20">
          Nurffe
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          high priority
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triage review
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: cuda
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          oncall: transformer/mha
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">### 🐛 Describe the bug

Similar to https://github.com/pytorch/pytorch/issues/9<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">4883

I'm trying to run textual inversion training using stable-diffusion with pytorch 2.0 using RTX 4090 and seeing `Expected is_sm80 to be true, but got false` which I believe should not trigger here?

`git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git &amp;&amp; git checkout a9fed7c364061ae6efb37f797b6b522cb3cf7aa2 `

modify webui-user.bat `COMMANDLINE_ARGS` and `TORCH_COMMAND` to use 2.0.0

```
set COMMANDLINE_ARGS= --opt-sdp-attention
set TORCH_COMMAND=pip install torch==2.0.0 torchvision --extra-index-url https://download.pytorch.org/whl/cu118
```

Run `webui-user.bat`

`Train -&gt; Train -&gt; Train embedding`

```
traceback (most recent call last):
  File "E:\git\stable-diffusion-webui\modules\textual_inversion\textual_inversion.py", line 503, in train_embedding
    scaler.scale(loss).backward()
  File "E:\git\stable-diffusion-webui\venv\lib\site-packages\torch\_tensor.py", line 487, in backward
    torch.autograd.backward(
  File "E:\git\stable-diffusion-webui\venv\lib\site-packages\torch\autograd\__init__.py", line 200, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "E:\git\stable-diffusion-webui\venv\lib\site-packages\torch\autograd\function.py", line 274, in apply
    return user_fn(self, *args)
  File "E:\git\stable-diffusion-webui\venv\lib\site-packages\torch\utils\checkpoint.py", line 157, in backward
    torch.autograd.backward(outputs_with_grad, args_with_grad)
  File "E:\git\stable-diffusion-webui\venv\lib\site-packages\torch\autograd\__init__.py", line 200, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Expected is_sm80 to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.)
```

```
python -c "import torch;print(torch.__config__.show(), torch.cuda.get_device_properties(0))"
PyTorch built with:
  - C++ Version: 199711
  - MSVC 193431937
  - Intel(R) Math Kernel Library Version 2020.0.2 Product Build 20200624 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v2.7.3 (Git Hash 6dbeffbae1f23cbbeae17adb7b5b13f1f37c080e)
  - OpenMP 2019
  - LAPACK is enabled (usually provided by MKL)
  - CPU capability usage: AVX2
  - CUDA Runtime 11.8
  - NVCC architecture flags: -gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_61,code=sm_61;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86;-gencode;arch=compute_90,code=sm_90;-gencode;arch=compute_37,code=compute_37
  - CuDNN 8.7
  - Magma 2.5.4
  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=11.8, CUDNN_VERSION=8.7.0, CXX_COMPILER=C:/actions-runner/_work/pytorch/pytorch/builder/windows/tmp_bin/sccache-cl.exe, CXX_FLAGS=/DWIN32 /D_WINDOWS /GR /EHsc /w /bigobj /FS -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOCUPTI -DLIBKINETO_NOROCTRACER -DUSE_FBGEMM -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_DISABLE_GPU_ASSERTS=OFF, TORCH_VERSION=2.0.0, USE_CUDA=ON, USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=OFF, USE_NNPACK=OFF, USE_OPENMP=ON, USE_ROCM=OFF,
 _CudaDeviceProperties(name='NVIDIA GeForce RTX 4090', major=8, minor=9, total_memory=24563MB, multi_processor_count=128)
```

I see that `CudaDeviceProperties` reports `major=8` and `minor=9` and NVCC architecture flags have `sml_80` so I assume check should not fail? Or am I misunderstanding something here?

### Versions

Collecting environment information...
PyTorch version: 2.0.0+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Pro
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.6 (tags/v3.10.6:9c7b4bd, Aug  1 2022, 21:53:49) [MSC v.1932 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19044-SP0
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4090
Nvidia driver version: 531.41
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=3900
DeviceID=CPU0
Family=107
L2CacheSize=4096
L2CacheSpeed=
Manufacturer=AuthenticAMD
MaxClockSpeed=3900
Name=AMD Ryzen 7 3800X 8-Core Processor
ProcessorType=3
Revision=28928

Versions of relevant libraries:
[pip3] numpy==1.23.3
[pip3] open-clip-torch==2.7.0
[pip3] pytorch-lightning==1.7.6
[pip3] torch==2.0.0+cu118
[pip3] torchdiffeq==0.2.3
[pip3] torchmetrics==0.11.4
[pip3] torchsde==0.2.5
[pip3] torchvision==0.15.1+cu118
[conda] Could not collect

cc @ezyang @gchanan @zou3519 @ngimel @jbschlosser @bhosmer @cpuhrsch @erichan1</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


eqy: <p>I believe that would qualify as an improvement to the implementation rather than relying on the current kernels.</p>