🤖 AI Summary
This work proposes TiledAttention, the first schedulable scaled dot-product attention (SDPA) forward operator built on cuTile Python (TileIR), which bridges the gap between high performance and kernel-level customization flexibility in PyTorch. TiledAttention enables direct manipulation of tiling shapes, staging strategies, and shared memory layouts at the Python level without requiring modifications to CUDA or CUTLASS templates. By integrating online softmax computation, streaming K/V tile loading, and NVIDIA’s CUDA Tile technology, it is seamlessly exposed through a PyTorch-callable interface. Experimental results on DGX GB10 demonstrate that TiledAttention consistently outperforms the standard eager execution path across diverse sequence lengths, head dimensions, and FP16/BF16 precisions, delivering both high performance and reproducibility—making it well-suited for real-world research and deployment scenarios.
📝 Abstract
TiledAttention is a scaled dot-product attention (SDPA) forward operator for SDPA research on NVIDIA GPUs. Implemented in cuTile Python (TileIR) and exposed as a PyTorch-callable function, it is easier to modify than low-level CUDA templates while retaining realistic behavior via online softmax and tiled $K,V$ streaming. The approach is both performant and directly editable at the schedule level from Python (tile shapes, staging, shared-memory layout), enabling rapid, reproducible kernel research without template-heavy CUDA/CUTLASS rewrites. We benchmark TiledAttention on an NVIDIA DGX GB10 node with a reproducible harness and compare against PyTorch SDPA (auto-dispatch) and explicit unfused baselines across sequence length, head dimension, and precision (FP16/BF16). While production fused baselines remain stronger overall, TiledAttention delivers large speedups over standard eager attention paths and is available for direct use within PyTorch workflows, providing a practical balance between performance and customizability.