Performance of torch.bmm with different CUDA kernels

seba-1511: <p>Good evening,</p>
<p>When using <code>torch.bmm()</code> to multiply many (&gt;10k) small 3x3 matrices, we hit a performance bottleneck apparently due to cuBLAS heuristics when choosing which kernel to call. For example, the colab notebook below shows that for 2^15 matrices the call takes 2s but only 0.5s for 2^16 matrices.</p>
<p>What’s the easiest way to fix this, keeping in mind that we’d like to keep the differentiability via autograd ?</p>
<p>Minimal Example Colaboratory: <a href="https://colab.research.google.com/drive/1BsqH2xYe61gzC8C4YuGCGmnCUzhSwBcS" rel="nofollow noopener">https://colab.research.google.com/drive/1BsqH2xYe61gzC8C4YuGCGmnCUzhSwBcS</a></p>
<p>StackOverflow discussion on batched GEMMs: <a href="https://stackoverflow.com/questions/48519861/strange-cublas-gemm-batched-performance" rel="nofollow noopener">https://stackoverflow.com/questions/48519861/strange-cublas-gemm-batched-performance</a></p>
<p>Thanks,<br>
Séb</p>

tom: <p>Quite likely, some sort of <code>doing your own</code> is required. If neither <code>bmm</code> nor spelling out the contraction (for 10k x 3x3 that might be an option, probably not for very large ones) works for you, you might allocate the result array and fee batches through <code>bmm_out</code>.<br>
Another alternative could be to see if CUTLASS works for you. Matrix multiplication should be simple to implement the backward for (which just is a couple of matrix multiplications itself).</p>
<p>Best regards</p>
<p>Thomas</p>

seba-1511: <p>Thanks for your answer!</p>
<p>By spelling of the contraction, do you mean writing the operation as an einsum ?</p>
<p>Using another library and writing the backward seems like a somewhat pain-free solution too, but I was wondering if there would be a way to indicate to cuBLAS which kernel to call. In the end, I suspect we might have to go that route.</p>
<p>Cheers,<br>
Seb</p>

tom: <aside class="quote no-group" data-username="seba-1511" data-post="3" data-topic="50042">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/seba-1511/48/13883_2.png" class="avatar"> seba-1511:</div>
<blockquote>
<p>By spelling of the contraction, do you mean writing the operation as an einsum ?</p>
</blockquote>
</aside>
<p>No, einsum will itself use bmm, I thought of materializing the elementwise product and .sum.<br>
(I do have a branch somewhere that uses TensorIterators for einsum instead. It’s so terrible on CPU (no AVX) that I didn’t look at GPU, but if you want to benchmark it on GPU, I can push it.)</p>
<p>Best regards</p>
<p>Thomas</p>

ngimel: <p>For an example of using CUTLASS for batched matrix multiply look here <a href="https://github.com/mlperf/training_results_v0.6/blob/master/NVIDIA/benchmarks/transformer/implementations/pytorch/fairseq/modules/strided_batched_gemm/strided_batched_gemm_cuda.cu" rel="nofollow noopener">https://github.com/mlperf/training_results_v0.6/blob/master/NVIDIA/benchmarks/transformer/implementations/pytorch/fairseq/modules/strided_batched_gemm/strided_batched_gemm_cuda.cu</a>. It’s hardcoded to half data type, so some changes are required.<br>
Cublas allows specifying algorithms, but pytorch uses “default” as the algo, in the hopes that it would be optimal. Apparently it’s not, in your case.<br>
There’s also a new “matmul” interface in cublas (not yet integrated into pytorch) that gives somewhat finer control over selected algorithms, <a href="https://docs.nvidia.com/cuda/cublas/index.html#using-the-cublasLt-api" rel="nofollow noopener">https://docs.nvidia.com/cuda/cublas/index.html#using-the-cublasLt-api</a>, it also might be useful, but integrating it into pytorch is some work.</p>