A simplified implementation of FlashAttention in PyTorch for educational purpose.