[控场AI]
· 8 分钟阅读· 4,191 字

让LLM更快更轻:重塑稀疏性以适配GPU的实战方案

让LLM更快更轻:重塑稀疏性以适配GPU的实战方案

核心问题:为什么"做更少的计算"反而更慢?

人类大脑的高效之处在于,它只激活完成特定思维所需的神经元。现代大语言模型(LLM)天然也具备这种特性——在前馈层中,对于任何给定的词,超过95%的神经元都处于沉默状态。然而,一个令人沮丧的悖论是:让模型做更少的数学运算,往往反而让它运行得更慢。

原因在于,非结构化稀疏性会引入不规则的内存访问模式,而GPU天生是为可预测的、密集的数据块运算而设计的。GPU的计算架构基于SIMT(单指令多线程)模型,其核心优势在于让数千个线程同时执行相同的操作。当数据以密集矩阵形式存在时,GPU可以通过合并内存访问(coalesced memory access)一次性读取连续的数据块,充分利用高带宽的HBM显存。然而,非结构化稀疏性意味着非零元素散布在矩阵的任意位置,GPU线程需要跳跃式地访问内存中不连续的地址,导致大量的缓存未命中和内存带宽浪费。更糟糕的是,稀疏矩阵还需要额外的索引数据结构来记录非零元素的位置,这些元数据本身也占用存储和带宽。因此,即使理论计算量减少了90%以上,实际执行时间可能仅减少10-20%甚至更差——稀疏矩阵虽然理论上计算量更少,但不规则的数据分布让GPU无法充分利用其并行计算能力,导致实际吞吐量大打折扣。

这正是Sakana AI与NVIDIA合作的最新ICML 2026论文——"Sparser, Faster, Lighter Transformer Language Models"——试图解决的核心问题。

论文概览

核心思路:让稀疏性适应GPU而非反过来

传统的稀疏加速方案通常试图修改GPU的执行逻辑来处理不规则的稀疏数据,但这往往事倍功半。本文的核心理念恰恰相反:重塑稀疏数据的格式,使其天然契合GPU的优化执行路径。

研究团队提出了一种"混合"格式,将稀疏性重新组织为GPU友好的结构。具体而言,该方案将99%的高度稀疏token通过一条快速路径处理,同时为那些罕见的"重型"token保留一个密集矩阵作为安全阀。在实际的LLM推理中,不同token的激活稀疏度存在显著差异——大多数token在前馈层中表现出极高的稀疏性,只有少量神经元被激活;但少数"重型"token(如句子开头的特殊标记、高信息密度的关键词等)可能激活大量神经元。如果强制所有token都走稀疏路径,这些重型token会因为格式中预留空间过大而导致严重的存储浪费和计算效率下降。混合策略通过设置稀疏度阈值,将token动态分流:高稀疏度token走快速路径享受稀疏加速,低稀疏度token走传统密集矩阵路径保证正确性和效率。这种设计类似于计算机体系结构中的"快速路径/慢速路径"模式,在保证鲁棒性的同时最大化平均性能,既充分利用了稀疏性带来的计算节省,又避免了不规则内存访问对GPU性能的拖累。

两大核心技术贡献

TwELL:面向GPU的新型稀疏打包格式

论文的第一个核心贡献是TwELL(Tile-wise ELLPACK),一种全新的稀疏数据打包格式。与传统稀疏格式不同,TwELL的设计目标是直接嵌入到已经高度优化的分块矩阵乘法(tiled matmul)内核中,而不会打断其执行流程。

ELLPACK本身是一种经典的稀疏矩阵存储格式,最早源于1980年代的ELLPACK数值计算软件包。其核心思想是:假设矩阵每行最多有K个非零元素,则用一个N×K的值数组和一个N×K的列索引数组来存储整个稀疏矩阵。相比更通用的CSR(压缩稀疏行)格式,ELLPACK的优势在于其规则的数据布局——每行占用固定的存储空间,这使得并行访问模式更加可预测,非常适合GPU的SIMD执行模型。然而,传统ELLPACK在行间非零元素数量差异较大时会造成严重的存储浪费。

TwELL在ELLPACK基础上进行了"分块化"改造——它按照GPU计算单元的tile大小(通常是16×16或32×32的子矩阵块)来组织稀疏数据,使得每个tile内部的稀疏模式独立管理。这样做的好处是,每个计算块都能以接近密集矩阵的效率被处理,既保留了ELLPACK的规则访问优势,又提高了存储效率。这种"以硬件为中心"的格式设计,是实现真正稀疏加速的关键。

要理解TwELL为何能无缝嵌入现有计算流程,需要了解分块矩阵乘法的工作原理。现代GPU上的高效矩阵运算将大矩阵分割为小的子块(tile),每个tile的大小恰好能装入GPU的共享内存或寄存器文件中。计算时,GPU的每个线程块负责一个输出tile的计算,通过反复从全局内存加载输入tile到共享内存,在片上完成乘累加操作,最后将结果写回全局内存。NVIDIA的Tensor Core进一步加速了这一过程,能在单个时钟周期内完成一个小矩阵的乘加运算。像cuBLAS和CUTLASS这样的库已经将分块矩阵乘法优化到接近硬件理论峰值的水平。TwELL的设计精妙之处在于,它将稀疏数据的组织方式与这些已经高度优化的tile计算流程对齐,使得稀疏内核可以复用密集内核的大部分优化策略。

