CHESS: Optimizing LLM Inference via Channel-Wise Thresholding and Selective Sparsification

TL;DR

CHESS tackles the challenge of edge-ready LLM inference by introducing a sparsity-aware activation pruning framework. It combines per-channel thresholding in FFNs with selective sparsification in attention modules and implements specialized sparse kernels to accelerate computation, explicitly modeling the sparsity-performance relationship. The approach yields lower accuracy degradation than prior methods while achieving up to 1.27x end-to-end speedups, demonstrated across multiple LLMs and downstream tasks. CHESS enables practical, training-free acceleration of LLMs on resource-constrained devices with broad applicability to edge deployment.

Abstract

Deploying large language models (LLMs) on edge devices presents significant challenges due to the substantial computational overhead and memory requirements. Activation sparsification can mitigate these resource challenges by reducing the number of activated neurons during inference. Existing methods typically employ thresholding-based sparsification based on the statistics of activation tensors. However, they do not model the impact of activation sparsification on performance, resulting in suboptimal performance degradation. To address the limitations, this paper reformulates the activation sparsification problem to explicitly capture the relationship between activation sparsity and model performance. Then, this paper proposes CHESS, a general activation sparsification approach via CHannel-wise thrEsholding and Selective Sparsification. First, channel-wise thresholding assigns a unique threshold to each activation channel in the feed-forward network (FFN) layers. Then, selective sparsification involves applying thresholding-based activation sparsification to specific layers within the attention modules. Finally, we detail the implementation of sparse kernels to accelerate LLM inference. Experimental results demonstrate that the proposed CHESS achieves lower performance degradation over eight downstream tasks while activating fewer parameters than existing methods, thus speeding up the LLM inference by up to 1.27x.

Figures (5)

• Figure 1: Distribution of absolute activation values $| a^{\text{up}}_i |$ across different inputs for various channels in the FFN of layer 16 of the Llama-3-8B model.
• Figure 2: Distribution of $\| W_{i,:} \|_2^2$ of different rows $i$ in attention projections of layer 16 of Llama-3-8B.
• Figure 3: End-to-end inference speedup
• Figure 4: Comparison between custom sparse kernels and PyTorch dense kernel on latency of linear projections
• Figure 5: Downstream performance and end-to-end speedups of each method under different sparsity levels.