TiledAttention提升注意力计算效率，支持快速定制和可复现测试。
TiledAttention: a CUDA Tile SDPA Kernel for PyTorch
- 基于TileIR的CUDA内核，可在Python中直接修改计算调度
- 相比标准注意力提速显著，在多种序列长度和精度下表现稳定
- 适合需要快速实验自定义注意力机制的研究者使用
TiledAttention 是一种用于 NVIDIA GPU 的缩放点积注意力（SDPA）前向算子，采用 cuTile Python（TileIR）实现并作为 PyTorch 可调用函数暴露。相较于低层 CUDA 模板，它更易修改，同时通过在线 softmax 与分块 K/V 流式传输保持真实行为。算法上遵循 FlashAttention 风格的在线 softmax 公式；创新点在于 cuTile/TileIR 实现策略、调度层级可修改性以及可复现的基准测试与性能分析流程。该方法在性能与可定制性间取得平衡，支持在 Python 中直接调整分块形状、数据搬运与共享内存布局，无需重写复杂 CUDA/CUTLASS 代码。我们在 NVIDIA DGX GB10 节点上进行可复现基准测试，对比 PyTorch SDPA（自动分派）、显式未融合基线（torch_sdpa_math、标准急切注意力）以及强制后端探测（FlashAttention2、EffecientAttention、CuDNN Attention），涵盖序列长度、头维度与精度（FP16/BF16）。尽管生产级融合基线仍更强，但 TiledAttention 在标准急切路径上获得显著加速，且可直接集成于 PyTorch 工作流，提供高效与灵活兼备的实验方案。
原文摘要 · Abstract (English)
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 (torch_sdpa_math, 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。