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

2025/05/25 by Xiao Liu, Lijun Zhang, Liu, Xiao +5
Computer Science · Engineering · #Neural Networks and Applications #Power Transformer Diagnostics and Insulation

paper · pdf · doi:10.48550/arxiv.2505.19342

Abstract

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.

Citations

• Eigen Attention: Attention in Low-Rank Space for KV Cache Compression
• The Llama 3 Herd of Models
• Galaxy: A Resource-Efficient Collaborative Edge AI System for In-situ Transformer Inference
• KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization
• CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving
• HiFi-Codec: Group-residual Vector quantization for High Fidelity Audio Codec
• DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale
• Knowledge Distillation via the Target-aware Transformer
• A Fast Post-Training Pruning Framework for Transformers
• Post-Training Quantization for Vision Transformer
• CoEdge: Cooperative DNN Inference with Adaptive Workload Partitioning over Heterogeneous Edge Devices
• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Q8BERT: Quantized 8Bit BERT
• ALBERT: A Lite BERT for Self-supervised Learning of Language Representations
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Natural Language Understanding with the Quora Question Pairs Dataset
• Are Sixteen Heads Really Better than One?
• Neural Network Acceptability Judgments
• Neural Discrete Representation Learning
• Pointer Sentinel Mixture Models
• Character-level Convolutional Networks for Text Classification

Related

• Kraken: Inherently Parallel Transformers For Efficient Multi-Device Inference
• Inference Optimization of Foundation Models on AI Accelerators
• Easy and Efficient Transformer : Scalable Inference Solution For large NLP model
• CoFormer: Collaborating with Heterogeneous Edge Devices for Scalable Transformer Inference
• Lean Attention: Hardware-Aware Scalable Attention Mechanism for the Decode-Phase of Transformers