I wrote a tiny FlashAttention kernel in CUDA C++: ~250 lines, up to 4.5x faster than naive PyTorch

I built a small educational FlashAttention-style forward pass in CUDA C++.

Repo: https://github.com/lavawolfiee/mini-flash-attention

The goal was to make something much easier to read than the official highly optimized kernels, but still fast enough to be interesting.

There are two implementations:

• flash_attn_wmma_cuda.cu: ~150 lines, mostly plain CUDA + WMMA. Tensor Cores for Q @ K^T, blockwise online softmax, simpler P @ V.
• flash_attn_cuda.cu: ~250 lines, CuTe/CUTLASS version. Tensor Core MMA for both Q @ K^T and P @ V, register-resident accumulators, and swizzled shared-memory layouts.

Current scope:

• forward only
• fp16
• head dim 64
• non-causal attention
• input layout [B x H, N, D]

Benchmarked on RTX A4000, B=1, H=8, D=64.

Median latency:

So the CuTe version is up to ~4.5x faster than naive PyTorch on this setup, while not materializing the full N x N attention matrix.

Official FlashAttention is still much faster, of course, but that is kind of the point: the code is small enough to read, understand and play with.

This is also my first project using CuTe, so I'd really love some feedback from people who have written CUDA/CuTe kernels!