Realizable N:M Sparse Transformer Inference via Search-Kernel Co-Design

2026-07-14

AI Summary

This work addresses the high inference latency of Vision Transformers (ViTs) and the limited practical acceleration achieved by existing N:M sparsity methods on modern GPUs. To overcome these challenges, the authors propose a hardware-software co-design approach. On the software side, they develop MD-SpMM, an adaptive parallel sparse kernel tailored for Tensor Cores that leverages micro-dense dataflow to enhance computational efficiency. On the algorithmic side, they introduce a three-stage heuristic search strategy to optimize layer-wise sparsity allocation under end-to-end latency constraints. This integrated method achieves the first significant real-world speedup for ViT and Swin inference without compromising model accuracy, delivering over 2.2× acceleration across multiple GPU platforms and attaining higher accuracy than prior approaches under identical latency budgets.

Abstract

Vision Transformers (ViTs) achieve strong accuracy but incur high inference latency. Semi-structured N:M sparsity can reduce arithmetic cost, yet its theoretical savings often fail to translate into proportional end-to-end speedups on modern GPUs. This mismatch arises because deployment latency depends not only on arithmetic reduction but also on execution regularity and hardware scheduling under sparsity. Achieving practical acceleration, therefore, requires coordinated design across sparse execution and sparsity configuration. To this end, we propose a hardware-software co-design framework for N:M sparse ViT inference. On the hardware side, we design MD-SpMM, an N:M sparse CUDA kernel that reorganizes sparse GEMM into micro-dense, Tensor-Core-aligned dataflow and uses inference-aware adaptive parallelism to sustain utilization. On the software side, we perform layer-wise sparsity search under explicit end-to-end latency budgets using a three-stage heuristic search with constraint relaxation to avoid premature convergence and enable deployment-aware sparsity allocation. Experiments on multiple ViT/Swin models and GPU platforms show that the framework achieves over 2.2x latency speedup while maintaining comparable accuracy and delivering superior accuracy under the same latency constraint. The source code is publicly available at https://github.com/liuganhuo/realizable-nm-sparse-transformer.

Problem

• N:M sparsity
• Vision Transformers
• inference latency
• sparse inference
• GPU acceleration

Innovation

• N:M sparsity
• hardware-software co-design
• sparse Transformer inference
• MD-SpMM
• deployment-aware sparsity search