Replacing single matmul with 2 einsums could be slower because of difference in speeds in GPU memory types. You have fast SRAM used for intermediate results and slow GPU memory for main storage. For low arithmetic intensity operations on V100, HBM2 can’t keep up with the compute units and becomes the bottleneck. For instance, you can’t do “torch.einsum(‘ni,nij->nj’, x, ww)” at nearly the same FLOPS as regular matmul or conv.