🤖 machine learning

Accelerating Sparse Transformer Inference on GPU

This paper presents STOF, a GPU framework that accelerates sparse Transformer inference by utilizing analytical modeling for efficient multi-head attention mapping and a two-stage search strategy to dynamically optimize operator fusion, achieving up to 1.6x and 1.4x speedups in MHA computation and end-to-end inference respectively.

Original authors: Wenhao Dai, Haodong Deng, Mengfei Rong, Xinyu Yang, Hongyu Liu, Fangxin Liu, Hailong Yang, Qianwen Cao, Qingxiao Sun

Published 2026-05-20

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Wenhao Dai, Haodong Deng, Mengfei Rong, Xinyu Yang, Hongyu Liu, Fangxin Liu, Hailong Yang, Qianwen Cao, Qingxiao Sun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to read a massive library of books (a Large Language Model) to answer a question. The library is organized into rooms called Transformers, and inside each room, there is a librarian (the Multi-Head Attention mechanism) who has to scan thousands of pages to find the specific sentences relevant to your question.

The problem is that for many questions, most of the pages are irrelevant. The librarian wastes time flipping through empty pages or pages that don't matter. This is where sparsity comes in: it's like putting "Do Not Read" stickers on the irrelevant pages.

However, the current librarians (existing software) are bad at using these stickers. They still walk past the "Do Not Read" pages, or they get confused when the stickers are placed in weird, random patterns. Furthermore, the library has other tasks (like summarizing or formatting) that are usually done separately, adding more walking time between tasks.

Enter STOF, a new system proposed by the researchers. Think of STOF as a super-efficient, smart library management system designed specifically for these "sparse" libraries. Here is how it works, broken down into simple parts:

1. The Smart Librarian (Unified MHA Kernels)

The researchers realized that different "Do Not Read" patterns need different strategies.

The Problem: Some patterns are neat rows of stickers (like a sliding window), while others are scattered randomly (like a lottery ticket). Old systems tried to use one "one-size-fits-all" method, which was slow.
The STOF Solution: STOF acts like a smart librarian who chooses the best tool for the job.
- If the stickers are in a neat, small cluster, the librarian uses a "Row-wise" approach: they grab a whole row of books at once and scan it quickly.
- If the stickers are scattered or the library is huge, they use a "Block-wise" approach: they break the books into small, manageable chunks and only open the specific chunks that have valid stickers.
The Result: By skipping the "Do Not Read" pages entirely instead of just ignoring them, the librarian works much faster.

2. The Assembly Line (Operator Fusion)

In a normal library, the librarian might finish reading, then walk to a different desk to summarize the text, then walk to another desk to format the answer. This walking (moving data between memory and the processor) is slow.

The Problem: Current systems often only combine simple tasks. They leave the heavy lifting (like complex math) for separate steps, causing traffic jams.
The STOF Solution: STOF builds a custom assembly line. It looks at the whole process and asks, "Can we combine these steps?"
- It doesn't just glue two simple tasks together; it figures out the perfect way to combine complex math tasks with formatting tasks.
- It uses a "search engine" to try out different ways of combining these tasks (like trying different assembly line layouts) to find the one that moves the fastest for the specific size of the library you are reading.

3. The Auto-Pilot (Hierarchical Search)

You can't manually design the perfect assembly line for every single book size and question type; there are too many combinations.

The STOF Solution: STOF has an Auto-Pilot that learns on the fly.
- Phase 1 (The Map): It looks at the library structure and draws a rough map of where the "Do Not Read" stickers are.
- Phase 2 (The Optimization): It runs a two-step search. First, it expands the assembly line boundaries to see how far it can go. Second, it fine-tunes the speed of the workers (parameters) based on how well the previous attempts worked.
- It remembers what worked (caching) so it doesn't waste time re-testing the same slow ideas.

The Results: How Much Faster?

The researchers tested STOF on powerful graphics cards (GPUs) using popular AI models (like BERT, GPT, and LLaMA).

Speed: Compared to the best existing methods, STOF made the core reading task (MHA) up to 1.6 times faster.
Overall Speed: When looking at the whole process of answering a question (end-to-end), it was up to 1.4 times faster.
Big Libraries: The bigger the library (longer text sequences), the more STOF shined, because it was skipping so much useless work.

Summary

Think of STOF as a system that stops the AI from wasting time reading pages it doesn't need to read, and stops it from walking back and forth between desks. It uses a smart, adaptive strategy to skip the junk and combine the useful steps into one smooth, fast motion. This makes AI models run significantly faster, especially when dealing with long or complex texts.

Technical Summary: Accelerating Sparse Transformer Inference on GPU

Problem Statement

Large Language Models (LLMs) rely heavily on the Transformer architecture, where Multi-Head Attention (MHA) is the core computational component. While sparsity-inducing mask layers (e.g., causal, sliding window, random) are widely used to reduce computational volume, they introduce significant performance challenges on GPUs:

Fragmented Computation: Sparse masks create fragmented data access patterns, exacerbating memory bandwidth bottlenecks and preventing efficient utilization of GPU resources.
Inflexible Optimization: Existing high-performance MHA implementations (e.g., FlashAttention) often support only specific, continuous masking patterns (like causal or sliding window). They struggle with discrete or unstructured sparsity (e.g., Bigbird, random attention) or fall back to inefficient methods like resetting score matrices post-computation.
Static Operator Fusion: Current operator fusion schemes in deep learning frameworks are largely static or rule-driven. They fail to adapt to diverse application scenarios, varying sequence lengths, and different model hyperparameters. Furthermore, the optimal configuration for fused operators often differs significantly from the sum of individually tuned operators, making sequential tuning suboptimal.

