Choice of torch.compile vs. triton

CDhere: <p>Hi, I’m new to <code>torch.compile</code> and the doc says</p>
<blockquote>
<p><code>torch.compile</code> is the latest method to speed up your PyTorch code! <code>torch.compile</code> makes PyTorch code run faster by JIT-compiling PyTorch code into optimized kernels, all while requiring minimal code changes.</p>
</blockquote>
<p>Does that mean if I use torch.compile on models/functions, it gives similar optimization of kernel fusion with <code>triton</code>?</p>
<p>What’s the difference between <code>torch.compile</code> and <code>triton</code>? Is either one sufficient for optimization, or can they be used orthogonally at the same time?</p>
<p>I’m new to all this kind of optimizations and super curious, any suggestions or references would be greatly appreciated! Thanks!</p>

marksaroufim: <p>On GPUs <code>torch.compile()</code> will apply various compiler optimizations with the most important ones being cuda graphs and fusions. Fusions specifically are done by code generating triton kernels so <code>torch.compile()</code> is actually a great starting point for writing triton kernels and you can always take that and try to write something faster yourself.</p>
<p>From a recent talk I gave</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/c/2cf008bd2a8bca80d6585c9570e7be3af197112c.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/2cf008bd2a8bca80d6585c9570e7be3af197112c" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/c/2cf008bd2a8bca80d6585c9570e7be3af197112c_2_690x392.jpeg" alt="image" data-base62-sha1="6pxhkgpfc3uv3VYTEpvFVm1QrlW" width="690" height="392" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/c/2cf008bd2a8bca80d6585c9570e7be3af197112c_2_690x392.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/c/2cf008bd2a8bca80d6585c9570e7be3af197112c_2_1035x588.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/c/2cf008bd2a8bca80d6585c9570e7be3af197112c.jpeg 2x" data-dominant-color="767777"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1119×637 91.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

CDhere: <p><a class="mention" href="/u/marksaroufim">@marksaroufim</a> Thank you! Would it be possible to share the talk?</p>
<blockquote>
<p><code>torch.compile()</code> is actually a great starting point for writing triton kernels and you can always take that and try to write something faster yourself.</p>
</blockquote>
<p>I don’t quite follow here: does that mean that <code>torch.compile()</code> can generate human-readable triton kernel files, and people can make manual tweaks from there and put them back to the model definition? My usage of <code>torch.compile</code> is simply wrapping a model with it during training.</p>

marksaroufim: <p>Yes that’s right, this is a link to the original talk <a href="https://www.youtube.com/watch?v=LuhJEEJQgUM&amp;t=1s">https://www.youtube.com/watch?v=LuhJEEJQgUM&amp;t=1s</a></p>
<p>It’s not that common for end users to modify the triton code but for pytorch compiler devs it’s quite useful</p>

CDhere: <p><a class="mention" href="/u/marksaroufim">@marksaroufim</a> May I ask to what extent <code>torch.compile</code> optimizes the computation for users? Given how powerful <code>torch.compile</code> is, should we always use it first before writing triton/CUDA kernels for extra efficiency?</p>
<p>I found many tutorials (triton, torch.compile, CUDA, etc) about making the model more efficient but just not sure if they are interchangeable.</p>

marksaroufim: <p>The main optimizations <code>torch.compile</code> will perform are</p>
<ol>
<li>Fusions</li>
<li>CUDA graphs</li>
<li>Template matching for SDPA (i.e: flash attention)</li>
</ol>
<p>It’s also an active project so I’m sure there’s more in flux. Can’t say what you should always do but personally yes I always reach for <code>torch.compile</code> first before attempting to write a custom kernel</p>

CDhere: <p>Thank you very much!</p>

vivy: <p>hi <a class="mention" href="/u/marksaroufim">@marksaroufim</a>. Why can’t <code>torch.compile</code> figure out the flash attention optimization part on it’s own? we have to use torch’s custom attention to get the flash attention speedup…</p>