Are the requirements for using `torch.utils.cpp_extension` with CUDA documented anywhere?

Pyra: <p>When I install PyTorch via Conda with <code>conda install pytorch torchvision torchaudio pytorch-cuda=12.4 -c pytorch -c nvidia</code>, I am able to use CUDA devices out of the box, without installing anything at a system level. However, it seems that <code>torch.utils.cpp_extension.load_inline</code> will not work without some further setup. Through trial and error, I have found that I need to (A) install CUDA on the system outside of Conda, and (B) install <code>gcc_linux-64</code> and <code>gxx_linux-64</code> through Conda.</p>
<p>Without (A), I get various errors  when attempting to compile CUDA code with <code>cpp_extension.load_inline</code>, e.g. <code>No such file or directory &lt;nv/target&gt;</code> or missing header files when including Cub/Thrust. Without (B), I get the error <code>GLIBCXX_3.4.32' not found</code>.</p>
<p>Are these requirements correct, and if so are they documented anywhere (I couldn’t see anything specific in the docs besides needing to install <code>ninja</code>)? If it’s correct that CUDA needs to be installed at the system level, then this could be quite annoying in terms of version mismatches, and I don’t really see why this would be necessary for <code>cpp_extension</code> but not anything else. I did try installing <code>cuda-toolkit</code> and <code>cuda-cccl</code> via Conda, but to no avail.</p>

Pyra: <p>In case anyone comes across this in the future, I did eventually find out how to make <code>cpp_extension.load_inline</code> use the version of CUDA installed through Conda:</p>
<p>The issue is that <code>nvidia::cuda-toolkit</code> doesn’t have all the necessary headers, but <code>conda-forge::cuda-toolkit</code> does, so you need to install the conda-forge version instead (I have no idea why these are different). The conda-forge version will also result in <code>which nvcc</code> returning the version in your conda env, which enables <a href="https://github.com/pytorch/pytorch/blob/9425c0767d358e1db51176e8a66d6b63ac6b01be/torch/utils/cpp_extension.py#L94-L119" rel="noopener nofollow ugc">cpp_extension.py</a> to find it (so you don’t need to manually set the <code>CUDA_HOME</code> environment variable like you do with nvidia::cuda-toolkit).</p>

safelix: <p>My understanding from the <a href="https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#conda-installation" rel="noopener nofollow ugc">Conda Section in the CUDA Installation Guide</a> is that <a href="https://anaconda.org/nvidia/cuda" rel="noopener nofollow ugc"><code>nvidia::cuda</code></a> is the preferred way to install all available packages for native CUDA development.</p>
<p>I also had issues with missing <code>cuda_runtime.h</code> and <code>&lt;nv/target&gt;</code> for <code>nvidia::cuda==12.4</code>. But the newest <code>nvidia::cuda</code> also includes <code>nvidia::cuda-runtime</code> which solved all issues for me.</p>
<p>To compile against the same runtime as the pytorch installation, you can explicitly include its path:</p>
<pre data-code-wrap="python"><code class="lang-python">from nvidia.cuda_runtime.include import __path__ as CUDA_RUNTIME_INCLUDES # .../site-packages/nvidia/cuda_runtime/include/
</code></pre>
<p>This should work for <code>nvidia-cuda-runtime-cu12</code>, and other dependencies of PyTorch which are distributed as pip wheels (see Sec. 7 of CUDA Installation Guide).</p>

Pyra: <aside class="quote no-group" data-username="safelix" data-post="4" data-topic="211222">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/96bed5/48.png" class="avatar"> safelix:</div>
<blockquote>
<p><code>from nvidia.cuda_runtime.include import __path__ as CUDA_RUNTIME_INCLUDES # .../site-packages/nvidia/cuda_runtime/include/</code></p>
</blockquote>
</aside>
<p>Wow, thank you, I didn’t know about this!</p>