Dispatch-Aware Ragged Attention for Pruned Vision Transformers

Seifeldin Abdellatif; Ahmad Almasri

cs.LGApr 16, 2026

Dispatch-Aware Ragged Attention for Pruned Vision Transformers

Authors: Seifeldin Abdellatif, Ahmad Almasri

Abstract

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.

Explore similar work

May 21, 2026cs.LG

ASAP: Attention Sink Anchored Pruning

Vision Transformers (ViTs) face severe computational bottlenecks due to the quadratic complexity of self-attention at high resolutions. Existing token reduction methods rely on local metrics - such as single-layer attention scores - that are inherently vulnerable to the attention sink phenomenon, where uninformative tokens are paradoxically preserved over salient foreground objects. We propose ASAP (Attention Sink Anchored Pruning), a training-free framework that recasts this sink as a feature. Modeling ViT information flow as a Lazy Random Walk, ASAP identifies the sink as a dominant accumulator of probability mass. By computing the diffusion distance to the sink within the cumulative transition matrix, ASAP partitions tokens via Radial Diffusion Clustering and compresses background redundancy through Transition Weight Pooling in a single shot. Extensive experiments across image, video, and vision-language tasks demonstrate ASAP outperforms state-of-the-art methods, accelerating throughput by up to 48% while maintaining - or even exceeding - baseline accuracy.

Jaehyuk Lee, Hanyoung Kim, Yanggee Kim +1

Visual Token Pruning Vision Transformer Pruning

Ask PDF

Aug 11, 2026cs.CV

Putting Registers to Work: Task Registers for Token Pruning in Vision Transformers

Token-pruning policies are usually designed for a single recognition pipeline, but pretrained Vision Transformers are reused across tasks with different spatial demands. We ask which parts of a pruning policy transfer across image classification, semantic segmentation, and object detection. For each pipeline, controlled probes freeze the no-pruning checkpoint and apply a series of parameter-free reduction criteria at one eligible layer at a time without retraining. The probes reveal three differences: segmentation and detection rank the criteria differently, classification is especially sensitive to attention-based pruning in the earliest layers, and the dense tasks prefer opposite recovery endpoints. These findings motivate Task-Adaptive Pruning (TAP). Existing register tokens serve as task-agnostic storage for feature artifacts. TAP instead introduces one task register per task and activates only the current one. Its evolving state ranks tokens, distributes an exact removal budget over depth, and sets the recovery scale for dense features. At a final keep rate of

ρ=0.5

, our jointly adapted model, TAP-J, reaches

47.0

mIoU at

1.30\times

encoder throughput on ADE20K and

53.7

box AP at

1.32\times

encoder throughput on COCO while remaining competitive on ImageNet-1K.

Hongsen Cao, Mona Jaber, Shanxin Yuan +1

Visual Token Pruning Self-Supervised Vision Transformers Vision Transformer Pruning Object Detection

Ask PDF

Jul 4, 2026cs.CV

Rethinking Depth Pruning for Vision Transformers: A Heterogeneity-Aware Perspective

While prior studies have successfully compressed vision Transformers (ViTs) through various pruning techniques, most have concentrated on width pruning to achieve significant reductions in model size. Depth pruning, which removes entire layers from a ViT, is notoriously difficult for accuracy recovery despite its potential to deliver higher speedups, limiting the acceleration achieved by existing joint width-and-depth pruning methods. In this work, we reveal that the failure of existing depth pruning methods lies in their neglect of heterogeneity between different layers, and we introduce HetDPT, a heterogeneity-aware depth pruning method that avoids dimension mismatch. Comprehensive experiments on ImageNet-1K, CIFAR-100, COCO, and ADE20K validate our method: HetDPT achieves a 1.58

\times

speedup for DeiT-B while maintaining accuracy and a 1.39

\times

speedup for DeiT-S with nearly no accuracy degradation. Furthermore, when combined with width pruning, HetDPT+ sets a new state-of-the-art record in extreme ViT pruning, enhancing the acceleration ratio from 4.24

\times

to 5.19

\times

for the Isomorphic-Pruning-2.6G configuration while maintaining near-lossless accuracy; our code is available at https://github.com/Efficient-AI-for-All/HetDPT.

Zhenfeng Su, Kang Zhao, Han Bao +4

Self-Supervised Vision Transformers Vision Transformer

Ask PDF