LV-XAttn: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models

TL;DR

LV-XAttn is a distributed cross-attention mechanism designed for large visual inputs in multimodal large language models, significantly reducing communication overhead and improving training and inference efficiency.

Contribution

We introduce LV-XAttn, a novel distributed cross-attention method that minimizes communication overhead by keeping large key-value blocks local and exchanging smaller query blocks, enabling efficient processing of large visual inputs.

Findings

Achieves up to 10.62× speedup over existing methods.

Reduces communication overhead in distributed cross-attention.

Supports longer visual context with activation recomputation.

Abstract

Cross-attention is commonly adopted in multimodal large language models (MLLMs) for integrating visual information into the language backbone. However, in applications with large visual inputs, such as video understanding, processing a large number of visual tokens in cross-attention layers leads to high memory demands and often necessitates distributed computation across multiple GPUs. Existing distributed attention mechanisms face significant communication overheads, making cross-attention layers a critical bottleneck for efficient training and inference of MLLMs. To address this, we propose LV-XAttn, a distributed, exact cross-attention mechanism with minimal communication overhead. We observe that in applications involving large visual inputs, the size of the query block is typically much smaller than that of the key-value blocks. Thus, in LV-XAttn we keep the large key-value blocks locally on each GPU and exchange smaller query blocks across GPUs. We also introduce an efficient activation recomputation technique to support longer visual context. We theoretically analyze the communication benefits of LV-XAttn and show that it can achieve speedups for a wide range of models. Our evaluations with Llama 3-V, mPLUG-Owl3 and OpenFlamingo models find that LV-XAttn achieves up to 10.62$\times$ end-to-end speedup compared to existing approaches.