 
Alternatively, the original mask can be passed into the attn_mask field however due to the mentioned kernel constraints that would limit the implementation to only support the generic sdpa_math kernel.
Step 3 (Bonus): Faster matmuls with padding
On top of the performance improvements from SDPA, our analysis yielded a nice ancillary win.  In Andrej's words "The most dramatic optimization to nanoGPT so far (~25% speedup) is to simply increase the vocab size from 50257 to 50304 (nearest multiple of 64)."
