Shfl-BW：基于感知 Tensor Core 的权重剪枝加速深度神经网络推理

深度神经网络（DNNs）中的权重剪枝可以降低存储和计算成本，但难以在模型推理时间上带来实际的加速。Tensor cores 能够显著提升 GPU 在密集计算上的吞吐量，但将 tensor cores 用于稀疏 DNNs 非常具有挑战性。与现有的 CUDA cores 相比，tensor cores 需要更高的数据复用率和矩阵形状的指令粒度，这两者都很难从稀疏 DNN 内核中获得。现有的剪枝方法无法平衡精度和效率的需求：随机稀疏性能够很好地保持模型质量，但无法利用 tensor core 加速；而高度结构化的块状稀疏性可以利用 tensor cores，但会遭受严重的精度损失。在这项工作中，我们提出了一种新的稀疏模式，即置换块状稀疏性（Shfl-BW），旨在高效利用 tensor cores 的同时最小化对权重结构的约束。我们的见解是，行和列置换为权重结构提供了丰富的灵活性，而使用我们的 GPU 内核设计引入的开销可以忽略不计。我们优化了线性层和卷积层中 Shfl-BW 的 GPU 内核。评估结果表明，我们的技术可以在 GPU 上实现最先进的速度-精度权衡。例如，在精度损失较小的情况下，我们可以在 75% 的稀疏度下，分别在 NVIDIA V100、T4 和 A100 GPU 上将 Transformer 的计算密集层加速 1.81、4.18 和 1.90 倍。