Tiny FlashAttention

WIP

A tiny flash attention implement in python, rust, cuda and c for learning purpose.

• python version
  • [x] naive pure python code
• triton version
  • [x] triton code
• [c version]
  • TODO: [naive pure c code]()
  • [x] naive cuda code standalone
  • [x] naive cuda code python binding
  • [x] cutlass cuda code
• [rust version]

cutlass cute flash attention in action

my env: cutlass v3.4, torch 1.14, cuda 12.4

• en tutorial
• zh tutorial