Build a tiled, IO-aware Flash Attention implementation in CUDA, starting from elementary GPU primitives and progressing to a fused online-softmax attention kernel. Along the way you implement a naive attention baseline, the online softmax math, and finish with a causal variant suitable for autoregressive models.
python scaffold.py- 1. vector_add
- 2. scale_array
- 3. elementwise_exp
- 4. row_max
- 5. row_sum
- 6. dot_product
- 7. matmul
- 8. transpose
- 9. qk_scores
- 10. softmax_rows
- 11. pv_matmul
- 12. naive_attention
- 13. online_max
- 14. correction_factor
- 15. update_running_sum
- 16. rescale_output
- 17. load_tile
- 18. tile_scores
- 19. tile_rowmax
- 20. tile_exp
- 21. tile_rowsum
- 22. accumulate_pv
- 23. flash_attention_kernel
- 24. flash_attention_launcher
- 25. causal_mask
- 26. flash_attention_causal_kernel
Built on Deep-ML.