修剪视觉的调度感知参差不齐的注意力 Transformer

Vision Transformer (ViTs) 的词元修剪方法承诺通过删除无信息的补丁来二次减少注意力失败。然而，标准的可变长度注意力 API（包括 FlashAttention-2 的 varlen 和 PyTorch 的 NestedTensor SDPA）无法将这些节省转化为 ViT 典型的短剪枝后序列长度（$\leq$197 词元）的比例挂钟增益。我们发现了一个调度开销瓶颈：在这些长度下，无论工作负载如何，主机端内核调度都会消耗 ${\sim}$50\,$μ$s，超过中等到高修剪率下的实际 GPU 计算时间。我们提出了一个轻量级双向 Triton 注意力内核，其调度底层为 ${\sim}$24\,$μ$s ——大约比 FlashAttention-2 varlen 低 2.17$\times$ ——允许修剪节省在挂钟时间内变得可见。集成到完整的打包-参与-解包管道中，并在 NVIDIA RTX 4000 Ada Generation GPU 上进行评估，我们的系统在标准 224$\times$224 输入下比填充的 PyTorch SDPA 实现了 1.88$\times$ 端到端吞吐量，在 384$\times$384 输入下扩展到 2.51$\times$。与 FlashAttention-2 varlen（最强基线）相比，我们的内核在服务批量大小 (BS=1-4) 下提供了 9-12\% 的更高吞吐量，并且在 80\% 词元修剪时将内核延迟降低了 2.17$\times$。数值正确性通过最大绝对 logits 差异 $<$0.004 和位精确的 top-1 预测进行验证。