Reimplementing FlashAttention for performance and giggles