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

2025-02-04

AI Summary

To address memory explosion and cross-GPU communication bottlenecks induced by cross-attention in multimodal large language models (MLLMs) processing long visual sequences (e.g., videos), this paper proposes a low-communication-overhead distributed exact cross-attention mechanism. Our method introduces a novel heterogeneous distribution strategy—“query blocks exchanged across GPUs while key-value blocks remain local”—combined with efficient activation recomputation to substantially alleviate GPU memory pressure and enable longer visual context. Theoretical analysis proves its communication complexity is asymptotically optimal. Experiments on mPLUG-Owl3 and OpenFlamingo demonstrate end-to-end training/inference speedups of up to 5.58×, alongside significant reductions in inter-GPU communication volume. This work establishes a scalable new paradigm for efficient MLLM training and deployment with extended visual inputs.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Large Vision Models
• Search and Optimization: Distributed Search

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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 enabling support for 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 mPLUG-Owl3 and OpenFlamingo models find that LV-XAttn achieves up to 5.58$ imes$ end-to-end speedup compared to existing approaches.

Problem

• Distributed cross-attention for large visual inputs
• Reduces memory demands in multimodal models
• Minimizes communication overhead in GPU computations

Innovation

• Distributed cross-attention minimizes communication
• Local GPU key-value blocks reduce overhead
• Efficient activation recomputation supports longer contexts