Weired cublas gemm kernel calling

nahso: <p>Hi,</p>
<p>The code fragment below is extracted from my project. When I profile the while project(the model has been export as a TorchScript, and I called it with libtorch) with <code>nsys</code>, it show me that the <code>lin1</code> here called <code>cublasSgemmStridedBatched</code> and  <code>ampere_sgemm_128x128_tt</code> kernel to perform gemm, which has a huge performance impact.</p>
<p>But when I ran this code fragment directly, <code>nsys</code> show me <code>cublasSgemm_v2</code> and <code>ampere_sgemm_128x64_tn</code> kernel was called. This kernel is much faster than <code>ampere_sgemm_128x128_tt</code>.</p>
<pre><code class="lang-auto">lin1 = Linear(128, 128, bias=False, device='cuda')
node_vector = torch.randn((27648, 128, 3), device='cuda') # pseudo input, same size as real input
v_tranpose = node_vector.transpose(1, 2)
v_u = lin1(v_tranpose).transpose(1, 2)
</code></pre>
<p>I wonder how does pytorch determine the kernel type? How to call the right kernel in my project?</p>

ptrblck: <p>Your current code will fail if the transposes are removed so could you show the code calling into different kernels using the “same” workload?</p>

nahso: <pre><code class="lang-auto">v_tranpose = node_vector.transpose(1, 2)
v_u = lin1(v_tranpose).transpose(1, 2)
</code></pre>
<p>is exactly the code calling into different kernels.</p>

ptrblck: <p>Your code shows a single linear layer, calling into a single kernel, which is <code>ampere_sgemm_128x64_tn</code> in my setup.</p>

nahso: <p>It will be <code>128x128_tt</code> kernel in my project, it is the code snippet in my project, where <code>lin1</code> and <code>lin2</code> are <code>Linear(128, 128, bias=False)</code>, <code>pu</code> and <code>po</code> are the wrapper of <code>nvtx_range_push</code> and <code>nvtx_range_pop</code></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/6/4643258b16b5b9f65ac1a7889b937e1fa0289239.png" data-download-href="https://discuss.pytorch.org/uploads/default/4643258b16b5b9f65ac1a7889b937e1fa0289239" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/6/4643258b16b5b9f65ac1a7889b937e1fa0289239_2_689x96.png" alt="image" data-base62-sha1="a1zhBALZrdti38h6mZQRk7EvC65" width="689" height="96" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/6/4643258b16b5b9f65ac1a7889b937e1fa0289239_2_689x96.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/6/4643258b16b5b9f65ac1a7889b937e1fa0289239_2_1033x144.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/6/4643258b16b5b9f65ac1a7889b937e1fa0289239.png 2x" data-dominant-color="F8F7F8"><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">1196×167 26.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>
<p>The nsys result shows:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/d/ed5c52a52c68d59df41ab2e25ab4ec857a42d943.png" data-download-href="https://discuss.pytorch.org/uploads/default/ed5c52a52c68d59df41ab2e25ab4ec857a42d943" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/d/ed5c52a52c68d59df41ab2e25ab4ec857a42d943_2_690x388.png" alt="image" data-base62-sha1="xRMXB3Du7mKL0VfruDipJ4E7Etl" width="690" height="388" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/d/ed5c52a52c68d59df41ab2e25ab4ec857a42d943_2_690x388.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/d/ed5c52a52c68d59df41ab2e25ab4ec857a42d943_2_1035x582.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/d/ed5c52a52c68d59df41ab2e25ab4ec857a42d943_2_1380x776.png 2x" data-dominant-color="B4CEBB"><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">1482×835 45.2 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>
<p>maybe there are too many complex operations before this Linear?</p>