Step 1: Identify the existing SDPA implementation
In the case of nanoGPT, SDPA is implemented in the model’s CausalSelfAttention class. The original implementation at time of writing is adapted below for this post.
 
Step 2: Replace with Torch’s scaled_dot_product_attention
At this point we can note the following:

Lines 36 - 42 define the mathematical implementation of SDPA which we are replacing
The mask applied on line 39 is no longer relevant since we are using scaled_dot_product_attention’s is_causal flag.
The dropout layer used in line 41 is also now unnecessary. 

Swapping out the SDPA implementation for torch’s scaled_dot_product_attention and removing the now redundant code yields the following implementation.