一句话：提出 ACSR 稀疏格式与 SPLAT 代码生成框架，为多样稀疏注意力模式在 GPU 上生成高性能实现。
多头自注意力（MHSA）在 NLP 与视觉任务上表现优异，但其对序列长度的二次依赖成为推理速度瓶颈。稀疏 MHSA 通过只计算部分注意力来缓解该问题，然而现有稀疏库与编译器受底层稀疏格式限制，难以兼顾通用性与性能：专用格式覆盖面窄，通用格式在稀疏 MHSA 常见的中等稀疏度（10%–50% 非零元）矩阵上元数据开销过大。本文提出仿射压缩稀疏行（affine-compressed-sparse-row, ACSR）格式及配套代码生成方案 SPLAT。其核心观察是常见稀疏 MHSA 模式具有独特的规则几何性质，可在即时（just-in-time）分析中提取，从而暴露新的优化与分块策略，SPLAT 据此为多样稀疏模式生成高性能 GPU 实现。在 A100 单精度下，SPLAT 为多种稀疏 MHSA 模型生成代码，相对 Triton 手写 kernel 最高加速 2.05 倍，相对 TVM 最高加速 4.05 倍。
Multi-head-self-attention (MHSA) mechanisms achieve state-of-the-art (SOTA) performance across natural language processing and vision tasks. However, their quadratic dependence on sequence lengths has bottlenecked inference speeds. To circumvent this bottleneck, researchers have proposed various sparse-MHSA models, where a subset of full attention is computed. Despite their promise, current sparse libraries and compilers do not support high-performance implementations for diverse sparse-MHSA patterns due to the underlying sparse formats they operate on. These formats are either too specialised, failing to cover a wide-range of sparse patterns, or too general, incurring high metadata overhead when computing on the moderately sparse (10-50% non-zeros) matrices present in sparse-MHSA. We bridge this gap, achieving both generality and performance, by proposing a novel sparse format: affine-compressed-sparse-row (ACSR) and supporting code-generation scheme, SPLAT, that generates highperformance implementations for diverse sparse-MHSA patterns on GPUs. Core to our proposed format and code generation algorithm is the observation that common sparse-MHSA patterns have uniquely regular geometric properties. These properties, which can be analyzed just-in-time, expose novel optimizations and tiling strategies that SPLAT exploits to generate high-performance implementations for diverse patterns. To demonstrate SPLAT’s efficacy, we use it to generate code for various sparse-MHSA models, achieving speedups of up-to 2.05x and 4.05x over hand-written kernels written in Triton and TVM respectively on A100 GPUs in single-precision.
@article{paperbot3192,
title = {SPLAT: A Framework for Optimised GPU Code-Generation for SParse reguLar ATtention},
author = {Ahan Gupta and Yueming Yuan and Devansh Jain and Yuhao Ge and David Aponte and Yanqi Zhou and Charith Mendis},
journal = {Proceedings of the ACM on Programming Languages},
volume = {9},
number = {OOPSLA1},
year = {2025},
doi = {10.1145/3720503}
}