优化剪枝视觉变压器的注意力计算，显著提升推理速度。
Dispatch-Aware Ragged Attention for Pruned Vision Transformers
- 设计轻量级双向Triton注意力核，降低调度开销
- 在80%剪枝率下实现2.17倍更低的内核延迟
- 适合高吞吐量视觉模型部署与实时推理场景
视觉变换器（ViTs）的标记剪枝方法可通过丢弃无信息块，将注意力计算的浮点运算量减少近二次方。然而，标准可变长度注意力接口——包括FlashAttention-2的varlen和PyTorch的NestedTensor SDPA——在典型剪枝后序列长度（≤197个标记）下未能带来相应的时间收益。我们识别出调度开销是瓶颈：在该长度范围内，主机侧内核调度耗时约50μs，超过中高剪枝率下的实际GPU计算时间。为此，提出一种轻量级双向Triton注意力核，其调度下限约为24μs，较FlashAttention-2 varlen降低约2.17倍，使剪枝带来的性能优势可在实际运行时间中体现。集成至完整的打包-注意力-解包流水线，在NVIDIA RTX 4000 Ada Generation GPU上测试，相比填充后的PyTorch SDPA，在标准224×224输入下实现1.88倍端到端吞吐量，384×384下达2.51倍；相较最强基线FlashAttention-2 varlen，服务批量（BS=1-4）下吞吐量提高9-12%，80%剪枝率下内核延迟降低2.17倍。数值正确性通过最大绝对logit差异<0.004及比特精确的top-1预测验证。
原文摘要 · Abstract (English)
Token pruning methods for Vision Transformers (ViTs) promise quadratic reductions in attention FLOPs by dropping uninformative patches. Yet standard variable-length attention APIs -- including FlashAttention-2's varlen and PyTorch's NestedTensor SDPA -- fail to translate these savings into proportional wall-clock gains at the short post-pruning sequence lengths typical of ViTs ($\leq$197 tokens). We identify a dispatch-overhead bottleneck: at these lengths, host-side kernel dispatch consumes ${\sim}$50\,$μ$s regardless of workload, exceeding the actual GPU compute time at moderate-to-high pruning rates. We present a lightweight bidirectional Triton attention kernel whose dispatch floor is ${\sim}$24\,$μ$s -- roughly 2.17$\times$ lower than FlashAttention-2 varlen -- allowing pruning savings to become visible in wall-clock time. Integrated into a complete pack-attend-unpack pipeline and evaluated on an NVIDIA RTX 4000 Ada Generation GPU, our system achieves 1.88$\times$ end-to-end throughput over padded PyTorch SDPA at standard 224$\times$224 inputs, scaling to 2.51$\times$ at 384$\times$384. Against FlashAttention-2 varlen -- the strongest baseline -- our kernel delivers 9-12\% higher throughput at serving batch sizes (BS=1-4), and 2.17$\times$ lower kernel latency at 80\% token pruning. Numerical correctness is verified with max absolute logit difference $<$0.004 and bit-exact top-1 predictions.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。