Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fused3S: Fast Sparse Attention on Tensor Cores

Zitong Li, Aparna Chandramowlishwaran

TL;DR

This work targets the bottleneck of sparse attention by presenting Fused3S, a fully fused 3S kernel optimized for tensor cores. By introducing a Binary Sparse Block format and multi-level tiling, it fuses SDDMM, softmax, and SpMM into a single on-chip workflow with mixed-precision arithmetic and online softmax for numerical stability. The approach yields large end-to-end speedups on real-world graphs and diverse GPU architectures, significantly outperforming state-of-the-art 3S kernels and CPU/CUDA-core baselines. The results demonstrate practical impact for graph neural networks and graph transformers, enabling faster inference on large, irregular sparsity patterns while reducing memory traffic.

Abstract

Sparse attention is a core building block in many leading neural network models, from graph-structured learning to sparse sequence modeling. It can be decomposed into a sequence of three sparse matrix operations (3S): sampled dense-dense matrix multiplication (SDDMM), softmax normalization, and sparse matrix multiplication (SpMM). Efficiently executing the 3S computational pattern on modern GPUs remains challenging due to (a) the mismatch between unstructured sparsity and tensor cores optimized for dense operations, and (b) the high cost of data movement. Previous works have optimized these sparse operations individually or addressed one of these challenges. This paper introduces Fused3S, the first fused 3S algorithm that jointly maximizes tensor core utilization and minimizes data movement. Across real-world graph datasets, Fused3S achieves $1.6- 16.3\times$ and $1.5-14\times$ speedup over state-of-the-art on H100 and A30 GPUs. Furthermore, integrating Fused3S into Graph Transformer inference accelerates end-to-end performance by $1.05-5.36\times$, consistently outperforming all 3S baselines across diverse datasets (single and batched graphs) and GPU architectures.

Fused3S: Fast Sparse Attention on Tensor Cores

TL;DR

This work targets the bottleneck of sparse attention by presenting Fused3S, a fully fused 3S kernel optimized for tensor cores. By introducing a Binary Sparse Block format and multi-level tiling, it fuses SDDMM, softmax, and SpMM into a single on-chip workflow with mixed-precision arithmetic and online softmax for numerical stability. The approach yields large end-to-end speedups on real-world graphs and diverse GPU architectures, significantly outperforming state-of-the-art 3S kernels and CPU/CUDA-core baselines. The results demonstrate practical impact for graph neural networks and graph transformers, enabling faster inference on large, irregular sparsity patterns while reducing memory traffic.

Abstract

Sparse attention is a core building block in many leading neural network models, from graph-structured learning to sparse sequence modeling. It can be decomposed into a sequence of three sparse matrix operations (3S): sampled dense-dense matrix multiplication (SDDMM), softmax normalization, and sparse matrix multiplication (SpMM). Efficiently executing the 3S computational pattern on modern GPUs remains challenging due to (a) the mismatch between unstructured sparsity and tensor cores optimized for dense operations, and (b) the high cost of data movement. Previous works have optimized these sparse operations individually or addressed one of these challenges. This paper introduces Fused3S, the first fused 3S algorithm that jointly maximizes tensor core utilization and minimizes data movement. Across real-world graph datasets, Fused3S achieves and speedup over state-of-the-art on H100 and A30 GPUs. Furthermore, integrating Fused3S into Graph Transformer inference accelerates end-to-end performance by , consistently outperforming all 3S baselines across diverse datasets (single and batched graphs) and GPU architectures.
Paper Structure (17 sections, 7 equations, 8 figures, 7 tables, 2 algorithms)

This paper contains 17 sections, 7 equations, 8 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Computational Pattern in Sparse Attention
  4. Tensor Core and Operand Shapes
  5. Fused3S
  6. Sparse Format for Tensor Cores
  7. Fusion and Thread-block Parallelization
  8. Warp Partitioning Strategies
  9. Data Layout and Memory Accesses
  10. Mixed Precision and Stability
  11. Results
  12. Setup
  13. 3S Kernel Performance
  14. Fused3S Performance Breakdown
  15. End-to-end Model Performance
  16. ...and 2 more sections

Figures (8)

  • Figure 1: Binary Sparse Block (BSB) format. The TCB size in this example is $4\times2$ while in practice the size is larger (i.e., $16\times8$). Red boxes highlight how the first row window in compacted, tiled, and stored in BSB format.
  • Figure 2: Comparison of node-parallel (top) and edge-parallel (bottom) strategies. Different colored blocks are owned by different thread blocks. In edge-parallel, blocks of ${\bm{\mathbf{Q}}}$ or ${\bm{\mathbf{O}}}$ are shaded with multiple colors if shared by multiple thread blocks. The figure is divided vertically into three stages: (1) SDDMM, (2) data distribution of ${\bm{\mathbf{S}}}$ and ${\bm{\mathbf{E}}}$ for softmax, and (3) SpMM.
  • Figure 3: Work partitioning among warps within a thread block. Top: split-column (column blocks of ${\bm{\mathbf{K}}}^T$ and ${\bm{\mathbf{V}}}$ are divided among warps). Each warp independently computes a $r\times c$ tile of ${\bm{\mathbf{S}}}$ and ${\bm{\mathbf{O}}}$. Bottom: split-row (row blocks of ${\bm{\mathbf{K}}}^T$ and ${\bm{\mathbf{V}}}$ are divided among warps). All warps collaborate to compute each $r\times c$ tile of ${\bm{\mathbf{S}}}$ and ${\bm{\mathbf{O}}}$.
  • Figure 4: Register remapping in SDDMM (left) and SpMM (right). Top: original layouts. Bottom: permuted layouts.
  • Figure 5: 3S kernel performance on single graph datasets in Table \ref{['tab:datasets']}. Graphs are ordered by increasing number of edges (left to right). Y-axis is in log-scale.
  • ...and 3 more figures