A CUDA kernel built on
Deep-ML and published to the Kernel Hub.
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.
1from kernels import get_kernel
2
3kernel = get_kernel("Deep-ML/flash-attention-in-cuda-from-scratch")
This repo follows the
kernel-builder
layout. Build the binaries with
nix build . (or the kernel-builder Docker image) before loading.
