Accelerating GPU Inference of Large Language Models with Moderately Unstructured Sparse Weight Matrices

技术摘要：加速具有中度非结构化稀疏权重矩阵的大语言模型 GPU 推理

问题陈述
大语言模型（LLM）的部署面临着推理成本的关键瓶颈。虽然模型剪枝技术（如 Wanda、RIA）可以通过引入权重矩阵的稀疏性来降低计算负载，但这些技术受限于模型质量的要求。与小规模模型不同，LLM 对高稀疏度非常敏感；为了维持性能，剪枝通常被限制在中度非结构化稀疏（约 50%）水平。

在这一特定稀疏度水平下，现有的稀疏矩阵乘法（SpMM）GPU 内核无法超越稠密矩阵乘法库（如 cuBLAS）。这种性能差距源于三个根本性的挑战：

1. Tensor Core 不兼容性：现代 GPU 提供高效的稀疏 Tensor Core，但它们严格支持 2:4 结构化稀疏。它们无法直接加速不规则的、非结构化的剪枝模式。
2. 元数据开销：传统的存储格式（如 CSR）需要为每个非零元素存储完整的位次元数据。在 50% 稀疏度下，元数据占用的空间可能与数值本身一样多，从而抵消了带宽收益。
3. 片上解码瓶颈：旨在减少元数据的压缩技术（如 bitmap 编码）会将工作负载转移到片上解码。在拥有高带宽内存（HBM）的现代 GPU 上，全局内存访问无法与执行较慢的 CUDA 核心进行的解码工作完全重叠，从而产生性能停顿。

方法论
作者提出了一个整体解决方案，包括一种新型的三层矩阵存储格式，以及一个针对中度非结构化稀疏性进行协同优化的 SpMM 内核。

1. 三层存储格式
所提出的格式将稀疏矩阵分解为三个不同的层，以平衡硬件兼容性和存储效率：

• Sparse-TC 层：该层从不规则分布中提取 2:4 结构化稀疏模式。它保留每四个连续权重中的两个非零值。由于该模式是固定的，因此每个值仅需 2 bit 的元数据，确保了与稀疏 Tensor Core 的直接兼容性且无需解码开销。
• 插槽填充层（Slot-Filling Layer）：无法适应 2:4 模式的剩余非零元素被重新分配到 Sparse-TC 层留下的空闲插槽中。为了最小化这部分约 17.7% 元素的元数据，作者引入了**并行差分距离（PDD）**编码。PDD 不存储绝对索引，而是对重新分配元素的相对距离进行编码。大多数距离足够小，可以用 4 bit 进行编码。对于较大的间距，通过插入虚拟点（dummy points）来打破跨度，确保没有任何单个距离超过编码限制。
• 残差层（Residual Layer）：剩余的非零元素（通常小于总量的 1%）无法被前两层高效处理，因此采用传统的压缩稀疏行（CSR）格式存储。由于数据量极小，其元数据开销可以忽略不计。

2. 协同优化 SpMM 内核
基于此格式，作者设计了一个 GPU 内核，通过流水线执行模型共同利用稀疏 Tensor Core 和 CUDA 核心：

• 并行解码：CUDA 核心负责并行解码来自插槽填充层的 PDD 元数据。由于该层仅占总数据的极小部分，其解码开销得到了有效控制。
• 结构化加速：稀疏 Tensor Core 使用 wgmma 指令（Hopper 架构）处理 Sparse-TC 层，利用 2:4 结构实现高吞吐量。
• 残差处理：CUDA 核心还处理稀疏的残差层，并将结果写入一个单独的片上缓冲区，以便在流水线的末端与 Tensor Core 的输出进行归约（reduce）。
• 流水线重叠：该内核实现了一个细粒度的流水线，使三个关键操作相互重叠：
  • 加载稠密输入矩阵（缓存于 L2）和稀疏权重块（从全局内存读取）。
  • 在 CUDA 核心上解码 PDD 元数据。
  • 在稀疏 Tensor Core 上执行矩阵乘法。
    这种“按稀疏加载，按稀疏计算”的策略确保了片上计算与全局内存访问相重叠，从而最大化了 HBM 的利用率。

核心贡献

• 新型存储格式：专为中度非结构化稀疏性设计的三层格式，能够实现对稀疏 Tensor Core 的兼容性，通过 PDD 编码大幅减少元数据开销，并最小化片上解码成本。
• 协同优化内核：一个能高效协调稀疏 Tensor Core 和 CUDA 核心的 SpMM 内核，利用流水线技术实现内存访问与计算的重叠。
• 首次超越稠密计算：该实现首次证明，在配备 HBM 的现代 GPU 上，中度非结构化稀疏 SpMM 可以超越稠密矩阵乘法的性能。

结果
评估是在配备 NVIDIA H100 SXM5 GPU 的 OPT-30B 和 OPT-66B 模型上进行的，稀疏度水平分别为 50%、60% 和 70%。

• 内核级性能：在 50% 稀疏度下，该方法实现了高达 1.64× 于 SpInfer (EuroSys'25 Best Paper) 的加速，并显著优于 SparTA 和 FlashLLM 等其他基准。它也超越了稠密 cuBLAS 的性能。
• 端到端推理：相比 FlashLLM (VLDB'24)，该方法实现了高达 1.41× 的端到端加速；在 Batch Size 为 16 的 OPT-30B 模型上，吞吐量比 FlashLLM 高出 41.4%。
• 内存效率：在 50% 稀疏度下，该方法比稠密基准降低了约 21.4% 的峰值全局内存使用量，尽管由于缺乏用于元数据的 bitmap 编码，其内存使用量略高于 SpInfer。
• 可扩展性：该方法在 50% 到 70% 的稀疏度水平下均保持优势。

意义与主张
本文声称解决了 LLM 推理优化中的一个关键空白：现有稀疏内核无法有效利用中度非结构化稀疏性的问题。通过引入一种能够弥合不规则剪枝模式与结构化硬件加速之间鸿沟的存储格式，作者证明了剪枝可以在不牺牲模型质量的前提下产生实际的推理加速。

这项工作强调，虽然高稀疏度（>70%）会降低模型质量，但中度稀疏度（~50%）是 LLM 的“黄金分割点”。所提出的解决方案释放了这一黄金分割点的性能潜力，而此前由于现有稀疏内核的低效，这一潜力是无法实现的。作者指出，其优化专门针对 LLM 推理的解码阶段（该阶段 Batch Size 较小且受内存带宽限制），这使其非常适合在 Prefill（预填充）和 Decode（解码）阶段经常解耦的生产环境中进行大规模部署。源代码已公开，以促进进一步的研究与应用。