It seems a bad kernel is selected in the default setup by cudnn and you can use torch.backends.cudnn.benchmark = True to use the cudnn benchmark mode to select the fastest kernel.
In this mode the first iteration will be slower, as multiple algorithms will be executed to select the fastest one.
After setting torch.backends.cudnn.benchmark = True at the beginning of the script, I’ll get ~11.44ms.
The fast kernel seems to be already selected in the upcoming cudnn8 release.