LRAgent: Efficient KV Cache Sharing for Multi-LoRA LLM Agents

AI summary

Overview

• Research area: Efficient inference for multi-LLM agent systems, specifically KV cache management in multi-LoRA architectures (machine learning / systems for LLM serving).
• Technical level: Advanced. The paper assumes familiarity with transformer attention, KV caching, LoRA low-rank decomposition, FlashAttention-style tiling, and multi-agent agentic pipelines.
• Scope: The paper proposes and evaluates LRAgent, a KV cache sharing framework that decomposes the value cache into a shared pretrained-weight component and a low-rank adapter component for role-specialized multi-LoRA agents, tested on HotpotQA and ScienceQA with two 8B-scale models.

What This Paper Is About

Multi-LLM agent systems often specialize each agent by fine-tuning LoRA adapters on top of one shared pretrained backbone. Even though the backbone weights are shared, every agent still builds and stores its own KV cache for the same long, tool-augmented trajectories, which wastes memory and repeats computation. LRAgent's goal is to share as much of that cache as possible across agents while keeping the small, agent-specific adapter contributions that preserve accuracy, and to do so with minimal runtime overhead.

Key Contributions

1. A cache-decomposition framework for multi-LoRA agents. LRAgent decouples the value cache into a shared base cache computed from the pretrained weights and an adapter-dependent component computed from the LoRA weights, based on the observation that base cache similarity across agents is high while adapter outputs are largely decorrelated.
2. Two sharing schemes: BaseShared and BaseLRShared. BaseShared shares the base cache and stores each agent's adapter output in its low-rank form (the LR cache); BaseLRShared additionally shares the LR cache by exploiting a shared-A multi-LoRA architecture, which removes redundant prefill computation for contexts already processed by other agents.
3. Flash-LoRA-Attention (FLA), a reordered attention kernel. It applies the attention weights directly to the low-rank cache before the LoRA up-projection, avoiding materialization of the LR cache to full dimension, and is implemented on top of FlashAttention.
4. Empirical validation across two models and two agentic QA benchmarks. The paper reports accuracy, throughput, TTFT, and memory comparisons against Non-Shared, FullShared, and DroidSpeak baselines.

Main Findings

• Base cache is highly similar across agents; adapter output is not. On 128 samples of 2k tokens from HotpotQA, average pairwise cosine similarity for LLaMA-3.1-8B-Instruct was 0.9576 for the full cache, 0.9726 for the base cache, and 0.0538 for the adapter output; for Ministral-8B-Instruct the values were 0.9200, 0.9530, and 0.0225. The authors state the full cache similarity is lower than the base cache "empirically by about 3% on average."
• Key cache similarity is very high. The paper reports that the cosine similarity of the key cache remains above 0.98 on average, indicating value cache management is the key factor for preserving accuracy in multi-agent inference.
• LR caches are also shareable under shared-A multi-LoRA. Table 2 reports LLaMA-3.1-8B-Instruct LR cache cosine similarity of 0.9486 (Plan, Action), 0.9634 (Action, Reflect), and 0.9607 (Reflect, Plan); Ministral-8B-Instruct reports 0.9473, 0.9526, and 0.9498 for the same pairs.
• Accuracy is preserved better than prior sharing baselines. BaseShared's average accuracy drop was at most 0.7% and BaseLRShared's at most 1.5%, versus up to 5.3% for FullShared and up to 2.6% for DroidSpeak. On LLaMA-3.1-8B-Instruct HotpotQA, Non-Shared averaged 38.88 and BaseShared 38.60; on ScienceQA, Non-Shared averaged 69.31 and BaseShared 69.24.
• Throughput approaches full sharing. With Flash-LoRA-Attention enabled, BaseShared achieved up to a 1.42x throughput gain and BaseLRShared up to a 2.46x gain, approaching the FullShared upper bound; DroidSpeak reached up to 1.36x.
• Flash-LoRA-Attention itself provides a measurable speedup. It yielded up to a 1.24x throughput gain for BaseShared and up to a 1.35x gain for BaseLRShared.
• TTFT reductions. BaseShared and BaseLRShared provided up to 1.63x and 4.44x TTFT reductions respectively, both exceeding DroidSpeak's up to 1.56x; BaseLRShared's TTFT was close to FullShared.
• Memory savings. BaseShared and BaseLRShared reduce KV cache memory by nearly 1/3 compared to the Non-Shared baseline, comparable to other cache-sharing baselines and only marginally higher than FullShared within 1GB.
• Non-Shared runs out of memory on long traces. In the throughput and TTFT tables, Non-Shared is marked OOM at the 66.4k sequence length for both models, and DroidSpeak is marked OOM at 66.4k for Ministral-8B-Instruct.

Methodology in Plain English

The researchers start from a measurement: when several LoRA-specialized agents process the same context, the part of the value cache produced by the frozen pretrained weights looks almost the same across agents, while the part produced by the LoRA adapters is small and nearly decorrelated between agents. That suggests sharing the large, similar part and keeping only the small, agent-specific part.

