Performance Issue: torch.matmul selecting cutlass sm75 kernel for A100

torchlearner: <p>I have an application which does torch.matmul on large tensors. Typical dimensions for my use cases are like A (32, 3072), B (3072, 4_000_000), where 32 is the batch size M, 3072 is the embedding dimension K, and 4_000_000 is N. Inputs are all fp16 (half). Other dimensions are M 32, K 512, N 16M etc.</p>
<p>I am using torch 2.7. I see from a profile when run on a A100, that it is using cutlass sm75 kernels like cutlass_tensorop_f16_s1688gemm_f16_256x128_32x2_nn_align1 for ~65% of the profile. Remaining 35% of the time it is using cutlass sm80 kernels like cutlass_tensorop_f16_s16816gemm_f16_128x64_64x3_nt_align8. The profile is taken during a benchmark run of my app with 1000 requests. For a single request it either picks a sm75 kernel or a sm80 kernel. We choose the request tensors using torch.rand.</p>
<p>I am looking for some insights on the kernel selection logic for torch. Specifically, I see significant perf difference (latency, TFLOPS) by running the cutlass sm80 kernels directly. I am looking for ways in which I can influence torch to always select sm80 kernels on A100 and sm90 kernels on H100.</p>

ptrblck: <p>The kernels are selected by cuBLAS using their heuristics and not defined by PyTorch.</p>

torchlearner: <p>Thanks. Does torch control when cuBLAS is used vs cutlass ? In my profile with torch 2.7 I see that all the kernels are cutlass and not cublas. In a H100 node, I saw a combination of cutlass sm75, cutlass sm80 and a cuBlas sm90 kernel.</p>
<ol>
<li>Does torch control whether the backend is cuBLAS or cutlass or is cuBLAS controlling that ?</li>
<li>If torch is controlling that, any pointers to how torch chooses between cuBLAS and cutlass ?<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/6/a67fb7c04b461500648d1c05d050c524d2b75538.png" data-download-href="https://discuss.pytorch.org/uploads/default/a67fb7c04b461500648d1c05d050c524d2b75538" title="Screenshot 2025-06-09 at 1.07.07 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a67fb7c04b461500648d1c05d050c524d2b75538_2_690x46.png" alt="Screenshot 2025-06-09 at 1.07.07 PM" data-base62-sha1="nKUTrcalo9HdTZiiM1xMOYMv9cA" width="690" height="46" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a67fb7c04b461500648d1c05d050c524d2b75538_2_690x46.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a67fb7c04b461500648d1c05d050c524d2b75538_2_1035x69.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a67fb7c04b461500648d1c05d050c524d2b75538_2_1380x92.png 2x" data-dominant-color="E9E9EA"><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">Screenshot 2025-06-09 at 1.07.07 PM</span><span class="informations">2560×174 41.9 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><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/7/d77d92fb8b1412bb48a9d064f329baa21b0657eb.png" data-download-href="https://discuss.pytorch.org/uploads/default/d77d92fb8b1412bb48a9d064f329baa21b0657eb" title="Screenshot 2025-06-09 at 1.12.24 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d77d92fb8b1412bb48a9d064f329baa21b0657eb_2_690x56.png" alt="Screenshot 2025-06-09 at 1.12.24 PM" data-base62-sha1="uKjHfhqLbTkPm9HNJGtMts3bTir" width="690" height="56" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d77d92fb8b1412bb48a9d064f329baa21b0657eb_2_690x56.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d77d92fb8b1412bb48a9d064f329baa21b0657eb_2_1035x84.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/7/d77d92fb8b1412bb48a9d064f329baa21b0657eb_2_1380x112.png 2x" data-dominant-color="ECEDED"><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">Screenshot 2025-06-09 at 1.12.24 PM</span><span class="informations">2848×234 59.7 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></li>
</ol>

ptrblck: <p>cuBLAS is able to call kernels from different backends including CUTLASS.</p>

torchlearner: <p>Thanks for your reply.<br>
In this case, I figured out that the sm75 kernel was getting selected when the input tensor was not contiguous. If the tensor was contiguous it was always choosing the sm80 kernel.</p>