For an example of using CUTLASS for batched matrix multiply look here 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 54. It’s hardcoded to half data type, so some changes are required.
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.
There’s also a new “matmul” interface in cublas (not yet integrated into pytorch) that gives somewhat finer control over selected algorithms, https://docs.nvidia.com/cuda/cublas/index.html#using-the-cublasLt-api 13, it also might be useful, but integrating it into pytorch is some work.