mlsanityData to 7 Oct 2026

TiledAttention: a CUDA Tile SDPA Kernel for PyTorch

arXiv 2603.01960 ↗First posted 2 Mar 2026cs.LG cs.AIcited by 0

Not on the map. Why?

Abstract (arXiv)

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. Algorithmically, TiledAttention follows the established FlashAttention-style online-softmax formulation; our novelty is the cuTile/TileIR implementation strategy, schedule-level modifiability, and reproducible benchmarking/profiling workflow. 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), explicit unfused baselines (torchsdpamath, standard eager attention), and forced backend probes (FlashAttention2, EffecientAttention, CuDNN Attention) 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.

Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 12 of its 21 references that are in the map. First posted on arXiv 2 Mar 2026.

Builds on

12 of 21 references.

References · 12

  1. 2017Attention Is All You Needcited by 59,932
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
  3. 2021Learning Transferable Visual Models From Natural Language Supervisioncited by 33,456
  4. 2020Language Models are Few-Shot Learnerscited by 32,199
  5. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  6. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 12,865
  7. 2022PaLM: Scaling Language Modeling with Pathwayscited by 5,304
  8. 2014Neural Machine Translation by Jointly Learning to Align and Translatecited by 4,530
  9. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
  10. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
  11. 2020Efficient Transformers: A Surveycited by 784
  12. 2021A Survey of Transformerscited by 399

Built on by

Not cited yet.