Toward Low-Latency and Memory-Efficient Deployment of Irregular Sparse Deep Learning Workloads
TL;DR
This work introduces a tile-aware scheduling framework for efficient sparse Vision Transformer execution on GPUs and introduces a training-aware extension that reuses the inference tile schedule and augments it with backward computation and activation-memory strategies.
Abstract
Our work introduces a tile-aware scheduling framework for efficient sparse Vision Transformer execution on GPUs. Sparse attention reduces the cost of high-resolution Vision Transformers, but its irregular masks produce blocks with diverse sizes, densities, and locations. Existing FlashAttention-style kernels rely on fixed tile configurations and cannot fully exploit these sparse patterns, leading to wasted computation and underutilized GPU resources. Our framework bridges this gap through four key techniques: sparse attention is represented as an adjacency matrix, structure-aware reordering algorithms improve locality, locally dense blocks are extracted as scheduling units, and offline profiling with integer linear programming (ILP) selects hardware-feasible tile assignments. Results show that our inference scheduler achieves up to 2.13 × end-to-end speedup over fixed-tile FlashAttention and up to 4.6 × speedup in high-resolution images. We further introduce a training-aware extension that reuses the inference tile schedule and augments it with backward computation and activation-memory strategies.