MACKO: Sparse Matrix-Vector Multiplication for Low Sparsity

AI summary

Overview

Research area: Machine learning systems and GPU kernel design — specifically sparse matrix-vector multiplication (SpMV) for compressed, pruned Large Language Models.

Technical level: Intermediate. The motivation (making pruned LLMs fit and run faster on consumer GPUs) is accessible, but the paper's core content assumes familiarity with GPU memory hierarchies, warp-level programming, and sparse matrix storage formats such as CSR.

Scope: The paper introduces MACKO-SpMV, a new sparse storage format paired with a co-designed GPU kernel that delivers memory reduction and speedup over dense and existing sparse baselines at the low, unstructured sparsity levels (30–90%) typical of pruned LLMs.

What This Paper Is About

Pruning a neural network removes individual weights, leaving an irregular pattern of zeros (unstructured sparsity). In theory this should shrink models and speed up inference, but in practice existing sparse matrix-vector multiplication (SpMV) routines perform so poorly at the low sparsity levels real pruned LLMs exhibit (30–90%) that unstructured pruning yields only limited memory savings and speedup. MACKO addresses this by designing a storage format and a matching GPU kernel together, so that sparse computation can beat dense computation at sparsity as low as 20–30% without special hardware units such as tensor cores or format-specific precomputation.

Key Contributions

1. The MACKO storage format. A coordinate-compression format for sparse matrices targeting 30–90% sparsity, in which column indices are delta-encoded and constrained to a fixed bit width, with strategic zero-padding of values so that compressed coordinates and values stay mutually aligned. The paper provides a formal analysis of the format's storage overhead in the best, expected, and worst cases.

2. An optimized SpMV GPU kernel. A SplitK-style kernel adapted for matrix-vector multiplication, in which each matrix row is handled by one warp of 32 threads, with warp-shuffle primitives used for prefix sums and reductions, and reverse offset memory alignment (ROMA, borrowed from Sputnik) for vectorized loads.

3. Integration into PyTorch. The kernel is wrapped into the PyTorch framework with torch compilation enabled, and released open-source at github.com/vlejd/macko_spmv along with matrix conversion utilities.

4. Extensive empirical evaluation. SpMV runtime is measured across 31 matrix shapes derived from Llama2, Llama3, OPT, Qwen2, and Mixtral-8×7B on three consumer GPUs (RTX 2080 SUPER, RTX 3090, RTX 4090), plus an end-to-end Llama2-7B inference experiment pruned with Wanda.

Main Findings

• Beating dense at low sparsity: At 50% sparsity, MACKO is described as the first approach with significant 1.5× memory reduction and 1.2–1.5× speedup over the dense representation. On the RTX 4090 for a 12288×12288 matrix (cuBLAS runtime 336 μs), MACKO reaches 1.3× speedup at 50% sparsity, 1.96× at 70%, and 5.62× at 95%. Figure 1 reports sparse computation exceeding dense performance at sparsity as low as 25% for a 36864×12288 matrix, and the text states MACKO matches cuBLAS performance "often as low as 20% and consistently at 30%."

• Speedups over sparse baselines: 2.8–13.0× over cuSPARSE, 1.9–2.6× over Sputnik, and 2.2–2.5× over DASP. MACKO outperforms all baselines on all sparsity levels from 30% to 90% on all three tested GPUs. At 95% sparsity, MACKO is outperformed by Sputnik (and by cuSPARSE).

• Effective density advantage: MACKO has lower effective density (memory footprint relative to dense) than CSR formats throughout the region of interest and beats bitmask encoding at low density d < 0.23. To match the compression rate of 8-bit quantization, it is sufficient to reach 60% model sparsity with MACKO versus 83% with CSR.

• Unstructured sparsity becomes practically viable at 50%: Thanks to MACKO, unstructured pruning at 50% sparsity is argued to be justified in real-world LLM workloads — a barrier the paper describes as broken given that existing libraries require 2.6× fewer non-zeros to achieve the same performance.

• End-to-end Llama2-7B results: On Llama2-7B pruned with Wanda in unstructured mode, generating 100 tokens from an empty prompt on an RTX 4090, MACKO gives 1.53× memory reduction and 1.4× speedup at 50% sparsity (dense size 13.59 GB versus MACKO size 8.87 GB; dense speed 66.53 toks/sec versus sparse 98.60 toks/sec). The abstract separately reports 1.5× memory reduction and 1.5× faster inference at fp16 for the 50% case.

• Diminishing returns at 90%: At 90% sparsity, speedup and memory reduction fall short of what pure SpMV numbers would predict, because Wanda (like most pruning algorithms) does not prune the input embedding layer or the output embedding head. For Llama2-7B these layers total 262.144 MB — only 2% of the original model but 10% of the model pruned to 90% sparsity.

• Storage overhead is bounded: In the expected case, the padding term z/(1−z) (with z = (1−d)^{2^{b_Δ}}) becomes noticeable only at low density d < 20% and dominant only below d < 4%. As a result, unstructured sparse matrices match the memory footprint of dense representation at density as high as 0.78 (sparsity as low as 22%).

Methodology in Plain English

The authors start from a bandwidth argument. Using the Roofline model, they show that both dense and sparse matrix-vector multiplication are memory-bandwidth-bound on modern GPUs, because ops-per-byte ratios are much larger than the compute intensity of the operation (80 for the RTX 4090, 38 for the RTX 3090, 45 for the RTX 2080). That means the way to speed up SpMV is to shrink the memory footprint of the format.

