
baseline (nanoGPT implementation):  ~143ms
sdpa_math (generic): ~134ms (6.71% faster)
mem_efficient kernel: ~119ms (20.16% faster)
flash_attention kernel: ~113ms (26.54% faster)
flash_attention + padded vocab:  ~87ms (64.37% faster)

All code was run on an 8 x NVIDIA Corporation A100 server with 80 GB HBM [A100 SXM4 80GB], and for the purpose of this experiment dropout was set to 0.
 
Figure 2: Using scaled dot product attention with custom kernels and torch.compile delivers significant speedups for training large language models, such as for nanoGPT shown here.
Enhancing Numerical Model Stability