TiledAttention: a CUDA Tile SDPA Kernel for PyTorch

TiledAttention is a new CUDA-based scaled dot-product attention kernel for PyTorch that enables easier modification of attention mechanisms for AI research. It provides a balance between performance and customizability, delivering significant speedups over standard attention implementations while remaining directly editable from Python.

• →TiledAttention offers a more accessible alternative to low-level CUDA templates for attention mechanism research on NVIDIA GPUs.
• →The implementation provides large speedups over standard eager attention paths while maintaining editability at the schedule level.
• →It supports online softmax and tiled K,V streaming for realistic behavior in attention computations.
• →The tool enables rapid, reproducible kernel research without requiring extensive CUDA/CUTLASS template rewrites.
• →Benchmarks show competitive performance against PyTorch SDPA auto-dispatch across various sequence lengths and precisions.