# Dispatch-Aware Ragged Attention for Pruned Vision Transformers

Source: [Dispatch-Aware Ragged Attention for Pruned Vision Transformers](<https://www.sota2.com/research/paper/69e5b3057d0172140ec7df43>)

Collection index: [llms.txt](<https://www.sota2.com/research/llms.txt>)

Data updated: 2026-05-14T05:00:13.174Z

Authors: Seifeldin Abdellatif, Ahmad Almasri

Published: 2026-04-16T15:48:44.000Z

## 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\\,$\\mu$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\\,$\\mu$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.

## Recorded benchmark results

A research result applies to its stated dataset, split, metric and evaluation configuration. Compare like-for-like settings, preserve the metric's unit and direction, and distinguish paper-reported results from independent reproduction. The best result recorded here is not a claim of exhaustive, current global coverage.

- [Image Classification on ImageNet 1K (val) (Speedup)](<https://www.sota2.com/research/sota/image-classification-on-imagenet-1k-val-speedup.md>): Image Classification; ImageNet-1k (val); Speedup: 2.13.

- [Image Classification on ImageNet-1K (val) (Peak Throughput)](<https://www.sota2.com/research/sota/image-classification-on-imagenet-1k-val-peak-throughput.md>): Image Classification; ImageNet-1k (val); Peak Throughput (images/s): 25348.

## Original sources and implementation

[Original paper](<https://arxiv.org/abs/2604.15408>)

[PDF](<https://arxiv.org/pdf/2604.15408.pdf>)