Methodology: The STOF Framework

The authors propose STOF, a framework designed to accelerate sparse Transformer inference on GPUs by integrating flexible masking support with adaptive operator fusion. STOF consists of two primary modules:

1. Unified MHA Module

This module addresses the challenge of efficiently computing sparse attention across diverse masking patterns.

Two-Level Storage Format: To handle arbitrary mask patterns (both structured and unstructured), STOF introduces a novel storage format combining Block Compressed Sparse Row (BSR) and bitmaps.
- Outer Tiles (OT): Represent global skipped blocks. "Full" OTs (all elements valid) and "Part" OTs (mixed valid/invalid) are tracked via pointer arrays.
- Inner Tiles (IT): Represent intra-block element distribution. Valid elements within "Part" OTs are encoded using 64-bit bitmaps (one uint64 per 8x8 block), allowing for precise masking without materializing the full sparse matrix.
Dual Kernel Strategy: STOF implements two distinct kernel types based on an analytical model:
- Row-wise Kernel: Optimized for highly sparse scenarios with short sequences. It slices the query ( $Q$ ) into rows, utilizes warp-level shuffles, and eliminates inter-warp synchronization to maximize data locality.
- Block-wise Kernel: A general-purpose kernel for broader scenarios. It partitions $Q$ , $K$ , and $V$ into sub-blocks, leveraging shared memory (SMEM) and asynchronous data copying to overlap memory transfers with computation.
Kernel Selection: An analytical model determines whether to use the row-wise or block-wise kernel based on mask sparsity and sequence length, ensuring the most efficient execution path.

2. Operator Fusion Module

This module addresses the need for adaptive fusion beyond the MHA layer.

Fusion Scheme Conversion: STOF converts fusion schemes into a numerical representation using hash encoding. Adjacent operators in the computational graph are mapped to binary arrays (representing fusion boundaries), which are then decoded into specific compilation templates (implemented via Triton and TileLang).
Hierarchical Search Engine: Instead of static rules, STOF employs a two-stage search mechanism to determine optimal fusion boundaries and kernel parameters:
1. Initialization: Uses neural hashing and predefined rules to identify frequent subgraph patterns and establish an initial fusion scheme.
2. Two-Stage Tuning:
  - Stage 1 (Fusion Expansion): Uses a reward-based algorithm (expand, seize, compete) to iteratively expand fusion boundaries, testing performance gains against a cache of previous results.
  - Stage 2 (Parameter Sampling): Once boundaries are fixed, it samples kernel parameters (e.g., block sizes, pipeline stages) to find the optimal configuration for the specific fusion scheme.

Key Contributions

Comprehensive Analysis: The authors analyze the impact of diverse masking patterns and inference configurations, exposing optimization opportunities in both MHA computation and downstream operator fusion.
Unified MHA Module: A novel module that implements row-wise and block-wise kernels with unique storage formats (BSR + bitmap) and optimizations, governed by an analytical selection model.
Adaptive Operator Fusion: A module that translates fusion schemes into compilation templates via numerical decoding and utilizes a two-stage search engine to determine optimal settings dynamically.
STOF Framework: A complete inference system that enables flexible masking patterns and optimal operator fusion on GPUs.

Experimental Results

The authors evaluated STOF on NVIDIA RTX 4090 (Ada) and A100 (Ampere) GPUs against state-of-the-art baselines including PyTorch Native, FlashAttention-2 (FA2), FlexAttention, ByteTransformer, MCFuser, and Bolt.

MHA Performance: STOF achieved maximum speedups of 1.6× over the best existing MHA implementations (e.g., FlexAttention) in MHA computation. It demonstrated consistent superiority across various masking patterns (causal, sliding window, Longformer, Bigbird) and sequence lengths (up to 4,096).
End-to-End Inference: STOF achieved maximum speedups of 1.4× in end-to-end inference across models like BERT, GPT-2, LLaMA, T5, and ViT.
Tuning Efficiency: The STOF tuning process was significantly faster than competitors (MCFuser and Bolt), being 6.7× to 6.9× faster in large-scale scenarios due to its reward-based sampling and caching mechanisms.
Scalability: In experiments with ultra-long sequences (up to 16k), STOF showed significant performance gains (up to 16.8× over PyTorch Compile) and better memory efficiency, avoiding Out-of-Memory (OOM) errors at sequence lengths where baselines failed.

Significance and Claims

The paper claims that STOF is the first system to simultaneously enable flexible masking patterns and diverse operator fusion schemes for sparse Transformer scenarios. Its significance lies in bridging the gap between hand-tuned, sparse-specific kernels and automatic, dense-operator fusion. By providing a complete stack that integrates custom MHA kernels with generative compilation templates, STOF establishes broader optimization opportunities that were previously unattainable with static or rule-driven approaches. The authors emphasize that their approach is extensible to emerging LLM architectures (e.g., Mixture-of-Experts) and newer GPU generations, offering a pathway to sustained performance improvements in large-scale model inference.