They therefore split the value cache in two. The base cache comes from the pretrained weight matrix and is computed once and reused by all agents. The adapter contribution is stored not at full output dimension but at the much smaller LoRA rank dimension — this is the LR cache, the activation right after the LoRA down-projection. When the full adapter contribution is needed, it is reconstructed on the fly by multiplying with the up-projection matrix.

Two variants follow. BaseShared shares only the base cache; each agent keeps its own LR cache, so memory drops but an agent switching roles still has to run a "LR prefill" over context it has not yet processed. BaseLRShared additionally shares the LR cache by using a shared-A multi-LoRA design where all agents use a common down-projection matrix and differ only in their up-projection matrices. Because both caches are then already available for all previously seen tokens, only newly appended tokens need computation.

To keep reconstruction cheap, they reorder the attention math: instead of expanding the LR cache to full dimension and then multiplying by attention weights, they multiply the attention weights by the low-rank cache first and apply the up-projection afterward. This makes the dominant cost scale with the small rank rather than the full output dimension. They implement this inside a FlashAttention-style tiled kernel called Flash-LoRA-Attention.

Evaluation follows the AutoAct multi-hop framework with three fine-tuned roles — plan, action, and reflect — using LLaMA-3.1-8B-Instruct and Ministral-8B-Instruct, LoRA rank r=8 on query and value projections, a split of 2.5k HotpotQA and 2.0k ScienceQA, and 20 iterations per difficulty level. Efficiency is measured on a controlled trace where retrieved context varies from 1k to 64k tokens, giving total sequence lengths of 2k to 66k, on a single NVIDIA A6000 48GB GPU. DroidSpeak is run with its Pareto-optimal configuration, recomputing the top 33% most sensitive layers.

Why This Matters

Impact on research. Most prior KV cache sharing work targets general fine-tuned models or positional misalignment from differing prefixes; the paper argues that cache sharing tailored to multi-LoRA systems is largely unexplored. It also positions its contribution as complementary to prior cache-sharing methods, and shows that a low-rank decomposition of the adapter contribution beats selective recomputation (DroidSpeak) on accuracy while matching or exceeding it on efficiency.

Real-world applications:

• Multi-agent question answering and research assistants that call search, Wikipedia, and other tools across long trajectories.
• Enterprise deployments serving several role-specialized agents from one backbone where GPU memory is the binding constraint.
• Long-context retrieval-augmented pipelines where the same retrieved documents are fed to multiple specialized models.
• Serving systems that must minimize time-to-first-token for interactive agents, where TTFT is summed over many agent steps.

Industry relevance. The approach keeps the standard multi-LoRA architecture intact (shared-A is implemented by duplicating the same A weights across agents, so it does not change efficiency comparisons), which lowers the barrier to adopting it in existing LoRA serving stacks. The reported memory reduction, throughput reaching close to full cache sharing, and TTFT reductions near FullShared are the metrics that determine serving cost and user-perceived latency in production agent products.

Future Directions

• Extending the sharing and low-rank reconstruction idea to the key projection, which the paper explicitly analyzes in an appendix rather than in the main implementation.
• Understanding when the observed base cache and LR cache similarity breaks down — the paper notes experiments use specific agent roles, models, and datasets, and that other configurations are examined in appendices.
• Combining LRAgent with prefix-aware cache reuse methods (KVLink, KVFlow, KVComm) and with selective recomputation methods such as DroidSpeak, which the authors describe as complementary.
• Broader evaluation beyond HotpotQA and ScienceQA, and studies of how accuracy and memory behave as the number of agents N grows, since the theoretical cache size scales as 1/N plus a low-rank term.

Target Audience

Researchers and engineers working on LLM inference systems, KV cache optimization, parameter-efficient fine-tuning, and multi-agent orchestration. It is most useful to readers already comfortable with attention internals, FlashAttention tiling, and LoRA mathematics; readers new to these topics will find the methodology sections dense, though the core insight about base versus adapter cache components is stated plainly.

Authors’ abstract

Role specialization in multi-LLM agent systems is often realized via multi-LoRA, where agents share a pretrained backbone and differ only by lightweight adapters. Despite sharing base model weights, each agent independently builds and stores its own KV cache for the same long, tool-augmented trajectories, incurring substantial memory and compute overhead. Existing KV cache sharing methods largely overlook this multi-LoRA setting. We observe that, cache differences across agents are dominated by adapter outputs, while activations from the shared pretrained backbone remain highly similar. Based on this observation, we propose LRAgent, a KV cache sharing framework for multi-LoRA agents. It decomposes the cache into two components, a shared base component derived from pretrained weights and an adapter-dependent component derived from LoRA weights. LRAgent reduces memory overhead by sharing the base component across agents and storing the adapter component in its inherent low-rank form. It also reduces computational overhead by sharing the low-rank cache, enabled by a shared-A multi-LoRA architecture. This avoids redundant computations for contexts that have already been processed by other agents. To efficiently reconstruct adapter contributions at runtime, we introduce Flash-LoRA-Attention, a kernel that reorders attention computation to avoid materializing the low-rank cache to full dimension. LRAgent achieves throughput and time-to-first-token latency close to fully shared caching, while preserving accuracy near the non-shared caching baseline across agentic question-answering benchmarks.