# Dispatch-Aware Ragged Attention for Pruned Vision Transformers

arXiv: https://arxiv.org/abs/2604.15408 · HTML: https://mlsanity.com/paper/2604.15408/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Token pruning and merging for ViTs, one of 12 directions in Efficient LLM inference and serving.
- Builds on 11 of its 13 references that are in the map.
- First posted on arXiv 16 Apr 2026.

## Placed in

- Efficient LLM inference and serving / [Token pruning and merging for ViTs](https://mlsanity.com/inference/token-pruning-and-merging-for-vits/)

## Builds on (most cited first, 11 of 11)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Token Merging: Your ViT But Faster](https://mlsanity.com/paper/2210.09461/) (2022)
- [DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification](https://mlsanity.com/paper/2106.02034/) (2021)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [Not All Patches are What You Need: Expediting Vision Transformers via Token Reorganizations](https://mlsanity.com/paper/2202.07800/) (2022)
- [Adaptive Token Sampling For Efficient Vision Transformers](https://mlsanity.com/paper/2111.15667/) (2021)
- [Self-attention Does Not Need O(n2) Memory](https://mlsanity.com/paper/2112.05682/) (2021)

## Built on by (most cited first, 1 of 1)

- [Variable-Granularity Tokenization for High-Resolution Object Detection](https://mlsanity.com/paper/2608.28706/) (2026)

## Abstract (arXiv)

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 (≤197 tokens). We identify a dispatch-overhead bottleneck: at these lengths, host-side kernel dispatch consumes ~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 ~24,μs -- roughly 2.17× 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× end-to-end throughput over padded PyTorch SDPA at standard 224×224 inputs, scaling to 2.51× at 384×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× lower kernel latency at 80% token pruning. Numerical correctness is verified with max absolute logit difference <0.004 and bit-exact top-1 predictions.