Accelerating GPU Inference of Large Language Models with Moderately Unstructured Sparse Weight Matrices

Technical Summary: Accelerating GPU Inference of Large Language Models with Moderately Unstructured Sparse Weight Matrices

Problem Statement
The deployment of Large Language Models (LLMs) faces a critical bottleneck in inference costs. While model pruning techniques (e.g., Wanda, RIA) can reduce computational load by introducing sparsity into weight matrices, they are constrained by model quality requirements. Unlike small-scale models, LLMs are highly sensitive to high sparsity levels; maintaining performance typically limits pruning to moderate unstructured sparsity (around 50%).

At this specific sparsity level, existing GPU kernels for Sparse Matrix Multiplication (SpMM) fail to outperform dense matrix multiplication libraries like cuBLAS. This performance gap stems from three fundamental challenges:

1. Tensor Core Incompatibility: Modern GPUs offer high-efficiency sparse tensor cores, but they strictly support structured 2:4 sparsity. They cannot directly accelerate irregular, unstructured pruning patterns.
2. Metadata Overhead: Traditional storage formats (e.g., CSR) require storing full positional metadata for every non-zero element. At 50% sparsity, this metadata can occupy as much space as the values themselves, negating memory bandwidth benefits.
3. On-Chip Decoding Bottlenecks: Compression techniques that reduce metadata (e.g., bitmap encoding) shift the workload to on-chip decoding. On modern GPUs with High-Bandwidth Memory (HBM), the global memory access cannot fully overlap with the decoding work performed by slower CUDA cores, creating a performance stall.

Methodology
The authors propose a holistic solution comprising a novel three-layer matrix storage format and a co-optimized SpMM kernel designed for moderate unstructured sparsity.

1. Three-Layer Storage Format
The proposed format decomposes the sparse matrix into three distinct layers to balance hardware compatibility and storage efficiency:

• Sparse-TC Layer: This layer extracts 2:4 structured sparsity patterns from the irregular distribution. It retains two non-zero values from every four consecutive weights. Because the pattern is fixed, it requires only 2 bits per value for metadata, ensuring direct compatibility with sparse tensor cores without decoding overhead.
• Slot-Filling Layer: Surplus non-zero elements that do not fit the 2:4 pattern are redistributed into the empty slots left by the Sparse-TC layer. To minimize metadata for these ~17.7% of elements, the authors introduce Parallel Differential Distance (PDD) encoding. Instead of storing absolute indices, PDD encodes the relative distances between redistributed elements. Most distances are small enough to be encoded in 4 bits. For larger gaps, dummy points are inserted to break the span, ensuring no single distance exceeds the encoding limit.
• Residual Layer: The remaining non-zero elements (typically <1% of the total) that cannot be efficiently handled by the first two layers are stored in a traditional Compressed Sparse Row (CSR) format. Due to the minimal volume of data, the metadata overhead here is negligible.

2. Co-Optimized SpMM Kernel
Building on this format, the authors design a GPU kernel that jointly utilizes sparse tensor cores and CUDA cores within a pipelined execution model:

• Parallel Decoding: CUDA cores are tasked with decoding the PDD metadata from the Slot-Filling layer in parallel. Since this layer represents a small fraction of the total data, the decoding overhead is managed effectively.
• Structured Acceleration: Sparse tensor cores process the Sparse-TC layer using wgmma instructions (Hopper architecture), leveraging the 2:4 structure for high throughput.
• Residual Handling: CUDA cores also process the sparse Residual layer, writing results to a separate on-chip buffer to be reduced with tensor core outputs at the end of the pipeline.
• Pipeline Overlap: The kernel implements a fine-grained pipeline that overlaps three critical operations:
  • Loading dense input matrices (cached in L2) and sparse weight tiles (from Global Memory).
  • Decoding PDD metadata on CUDA cores.
  • Executing matrix multiplication on sparse tensor cores.
    This "load-as-sparse, compute-as-sparse" strategy ensures that on-chip computation overlaps with global memory access, maximizing HBM utilization.

Key Contributions

• Novel Storage Format: A three-layer format tailored for moderate unstructured sparsity that enables sparse tensor core compatibility, drastically reduces metadata overhead via PDD encoding, and minimizes on-chip decoding costs.
• Co-Optimized Kernel: A SpMM kernel that efficiently coordinates sparse tensor cores and CUDA cores, utilizing pipelining to overlap memory access with computation.
• First-to-Beat-Dense: The implementation is the first to demonstrate that moderately unstructured sparse SpMM can outperform dense matrix multiplication on modern GPUs equipped with HBM.

Results
Evaluations were conducted on NVIDIA H100 SXM5 GPUs using OPT-30B and OPT-66B models with sparsity levels of 50%, 60%, and 70%.

• Kernel-Level Performance: At 50% sparsity, the proposed method achieves up to 1.64× speedup over SpInfer (EuroSys'25 Best Paper) and significantly outperforms other baselines like SparTA and FlashLLM. It also surpasses dense cuBLAS performance.
• End-to-End Inference: The method delivers up to 1.41× end-to-end speedup over FlashLLM (VLDB'24) and up to 41.4% higher throughput compared to FlashLLM on OPT-30B with a batch size of 16.
• Memory Efficiency: The approach reduces peak global memory usage by approximately 21.4% compared to dense baselines at 50% sparsity, though it incurs slightly higher memory usage than SpInfer due to the lack of bitmap encoding for metadata.
• Scalability: The method maintains advantages across 50% to 70% sparsity levels.

Significance and Claims
The paper claims to address a critical gap in LLM inference optimization: the inability of existing sparse kernels to leverage moderate unstructured sparsity effectively. By introducing a storage format that bridges the gap between irregular pruning patterns and structured hardware acceleration, the authors demonstrate that pruning can yield actual inference speedups without sacrificing model quality.

The work highlights that while high sparsity (>70%) degrades model quality, moderate sparsity (~50%) is the "sweet spot" for LLMs. The proposed solution unlocks the performance potential of this sweet spot, which was previously inaccessible due to the inefficiencies of existing sparse kernels. The authors note that their optimization is specifically targeted at the decode phase of LLM inference (where batch sizes are small and memory-bound), making it highly suitable for scalable deployment in production environments where prefill and decode phases are often decoupled. The source code is made available to facilitate further research and adoption.