提升稀疏Transformer在GPU上的推理速度，支持灵活掩码与自动融合优化。
Accelerating Sparse Transformer Inference on GPU
- 基于分析建模，动态选择行级或块级计算核与存储格式。
- 多头注意力计算加速1.6倍，端到端推理提速1.4倍。
- 适合需要高效推理的稀疏模型部署场景。
大型语言模型（LLMs）因其强大的理解能力广受欢迎。作为其核心组件，通过并行化加速Transformer逐渐成为研究热点。掩码层引入稀疏性以减少计算量，但现有工作很少关注稀疏Transformer的性能优化。此外，当前静态算子融合方案难以适应多样应用场景。为此，我们提出STOF框架，支持稀疏Transformer的灵活掩码与算子融合优化。针对多头注意力（MHA）结构，STOF根据分析建模结果，将计算映射至行级或块级核，并采用独特存储格式。对于下游算子，STOF将其映射至编译模板，并通过两阶段搜索确定最优运行配置。实验表明，相比最先进方法，STOF在MHA计算中最高提速1.6倍，在端到端推理中最高提速1.4倍。
原文摘要 · Abstract (English)
Large language models (LLMs) are popular around the world due to their powerful understanding capabilities. As the core component of LLMs, accelerating Transformer through parallelization has gradually become a hot research topic. Mask layers introduce sparsity into Transformer to reduce calculations. However, previous works rarely focus on the performance optimization of sparse Transformer. In addition, current static operator fusion schemes fail to adapt to diverse application scenarios. To address the above problems, we propose STOF, a framework that incorporates optimizations for Sparse Transformer that enables flexible masking and Operator Fusion on GPU. For multi-head attention (MHA) structure, STOF maps the computation to row-wise or blockwise kernels with unique storage formats according to analytical modeling. For downstream operators, STOF maps the fusion scheme to compilation templates and determines the optimal running configuration through two-stage searching. The experimental results show that compared to the stateof-the-art work, STOF achieves maximum speedups of 1.6x in MHA computation and 1.4x in end-to-end inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。