# ASTRA: Communication-Efficient Acceleration for Multi-Device Transformer Inference

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 40 references that are in the map.
- First posted on arXiv 25 May 2025.

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

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Neural Network Acceptability Judgments](https://mlsanity.com/paper/1805.12471/) (2018)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization](https://mlsanity.com/paper/2401.18079/) (2024)
- [DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale](https://mlsanity.com/paper/2207.00032/) (2022)
- [Q8BERT: Quantized 8Bit BERT](https://mlsanity.com/paper/1910.06188/) (2019)
- [Post-Training Quantization for Vision Transformer](https://mlsanity.com/paper/2106.14156/) (2021)
- [CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving](https://mlsanity.com/paper/2310.07240/) (2023)
- [A Fast Post-Training Pruning Framework for Transformers](https://mlsanity.com/paper/2204.09656/) (2022)
- [HiFi-Codec: Group-residual Vector quantization for High Fidelity Audio Codec](https://mlsanity.com/paper/2305.02765/) (2023)
- [Eigen Attention: Attention in Low-Rank Space for KV Cache Compression](https://mlsanity.com/paper/2408.05646/) (2024)
- [Knowledge Distillation via the Target-aware Transformer](https://mlsanity.com/paper/2205.10793/) (2022)
- [CoEdge: Cooperative DNN Inference with Adaptive Workload Partitioning over Heterogeneous Edge Devices](https://mlsanity.com/paper/2012.03257/) (2020)
- [Natural Language Understanding with the Quora Question Pairs Dataset](https://mlsanity.com/paper/1907.01041/) (2019)
- [Galaxy: A Resource-Efficient Collaborative Edge AI System for In-situ Transformer Inference](https://mlsanity.com/paper/2405.17245/) (2024)

## Abstract (arXiv)

Multi-device inference can reduce Transformer latency by parallelizing computation. However, existing methods require high inter-device bandwidth, making them impractical for bandwidth-constrained environments. We present ASTRA, a communication-efficient framework that integrates sequence parallelism with mixed-precision attention, where non-local token embeddings are transmitted as low-bit vector-quantized codes while local attention remains full precision. To preserve accuracy under aggressive compression, ASTRA introduces Noise-Augmented Quantization and Distributed Class Tokens. Across vision and language models (e.g., ViT and GPT2), ASTRA achieves up to 2.64× speedup over single-device inference and up to 15.25× over prior multi-device baselines while operating at bandwidths as low as 10 Mbps. ASTRA remains robust on large models (e.g., Llama-3-8B) even under non-ideal network conditions such as packet loss and dynamic networks.