They compare existing formats — CSR (as used by Sputnik), Tiled-CSL (Flash-LLM), and bitmask — and derive each one's "effective density" — how many bytes it takes relative to a dense matrix with the same values. CSR32/CSR16 wastes space because each nonzero carries a 32- or 16-bit column index; bitmask is compact but its access pattern is hostile to GPUs.

MACKO instead takes each row's column indices and delta-encodes them — storing the gaps between consecutive columns rather than the absolute positions — and fixes the delta width to b_Δ bits (the paper focuses on b_val = 16 and b_Δ = 4, though b_Δ ∈ {1, 2, 4, 8} are viable). When a gap exceeds the maximum representable delta (2^b_Δ), a padding zero is inserted so that every gap fits. Padding zeros cost value slots but save far more on the delta side. Values and deltas are packed into 8-bit words (8/b_Δ deltas per word) using bit shifts.

For execution, each matrix row is assigned to one warp of 32 threads, following the SplitK scheme. Each thread loads load_size = 8 elements per step; a warp loads 128 consecutive bytes of deltas (eight 4-bit deltas per thread) and 512 consecutive bytes of values (eight 16-bit values per thread), matching the GPU's 128-byte memory transactions. Reconstructing column indices requires an exclusive prefix sum across the warp, done with a binary reduction tree of depth log2(32) = 5 using __shfl_up_sync — register-only work that adds computation but no memory traffic, and since the kernel is bandwidth-bound this does not slow it down. Vectorized loading uses reverse offset memory alignment (ROMA) so each row's start is aligned to the vector width, with masking to avoid picking up values from the preceding row.

The runtime experiments use 1000 runs averaged with 100 warm-start iterations, invalidating the GPU cache before every run to model LLM inference; sparsity is swept from 0% to 95% "for completeness," though 30–90% is the region of interest. The end-to-end test compares dense cuBLAS against MACKO on Llama2-7B; because MACKO is loss-less, a fixed seed produces the same output up to floating-point rounding.

Why This Matters

Impact on research: The paper reframes the sparse-LLM bottleneck as one of effective density rather than raw nonzero count, and shows that a storage format can be designed to be simultaneously compact and hardware-friendly. It gives a concrete argument that 50% unstructured pruning — a regime widely treated as practically useless for speed — can now pay off, which may shift how pruning work evaluates itself. It also provides a bounded theoretical analysis of padding overhead, so the technique is not just an empirical trick.

Real-world applications:

• Running pruned LLMs locally on consumer GPUs, mobile devices, embedded systems, and IoT hardware, where memory capacity is the binding constraint — the paper notes Llama2-13B needs about 26 GB in FP16 while an RTX 4090 has only 24 GB.
• Local and on-device assistants where model weights must fit in a fixed memory budget.
• Combined compression pipelines, since MACKO targets weight pruning and is designed to be combinable with orthogonal system optimizations such as offloading (FlexGen, llama.cpp), attention kernels (FlashAttention), and graph restructuring (ByteTransformer, DeepSpeed).
• Serving stacks that decouple prefill and decode, since the decode stage is the one dominated by SpMV.

Industry relevance: The kernel is parameter-free and works out-of-the-box across matrix sizes, sparsity levels, and GPU models, which matters for deployment teams that cannot hand-tune per workload. It requires no tensor cores or vendor-specific hardware, is released open-source with PyTorch integration, and was benchmarked against three vendor and state-of-the-art libraries (cuBLAS, cuSPARSE, DASP, Sputnik) on consumer hardware rather than data-center accelerators. The paper frames its impact explicitly around democratization of LLMs on consumer hardware.

Future Directions

1. Hardware-specific optimization. The paper estimates a theoretical 1.6× speedup at 50% sparsity versus the 1.2–1.5× currently achieved, and notes that under an optimal setup for a specific GPU model it should be possible to reach cuBLAS-level memory bandwidth, yielding a speedup of 1/effd. Asynchronous memory loading, cache bypassing, explicit caching of V, and thread block size optimization are named as deferred, since they introduce tunable parameters.

2. Better performance above 95% sparsity. In this regime MACKO is outperformed by Sputnik, because value padding dominates effective density and bandwidth is further consumed by scattered access to the vector V. The paper notes b_Δ = 8 is a good alternative at very high sparsity, but that scattered V access remains the bottleneck.

3. Combining pruning with quantization. MACKO shows promising effective density at b_val = 8 and a viable density at b_val = 4, but the authors state more optimization is needed for efficient SpMV at these precisions.

4. Matrix-matrix multiplication with small batch dimensions and a wider range of GPU models. Both are listed as future work, which would extend the approach beyond the single-vector decode case to other inference regimes. The paper also notes that b_Δ = 2 is ideal for 15–45% sparsity in terms of storage but needs more kernel research, and that b_Δ = 1 shows limited practical applicability despite improving on dense memory at sparsity as low as 7%.

Target Audience

GPU kernel and systems researchers working on efficient LLM inference; practitioners who prune or compress language models and need the compression to actually translate into memory savings and speed; ML engineers deploying models on consumer or memory-constrained hardware; and anyone evaluating the practical value of unstructured versus structured sparsity. Readers without GPU programming background can follow the motivation, storage-format analysis, and benchmark results, but the kernel design section (SplitK, warp shuffling, prefix sums, ROMA) assumes comfort with CUDA execution-model concepts.

Authors’ abstract

Sparse Matrix-Vector Multiplication (SpMV) is a fundamental operation in the inference of sparse Large Language Models (LLMs). Because existing SpMV methods perform poorly under the low and unstructured sparsity (30-90%) commonly observed in pruned LLMs, unstructured pruning provided only limited memory reduction and speedup. We propose MACKO-SpMV, a GPU-optimized format and kernel co-designed to reduce storage overhead while preserving compatibility with the GPU's execution…