TiledAttention: a CUDA Tile SDPA Kernel for PyTorch
机器学习
2026-05-12 v2 人工智能
摘要
TiledAttention是面向NVIDIA GPU的缩放点积注意力(SDPA)前向算子实现。它基于cuTile Python(TileIR)实现，作为可调用的PyTorch函数暴露，比低级CUDA模板更易于修改，同时通过在线softmax和分块K,V流式处理保留真实行为。算法上，TiledAttention遵循既定的FlashAttention风格在线-softmax公式；我们的创新点在于cuTile/TileIR实现策略、调度层可修改性以及可重复的基准测试/分析工作流程。该方法在性能和可从Python(块状、存储、共享内存布局)调度层级直接编辑方面都具有优势，使其能够在无需模板密集型CUDA/CUTLASS重写的情况下，实现快速、可重复的内核研究。我们在NVIDIA DGX GB10节点上对TiledAttention进行基准测试，并使用可重复的测试平台与PyTorch SDPA(自动调度)、显式未融合基线(torch_sdpa_math、标准惰性注意力)以及强制后端探针(FlashAttention2、EfficientAttention、CuDNN Attention)进行比较。在序列长度、头维度和精度(FP16/BF16)方面，虽然生产融合基线总体更强，但TiledAttention在标准惰性注意力路径上实现了显著加速，并且可直接用于PyTorch工作流程中，提供了性能和可定制性之间的实用平衡。
引用
@article{arxiv.2603.01960,
title = {TiledAttention: a CUDA Tile SDPA Kernel for PyTorch},
author = {Taimur Khan},
journal= {arXiv preprint arXiv:2603.01960},
year = {2026}
}