🤖 AI Summary
This work addresses the challenge that, in pruned Vision Transformers (ViTs), attention latency fails to scale proportionally with reduced computation due to scheduling overhead dominating on short sequences. To overcome this, the authors propose the first low-overhead attention implementation tailored for pruned ViTs, featuring a lightweight bidirectional Triton attention kernel and a pack-attend-unpack pipeline compatible with diverse token pruning strategies (e.g., Threshold-L2, DynamicViT). The approach reduces scheduling latency to approximately 40 microseconds and achieves up to 2.24× end-to-end throughput improvement on DeiT-T/S/B models while preserving bit-level prediction consistency—demonstrated by a maximum logit discrepancy below 0.007—thereby closely approaching the theoretical acceleration limit.
📝 Abstract
Token pruning methods for Vision Transformers (ViTs) promise quadratic reductions in attention FLOPs by dropping uninformative patches. Yet when pruned sequences are executed with state-of-the-art variable-length attention APIs -- including FlashAttention-2's varlen and PyTorch's NestedTensor SDPA-the wall-clock attention latency doesn't scale accordingly. We trace this to a dispatch-overhead bottleneck: at the short, post-pruning sequence lengths typical of ViTs (<=197 tokens), actual matrix arithmetic completes in single-digit microseconds while the host-side dispatch path consumes 60-90 us. We present a lightweight, bidirectional Triton attention kernel whose dispatch floor is 40 us roughly 1.5x lower than FlashAttention-2 varlen-allowing pruning savings to become more visible in wall-clock time. Integrated into a complete pack-attend-unpack pipeline, our system achieves up to 2.24x end-to-end throughput over padded PyTorch SDPA consistently across four pruning algorithms (Threshold-L2, DynamicViT, EViT, ATS), scales across DeiT-T/S/B, and maintains bit-exact classification predictions with <0.007 max absolute logit difference.