metalSPARSE：面向 Apple Silicon GPU 的 2:4 半结构化稀疏 GEMV 内核

View PDF HTML (experimental)

Abstract:Autoregressive Large Language Models (LLMs) generate text one token at a time, and each token reads every weight from memory, so on edge devices, decoding is limited by memory bandwidth. 2:4 semi-structured sparsity keeps two of every four weights, but Apple-silicon GPUs have no sparse matrix hardware, so pruning helps only if a kernel actually reads fewer bytes. We present metalSPARSE, to our knowledge the first 2:4 sparse kernels for Apple GPUs, together with a fused decoder written in Metal. The kernels stream weights in compressed form and expand them in registers. On Qwen2.5-1.5B with 2:4 sparsity on every decoder linear layer, decoding runs 3.6x faster than HuggingFace Transformers on PyTorch-MPS on an 8 GB M2 MacBook Air and 1.52x faster than MLX on a 16 GB M4. The same engine on dense weights matches MLX, so the gain comes from the compressed weights, and sparsity delivers 97% of the speedup its byte reduction allows. Prompt processing, which is compute-bound, does not speed up. To preserve quality, we combine a Wanda mask with block-wise reconstruction and low-rank adapters trained jointly with the surviving weights and quantized to 4 bits. On WikiText-2 this brings perplexity from 3.01x dense after one-shot pruning (Wanda with a closed-form refit) to 1.34x, and the whole pipeline runs on the same laptop/hardware.

Submission history

From: Tek Raj Awasthi [view email]
[v1] Tue, 27 Oct 2026 10:46:56 UTC (145 KB)