WIDE：通过token级动态宽度剪枝提升自适应LLM推理

剪枝是提升大语言模型（LLM）效率的一种有前景的方法。现有的静态结构化剪枝方法对硬件友好并能带来实际的吞吐量提升，但其与输入无关的计算分配在高稀疏度下往往导致显著的精度下降。近期的动态稀疏性方法通过使计算适应单个输入来改善质量保持，但它们在很大程度上仍局限于粗粒度的结构决策，其在真实推理场景下的实际加速仍然具有挑战性。为应对这些挑战，我们提出了WIDE，这是首个端到端可微、同时面向prefill与decode两种场景设计的token级动态宽度剪枝框架。WIDE通过允许每个token动态选择注意力头组与FFN通道组实现细粒度计算分配，将动态剪枝从层级决策扩展到神经元块级粒度。通过两阶段训练流程，WIDE学习到有效的按token稀疏执行模式，并取得了远优于现有方法的质量保持效果。为使这种细粒度动态剪枝切实可行，我们进一步提出剪枝—内核协同设计框架，将动态稀疏加速分解为掩码重排序、硬件无关的块级跳过与硬件相关的块内跳过，从而在不同粒度上实现高效执行。在50%稀疏度下，在仅校准设置中与最先进的动态深度剪枝相比，WIDE带来55.1%的性能提升。在prefill与解码推理负载下，WIDE实现了接近理论值的内核级加速，prefill最高达1.98倍、解码最高达4.95倍，端到端加速分别为1.68倍与1.55倍。我们的代码已发布于 https://github.com/EIT-NLP/LLM-Pruning/tree/main/WIDE。