利用適度非結構化稀疏權重矩陣加速大語言模型的GPU推理
本文提出了一種針對大語言模型推理的高效GPU方法，採用三層矩陣存儲格式，包括稀疏張量核心層、插槽填充層和殘差層，實現稀疏矩陣乘法，在50%稀疏度下首次超越密集矩陣乘法，最高可達1.64倍內核加速和1.41倍端到端加速。
來源arXiv Machine Learning作者: Tao Lu, Haoyu Wang, Zonghui Wang, Keshen Xiang, Jiaheng Zhang, Wenzhi Chen
隨着大語言模型（LLM）的廣泛部署，推理成本已成為關鍵挑戰。剪枝技術通過引入權重矩陣的稀疏性來加速推理，但為了保持模型質量，通常只能將剪枝限制在適度非結構化稀疏度（約50%）範圍內。然而，在現有的GPU上，針對這種稀疏度的稀疏矩陣乘法（SpMM）內核均無法超越密集矩陣乘法的性能。本文提出了一種高效的GPU推理方法，專門針對具有適度非結構化稀疏度的LLM。
該方法的核心是一種創新的三層矩陣存儲格式：第一層是稀疏張量核心層（Sparse-TC），利用NVIDIA的稀疏張量核心加速SpMM；第二層是插槽填充層（Slot-Filling），採用並行差分距離對矩陣進行壓縮，並支持低成本的片上解碼；第三層是輕量級殘差層（Residual），確保SpMM計算的正確性。基於這一格式，作者設計了一個聯合利用稀疏張量核心和CUDA核心的SpMM內核，實現了高效的執行流水線，並重疊了片上計算與內存訪問。
在配備高帶寬內存（HBM）的現代GPU上進行的評估表明，該工作是首個在50%稀疏度下超越密集矩陣乘法的工作。與現有最佳方法相比，內核算加速比最高達到SpInfer（EuroSys'25最佳論文）的1.64倍，端到端加速比最高達到FlashLLM（VLDB'24）的1.41倍。該論文已被DAC 2026接收，源代碼已在GitHub上公開。這一成果有望顯著降低LLM推理成本，推動剪枝技術在實際部署中的應用。