Skip to main content

Mini FlashAttention

Key Insight

Writing a simplified FlashAttention kernel in Triton demonstrates how online softmax computation enables tiling of the attention mechanism. By calculating softmax scaling factors incrementally, we avoid materializing the massive intermediate attention matrix in global memory, keeping all intermediate data within the SM's fast SRAM. This project illustrates the fundamental design pattern of modern high-performance deep learning: transforming a memory-bound operation into a compute-bound one.