LLM decoding sped up by pruning tuned to generated tokens
SparseDecoding: Decoding-Aware Pruning for Accurate and Efficient LLM Inference
Machine LearningComputation and Language
Summary
Decoding large language models takes time because the computer has to read many numbers from memory. Pruning, which means cutting out less important parts of the model, can speed this up, but usually it’s done using sample text different from what the model actually produces when guessing the next word. The authors found that this difference makes pruning less effective. They created a new method called SparseDecoding that prunes the model based on the data it generates itself, resulting in more accurate and faster language generation. They also built a new computer program to better handle the model’s math during decoding on GPUs.
What this means in practice
- •For ai infrastructure teams: Deploy faster large language model inference by pruning models based on their self-generated tokens to reduce memory bottlenecks during decoding.
- •For gpu software engineers: Implement efficient sparse matrix-vector kernels specialized for large language model decoding to gain practical speedups on GPU hardware.
Authors
Qitong Wang, Xinwei Niu, Mingluo Su, Shanwei Zhao, Shiai Zhu, Huan Wang
Abstract
The memory-bound nature of the decoding stage of large language model (LLM) inference incurs significant latency. Layer-wise training-free network pruning approaches guided by the Hessian have been a prominent solution to this problem, as pruning reduces the number of nonzero parameters read from memory during decoding. Nevertheless, typical methods in this line compute the Hessian using pre-collected natural sequences, whereas the model is fed self-generated tokens during decoding, creating a distribution shift between the two sequences. The Hessian calculated on the natural sequence is different from that calculated on the generated sequence. We observe that this discrepancy causes the activation distribution during generation to deviate from that used for pruning, further hurting the pruned model performance. Moreover, most existing LLM pruning methods that bring actual speedup primarily target the sparse matrix-matrix (SpMM) multiplication, providing limited support for the sparse matrix-vector (SpMV) operations, which dominate decoding. To solve these problems, we introduce SparseDecoding, a principled decoding-aware pruning framework tailored for accurate and efficient LLM decoding. Specifically, at the algorithmic axis, SparseDecoding constructs calibration matrices from layer-wise activations collected during the dense-model autoregressive generation, excluding prefill, thereby aligning the pruning objective with the decoding activations. At the system axis, we develop an optimized N:M sparse matrix-vector kernel with bitmask indexing and fixed-step traversal. Substantial empirical results on representative LLMs (Llama-3.1-8B, Llama-3.3-70B, Qwen3-14B / 32B) demonstrate that our method consistently outperforms standard fixed-text calibration on the long-form generation benchmarks while achieving up to 1.48x end-to-end wall-clock decoding speedup on A100 GPUs.