LV-XAttn: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
Authors: Tzu-Tao Chang
Shivaram Venkataraman
Venue: ICML 2025 poster (2025)
Keywords: Distributed System
Cross-Attention
Long Context
Multimodal Large Language Model
Primary area: deep_learning->attention_mechanisms
TL;DR: LV-XAttn is a distributed cross-attention mechanism with minimal communication overhead for multimodal large language models.

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.

LV- LV-XATTN: Distributed Cross- for Long Visual Inputs in Multimodal Large Language Models
Tzu-Tao Chang'
Shivaram Venkataraman
Abstract
pletion (Shrivastava al., 2023; Zhang al., 2023). Their strong reasoning capabiliitess have motivated efforts to ex- pand beyond language inputs, giving rise to multimodal large language models (MLLMs). These models can pro- cess and reason about other modalities, such as visual inputs, enabling applications like video understanding (Qian et al., 2024; Islam et al., 2024; He et al., 2024) and image process- ing (Guo et al., 2024; Li et al., 2023b).
Cross-attention is commonly adopted multi- modal large language models (MLLMs) for inte- grating visual information into the language back- bone. 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 multi- ple GPUs. Existing distributed attention mecha- nisms face significant communication overheads, making tross-attention layers critical bottleneck for efficient training and inference of MILMs. To address this, we propose LV-XAttn, a distributed, exact cross-attention mechanism with minimal communication overhead. We observe that in ap- plications 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 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 commu- nication benefits of LV-XAttn and show that it can achieve speedups for wide range of models. Our evaluations with Llama 3-V, mPLUG-OW13 and OpenFlamingo models find that LV-XAttn achieves up to 10.62x end-to-end speedup com- pared to existing approaches.
A common approach to integrating visual inputs into LLMs is through cross-attention (Alayrac et al., 2022; Laurençon et al., 2024; Li et al., 2023a; Ye et al., 2024; Grattafiori et al., 2024; Dai et al., 2024), where queries derived from the text input interact with keys and values derived from the visual inputs. This enables effective fusion of multi- modal information. In MLLMs, cross-attention layers are inserted between language model blocks, enabling the LLM to process intermediate representations that are integrated with visual information.
However, the memory requirement of cross-attention layers is limiting factor for applications involving large visual inputs, such long video understanding. For example, in Ilama 3-V (Grattafiori' et al., 2024), applying cross-attention to a text sequence of length 2048 and a 20-minute video sampled at frame per second (fps) requires over 234 GB of memory, even with a memory-efficieee attention imple- mentation (Dao, 2024). This exceeds the memory capacity of existing accelerators, necessitating the distributed com- putation of the attention operation across multiple workers.
Existing distributed attention approaches can be categorized into two classes: head sead-parallelism and sequence-pane sequence-parallelism. parallelism. Head-parallelism methods such as Deepspeed-UUyssee (Ja- cobs et al., 2024) and Megatron-LM (Korthikanti et al., 2023a) partition the computation along the head dimension of muat attention. Consequently, maximum degree of parallelism is capped by the number of heads used in multi- head attention. This translates to an upper bound in terms of memory capacity, preventing them from processing longer visual inputs which have memory demands beyond this. In addition, the number of workers must be divisible by the total number of heads to ensure balanced load across work- ers. Otherwise, resource underutilization may occur due stragglers. On the other hand, sequence parallel methods
1. Introduction
Large language models (LLMs) have shown exceptional performance in language processing tasks that involves long context, such as long document understanding (Sun al., 2024; Bertsch et al., 2023) and repository-level code com-
Universiity of Wisconsin-Madde Tao Chang <tchang85 @ wisc.edu>.
Correspondence to: Tzu-
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
such Ring Attention (Liu al., 2024a) partition the com- putation along the input sequence dimension, overcoming the limitation of head- pardlpaalllllsm methods. However, when applied to cross-attention with large visual inputs, these ap- proaches suffer from large communication overheads even after overlapping computation and communication. Figure 2 shows that cross tross-attention operations distributed with Ring Attention can account for up 88% of the iteration time, despite comprising only 3% of the total parameters.
The man chooses the red pill.
LM Block
Block Cross-Arth
LM Block
LM Block
Block
Cross Attention
LM Block w/ Cross-Attn
LM Block
Block
In this work, we present LV-XAttn, distributed, exact cross- attention mechanism that employs sequence- pareencllaallllllss with minimal communication overhead. Our main obser- vation is that while keys and values derived from visual inputs are large, the queries derived from text input are typically small in MLLMs. For example, in the video un- derstanding benchmark Video-MME (Fu et al., 2024), input processed with Llama 3-V results in an average se- quence length of 15,279,944 for keys and values and 5,514 for queries, when frames are sampled at I fps. Based on this, LV-XAttn organizes each worker to locally store a partition of the large key and value blocks, while small query blocks are transmitted between workers compute the attention output in a blockwise fashion. This significantly reduces the communication volume compared to Ring Attention. For instance, on the Video-MME benchmark, LV-XAttn reduces communication volume to just 0.04% of that required by Ring Attention. Furthermore, the reduced communication can be effectively overlapped by computation, allowing dis- tributed cross-attentie to be performed without incurring any communication overhead.
1MM Block w/ Cross-
LM Block
Projection
Visual Encoder
Word Embedding
Vidde What color is pill chosen
Figure I. MLLM with cross-attention.
incurs less than 0.42% overhead compared to the theoretical no-communication baseline. LV-XAttn is available at
https://giie tig...hh........ com/uw /ww saan//wwwd-----e mad- deddddass/V-X---e
2. Background
(Vaswani et al., 2017) is a variant of self-attention to model interactions between different sequences. The input to cross-attention consists of two sequences x E x and y E SxV Xohantes where SQ and Skv denote the sequence lengths of and y, respectively, and dembed the embedding dimension. The input sequence is multiplied with the projection matrices Wq E IR xd to obtain the queries Q E 5&x while the input sequence y is multiplied with the projection ma- trices WK. Wy E xd to obtain the keys and values K, V E RKv xd where is the hidden dimension. The attention output O R. xd is then computed as:
Cross-attention
To further enable the processing of longer visual inputs, we employ an activation recomputation technique that is specific to MI..Ms. In standard attention implementations, activations including queries, keys, and values need to be saved for backward pass (Korthikanti et al., 2023b). Storing the large key and value tensors for every cross-attention layer introduces additional memory pressure. We observe that since cross-attention layers in MLLMs share the same input visual tokens, we can maintain single copy of the visual tokens accessible to all cross-attention layers and re- compute activations during the backward pass. This allows us to process up to 1.6x longer visual inputs with just less than 8% overhead in terms of iteration time.
QKT
0 = softmax
V
Multimodal Large Language Models As LLMs continue to evolve, researchers are investigating how to incorporate vision and other modalities into these models. One com- mon way to embed cross-altention layers into the lan- guage model. This design has been adopted by a num- ber of models including Flamingo (Alayrac et al., 2022), Otter (Li et al., 2023a), mPLUG-Ow IDEFICS (Laurençon et al., 2024), Llama 3-V (Grattafiori et al., 2024), NVLM-X (Dai et al., 2024). and NVLM- H (Dai et al., 2024). Broadly, these models follow the archi- tecture illustrated in Figure 1. They include a visual encoder, a projection layer, and an L.M.M Cross-attention layers are
We perform comprehensive evaluation of LV-XAttn Llama 3-V, mPL.UG-O113 (Ye et al., 2024) models and OpenFlamingo (Awadalla et al., 2023) models across multiple cluster configurations, including setups with A100 and A30 GPUs LV-XAttn speeds up the cross-attention op- eration by up to 45.85x and overall model iteration time by up to 10.62x compared Ring Attention. By minimizing communication volume and further overlapping communi- cation with computation, we demonstrate that LV-XAttn
(Ye et al., 2024),
2

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
Fwd Vision
Pwd Non-CA
Bwd Non-CA
Fwd CA
Bodd CA
To handle large visual inputs, distributed attention ap- proaches have been proposed. These methods can be cat- egorized into two classes: head-paralle and sequence- parallelism. Head headdparallelism methods such as Deepspeed- Ulysses (Jacobs et al., 2024) and Megatron-LM (Korthikanti et al., 2023a) distribute the computation of different atten- tion heads across multiple workers. However, their scala- bility is limited by the number of attention heads, which imposes upper bound on memory capacity and thus max- imum sequence length.
Time (seconds)
150
1.00
0
Ring Attn 1V-XAI5n
Ring Attn LV-XAmt mPLIG-OW13
Ring LV-Xattn OpenFlamingo-33
Llama 3-V
Figure 2. Runtime breakdown for single iteration Llama 3-V, mPLUG-OWI377b and OpenFlamingo-336 using Ring Attention and LV-XAttn on 16 A100 GPUs LV-XAttn reduces the time spent on cross-attention computation by 96%, 93%, and 53% for the three models, respectively, compared to Ring Attention. "FWD Vision" refers to the forward pass through the vision encoder and the projection layer: "FWD CA" and "BWD CA" refer to the forward and backward passes through the cross- coss-attention layers in the LLM; and "FWD Non-CA" and "BWD Non-CA" refer to the forward and backward passes through the non-cross-attention lay- ers in the LLM. Llama 3-V was evaluated with text length of 1K and a frame count of 192 (Sq 1K, Skv 1200K); mPL.G- Dw was evaluated with a text length of 4K and frame count of 2K (Sq = 4K, SKY 1458K); and OpenFlamingo- 3b was evaluated with a text length and frame count of 32K (Sq² = 32K, Skv 2018K).
To overcome this limitation, sequence- peuuence-parallelim methods such as Ring Attention (Liu et al., 2024a) propose distribut- ing the attention operation across multiple workers along the sequence dimension. Specifically, with n workers, each worker is responsible for storing one block of query, key
and value Qi E R 22 K; E R "gy xd V, E R Sgg xd
sg xd Computa-
and computing the attention block O₂ E R tion of O can be decomposed to
QilKe: ..., Kn- 1 Vo, Vn-1]
O, softmax
Vd
Og is computed iteratively by transmitting key-value blocks among workers in ring-liks fashion. To facilitate the block- wise computation of Os, worker i has to maintain necessary Sq
softmax statistics L3 R putes partial attention using the blocks Q; ((i-r) mod n- and V(t-r) mod n and updates O and L, The key-value block is then sent to worker + 1 while a new key-value block is received from worker i 1-1.
During round worker com-
interleaved between the layers of the LLM. Language inputs are fed directly into the LLM and the resulting intermediate representations are passed to the cross-attention layers as x. Visual inputs are processed by the visual encoder and the projection layer to produce visual tokens, which are then passed to the cross-attention layers y, enabling the incorporation of visual information.
While Ring Attention can be used to distribute the cross- attention operation, the presence of large key-value blocks makes Ring Attention communication-bount resulting in highly inefficient cross-attention operations. As illustrated in Figure 2, despite comprising only 3% of the total param- eters, cross-attention operations can account for up to 88% of the iteration time when using Ring Attention.
Challenges with Large Visual Inputs When applied to scenarios with large visual inputs, cross-attention requires significant memory resources and therefore presents a scal- ing challenge. The standard implementation of attention in- volves materializing the matrix product QKT 5 x SKv resulting in memory complexity that scales with the prod- uct of text sequence and the visual sequence length. The large amount of visual tokens from long videos inputs thus causes a memory bottleneck. For example, in Llama 3-V, a 20-minute video sampled at 1 fps is encoded to a visual input with more than 7 million tokens.
3. LV-XAttn: Distributed Cross-Attention with Minimal Communication Overhead
In this section, we introduce LV-XAItn, our method for efficiently distributing the cross-attention operation with minimal communication overhead. We also present ac- tivation recomputation technique specific to MLLMs to re- duce memory pressure, enabling the processing of longer visual contexts.
While memory-efficciee methods like FlashAttention (Dao al., 2022; Dao, 2024) reduce the memory footprint of at- tention operations to enable handling longer context lengths, the amount of memory required still often surpasses the capacity of single worker. For example, processing a 20- minute long video with a language input 2048 tokens in Llama 3-V demands over 234 GB of memory for the cross-attention operation, even with PtastAntention.
3.1. Method
The primary observation that motivates our work is that, in applications involving large visual inputs, the size of the query block is typically much smaller than that of the key-
3

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
Round
Round
Round
Round 3
Og
Qo
K₂
Vi
Qa
Ko
Vi
O2 +=
Q2
Kn
Vo
O
+=
Q1
Kg
Vi
Q1
K₁
K
Os
Qa
Vi
O
Q2
K₁
Vi
Q2 +=
Q2
K₂
V/2
0 += Q1
K₂
V/2
On +=
Qo
K₂
V2
Og +=
Q5
K₃
V2
Ks
O
Q2
K
O1
Q
V/
Oo += Qo
Ka
Figure 3. LV-XAtt with workers. We partition the KV blocks and each worker stores their respective large key-value blocks K,, Vi We also partition the query (Q.), output (O,), and softmax statistics (m, and l omitted in the figure). The query and output are rotated among workers compute the attention.
Algorithm LV-XAttn Forward Pass for Worker i
received attention block O, and softmax statistics L, by rescaling them using AO and AT. The worker then sends Qj. Oj and L, to the next worker in the ring topology and receives Qj-1, Oj-1 and Lj-1 from the previous worker. After rounds, the computed attention block O and soft- max statistics I; are returned to worker i.
Input: data Qú, Kú, V; Initialize O,, L; 0
for round = 0 to n do Jorev & (i round + 1) mod n
f (i - round) mod n
Jnext < (i - round - 1) mod n
Overlapping Computation and Communication To fur- ther reduce communication overhead, we can overlap the attention computation with data transmission between work- ers. While performing attention computation with Qj. K; and V4 worker also does the following in parallel
do in parallel:
Send Jeen Lin Qj to worker (i + 1) mod n Recv Og. Ly. hiee from worker (i - 1) mod n
20, AL <- FlashAttention(() Qj, K,, V() Oj. Lj + Rescale(O), AO, AL)
end for
Receive Oj and Ly from worker i-1, which are needed for rescaling in this round.
which is needed for attention computation in the next round.
Receive Q) from worker
value blocks. For instance, in the widely-used vidco under- standing benchmark Video-MME (Fu al., 2024), videos have an average duration of 2,386 seconds. Each frame is encoded by the visual encoder and projection layer in MILMM into multiple visual tokens. For example, Llama 3- V generates 6404 visual tokens per frame. With a sampling rate of 1 fps, each video results in a key-value sequence length SKY = 15279944. On the other hand, the average text prompt for long videos consists of 3128 words, includ- ing the question, options, answer, and subtitles, resulting in query sequence length of SQ 2386 3128 5514. As result, distributed attention mechanisms that involves movement of key-value blocks incur substantial communi- cation overhead.
Send Oj+1 and L computed in the previous round worker i+
Send the already present Q2 to worker i + 1.
After receiving O, and L,j, and computing AO and AL, we can perform rescaling to update O,, and Lj. We de- scribe our distributed attention procedure in Algorithm 1. As demonstrated in Section 4.4, the substantial reduction in communication volume enables complete overlap with computation, effectively eliminating any communication overhead.
To address this, we propose LV-XAutn, which keeps the large key-value blocks locally on each worker, while smaller query blocks, attention blocks, and necessary softmax statis- ties are exchanged among workers in a ring-style fashion. This is illustrated in Figure 3. During each round, each worker computes attention using its local key-value blocks K; and V; and query blocks Q; received from peers. This computation generates partial attention blocks AO and par- tial softmax statistics A/. The worker then updates the
Runtime Analysis Let f(query size, key-value size) repre- sent the time required to perform the forward-past attention computation, and let commitenso size) denote the time
to transmit a tensor. LV-XAttn, Qi O2 R sq xd and L, E IR are transmitted during each round. This results in a per-round runtime of:
max (sc Sed Skvd n comm SQd
SQ 5a)) (1)

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
16 x107
Table 2. Evaluated models.
Ulura Iffe
Num. of CA Layers
Num. of LM Blocks
1.4
Llana 3 i nafee
Model
1.2
GPUUG Owtt & aabe
Llama 3-V-1
8
40
Qlemiiig fip-
1.0
Pheniig ds 0.
mPLUG-OWI3377
4
28
Speedup
HPPLUG-OW1--
4
28
05
0.8
mPLUG O0-3--1
4
24
0.6
Open Openinoamiggo99
8
32
0.4
OpenFlamingo-3b
24
24
3.2. Activation Recomputation for MILM
3000
10000
15000 20000
Sig
In standard attention implementation, during forward pass computation, input tensors Q; K; Vi and output tensors O, and Lu are saved for backward pass. However, storing large key-value blocks K; and V, increases memory usage, thereby limiting the maximum number of visual inputs that can be processed. For instance, in the case of mPLUG- Ow13-7b with 3600-frame video, storing K; and Vi takes 70.08GB per cross-attention layer.
Figure 4. The thcoretical speedup of LV-XAttn over Ring Attention for cross-attention on a 4-node cluster. Each node is equipped with 4 A100 GPUs, and nodes are interconnected by a 25 GB/s network. The markers represent processing a 2,386-second video and a 3, 1ord text prompt- - -vverge values for long videos in Video-MME (Fu et al., 2024) using LLama-3V, mPLUG-Ow13 and OpenFlamingo models at different frame rates. Note that for cach frame, special token <image> have to be added to the text-prompt, resulting in SQ = 2386 + 3128 = 5514.
To address this, we observe that while language features x differ across cross-attention layers as they pass through various LM blocks, the visual features y remain unchanged throughout all cross-attention layers, as they are only fed into the cross- coos-att layers. Thus, instead of storing key- value blocks for each cross coos-attention layer, we propose to keep single copy of visual features y that can be accessed by all cross-attention layers. During the backward pass, y/ is projected to recompute key-value blocks K; and V,. With Q, also being recomputed, we only need to save x, O,, and La during each cross-attention forward pass.
Table 1. Runtime analysis for Ring Attention and LV XAttn. h rep- resents the number of heads in multi-head attention. The attention FLOPs are calculated similar to previous work (Dao, 2024). Forward
Backward
2. SK -t-d
4. 5KK. -tdd Net Bandwidth
Ring Attention
Not Bandwidth
60 Sky...d ..dd
10. 3q Sgy. had
LV-XAttn
FLOPS
GLOPS
Speedup
GPU FLOPS
FLOPS
2 2q Net Bandwidth
5 S
Net Bandwadth
As demonstrated in the ablation study in Section 4.4, this approach incurs a runtime overhead of less than 8% while enabling the system to handle 1.6x more visual inputs.
In contrast, Ring Attention transmits key-value blocks
K2, V; E R SKx xd in each round, leading per-round runtime of:
4. Evaluation
max f( Sod Skvd comm Skvd
4.1. Experimental Setup
(2)
n
Implementation LV-XAttn is implemented using Py- Torch and Triton (Tillet et al., 2019). It uses torch. distributed for distributed communication. while the modified FlashAttention kernels to support rescal- ing operations are implemented with Triton. We validated the correctness of LV-XAttn by confirming that its output matches that of Py PyTorch's 's scaled dot- dot-product attention im- plementation.
Figure 4 shows the theoretical speedup of LV-XAttn over Ring Attention across different Sq and SKV. For MLLM with large visual inputs, where SKY > Sg. and slow cross-node interconnect an unavoidable constraint for in- puts exceeding single-node memory capacity- LV-XAttn is compute- bound, while Ring Attention is communication- bound. Consequently, the runtime for LV-XAtttn in Equa- tion reduces to f (Sad, Skyd) while the runtime for Sgd
Model Setup We evaluate our methods and baselines on 6 models shown in Table 2. Using their publicly available model checkpoints, we replace the cross-attention opera- tions with LV-XAttn and other distributed attention base- lines. Following their default configuration, each frame is encoded into 6404 visual tokens for Llama 3-V, 729 vi-
Ring Attention in Equation 2 becomes comm (2 Sxvd) A similar analysis applies to the backward pass. Table summarizes the runtime and corresponding speedup for multi-head cross-attention More in-depth discussion in Appendix A.
5

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
sual tokens for the mPLUG-OW11 models, and 64 visual tokens for the xOpenFlamingo models. This implies that given the same amount of memory capacity, OpenFlamingo models can accommodate more frames than Llama 3-V and mPLUG-Owl models.
larger number of frames (as each frame is encoded into fewer visual tokens compared to Llama 3-V) and thus have longer text lengths (due to the inclusion of a special token < <image> per frame) and larger Sq. the speedup of cross- attention remains significant, though less pronounced: 7.04 15.32x for the mPLUG-OWL models and 1.47 - 2.14x for the OpenFlamingo models. Notably, OpenFllamingo-Sh with denser cross- cross-attention layers, spends a larger portion of its time in cross-attention compared to OpenFlamingo-99 when using Ring Attention. Consequently, the speedup in cross-attention translates to more substantial end-to-end speedup for
For all models, a special token <image> must be included in the text prompt for each frame. Consequently, the length of the text prompt must be least equal the number of frames in the visual imput.
We use a batch size of and fully sharded tensor parallelism for all models to enable a larger context length.
Cluster Setup We evaluate our method and baselines on the following configurations: (1) A 16-GPU cluster, each node equipped with 4 A100 80GB GPUs, with the GPUs within a node interconnected via NVLink and cross-node bandwidth of 25 GB/s, representing typical setting for cross-node training of up to millions of tokens. (2) An 8- GPU cluster, each node equipped with A30 24GB GPU, with a cross- node bandwidth of 1.25 GB/s, representing a more resource-constraine
Table shows the same experiment on 8 A30 24GB GPUs. In this setup, we have smaller text lengths and fewer frames due to the smaller memory capacity. The speedup for cross-attention operation is greater than that on 16 A100 GPUs: 22.43 - 33.77 x for Llama 3-V, 20.6 - 455,5 for the mPLUG-OW models, and 3.97 7.22x for the Open- Flamingo models. This is due to smaller query block sizes (shorter computations favors computation-hound LV- n X'ttn) and slower interconnect bandwidth (longer com- munication hurts communication-boun Ring Attention), as shown in Table 1. However, the larger cross-attention speedups do not translate into a larger total speedup, as the portion of time spent on cross-attention layers decreases due slower self-attention layers in LM blocks (caused by lower GPU FLOPS and slower interconnect). Despite this, the total speedup remains 2.2 - 3.13x for Llama 3-V, 1.65 3.45 x for the mPLUG-OWI3 models, and 1.04 - 2.22x x for the OpenFlamingo models.
Se
setup with slower interconnect bandwidth. (3) A 12-GPU cluster, each node equipped with 3 A100 40GB GPUs, with the GPUs interconnected via 64 GB/s PCIe and a cross-node bandwidth of 25 GB/s, used for smaller-scale case studies and ablation studies.
Baselines For our method, we use LV-XAttn for the cross- attention layers and Ring Attention for the LM blocks. Our primary baseline is the setup where Ring Attention is used for both the cross-attention layers and LM blocks. We apply our activation recomputation technique to both of these settings for enabling longer context length. We also compare against Deepspeed- Ulysses (Jacobs et al., 2024), which employs sequence parallelism for non-attention layers and head parallelism for attention layers. All methods use Flash Attentionn....
4.3. Comparison with DeepSpeed-Ulsses
For Deepspeed- each attention operation involves two all-to-all communications: one before the computation to gather input query, key and value blocks, and another afterward distribute attention output along the sequence dimension. The first all-to-all is expensive as it involves communicating the large key-value blocks. To see this, we compare Prepspeed-Ulyseeee with LV-XAttn on mPLUG- Ow13-2b using the cluster with A100 80GB GPUs. As shown in Table 5, LV-XAttn achieves 1.34 - 1.55x speedup compared to Deepspeed- Ulysses.
Our evaluation focuses on comparing the runtime of dis- tributed cross-attention mechanisms. We follow the same benchmarking methodology used in previous work (Dao, 2024; Li et al., 2024), measuring runtime on randomly gen- erated inputs of specific sizes to compare our method with the baselines across different scales. The reported runtime averaged over 5 trials, following warmup runs.
In addition, without activation recomputation, the larger memory footprint of process large visual inputs. When running OpenFlamingoo- 3b on the cluster with A30 24GB GPUs, Table 6 shows that Attn is able to process more than 4x longer text and visual inputs compared to Deepspeed-Ulyssee
4.2. Comparison with Ring Attention
limits its ability to
Table shows the per iteration time of six models using LV-XAttn and Ring Attention on 16 A 100 80GB GPUs. For Llama 3-V, LV-XAttn speeds up the cross-attention operation by 17.28 27.59x. Since the cross-attention operation accounts for the majority of the total iteration time when using Ring Attention, this reduction results in a significant total iteration speedup of 7.51 - 10.62x. For the mPLUG-OWL and OpenFlamingo models, which process
Notably, the head parallelism in Deepspeed- Ulysses restricts both its scalability and fiexibility: the maximum degree of parallelism is limited by the number of heads, and the num- ber of heads has to be divisible by the number of workers.
6

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
Table Per iteration wall-clock time (in seconds) 16 100 80GB GPUs with Ring Attention and IV-XAttm. "CA" represents the time spent on cross-attention operations. As So doubles, the gross-attentie speedup nearly halves because the runtime for Ring Attention,
which is communication-on bound. remains constant, while runtime LV-XAtta, which computation boundd, doubles. On the other hand, as SKv doubles, both communication and computation also double, the speedup remains roughly the same.
Text
Frame
Ring Attention
LV-XAttn
Speedup
Model
SQ
SKv
length
count
CA (s)
Total (s)
CA (s)
Total (s)
Total
2K
384
2K
2401K
301.81
333.29
14.22
31.38
21.22x
10.62
Llama 3-b-11b
2K
192
2K
1200K
144.5
167.2
8.36
22.27
17.28x
7.51
IK
192
1K
1200K
151.83
176.16
5.5
19.93
27.59x
8.84x
8K
4K
8K
2916K
174.73
202.84
24.08
42.79
7.26x
4.74x
mPLUG-OWl3-7b
8K
2K
8K
1458K
89.88
112.28
12.14
32.72
7.41x
3.43x
4K
2K
4K
1458K
92.48
107.01
6.45
19.5
14.33x
5,49x
8K
4K
8K
2916K
83.41
90.1
10.33
17.39
8.07
5.18x
mPLUGG-W13-22
8K
2K
8K
1458K
36.66
45,42
5.21
11.28
7.04
4.03x
4K
2K
4K
1458K
37.78
44.8
2.79
8.25
13,52x
5,43x
8K
4K
8K
2916K
47.12
55.1
5.17
12.69
9.12x
4.34x
mPLUG-OWI3-Ib
SK
2K
8K
1458K
22.63
28.81
2.62
7.99
8.64
3.6x
4K
2K
4K
1458K
23.26
29.24
1.52
5.24
15.32
5.58x
64K
64K
64K
4096K
95.13
165.17
62.4
126.71
1.52x
1.3x
OpenFllamingo-91
64K
32K
64K
2048K
47.86
101.01
31.44
86.89
1.52x
1.16x
32K
32K
32K
2048K
33.58
69.53
16.0
49.5
2.1x
1,4x
64K
64K
64K
4096K
276.69
306.51
187.45
226.04
1.48x
1.36x
OpenFlamingo-3b
64K
32K
64K
2048K
138.05
166.36
94.1
120.09
1.47
1.39x
32K
32K
32K
2048K
102.82
118.01
47.98
61.57
2.14x
1.92x
Por Beeraion CA Time (sec)
time by transmitting significantly smaller query, output, and softmax statistics blocks. The overlapping further hides the communication time, enabling distributed attention with no communication overhead.
100
Ring w/o Overlap
75.
Ring w/ Overlap
50
LVXAtto w/o Overlap
DV-Xttto w/ Overlap
25
communication
Activation Recomputation Figures 6a and 6b show the iter- ation time for rumning mPLUG-Owl-76 and OpenFlamingo- 3b on a single node with 3 A 1100 40GB GPUs, with and with- out employing activation recomputation for cross-attention layers. By omitting the saving of large key-value blocks, the reduced memory consumption enables the processing of larger number frames, increasing by 1.5x and 1.6x for mPL.UG-OWl-7b and OpenFlamingo-3b. respectively, with runtime overhead of less than 8%.
4096 Text Length per Worker
Figure Ablation study on the effect overlapping communi- cation and computation with 6 A100 40GB GPUs. The frame count is set to 2048 per worker. Since processing the same total single GPU not feasible due memory number of frames constraints, the "no communication" runtime is derived by running the same per- worker input size single GPU and then scaling the result by 6. IV-XAttn incurs an overhead of less than 0.42% compared the no-communication baseline.
5. Discussion
In this section, we discuss the applicability of LV-XAttn to MLLMs with different architectures.
4.4. Ablation Study
There two main classes of MILM architectures: concutenation--usel and cross-altention-baeell The first design concatenates tokenized visual inputs with text tokens and feeds them into the LLM backbone. Models such as LA (Liu et al., 2023), InternVL (Chen et al., 2024), NVILA (Liu al., 2024b), and MiniGPT-4 (Zhu al., 2024) follow this paradigm. Concatenation- els have shown strong unified multimodal reasoning capa- bilities. However, the large number of tokens fed to the
Overlapping Communication and Computation Figure 5 shows the time spent on cross-attention in OpenFlamingo- 3b using Ring Attention and LV-XAttn, with and without overlapping communication and computation, on 6 A100 40GB GPUs. While overlapping reduces the runtime for Ring Attention, its effect is limited as the large communica- tion overhead of key-value blocks cannot be fully hidden by computation. In contrast, LV- Attn reduces communication
mod-

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
Table Per iteration wall-clock time (in seconds) A30 24GB GPUs with Ring Attention and LV-XAttm. "CA" represents the time
spent on cross aoss--atention operations.
Text
Frame
Ring Attention
LV-XAttn
Speedup
Model
Sq
Skv
length
count
CA
Total
CA
Total
CA
Total
512
64
512
400K
109.07
154.18
3.35
49.24
32.56x
3.13x
Llama 3-V-11b
512
32
512
200K
55.8
101.13
2.49
45.9
22.43x
2.2x
256
32
256
200K
57.5
100.75
1.7
44.93
33.77x
2.24x
IK
512
1K
364K
42.41
74.28
1.42
33.32
29.96x
2.23x
mPLUG-OW13-76
IK
256
IK
182K
20.81
50.61
0.66
30.63
31.31x
1.65x
512
256
512
182K
20.84
49.89
0.45
28.95
45,85
1.72x
IK
512
IK
364K
17.85
25.94
0.78
8.71
22.89
2.98x
mPLUG-OW13-2b
IK
256
IK
182K
9.06
16.56
0.44
7.86
20.6x
2.11x
512
256
512
182K
9.15
16.34
0.3
7.44
30.39
2.19x
IK
512
IK
364K
10.6
14.41
0.44
4.18
24.25x
3.45x
mPLUG-OWI3-1b
IK
256
1K
182K
5.38
8.4
0.25
3.36
21.19x
2.5x
512
256
512
182K
5.31
8.22
0.18
3.03
29.44
2.71x
8K
8K
8K
512K
17.28
65.75
3.99
53.22
4.33x
1.24x
OpenFlamingo-9b
8K
4K
8K
256K
8.74
54.09
2.2
52.17
3.97x
1.04
4K
4K
4K
256K
8.87
52.04
1.23
44.18
7.2x
1.18x
8K
8K
8K
512K
52.26
69.45
12.25
32.71
4.27x
2.12x
OpenFlamingo-3b
8K
4K
8K
256K
26.09
41.73
6.43
22.22
4.06x
1.88x
4K
4K
4K
256K
25.84
40.59
3.62
18.28
7.14x
2.22
Table Per iteration wall-lock time (in seconds) of mPLUG-OW13- 2b ran A100 80GB GPUs. The model uses multi-hean attention with 12 heads.
Table 6. Per iteratioon wall-lock time (in seconds) of OpenFlamingo- 3b ran A30 24GB GPUs. The model uses multi-head attention with 8 heads. Cluster
Cluster
Text/
Frame
Text/
Frame
DS (s)
LV-XAttn (s)
DS (s)
LV-XAttn (s)
Config.
worker
worker
Config.
worker
worker
512
256
OOM
13.38
IK
IK
OOM
32.71
12 GPUs
512
128
12.15
8.71
8 GPUs
512
512
OOM
18.28
256
128
9.32
6.1
256
256
OOM
14.46
512
256
16.36
10.58
IK
1K
OOM
19.71
6 GPUs
512
128
10.41
7.09
4 GPUs
512
512
OOM
13.34
256
128
8.81
5.83
256
256
11.65
11.29
512
256
15.64
10.11
1K
1K
OOM
13.75
3 GPUs
512
128
10.61
7.8
2 GPUs
512
512
10.19
9.24
256
128
9.91
7.37
256
256
8.04
7.87
LLM significantly slows down both training and inference, limiting their scalability and efficiency (Ye et al., 2024; Grattafiori et al., 2024).
handling large visual inputs (Dai et al., 2024). For example, cross-attention based MLLMs exhibit 10x better prefill la- tency than similar-sized concatenation--ase MLLMs (Qiu et al., 2025).
The second design relies on cross-attention, where cross- attention layers are inserted between LLM layers to in- corporate visual features into its intermediate representa- tions. Models such Flamingo (Alayrac al., 2022), IDEFICS (Laurençon al., 2024), Otter (Li et al., 2023a), mPL.UGOW13 (Ye et al., 2024), Llama 3-V (Grattafiori et al., 2024), and NVLM-X (Dai et al., 2024) adopt this strategy. By bypassing the need to process a large number of visual tokens through the LLM backbone, tross-attention based ML.LMs offer superior computational efficiency when
However, as the size of visual inputs increase, even cross- attention- based MLLMs encounter efficiency bottlenecks due to communication overhead. To address this, pro- posed LV-XAttn, which substantially reduces the communi- cation volume to achieve significant speedup.
To achieve both high accuracy and computational efficiency, recent work have explored hybrid architecture that uses both concatenation and cross-attention (Dai et al., 2024). We
8

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
Por Iteration Time (sec)
parallelism (Shoeybi al., 2019), and fully sharded data parallelism (Rajbhandari et al., 2020), can be combined with our approach.
30
Not save QKV Save QKV
20
10-
7. Conclusion
DOM 480
OOM
192
Frame Count
We introduced LV-XAttn, a distributed, exact cross-attention mechanism for MLLMs with minimal communication over- head. By storing large key-value blocks locally each worker and transmitting only smaller query blocks, LV- XAttn significantly reduces communication volume, which can be fully hidden by computation. Additionally, the ac- tivation recomputation technique reduces memory usage, enabling the processing of longer visual inputs with mini- mal overhead. Our evaluation demonstrates that LV-XAttn speeds up MLLM iteration by up 10.62x x and enables the processing of visual inputs up 1.6x longer.
(a) mPLUG-Owl-7b
Per Deration Time (sec)
60
Not save QKV
Save OKV
40-
20-
10M
DOM
DOM
2304
3072
3840
4608
5376
6144
(b) OpenFlamingo-3b
Figure 6. Ablation study the effect of activation recomputation for cross-attention layers with 3 A30 24GB GPUs. Text length is set 2K and 8K for mPL.UG-WWl-7b and OpenFlamingo-3h, respectively.
Acknowledgements
The authors were supported by NSF award CNS 2237306. We thank Zhao Zhang from Rutgers University for provid- ing us access computing resources. The authors acknowl- edge the Texas Advanced Computing Center (TACC) The University of Texas at Austin for providing computational resources that have contributed to the research results re- ported within this paper. This research used resources of the National Energy Research Scientific Computing Center (NERSC), a Department of Energy Office Science User Facility using NERSC award DDR-ERCAP0029980. This research also used computational resources from the NSF Cloudlab (Duplyakin et al., 2019) facility. We also thank Wenxuan Tan for his insightful suggestions on improving the paper.
note that LV-XAttn can also be applied to such architecture to address communication bottlenecks.
6. Related Work
Memory-efficient Attention The attention operation has memory complexity that scales quadratically with sequence length, limiting its scalability for longer contexts. Approx- imate methods (Kitaev et al., 2020; Zaheer et al., 2020; Beltagy al., 2020; Choromanski et al., 2021; Ding et al., 2023) and compression techniques (Chevalier et al., 2023; Munkhdalai et al., 2024) reduce memory requirements by sacrificing some model quality. For exact attention, FlashAt- tention (Dao et al., 2022: Dao, 2024) proposes block-wise computation, which reduces memory complexity to linear while providing runtime speedups by minimizing I/O be- tween GPU HBM and SRAM. However, for applications like long video understanding, which require context lengths exceeding single GPU's memory capacity, distributed setup is necessary.
Impact Statement
This paper aims improve the efficiency of MLLMs by alle- viating communication bottlenecks when processing long vi- sual inputs. By reducing end-to-end runtime, it helps lower the computational cost of training and deploying MLLMs, which in turn can reduce energy consumption and environ- mental impact. The approach does not alter model outputs, and thus does not introduce additional risks beyond those already associated with MLLMs.
Parallelism Distributed attention can be classified into head- parallelism approaches, such as Megatron-LM (Korthikanti et al., 2023a) and Deepspeed-Ulyssee (Jacobs et al., 2024), and sequence-parallelinn
approaches like Ring Attention, its variants DistFlashAttn (Li al., 2024), and Striped At- tention (Brandon et al., 2023). As discussed in Section 2, these methods face significant communication overhead. Our work proposes a solution with minimal communication overhead for cross-attention MLLMs. General paral- lelism approaches, such as data parallelism (Dean et al., 2012), pipeline parallelism (Narayanan al., 2019), tensor
References
Alayrac, J.-B., Donahue, J., Luc, P., Miech, A., Barr, I., Has- son, Y., Lenc, K., Mensch, A., Millican, K., Reynolds, M.. Ring. R., Rutherford, E., Cabi, S., Han, T., Gong, Z., Samangooei, S., Monteiro, M., Menick, J., Borgeaud, S., Brock, A., Nematzadeh, A., Sharifzadeh, S., Binkowski, M., Barreira, R., Vinyals, O., Zisserman, A., and Si-

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
monyan, K. Flamingo: visual language model for few-shot learning. In Oh, A. H., Agarwal, A., Bel- grave, D., and Cho, K. (eds.), Advances in Neural In- formation Processing Systems, 2022. URL https: //openreviews wwtrriiwwweeeeeeeeeeee /oet/fooo forum? wtw.rrummmddiuuiiiiiioi w...EbMuimbPPo..
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. FlashAt- tention: Fast and memory-efficiet exact attention with IO-awareness. In Advances in Neural Information Pro- cessing Systems (NeurIPS), 2022.
Dean, J., Corrado, G. S., Monga, R., Chen, K., Devin, M.. Le, Q. V., Mao, M. Z.,, Ranzato, M., Senior, A., Tucker, P., Yang, K., and Ng, A. Y. Large scale distributed deep networks. In Proceedings of the 26th International Con- ference on Neural Information Processing Systems Vol- une 1, NIPS' 12, pp. 1223-1231, Red Hook, NY, USA, 2012. Curran Associates Inc.
Awadalla, A., Gao, I., Gardner, J., Hessel, J., Hanafy, Y., Zhu, W., Marathe, K., Bitton, Y., Gadre, S., Sagawa, S., Jitsev, J., Kornblith, S., Koh, P. W. Ilharco, G., Worts- man, M., and Schmidt, L. Openflamingo: An open- source framework for training large autoregressive vision- language models, 2023. URL httph::p////// http///xxxvvoorog abs/ 2308/230.02 01390.
Ding, J., Ma, S., Dong, L.,, Zhang, X., Huang, S., Wang, W., Zheng, N., and Wei, F. Longnet: Scaling transformers to 1,000,000,000 tokens, 2023. URL https: //arxiv. org/ /200 02486.
Beltagy, I,, Peters, M. E.,, and Cohan, A. Longformer: The long- doc-mentnnnt transformer. arXiv:2004.055 150, 2020.
Bertsch, A., Alon, U., Neubig, G., and Gormley, M. R. Unlimiformer: Long-range transformers with unlimited length input. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https: //openrevien wwt.ileweee/eeeeeeeme wtt::/orommmddddtoooe
Duplyakin, D., Ricci, R., Maricq, A., Wong, G., Duerig. J., Eide, E., Stoller, L.,, Hibler, M., Johnson, D.. Webb, K., Akella, A., Wang. K., Ricart, G., Landweber, L., Elliott, C., Zink, M., Cecchet, E., Kar, S., and Mishra, The design and operation of Cloudl ab. In Proceedings of the USENIX Annual Technical Conference (ATC), pp. 1-14, July 2019. URL https: www. www..utah.ee.. edu/ paper/ htt.//paaaaaaiaatlll.
Brandon, W., Nrusimha, A., Qian, K., Ankner, Z., Jin, T., Striped attention: Faster
Song, Z., and ring attention for causal transformers. arXiv preprint arXiv:231 1.09431, 2023.
Chen, Z., Wu, J., Wang, W., Su, W., Chen, G., Xing, S., Zhong, M., Zhang, Q., Zhu, X., Lu, L.,, et al. Internvl: Scaling up vision foundation models and aligning for generic visual-linguistie tasks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24185 24185-24198, 2024.
Fu, C., Dai, Y., Luo, Y., Li, L., Ren, S., Zhang, R., Wang, Z., Zhou, C., Shen, Y., Zhang, M., et al. Video-mme: The first-ever comprehensive evaluation benchmark of multi-modal Ilms in video analysis. arXiv preprint arXiv:2405.21075. 2024.
Grattafiori, A., Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Vaughan, A., Yang, A., and et al, A. F. The llama 3 herd models, 2024. URL https: /ttp/aiivvv./ org org//bs/ /ba/ 2407.
Chevalier, A., Wettig, A., Ajith, A., and Chen, D. Adapting language models to compress contexts. The 2023 Con- ference on Empirical Methods Natural Language Pro- cessing, 2023. URL https:: http://opnneeeeiieeeeeee forum? d= wtt.dddpppppuppx.....
Guo, Z., Xu, R., Yao, Y., Cui, J., Ni, Z., Ge, C., Chua, T.-S., Liu, Z.,, and Huang, G. Llava-uhd: An Imm per- ceiving any aspect ratio and high-resolution images. In Computer Vision ECCV 2024: 18th European Confer- ence, Milan, Italy, September 29-October 4, 2024, Pro- ceedings, Part LXXXIII, pp. 390-406, Berlin, Heidelberg, 2024. Springer- Verlag. ISBN 978-3-031-73009- doi: 10. 10.10077998833331113810----- org 1001..1.0000 1007 /900//788--22
Choromanski, K. M., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J. Q.. Mo- hiuddin, A., Kaiser, L.., Belanger, D. B., Colwell, L. J., and Weller, A. Rethinking attention with performers. In International Conference Learning Representations, 2021. URL https: http::://peeneeeeee/e (d=Ua6zukOMES
URL ht tps: //oi.. 73010-
Dai, W., Lee, N., Wang, B., Yang, Z., Liu, Z., Barker, J., Rintamaki, T., Shoeybi, M., Catanzaro, B., and Ping. W. Nvim: Open frontier-class multimodal llms. arXiv preprint, 2024.
He, B., Li, H., Jang, Y. K., Jia, M., Cao, X., Shah, A., Shrivastava, A., and Lim, S.-N. Ma-lmm: Memory- augmented large multimodal model for long-term video understanding. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition (CVPR), 2024.
Dao, MtshAtttntionn22. Faster attention with better paral- lelism and work partitioning. In International Conference on Learning Representations (ICLR), 2024.
10

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
Isla,, M. M., Ho, N., Yang, X., Nagarajan, T., Torresani, L., and Bertasius, G. Video recap: Recursive captioning of hour-long videos. arXiv preprint 1rX1v:2402.13250, 2024.
and large language models. In Proceedings of the 40th In- ternational Conference Machine Learning, ICML223 JMLR.org, 2023b.
Liu, H., Li, C., Wu, Q. and Lee, Y. J. Visual instruction tuning, 2023.
Jacobs, S. A.. Tanaka, M., Zhang, C., Zhang. M., Am- inadabi, R. Y. Song, S. L.., Rajbhandari, S., and He, Y. System optimizations for enabling training of ex- treme long sequence transformer models. In Pro- ceedings of the 43rd ACM Symposium on Principles of Distributed Computing, PODC '24 pp. 121-130, New York, NY, USA, 2024. Association for Comput- ing Machinery. ISBN 9798400706684. doi: 10.1145/
Liu, H., Zaharia, M., and Abbeel, P. Ringattention with blockwise transformers for near-infinite context. The Twelfth International Conference on Learning Represen- tations, 2024a. URL https: http///ppneeeeeeeeeeeeeeeeoeeeep net/ forum? Brrum?dd=sRRRHHa450
3662158.3662888 URL https: pattss///ooi http://oo...ggg//// 1145/3662158 36628066.
Liu, Z., Zhu, L.. Shi, B., Zhang, Z., Lou, Y., Yang. S., Xi, H., Cao. S., Gu, Y., Li, D., Li, X., Fang, Y., Chen, Y., Hsieh, C.-Y., Huang, D.-.,, Cheng, A.-C., Nath, V., Hu, J., Liu, S., Krishna, R., Xu, D., Wang. X., Molchanov, P., Kautz, J., Yin, H., Han, S., and Lu, Y. Nvila: Efficient frontier visual language models, 2024b. URL https: / //arxiv. www.orr/gaaooo. abbs//b/////// 221/241222222 04468.
Kitaev, N., Kaiser, L., and Levskaya, The efficient transformer. In International Conference on Learning Representations, 2020, URL https:// openreview weeet//oormmiddzkg/liie
Reformer:
Korthikanti, V. A., Casper. J., Lym, S., McAfee, L., Ander- sch, M., Shoeybi, M., and Catanzaro, B. Reducing acti- vation recomputation in large transformer models. Pro- ceedings of Machine Learning and Systems, 5:341-353, 2023a.
Munkhdalai, T., Faruqui, M., and Gopal, S. Leave text behind: Efficient infinite context transformers with infini-attention, 2024. URL https: abs/2404.07145
Korthikanti, V. A., Casper, J., Lym, S., McAfee, L.,, Ander- sch, M., Shoeybi, M., and Catanzaro, B. Reducing acti- vation recomputation in large transforeer models. Pro- ceedings of Machine Learning and Systems, 5:341-353, 2023b.
Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, P. B., and Zaharia, M. Pipedream: generalized pipeline paral- lelism for dnn training. In Proceedings of the 27th ACM Symposium Operating Systems Principles, SOSP '19, pp. 1-15, New York, NY, USA, 2019. Associa- tion for Computing Machinery. ISBN 9781450368735. doi: 10.1145/3341301 org http:11.1111/3333
Laurençon, H., Saulnier, L., Tronchon, L., Bekman, S., Singh, A., Lozhkov, A., Wang, T., Karamcheti, S., Rush, A. M., Kiela, D., Cord, M., and Sanh, V. Obelics: an open web-scale filtered dataset of interleaved image-text documents. In Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS 23, Red Hook, NY, USA, 2024. Curran Associates Inc.
URL https:: /doi. 3359646.
Qian, L., Li, J., Wu, Y., Ye, Y., Fei, H., Chua, T.-S., Zhuang, Y., and Tang, S. Momentor: advancing video large lan- guage model with fine fine-grained temporal reasoning. In Proceedings of the 41st International Conference Ma- chine Learning, ICML'24. JMLR.org, 2024.
Li, B., Zhang, Y., Chen, L., Wang, J., Pu, F., Yang. J., Li, C., and Liu, Z. Mimic-it: Multi-modal in in-context instruction tuning. CoRR, abs/2306.0542 2023a. URL https: (addi.oor///ee
Qiu, H., Biswas, A., Zhao, Z., Mohan, J., Khare, A.. Choukse, E., fñigo Goiri, Zhang, Z., Shen, H., Bansal, C., Ramjee, R., and Fonseca, R. ModServe: Scalable and resource-efficient large multimodal model serving, 2025. URL https:ttp///mm aawiiv.vvvv../////v//aaaaaa/avav woggorrgsss... /2002 00937.
(18555//XXiii 2306.000000 05425.
Li, D., Shao, R., Xie, A., Xing, Ma, X., Stoica, I., Gonzalez, J. E., and Zhang, H. DISTELASHATTN: Distributed memory-efficient attention for long-context LLMs training. In First Conference on Language Mod- eling, 2024. URL https:/ /ttp///pnweveeeeeeeeee Forum?id=pUED<20
Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. Zero: memory optimizations toward training trillion parame- ter models. In Proceedings of the International Con- ference for High Performance Computing, Networking. Storage and Analysis, SC '20. IEEE Press, 2020. ISBN 9781728199 1999999..
Ji, Li, D., Savarese, S., and Hoi, S. Blip-2: bootstrapping ianguage-image pre-training with frozen image encoders

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-Im: Training multi- billion parameter language models using model paral- lelism. arXiv preprint arXhv:1909.08053. 2019.
Zhang, E, Chen, B., Zhang, Y., Keung, J., Liu, J., Zan, D., Mao, Y., Lou, J.-G., and Chen, W. RepoCoder: Repository-Ievel code completion through iterative re- trieval and generation. In Bouamor, H.. Pino, J., and Bali, K. (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 2471-2484, Singapore, December 2023. Association for Computational Linguisties. doi: (0.18653/v1/2022 emnlp-main. 151. URL https: http://aaaaantoo.y....oo w022222...mpppp..mmmmmmmmm.m2... mamolpppmai...m 151/.
Shrivastava, D., Larochelle, H., and Tarlow, D. Repository- level prompt generation for large language models of code. In Krause, A., Brunskill, E., Cho, K., En- gelhardt, B., Sabato, S., and Scarlett, J. (eds.), Pro- ceedings of the 40th International Conference on Ma- chine Learning, volume 202 Proceedings of Machine Learning Research, pp. 31693-31715. PMLR, 23-29 Jul 2023. URL httpcttpps//proceeeddn ngs mlr m.r.ppress. v202 #202/shrivastava333 html.
Zhu, D., Chen, J., Shen, X., Li, X., and Elhoseiny, M. MiniGPT-4: Enhancing lasion-languag understanding with advanced large language models. In The Twelfth International Conference on Learning Representations, 2024. URL h /oppnreviewwwe.eeeoooeoeeerreewe net forum?? i id-1t2bq88E27.
Sun, S., Liu, Y., Wang, S., Iter, D., Zhu, C.,, and lyyer, M. PEARL: Prompting large language models to plan and execute actions over long documents. In Graham, Y. and Purver, M. (eds.), Proceedings of the 18th Con- ference of the European Chapter of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 469-486, St. Julian's s, Malta, March 2024. Asso- ciation for Computational Linguistics. URL https: / //aclanthology. org/2024 eacl eacl-llng..222 long.ong..2.. 29/.
Tillet, P. Kung, H. T., and Cox, D. Triton: an inter- mediate language and compiler for tiled neural net- work computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learn- ing and Programming Languages, MAPL 2019, pp- 10-19, New York, NY, USA, 2019. Association for Com- puting Machinery. ISBN 978145 9781450367196. doi: 10. 1145/331 URL. https:////// http://oo.orr//// 3329973.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, I.,, Gomez, A. N., Kaiser, I.., and Polosukhin, I. Attention is all you need. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS' 17, pp. 1000-6010, Red Hook, NY. USA, 2017. Curran Associates Inc ISBN 9781510860964
Ye, J.. Xu, H., Liu, H., Hu, A., Yan, M., Qian, Q.. Zhang, J.. Huang, F., and Zhou, mplug-ows owl3: Towards long image-sequence understanding in multi-modal large lan- guage models, 2024. URL https: http://aaxvvvv.oo org/ abs/2408 04840.
Zaheer, M., Guruganesh, G., Dubey, A., Ainslie, J., Alberti, C. Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., and Ahmed, A. Big bird: transformers for longer sequences. In Proceedings of the 34th International Con- ference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc.
ISBN
12

---

LV-XATTN: Distributed Cross-Attention for Long Visual Inputs in Multimodal Large Language Models
256K
196K
13 Log2 speedup
LV-XAtta Comp Bound
SKV
128K
64K
Ring Attention Comm Bound
64K
128K
196K
256K
SQ
Figure The theoretical speedup of LV-XAttn over Ring Attention for on cluster with nodes, each equipped with 4 A100 GPUs. The uncolored region indicates a speedup of less than 1, meaning LV-XAttn performs slower than Ring Attention. Top-left and bottom-left quadrant represents typical use case MLLM with large visual inputs: SKY > Sq.
A. Comparison of LV-XAttn and Ring Attention for General Use Case
We have shown that for applications with large Sxv and small Sq such as long video understanding. LV-XAttn achieves significant speedup over Ring Attention. Here, we provide more in-depth analysis that generalizes to broader range of cases.
Figure 7 plots the theoretical speedup LV- Attn over Ring Attention for general So and SKV. When SKY large enough (above the horizontal bright red line), the transmission of Q. O and L; LV-XAttn is hidden by computation. making LV-XAttn compute-bound. On the other hand, when Sq is small (to the left of the vertical dark red line), the transmission of K; and V; are too large be hidden by computation, making Ring Attention compute-bound. Their intersection (the top-left quadrant) represents the typical MILMM use case with large visual inputs and small text prompt.
For smaller visual inputs, the reduced SK causes LV-XAttn to also become communication. bound (the bottom-left quad- rant). When both LV-XAttn and Ring Attention are communication bound, their relative speed depends on communication
volume: LV-XAttn sends 2 Sod + Se, while Ring Attention sends 2 Skvd Roughly, when SKY > Sq. LV-XAttn still remains faster than Ring Attention. For MLLMs, each image is encoded into a large number of visual tokens e.g.. 6404 for Llama 3-V, 729 for mPLUG-Owls and 64 for OpenFlamingo so this condition typically satisfied, making LV-XAttt faster for MILMS in general.
This also suggests that for self-attention, where SQ kv, Ring Attention is preferable. This is why in experiments, apply Ring Attention LM blocks for all baselines. However, when the context length is very large (top-right quadrant), both LV-XAttn and Ring Attention become compute-bour bound, resulting in identical iteration times, making the choice between them effectively irrelevant.
13


Reviews:
[{"review_id": "DGV9wB5Cgq", "reviewer": "Reviewer_zbCn", "summary": "This paper presents LV-XAttn, a distributed cross-attention mechanism with low communication overhead that can significantly reduce the inference time of MLLMs adopting cross-attn strategy. The authors also introduce the activation recomputation technique enabling supporting longer visual inputs. The proposed LV-XAttn is more efficient compared to Ring-Attention.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive feedback and questions!\n\n**Q1: Llama-3V is not discussed in the paper.**\n\nA1: Thank you for pointing this out! Llama-3V [1] also utilizes a cross-attention architecture, allowing LV-XAttn to be applied to it for large visual inputs. We conducted additional experiments on Llama-3V (`Llama-3.2-11B-Vision-Instruct`), comparing LV-XAttn and Ring Attention, following a similar setup to Section 4.2. The experiments were performed on a 6-GPU cluster, with each node equipped with three A100 40GB GPUs. The GPUs inside a node were interconnected via a 64GB/s PCIe connection, with a cross-node bandwidth of 25GB/s. The results are presented in the table below.\n\n| Model              | Text Length | Frame Count | $S_Q$  | $S_{KV}$   | Ring Attention CA (s) | Ring Attention Total (s) | LV-XAttn CA (s) | LV-XAttn Total (s) | Speedup CA | Speedup Total |\n|--------------------|------------|-------------|-----|-------|-----------------------|-------------------------|-----------------|-------------------|------------|---------------|\n| Llama-3V      | 384         | 120          | 384  | 750K | 23.30                | 37.54                  | 0.43           | 14.45             | 54.19$\\times$      | 2.60$\\times$        |\n| Llama-3V      | 192         | 120          | 192  | 750K | 23.01                | 37.21                  | 0.40           | 14.40             | 57.53$\\times$      | 2.58$\\times$        |\n| Llama-3V      | 192         | 60          | 192  | 375K | 11.69               | 22.92                  | 0.22           | 11.61             | 53.14$\\times$      | 1.97$\\times$        |\n\nFor Llama-3V, each frame is encoded into 6,404 visual tokens. In comparison, OpenFlamingo and mPLUG-Owl3 models encode each frame into 64 and 729 visual tokens, respectively. This results in an even greater imbalance between query block size and key-value block size for Llama-3V. Since LV-XAttn transmits the smaller query blocks while Ring Attention transmits the larger key-value blocks, LV-XAttn achieves a significant speedup over Ring Attention.\n\n**Q2: Applicability to MLLMs with alternative architecture.**\n\nA2: Recent MLLMs adopt two main classes of architecture: cross-attention based, where visual information is fused with text information through cross-attention, and concatenation based, where visual tokens are directly concatenated with text tokens and fed into the LLM backbone as inputs. While LV-XAttn is designed to address the communication overhead of the former class of MLLMs, it might still be useful for concatenation based MLLMs for the following reasons. \n\nFirst, to generate visual tokens, concatenation based MLLMs often use a visual adapter to align visual features generated by the vision encoder with textual prompt. Cross-attention is widely adopted in these visual adapters (e.g. BLIP-2 [2], Qwen-VL [3], Video-Llama [4], InternVL [5]). LV-XAttn can be used in the visual adapters when processing large visual inputs.\n\nSecond, with large amounts of visual tokens fed into the LLM backbones, concatenation based MLLMs also require distributing the attention operation. As discussed in Appendix A, with large context size, LV-XAttn can be used for distributed attention without slowdown when compared to Ring Attention. \n\nFinally, we would also like to note that cross-attention-based architectures are widely used in recent models, such as Llama-3V [1] and NVLM-H [6], due to their computational efficiency during training and inference, and due to their ability to prevent degradation for text-only tasks [1, 6]. As a result, addressing the communication bottleneck in these models with LV-XAttn is important for improving their overall efficiency.\n\n**Q3: Figure 2 setup clarification.**\n\nA3: For Figure 2, mPLUG-Owl-7b was run with a text length of 4K and a frame count of 2K ($S_Q=4K$, $S_{KV}=1458K$), while OpenFlamingo-3b was run with a text length and frame count of 32K ($S_Q=32K, $S_{KV}=2048K$). These results come from the same experiment presented in Table 3. We will ensure these clarifications are included in the caption in our revision.\n\n[1] The Llama 3 Herd of Models: https://arxiv.org/abs/2407.21783\n\n[2] BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models: https://proceedings.mlr.press/v202/li23q/li23q.pdf\n\n[3] Qwen-VL: A Versatile Vision-Language Model for Understanding, Localization, Text Reading, and Beyond: https://arxiv.org/pdf/2308.12966\n\n[4] Video-LLaMA: An Instruction-tuned Audio-Visual Language Model for Video Understanding: https://aclanthology.org/2023.emnlp-demo.49/\n\n[5] InternVL: Scaling up Vision Foundation Models and Aligning for Generic Visual-Linguistic Tasks: https://openaccess.thecvf.com/content/CVPR2024/papers/Chen_InternVL_Scaling_up_Vision_Foundation_Models_and_Aligning_for_Generic_CVPR_2024_paper.pdf\n\n[6] NVLM: Open Frontier-Class Multimodal LLM: https://arxiv.org/pdf/2409.11402", "claims_and_evidence": "The claims are clear. For example in Figure 2, the comparison of the overheads is well illustrated and evidenced.", "methods_and_evaluation_criteria": "The evaluation adopted in this paper is sense-making.", "theoretical_claims": "Theoretical claims presented in this paper are reasoning. For example the Figure 4 depicting the theoretical speedup of LV-XAttn over Ring Attention.", "experimental_designs_or_analyses": "I have checked the soundness and validity of the experimental designs. For example the Model Setup (i.e., using mPLUG-Owl3, OpenFlamingo), the Cluster Setup (i.e., using A100 80G GPUs), and the Baselines (i.e., Ring Attention, and Deepspeed-Ulysses). There are no issues.", "relation_to_broader_scientific_literature": "The technique proposed in this paper, which speedup the MLLMs on long videos inputs are useful.", "essential_references_not_discussed": "The references are sufficient. However, I noticed that the Llama 3V, a MLLM which is typical and also adopt the cross-attention paradigm to  perceive visual inputs is not discussed in this paper.", "other_strengths_and_weaknesses": "Strengths\n- The strength of this paper is very clear: the proposed cross-attention mechanism can reduce the communication overhead for MLLMs leveraging cross-attention paradigm on video scenario.\n\nWeakness\n- I think the proposed method is an extend version of the Ring-Attention on videos scenario, which is specifically effective for MLLMs who adopt the cross-attention paradigm to perceive visual inputs. However, the mainstream architecture adopted in current MLLMs is the concatenation paradigm, which prepends the visual tokens to prompt tokens directly and feeds them to LLM layers. Therefore I think the applicable scenarios are limited for the LV-XAttn.", "other_comments_or_suggestions": "Suggestions\n- The data with statistics characteristics used in the evaluation of Figure 2 should be included in the caption.", "questions_for_authors": "Could the proposed LV-XAtten also improve the efficiency for MLLMs using concatenation paradigm?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aRvuscRhUe", "reviewer": "Reviewer_EkKR", "summary": "The paper introduces ​LV-XAttn, a distributed cross-attention mechanism designed to handle large visual inputs in ​Multimodal Large Language Models (MLLMs) with minimal communication overhead. Cross-attention is commonly used in MLLMs to integrate visual information into the language backbone, but processing large visual inputs (e.g., long videos) leads to high memory demands and significant communication costs in distributed setups. Existing distributed attention mechanisms, such as ​head-parallelism and ​sequence-parallelism, face scalability and efficiency challenges, particularly with large key-value blocks. The key contributions of the paper:\n1. The proposed mechanism keeps large key-value blocks locally on each GPU and exchanges smaller query blocks across GPUs, significantly reducing communication volume. This approach leverages the observation that in applications with large visual inputs, the query block size is typically much smaller than the key-value blocks.\n​2. To further reduce memory pressure, LV-XAttn introduces a technique where visual tokens are shared across cross-attention layers, and activations are recomputed during the backward pass. This allows processing longer visual inputs with minimal overhead (less than 8%).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful questions and feedback!\n\n**Q1: Details on implementation.**\n\nA1: LV-XAttn is implemented using PyTorch and Triton. It uses `torch.distributed` for distributed communication, while the modified FlashAttention kernels to account for rescaling operations are implemented with Triton. We will add these details to the paper. We plan to open-source LV-XAttn.\n\n**Q2: Limited applicability to certain MLLM architectures.**\n\nA2: Please refer to Q2 of Reviewer zbCn.\n\n**Q3: ​Lack of analysis on single image processing.**\n\nA3: LV-XAttn addresses the communication overhead caused by the large number of visual tokens. While this issue is common in long video inputs due to the high frame count, it can also arise in single-image inputs. This can be due to either a large image size (for example, in mPLUG-Owl3, processing a single 3,840×3,840 high-resolution image is equivalent to processing a 100-frame video of standard frame size) or a high number of visual tokens per image (for instance, Llama-3V encodes each image into 6,404 visual tokens, whereas OpenFlamingo uses only 64, making a single image in Llama-3v equivalent to a 100-frame video in OpenFlamingo). \n\nIn both video and single-image inputs, these scenarios result in large key-value blocks, which are the primary source of communication overhead in existing distributed attention mechanisms. LV-XAttn effectively addresses this challenge.\n\nFor single images with small image size or low visual token count, we note that distributed attention may not be necessary. Even if LV-XAttn is applied in such scenarios, we expect its performance to be comparable to Ring Attention, as discussed in Appendix A.", "claims_and_evidence": "The claims made in the submission are generally supported by clear and convincing evidence, as the paper provides theoretical analysis, empirical evaluations, and comparisons with existing methods.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are well-designed and appropriate for the problem at hand. LV-XAttn effectively addresses the bottlenecks of cross-attention in MLLMs, and the evaluation demonstrates its advantages in terms of efficiency, scalability, and practicality for real-world applications.", "theoretical_claims": "The paper presents theoretical claims and analyses, particularly regarding the communication benefits and runtime efficiency of ​LV-XAttn compared to ​Ring Attention. I have checked the runtime analysis and don't find problems, and the empirical results support the claims.", "experimental_designs_or_analyses": "The paper provides a detailed algorithm (Algorithm 1) for the forward pass of LV-XAttn, which is helpful for reproducibility. However, it would be beneficial to include more details on the implementation, such as the specific libraries and frameworks used, to facilitate replication of the results.", "relation_to_broader_scientific_literature": "None", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Weakness：\n​Limited Applicability to Certain MLLM Architectures: The proposed method is specifically designed for MLLMs that utilize cross-attention mechanisms to process visual tokens. Consequently, it may not be directly applicable to other MLLM architectures, such as ​LLaVA, which employ different approaches for integrating visual information. This limitation restricts the generalizability of LV-XAttn to a broader range of multimodal models.\n\n​Lack of Analysis on Single Image Processing: In real-world deployment scenarios, MLLMs are often required to process not only long videos but also individual images. The paper does not provide an analysis of the runtime performance or potential trade-offs when applying LV-XAttn to single-image inputs. This omission raises concerns about whether the proposed method might introduce inefficiencies or negative effects in such cases, which are critical for practical applications. A more comprehensive evaluation encompassing both long videos and single images would strengthen the paper's relevance to real-world use cases.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wxUBFF8NYa", "reviewer": "Reviewer_c9FC", "summary": "Cross-attention layers takes a significant memory for applications involving large visual inputs, such as long video understanding, making scaling difficult due to high memory requirements. To address this, the authors present LV-XAttn, a distributed and exact cross-attention mechanism that leverages sequence-parallelism with minimal communication overhead. LV-XAttn partitions large key and value blocks across workers while transmitting only small query blocks, enabling blockwise attention computation. This approach significantly reduces communication volume compared to Ring Attention, improving scalability for large-scale visual tasks.\n\n## update after rebuttal\nThank you to the authors for their responses. While some concerns have been addressed, the contributions appear limited to cross-attention operations, making the scenario less compelling. Therefore, I will maintain my current score.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed questions and useful feedback!\n\n**Q1: Which datasets were used for training and evaluation?**\n\nA1: For our experiments, we do not pretrain or finetune the models. Instead, we use the checkpoints from pretrained models and replace their cross-attention operations with LV-XAttn and other distributed attention baselines.\n\nSince LV-XAttn does not affect accuracy, our evaluation focuses on its runtime. We follow the same benchmarking methodology used in Flash Attention and DistFlashAttn to measure the runtime with randomly generated inputs of specific sizes, validating its effectiveness across different scales. It is important to note that our performance benefits are independent of the datasets used.\n\n**Q2: How does LV-XAttn affect the model’s performance (e.g. accuracy)?**\n\nA2: LV-XAttn is an exact distributed attention mechanism and does not change the output of the cross-attention operator. Hence LV-XAttn does not affect the accuracy. In addition, we have validated the correctness of LV-XAttn by confirming that its output matches that of the PyTorch attention implementation. These tests are the same ones used by Flash Attention and DistFlashAttn.\n\n**Q3: Why is Ring Attention used for the LM blocks?**\n\nA3: For self-attention in LM blocks, the short text length and the equal size of query blocks and key-value blocks result in similar performance for LV-XAttn and Ring Attention, making the choice between them less impactful. To confirm this, we conducted additional experiments with mPLUG-Owl3-1b, maintaining a fixed frame count and varying the text length (since LM blocks are only affected by text length). The experiments were performed on a 6-GPU cluster, with each node equipped with three A100 40GB GPUs. The GPUs were interconnected via a 64GB/s PCIe connection, with a cross-node bandwidth of 25GB/s. The results are presented in the table below.\n\n| Text Length | Frame Count | $S_Q$   | $S_{KV}$   | Ring LM + Ring CA (s) | LV-XAttn LM + Ring CA (s) | Ring LM + LV-XAttn CA (s) | LV-XAttn LM + LV-XAttn CA (s) |\n|-------------|-------------|-------|----------|-------------------|-----------------------|-----------------------|---------------------------|\n| 24K        | 1152        | 24K | 820K   | 20.34             | 19.81                 | 16.7                  | 16.06                     |\n| 18K        | 1152        | 18K | 820K   | 17.64             | 16.81                 | 12.91                 | 12.86                     |\n| 12K        | 1152        | 12K | 820K   | 14.37             | 14.35                 | 10.11                 | 9.87                      |\n| 6K        | 1152        | 6K  | 820K   | 13.08             | 12.81                 | 6.68                  | 6.47                      |\n\nHere, “X LM + Y CA” indicates that the LM blocks are implemented with X, and the cross-attention layers are implemented with Y. As shown in the table, there is minimal difference between LM blocks implemented with Ring Attention or LV-XAttn. The speedup observed is primarily due to how the cross-attention layers are implemented.\n\n**Q4: Why is the end-to-end speedup much smaller than the speedup of cross-attention?**\n\nA4: The total iteration time includes the time spent on both cross-attention and non-cross-attention operations within the MLLM. LV-XAttn significantly accelerates the cross-attention layers, but the time spent on the remaining components of the MLLM remains unchanged. This is illustrated in Figure 2, where LV-XAttn reduces both the forward and backward cross-attention times (FWD CA and BWD CA), while the forward and backward times for vision and non-cross-attention layers (FWD Vision, FWD Non-CA, and BWD Non-CA) stay the same. As a result, the overall end-to-end speedup is limited by the time spent on non-cross-attention layers, meaning the speedup is smaller than that observed in the cross-attention layers. This is particularly true for mPLUG-Owl3 models, which only have 4 cross-attention layers out of a total of 24 or 28 LM blocks.\n\n**Q5: Figure 2 needs additional clarification on the legends.**\n\nA5: The abbreviations used in the legend for Figure 2 are as follows:\n* FWD CA: Forward pass through the cross-attention layers.\n* BWD CA: Backward pass through the cross-attention layers.\n* FWD Vision: Forward pass through the vision encoder.\n* FWD Non-CA: Forward pass through the non-cross-attention layers.\n* BWD Non-CA: Backward pass through the non-Cross-attention layers.\n\nWe will ensure these clarifications are included in our revision.", "claims_and_evidence": "The claims sound reasonable to me. However, I have some concerns and questions about the experiments, as described below.", "methods_and_evaluation_criteria": "**Strengths:** \n- The authors evaluated LV-XAttn on various MLLM models and various Cluster Setup. \n\n**Weaknesses and Questions:**\n- Which datasets were used for training and evaluation?\n- Although the authors demonstrate a significant speedup, there is no discussion on how LV-XAttn affects performance (e.g. accuracy).", "theoretical_claims": "The theoretical claims written in Section 3 sound reasonable to me.", "experimental_designs_or_analyses": "**Weaknesses and Questions:** \n- Why is the proposed method implemented with Ring Attention for the LM blocks? Is LV-XAttn not very efficient for that? Section A in the Appendix provides further analysis (Lines 632–634), but a more detailed investigation is needed to clarify the reasons.\n- In the result tables, \"CA\" indicates that the proposed model achieves a significant speedup in cross-attention operations. Nevertheless, the total speedup is much smaller than that of CA. Where does the proposed model spend more time?", "relation_to_broader_scientific_literature": "The experimental results may be useful for running models in the specific scenario described in the paper.", "essential_references_not_discussed": "I believe that the paper covers most of the essential references.", "other_strengths_and_weaknesses": "Please check my comments above.", "other_comments_or_suggestions": "The manuscript requires proofreading. For instance, Figure 2 needs additional clarification on the legends, as they are written entirely in abbreviations, such as ``Fwd Non-CA``.", "questions_for_authors": "Please answer my questions above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MjRpekaUuE", "reviewer": "Reviewer_c1fA", "summary": "1. The paper presents LV-XAttn, a distributed cross-attention mechanism designed to address the high memory demands and communication overheads in multimodal large language models (MLLMs) when processing large visual inputs.\n\n2. LV-XAttn reduces communication volume by keeping large key-value blocks locally on each GPU and exchanging smaller query blocks across GPUs, while also introducing an efficient activation recomputation technique to support longer visual context. \n\n3. Evaluations demonstrate that LV-XAttn achieves significant speedups, up to 5.58× end-to-end speedup compared to existing approaches like Ring Attention, making it a more efficient solution for distributed training and inference of MLLMs.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and feedback.\n\n**Q1: Limited model scope.**\n\nA1: Please refer to Q2 of Reviewer zbCn.\n\n**Q2: Lack of statistical significance testing.**\n\nA2: All runtime results are averaged over five iterations, excluding the first two warm-up iterations.\n\n**Q3: Limited discussion of trade-offs.**\n\nA3: In Sections 3.2 and 4.4, we discussed and demonstrated how LV-XAttn’s activation recomputation trades off minimal runtime for reduced memory pressure. For smaller visual inputs, as noted in the Appendix (and discussed in Reviewer EkKR Q3), we expect LV-XAttn to perform similarly to Ring Attention.\n\n\n**Q4: No comparison with other distributed cross-attention methods.**\n\nA4: To the best of our knowledge, existing distributed attention mechanisms either use sequence-parallelism or head-parallelism. Ring Attention and DeepSpeed-Ulysses are two widely adopted and representative mechanisms for each approach, respectively. We believe comparing LV-XAttn to these two methods provides a comprehensive evaluation, as they cover the primary paradigms in distributed cross-attention.", "claims_and_evidence": "1. Theoretical Analysis: They present a detailed theoretical analysis of the communication benefits of LV-XAttn compared to existing methods like Ring Attention. This includes mathematical formulations of communication volumes and runtime analysis for different scenarios.\n\n2. Empirical Evaluations: The paper includes comprehensive empirical evaluations across multiple models (mPLUG-Owl3 and OpenFlamingo) and cluster configurations (16 A100 GPUs, 8 A30 GPUs, etc.). These evaluations demonstrate significant speedups for both cross-attention operations and overall model iteration time.\n\n3. Ablation Studies: The authors conduct ablation studies to isolate the effects of specific techniques, such as the activation recomputation method. These studies show that the proposed methods achieve their claimed benefits with minimal overhead.\n\n4. Comparison with Baselines: The paper compares LV-XAttn with existing distributed attention mechanisms like Ring Attention and DeepSpeed-Ulysses, demonstrating superior performance in terms of speed and memory efficiency.", "methods_and_evaluation_criteria": "1. The core idea of LV-XAttn—keeping large key-value blocks locally while exchanging smaller query blocks across GPUs—makes sense given the observation that query blocks are typically much smaller than key-value blocks in MLLMs. This approach directly addresses the communication bottleneck identified in existing distributed attention mechanisms.\n\n2. The activation recomputation technique specifically designed for MLLMs is logical, as it leverages the fact that visual features remain unchanged across cross-attention layers, allowing for memory savings without significant recomputation overhead.\n\n3. The theoretical analysis of communication benefits provides a clear rationale for why LV-XAttn should outperform existing methods like Ring Attention, especially in scenarios with large visual inputs.\n\n4. The use of mPLUG-Owl3 and OpenFlamingo models as test cases is appropriate, as these are representative MLLMs that incorporate cross-attention mechanisms and are known to have challenges with large visual inputs.\n\n5. The evaluation across different cluster configurations (16 A100 GPUs, 8 A30 GPUs, etc.) helps demonstrate the method's effectiveness in various distributed settings and resource constraints.\n\n6.The focus on both cross-attention operation speedup and overall model iteration time speedup provides a comprehensive view of the practical benefits, showing how improvements in one component translate to overall system performance.", "theoretical_claims": "The paper presents several theoretical claims about the communication benefits and speedups of LV-XAttn compared to existing methods like Ring Attention. These claims are supported by mathematical formulations of communication volumes and runtime analysis. The theoretical speedup analysis in Figure 4 and the accompanying equations seem reasonable and align with the intuition that reducing communication volume would lead to significant speedups in distributed settings, especially when dealing with large key-value blocks.", "experimental_designs_or_analyses": "1. Main Evaluations Comparing LV-XAttn with Ring Attention\n1.1 Design: The authors evaluated LV-XAttn on multiple models (mPLUG-Owl3 and OpenFlamingo) across different cluster configurations (16 A100 GPUs, 8 A30 GPUs, etc.). They measured both cross-attention operation speedup and overall model iteration time speedup.\n\n1.2 Validity: This design is appropriate for demonstrating the practical benefits of LV-XAttn in real-world distributed settings. The use of multiple models and cluster configurations helps establish generalizability.\n\n1.3 Issues: The paper doesn't provide statistical significance testing for the results, which would strengthen the claims. Additionally, while the results are impressive, more detailed analysis of how different parameters (like batch size or sequence length) affect performance could be beneficial.\n\n2. Ablation Studies\n2.1 Design: The authors conducted ablation studies to examine the effects of overlapping computation/communication and activation recomputation.\n2.2 Validity: These studies effectively isolate specific components of the proposed method and demonstrate their individual contributions. The results (Figures 5 and 6) clearly show the benefits of these techniques.\n2.3 Issues: The ablation studies are somewhat limited in scope and could benefit from more detailed analysis of how different parameters affect performance. For example, varying the number of workers or the size of visual inputs could provide additional insights.", "relation_to_broader_scientific_literature": "1. Cross-Attention in Multimodal Models: Cross-attention has been widely adopted in multimodal large language models (MLLMs) for integrating visual information into language backbones. The proposed LV-XAttn mechanism advances this area by specifically addressing the challenges of distributed computation for cross-attention with large visual inputs, which is a critical bottleneck for efficient training and inference.\n\n2. Distributed Attention Mechanisms: Existing distributed attention approaches, such as head-parallelism methods (e.g., Deepspeed-Ulysses) and sequence-parallelism methods (e.g., Ring Attention), face significant communication overheads. LV-XAttn introduces a novel distributed cross-attention mechanism that minimizes communication overhead by keeping large key-value blocks locally on each GPU and exchanging smaller query blocks. This approach is a significant improvement over previous methods, especially for applications with large visual inputs where the size of the query block is much smaller than that of the key-value blocks.\n\n3. Memory-Efficient Attention: The paper's introduction of an efficient activation recomputation technique is related to broader efforts in developing memory-efficient attention mechanisms. Techniques like Flash Attention aim to reduce the memory footprint of attention operations. The activation recomputation method in LV-XAttn specifically targets the memory pressures introduced by storing large key and value tensors for every cross-attention layer in MLLMs, allowing for the processing of longer visual contexts with minimal overhead.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "1. Limited Model Scope: The evaluations primarily focus on mPLUG-Owl3 and OpenFlamingo models. While these are representative MLLMs, the approach could benefit from validation on a wider range of multimodal models to demonstrate broader applicability.\n\n2. Lack of Statistical Significance Testing: While the results show substantial speedups, statistical significance testing isn't provided. This could strengthen the claims by quantifying the reliability of the observed performance improvements.\n\n3. Limited Discussion of Trade-offs: The paper could benefit from a more detailed discussion of potential trade-offs, such as increased computational complexity from activation recomputation or limitations in scenarios with very small visual inputs.\n\n4. No Comparison with Other Distributed Cross-Attention Methods: The paper primarily compares with Ring Attention and DeepSpeed-Ulysses but doesn't evaluate against other potential distributed cross-attention mechanisms that might exist in the literature.", "other_comments_or_suggestions": "1. Limited Model Scope: The evaluations primarily focus on mPLUG-Owl3 and OpenFlamingo models. While these are representative MLLMs, the approach could benefit from validation on a wider range of multimodal models to demonstrate broader applicability.\n\n2. Lack of Statistical Significance Testing: While the results show substantial speedups, statistical significance testing isn't provided. This could strengthen the claims by quantifying the reliability of the observed performance improvements.\n\n3. Limited Discussion of Trade-offs: The paper could benefit from a more detailed discussion of potential trade-offs, such as increased computational complexity from activation recomputation or limitations in scenarios with very small visual inputs.\n\n4. No Comparison with Other Distributed Cross-Attention Methods: The paper primarily compares with Ring Attention and DeepSpeed-Ulysses but doesn't evaluate against other potential distributed cross-attention mechanisms that might exist in the literature.", "questions_for_authors": "1. Limited Model Scope: The evaluations primarily focus on mPLUG-Owl3 and OpenFlamingo models. While these are representative MLLMs, the approach could benefit from validation on a wider range of multimodal models to demonstrate broader applicability.\n\n2. Lack of Statistical Significance Testing: While the results show substantial speedups, statistical significance testing isn't provided. This could strengthen the claims by quantifying the reliability of the observed performance improvements.\n\n3. Limited Discussion of Trade-offs: The paper could benefit from a more detailed discussion of potential trade-offs, such as increased computational complexity from activation recomputation or limitations in scenarios with very small visual inputs.\n\n4. No Comparison with Other Distributed Cross-Attention Methods: The paper primarily compares with Ring Attention and DeepSpeed-Ulysses but doesn't evaluate against other potential distributed cross-attention mechanisms that might exist in the literature.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents LV‑XAttn, which exchanges small query blocks while keeping large key/value blocks local and using activation recomputation to extend visual context. The paper presents theoretical rationale and empirical speedups (up to 5.6× over Ring Attention and DeepSpeed‑Ulysses on mPLUG‑Owl3 and OpenFlamingo across diverse GPU clusters). Reviewers highlight that LV‑XAttn’s benefits are confined to cross‑attention MLLM architectures, leaving out concatenation‑based models (e.g. LLaVA, Llama 3V), and that its performance on single‑image inputs, accuracy trade‑offs, and statistical significance are not evaluated. They also call for broader model coverage, more detailed implementation/runtime and accuracy analyses, and minor clarifications in figures, legends, and experimental descriptions. In light of its significant contributions but narrow applicability and the need for additional validation and clarity, I recommend a weak accept. And I strongly encourage the authors to revise the paper by following reviewers' suggestions.
Decision: Accept (poster)