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

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$\times$ speedup over single-device inference and up to 15.25$\times$ 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.

In recent years, many modern AI systems rely on large Transformer models. These models can be heavy to run on one device, so a natural idea to accelerate their execution is to distribute their computation across multiple devices. The challenge is that such distribution results in repeated information exchange across devices, which can be slow on wireless, edge, or consumer-grade networks. We developed ASTRA to solve this communication bottleneck. Instead of sending full internal data between devices, ASTRA sends compact codes for the exchanged information to other devices, while keeping each device’s local information precisely. This design enables the devices to cooperate even when the connection between them is limited. In experiments on image and language tasks, ASTRA provided faster inference under low-bandwidth settings while keeping model quality close to the original model. This can make large AI models more practical when fast responses are needed but high-end networking is unavailable.