定制CUDA内核:融合与压缩

第二个核心贡献是一套定制的CUDA内核,具备两个关键能力:

  • 融合多个稀疏矩阵乘法:通过将多个操作融合为单次内核调用,最大化GPU吞吐量,减少内核启动开销和中间数据的内存读写。
  • 压缩TwELL为混合表示:将稀疏格式进一步压缩为混合表示,最小化激活值的存储大小,从而显著降低内存占用。

在GPU编程中,每次内核启动都伴随着不可忽视的开销:CPU与GPU之间的同步、内核调度延迟、以及中间结果需要写回全局内存再被下一个内核读取的内存往返开销。内核融合是将多个逻辑上独立的操作合并为单个CUDA内核执行的优化技术——中间结果可以保留在寄存器或共享内存中,避免昂贵的全局内存读写;同时减少了内核启动次数,降低了调度开销。在Transformer模型中,前馈网络层通常包含多个矩阵乘法和激活函数,将这些操作融合可以显著提升吞吐量。本文的定制CUDA内核不仅融合了计算操作,还将稀疏格式的压缩和解压缩逻辑嵌入其中,实现了计算与数据管理的一体化。

这两项技术协同工作,使得稀疏Transformer在实际硬件上的运行效率得到了实质性提升。

实验结果:十亿参数规模的性能验证

研究团队使用这些内核在十亿参数规模的稀疏LLM上进行了训练和基准测试,结果令人振奋:

  • 推理和训练速度提升超过20%
  • 峰值内存占用显著降低
  • 能耗节省甚至高于速度提升的比例

在深度学习优化研究中,许多技术在小规模模型上表现优异,但在扩展到实际生产规模时会遇到各种问题:通信开销、内存碎片、数值稳定性等。本文选择在十亿参数(1B)规模进行验证具有重要意义——这一规模已经足够大,能够暴露出实际部署中的工程挑战,同时又是当前许多端侧和边缘部署场景的目标规模(如手机、嵌入式设备上的LLM)。20%以上的端到端加速意味着在相同硬件上可以服务更多用户请求,或者在相同延迟要求下使用更便宜的硬件,这对于LLM的商业化部署具有直接的经济价值。

有意思的是,这些加速不是在理论FLOPS层面的估算,而是在实际GPU硬件上测得的端到端性能提升。这意味着稀疏性的理论优势终于开始转化为真实的工程收益。

为什么这项工作值得关注?

打破稀疏性的"实用性困境"

长期以来,模型稀疏性一直是深度学习优化的"圣杯"之一。研究者们早就知道,训练好的神经网络中存在大量冗余参数和激活,但将这种稀疏性转化为实际的速度和内存收益,一直是一个未解难题。早在2018年,"彩票假说"(Lottery Ticket Hypothesis)就指出密集网络中存在稀疏子网络能够达到同等性能,但如何在不重新训练的情况下找到并高效执行这些子网络,始终缺乏实用的硬件支持方案。本文通过"格式-内核"协同设计的方式,为这个问题提供了一条可行路径。

软硬件协同设计的范式转变

这项工作体现了一个重要趋势:未来的AI加速不仅仅是算法层面的优化,更需要算法与硬件的深度协同。 Sakana AI与NVIDIA的合作本身就说明了这一点——只有理解GPU的底层执行模型,才能设计出真正高效的稀疏格式和内核。这种趋势在业界已有先例:Google的TPU针对矩阵运算进行了专门的架构设计,NVIDIA在Ampere架构中引入了2:4结构化稀疏支持(每4个元素中强制2个为零),而本文的工作则展示了如何在更高稀疏度(95%以上)下实现类似的硬件友好性。

开源生态贡献

论文配套发布了开源的GPU内核和数据格式实现(GitHub仓库),这为社区进一步探索稀疏LLM提供了坚实的基础设施。研究者和工程师可以直接在此基础上构建更高效的稀疏模型。

未来展望

随着LLM规模持续增长,推理成本和能耗问题日益突出。本文展示的"让稀疏性适应硬件"的思路,可能成为下一代高效LLM的重要技术方向。如果稀疏格式和内核能够进一步成熟,我们或许能在不牺牲模型能力的前提下,将LLM的运行成本降低一个数量级。

值得注意的是,这一方向与当前其他效率优化技术(如量化、知识蒸馏、推测解码等)并不冲突,而是可以叠加使用。例如,将TwELL格式与INT4量化结合,理论上可以同时获得稀疏性和低精度带来的双重加速。随着下一代GPU架构(如NVIDIA Blackwell及其后续产品)可能引入更灵活的稀疏计算支持,软件层面的稀疏格式创新将变得更加重要。

这项工作将在ICML 2026上正式发表,感兴趣的读者可以通过论文技术博客深入了解技术细节。

核心要点

分享:

相关推荐