PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Hyesung Jeon, Hyeongju Ha, Jae-Joon Kim

arXiv:2602.01053v2cs.LG

TL;DR

Multi-LoRA agents redundantly build KV caches for shared, long trajectories, while existing sharing methods largely overlook this architecture. LRAgent separates shared base and low-rank adapter caches, achieving accuracy near non-shared caching and throughput and time-to-first-token latency close to fully shared caching.

  • Problem

    Existing KV-cache sharing methods largely overlook multi-LoRA agents and may reduce memory without eliminating redundant computation for already processed contexts.

  • Method

    LRAgent shares a base cache across agents, stores adapter contributions as low-rank caches, and reconstructs them efficiently with Flash-LoRA-Attention.

  • Results

    Across agentic question-answering benchmarks, LRAGENT consistently outperforms prior cache-sharing baselines in both accuracy and efficiency.

  • Takeaways & Limitations

    LRAGENT preserves role-specific behaviors and accuracy near the non-shared baseline while delivering efficiency close to fully shared caching.

  • Takeaways & Limitations

    Under group-query attention, DroidSpeak’s hidden-state cache can approach out-of-memory behavior in some experimental cases.

Abstract

from arXiv · show

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.

1. Introduction

Multi-agent LLM systems incur severe KV-cache memory and computation overhead because agents maintain separate caches and redundantly prefill shared context. LRAgent addresses this largely unexplored multi-LoRA setting by separating shared backbone activations from agent-specific LoRA-induced cache differences.

  • Long agent trajectories amplify KV-cache memory and compute overhead because each agent maintains its own cache and redundantly prefills overlapping context.
  • Existing KV-cache sharing methods primarily address positional misalignment or selective recomputation, while still rebuilding hidden states for contexts processed by other agents.
  • KV-cache sharing schemes that explicitly exploit multi-LoRA architectures remain largely unexplored.
  • For the same context, cache discrepancies are dominated by task-specific LoRA outputs, whereas shared-backbone activations remain highly similar.
  • LRAgent decomposes each cache into a shared base component computed from pretrained weights and an agent-dependent component.

2. Background

LoRA adapts frozen pretrained weights with low-rank matrices, enabling parameter-efficient task specialization. Multi-LoRA extends this setup across agent roles, while shared-A designs and KV-cache reuse address redundancy in multi-agent inference.

  • LoRA: LoRA adds low-rank matrices A and B to frozen pretrained weights, reducing trainable parameters, memory usage, and computation versus full fine-tuning.The update uses rank r much smaller than the input and output dimensions; query and value projections provide strong accuracy for a given parameter count.
  • Multi-LoRA: Multi-LoRA augments one pretrained base weight with multiple task- or agent-role-specific low-rank weights for inference.Each role i uses its own matrices A_i and B_i.
  • Multi-LoRA with Shared-A: Shared-A multi-LoRA shares the down-projection matrix because task differences are primarily driven by B_i, while A_i encode highly similar intrinsic information.Sharing A can improve accuracy by making the down-projection more generalizable across tasks, with effectiveness reported across diverse domains and datasets.
  • Multi-LLM Agent: Multi-LLM agents may use role-specific prompts with one model or multiple models fine-tuned from the same backbone.Common systems include heterogeneous planning, retrieval, and other agent roles, although the supplied passage truncates the role list.
  • KV Cache Sharing: Long reasoning trajectories and repeated tool-retrieved contexts amplify memory and compute overhead because each agent stores and independently processes its own KV cache.The redundancy arises even when much of the context is shared across agents.
  • KV Cache Sharing: Prior methods reuse KV caches through selective recomputation, including critical tokens in CacheBlend and predefined critical layers in DroidSpeak.DroidSpeak targets fine-tuned LLMs sharing a backbone and reports higher accuracy than token-wise recomputation methods, though the passage truncates the comparison.

3. Methodology

LRAgent separates multi-LoRA value caches into shared base and adapter-dependent low-rank components, exploiting cross-agent similarity to reduce memory and redundant computation. BaseLRShared and Flash-LoRA-Attention further share low-rank caches and avoid full-dimensional expansion during attention.

  • Cache observations: Cache differences across agents mainly arise from small but agent-specific adapter outputs, while shared-backbone hidden states and base caches remain highly similar.Key-cache cosine similarity remains above 0.98 on average, making value-cache management central to preserving accuracy.
  • BaseShared: BaseShared shares one pretrained-weight base cache across agents and stores each adapter contribution as a low-rank LR cache reconstructed at runtime.The adapter output is represented as Ylr,i = XiAi and expanded through Ylr,iBi; key caches are fully shared.
  • BaseShared: BaseShared reduces total KV cache size to 1/N + r/dout ≃1/N of non-shared caching, but agent switching still requires LR prefill over unprocessed accumulated context.Per-agent LR caches are smaller than full caches by r/dout ≪1, while only the base cache is shared computationally.
  • BaseLRShared: BaseLRShared uses shared-A multi-LoRA to share the LR cache, eliminating redundant prefill computation while reducing memory to 1/N + r/(Ndout) ≃1/N.Previously available base and LR caches let agents compute only newly appended context rather than recomputing earlier tokens.
  • Flash-LoRA-Attention: Flash-LoRA-Attention reorders attention so PVlr is computed in Θ(Lr) before up-projection in Θ(rdout), replacing naive Θ(Lrdout) expansion overhead.The reordered total is Θ(Lr + rdout), avoiding materialization of VlrB across all L tokens.

4. Experiments

Experiments evaluate LRAgent’s accuracy and efficiency in multi-hop agent execution on HotpotQA and ScienceQA using role-specific agents and cache-sharing baselines. BaseShared and BaseLRShared preserve accuracy near Non-Shared while reducing memory and gaining throughput with Flash-LoRA-Attention.

  • Agent Setup: Experiments fine-tune plan, action, and reflect agents in a multi-hop execution framework with external tool use.The plan agent reasons and selects tools, while the action agent generates tool arguments and executes API calls.
  • Models and Datasets: Evaluations use LLaMA-3.1-8B-Instruct and Ministral-8B-Instruct on 2.5k HotpotQA and 2.0k ScienceQA examples across three difficulty levels.Training uses synthetic and filtered AutoAct trajectories, while evaluation runs multi-hop inference on test questions.
  • Accuracy: BaseShared’s average accuracy drop is at most 0.7%, and BaseLRShared’s is at most 1.5%, versus up to 5.3% for FullShared and 2.6% for DroidSpeak.These results show that the proposed cache-sharing schemes preserve accuracy more effectively than the prior baselines.
  • Efficiency: Flash-LoRA-Attention improves throughput by up to 1.24× for BaseShared and 1.35× for BaseLRShared.The reported gains result from reducing low-rank cache expansion overhead.
  • Memory: BaseShared and BaseLRShared reduce KV-cache memory by nearly 1/3 versus Non-Shared, while remaining within 1GB of FullShared.Their low-rank caches are negligible in size; under GQA, DroidSpeak’s hidden-state cache can cause near-OOM behavior.

5. Conclusion

LRAGENT shares KV-cache components across multi-LoRA agents by separating a shared base cache from an adapter-dependent low-rank cache. BaseShared reduces memory, BaseLRShared reduces computation, and Flash-LoRA-Attention enables efficiency close to fully shared caching while preserving near-baseline accuracy.

  • Conclusion: LRAGENT decouples the value cache into a shared base cache and an adapter-dependent LR cache for multi-LoRA agent systems.BaseShared shares the base cache, while BaseLRShared shares the LR cache under shared-A multi-LoRA variants while preserving role-specific behaviors.
  • Conclusion: The proposed methods preserve accuracy close to the non-shared baseline and consistently outperform prior cache sharing baselines in accuracy and efficiency.
  • Conclusion: Flash-LoRA-Attention avoids full-dimension materialization of the LR cache, enabling throughput and TTFT improvements close to fully shared caching.

Software and Data … A.3. Cosine Similarity Bound

The paper provides reproducibility materials and analyzes why base-cache sharing is effective: shared-backbone activations produce similar base caches, while adapter contributions are smaller and less correlated. Under approximate orthogonality assumptions, the resulting full-cache cosine similarity is lower than base-cache similarity.

  • Software and Data: A file upload reproduces the paper’s training, evaluation, and latency results under the same experimental settings.It includes environment setup, commands, and step-by-step experiment guidelines.
  • A. Base Cache and Adapter Output: Base-cache and adapter-output magnitudes follow similar layerwise trends, but their contributions differ in magnitude.The analysis uses a three-agent multi-LoRA system and treats pretrained-weight and LoRA contributions as separate components.
  • A.1. Cache L1 norm: The highly similar base cache and largely decorrelated adapter outputs motivate sharing the base cache instead of the full cache.The adapter output is modeled as a small but non-trivial, approximately random perturbation to the base cache.
  • A.2. Input Activation Cosine Similarity: Input activations are highly similar across agents, with similarity higher in earlier layers and gradually decreasing in deeper layers.This layerwise trend is consistent with other reported cache-similarity observations.
  • A.2. Input Activation Cosine Similarity: Shared pretrained projections preserve activation similarity in the base cache, whereas adapter weights introduce additional full-cache variation.Thus, identical input activations can still yield lower full-cache similarity when adapter contributions are included.
  • A.3. Cosine Similarity Bound: Under approximate orthogonality between adapter outputs and the base cache, plus decorrelation across agents, the base-cache cosine similarity exceeds full-cache cosine similarity.The orthogonality between base and adapter contributions is presented as a practical approximation, especially for heterogeneous agent roles.
  • A.3. Cosine Similarity Bound: The analysis defines the base cache as Ybase,i := XiW0, the adapter output as ∆Yi := Xi∆Wi, and the full cache as Yi := Ybase,i + ∆Yi.These definitions follow the multi-LoRA convention used in the paper.

A.4. Key Cache Cosine Similarity … C.1. Agent Prompts and Trajectory Templates

The appendix shows that key caches are highly similar across multi-LoRA agents, while adapter-dependent value differences motivate low-rank Flash-LoRA-Attention. It also describes heterogeneous-role evaluation limits and the shared prompts and tool-augmented trajectories used in implementation.

  • A.4. Key Cache Cosine Similarity: Because key caches remain highly similar, cross-agent differences primarily arise in the value cache through adapter-induced components.The paper therefore shares the entire key cache.
  • A.4. Key Cache Cosine Similarity: Key-cache similarity averages 0.9922 for LLaMA-3.1-8B-Instruct and 0.9840 for Ministral-8B-Instruct across agents.Minimum pairwise similarities are 0.9726 and 0.9530, respectively, exceeding corresponding average base-cache similarities.
  • A.5. Agent Role Heterogeneity and Cache Similarity: LRAgent depends on base caches being more similar than full KV caches, a relationship expected to persist as heterogeneous roles lower full-cache similarity more strongly.The evaluation includes strongly heterogeneous planning, action, and reflection roles, plus debate-style agents on MATH.
  • A.5. Agent Role Heterogeneity and Cache Similarity: Experiments are limited to AutoAct-based settings because open-sourced agent trajectories and evaluation frameworks are unavailable.The authors expect effectiveness in diverse heterogeneous scenarios and plan broader evaluation as resources become available.
  • B. Flash-LoRA-Attention: Flash-LoRA-Attention stores the adapter value cache V_lr ∈ R^L×r with rank r ≪ d_head and reconstructs contributions using the up-projection matrix B.The method analyzes a single attention head with accumulated length L = L_p + L_c.
  • B. Flash-LoRA-Attention: Reordering computes PV_lr in rank r before multiplying by B, replacing the dominant L_cL d_head expansion with L_cL r accumulation.The head-dimension multiplication is applied only once per query block.
  • B. Flash-LoRA-Attention: The kernel streams key/value blocks, accumulates base and low-rank outputs with online softmax, then applies one post multiplication O_lr,iB while preserving FlashAttention’s I/O pattern.The low-rank intermediate remains in rank r throughout the streaming pass.
  • C. Implementation Details / C.1. Agent Prompts and Trajectory Templates: Each trajectory combines a shared system prompt, user question, generated tokens, and rule-based inserts such as retrieved tool outputs.The system prompt is identical across agents, so prefix positional-alignment sharing methods reduce to FullShared in this setup.

C.2. Shared-A Multi-LoRA Architecture

Sharing the LoRA down-projection matrix A improves accuracy over conventional independent multi-LoRA training without changing inference efficiency or model structure. It also reduces trainable parameters by 33%, motivating shared-A weights for the main experiments.

  • Accuracy: Shared-A multi-LoRA improves accuracy over independent (A_i, B_i) pairs across settings.The reported comparison uses HotpotQA accuracy under identical training conditions and hyperparameters.
  • Inference efficiency: Duplicating shared A across agents preserves the conventional multi-LoRA implementation, leaving inference efficiency unchanged.The shared-A design introduces no model-structure change or inference overhead.
  • Parameter efficiency: Sharing A reduces the number of trainable parameters by 33%.This parameter reduction provides an efficiency benefit during training.

C.3. Shared-A on Multi-Domain Dataset … D. Ablative Experiments

The appendix validates shared-A across mixed HotpotQA and ScienceQA training, documents training and baseline configurations, and defines controlled traces for cache-sharing efficiency comparisons.

  • C.3. Shared-A on Multi-Domain Dataset: Mixed HotpotQA and ScienceQA trajectories were shuffled for training and evaluated separately under the same hyperparameters and setup.The study focuses on these two datasets because open-source agent trajectories are limited.
  • C.3. Shared-A on Multi-Domain Dataset: Mixed-dataset training does not degrade benchmark accuracy, remaining comparable to separate training within reported evaluation variance.The result supports shared-A generalization to mixed-task settings without benchmark-accuracy loss.
  • C.3. Shared-A on Multi-Domain Dataset: BaseLRShared may not directly apply to deployed multi-LoRA models whose weights cannot be modified because it requires shared-A.The authors nevertheless frame shared-A as an effective design opportunity for newly built multi-LoRA systems.
  • C.4. Hyperparameter and Loss Curve: Training across all agents takes 3.9 hours for HotpotQA and 3.4 hours for ScienceQA on a single 48GB NVIDIA A6000 GPU.Hyperparameters largely follow AutoAct, with grid searches over learning rates and training epochs.
  • C.5. DroidSpeak Recomputation Layer Selection: DroidSpeak selectively recomputes critical layers identified by probing accuracy drops when layer KV caches are reused.Its Pareto-optimal configuration balances accuracy degradation against inference-efficiency gains from cache sharing.
  • C.6. Emulated Trace for Efficiency Analysis: Efficiency comparisons use a fixed context-length trace and agent schedule to isolate cache-sharing effects under identical context lengths.The trace derives from profiled trajectories, including concatenated context lengths at agent steps and steps per iteration.
  • C.6. Emulated Trace for Efficiency Analysis: Retrieved context length Lctx varies from 0.25k to 16k, producing total trajectory lengths from 1.9k to 66.4k.Detailed agent trajectory templates are provided in Appendix C.1.

D.1. Ablation on LoRA Application

The ablation evaluates qkvo, which applies rank-4 LoRA to query, key, value, and output projections under matched parameter budget. Despite lower accuracy and throughput than qv, LRAgent’s decoupled cache remains effective, with BaseLRShared retaining an efficiency advantage.

  • Configuration: The qkvo ablation applies rank-4 LoRA to query, key, value, and output projections, matching qv’s parameter budget.It evaluates HotpotQA accuracy, system throughput, and TTFT under the same emulated trace.
  • Accuracy: The qkvo non-shared baseline has lower accuracy than qv, and this degradation extends across all cache-sharing methods.The evaluation uses average HotpotQA benchmark accuracy under qkvo.
  • Accuracy: LRAgent’s methods achieve the best accuracy among qkvo cache-sharing approaches, showing that base-cache and low-rank-cache decoupling remains effective with key-projection LoRA.The result specifically supports decoupling the shared base cache from the adapter-dependent LR cache.
  • Efficiency: qkvo lowers throughput because multiple projections add LoRA computation, while key-cache reconstruction must occur in head dimension before RoPE.Position-dependent RoPE rotations prevent the same associativity-based reordering used for the value cache.
  • Efficiency: BaseShared and BaseLRShared achieve lower qkvo throughput than their qv counterparts, but BaseLRShared remains more efficient than DroidSpeak.The comparison concerns system throughput across the total sequence lengths in the qkvo traces.
  • Efficiency: qkvo TTFT increases marginally overall because adapter-output reconstruction primarily affects generation rather than prefill.TTFT is reported in seconds across the qkvo trace sequence lengths.

D.2. Latency on HotpotQA Benchmark … D.5. Accuracy Score Deviation

The evaluation shows that latency and out-of-function behavior depend on both cache-sharing efficiency and generation quality, while BaseShared maintains strong accuracy across LoRA ranks. Accuracy comparisons remain reliable despite nondeterministic variation in external tool usage.

  • D.2. Latency on HotpotQA Benchmark: Latency is separated into model inference, end-to-end execution, and time-to-first-token measurements on HotpotQA.End-to-end latency additionally includes function calls and processing retrieved context.
  • D.2. Latency on HotpotQA Benchmark: Lower-accuracy methods such as FullShared and DroidSpeak can produce longer sequences and exceed Non-Shared latency despite strong cache-sharing efficiency.The results show that generation quality affects overall latency alongside cache-sharing efficiency.
  • D.3. Out-of-Function Ratio: Out-of-function incidents are separately reported because failing to answer within the maximum iterations can be more severe than returning an incorrect answer.OOF cases count as incorrect in benchmark accuracy but are distinguished for user-experience analysis.
  • D.3. Out-of-Function Ratio: Lower-accuracy methods generally have higher OOF ratios, while BaseShared and BaseLRShared achieve lower ratios among cache-sharing methods except for Ministral-8B-Instruct on ScienceQA.Table 18 reports OOF rates and their differences from the Non-Shared baseline.
  • D.4. Rank Ablations: BaseShared accuracy rises noticeably from rank 1 to rank 8 because planning and tool selection require precise trajectory-specific adaptation.The rank ablation evaluates BaseShared across ranks from 1 to 32.
  • D.4. Rank Ablations: Accuracy differences become marginal above rank 8 because the agent-trajectory training data are relatively small and easy to adapt.This supports avoiding unnecessarily high LoRA ranks once adaptation is sufficient.
  • D.5. Accuracy Score Deviation: Accuracy deviations are measured with 20 iterations despite using a single random seed, because external tool usage remains subtly nondeterministic.The reported standard deviations concern average benchmark accuracy.
  • D.5. Accuracy Score Deviation: Accuracy gaps between methods exceed the observed deviations, so the comparisons are considered reliable.The reliability assessment accompanies the standard-deviation analysis in Table 20.

D.6. Ablation on the Context Overlap Ratio · D.7. Memory Usage

D.6 shows that lower context overlap reduces reusable cache and narrows throughput gains, while D.7 examines memory usage as long trajectories accumulate KV-cache overhead. Across overlap settings, LRAgent methods outperform conventional sharing, and FullShared minimizes memory usage.

  • D.6. Ablation on the Context Overlap Ratio: Lower context overlap reduces reusable cache, causing cache-sharing methods to approach the non-shared setting.The analysis uses an emulation trajectory length of 33.7k; 100% overlap represents the fully shared setting.
  • D.6. Ablation on the Context Overlap Ratio: Throughput-sharing advantages decrease as overlap approaches 0%, converging toward the non-shared baseline.This trend is reported across the evaluated overlap settings.
  • D.6. Ablation on the Context Overlap Ratio: Across all overlap settings, the proposed methods achieve higher throughput than conventional sharing baselines such as DroidSpeak.BaseLRShared performs best overall even when reusable context is reduced.
  • D.6. Ablation on the Context Overlap Ratio: Table 21 reports throughput in tokens/s under varying context overlap ratios for LLaMA-3.1-8B-Instruct with trajectory length 33.7k.The table evaluates throughput across the overlap-ratio ablation.
  • D.7. Memory Usage: Memory usage is evaluated across diverse trajectory lengths on Ministral-8B-Instruct, where long-context retrieval makes KV-cache memory severe.The pretrained model weights consume 14.95 GB, while three LoRA weights add 0.11 GB.
  • D.7. Memory Usage: KV-cache sharing methods typically maintain one shared KV cache for three agents and recompute and overwrite it when needed.This memory behavior is described alongside the long-context memory analysis.
  • D.7. Memory Usage: FullShared has the lowest memory usage because it directly reuses the KV cache without additional components.DroidSpeak additionally maintains a hidden-state cache, which can become overhead in modern GQA models.
  • D.7. Memory Usage: Table 22 reports memory usage in GB for each total-sequence-length trace.The table summarizes memory consumption across the evaluated trace lengths.
Loading 2602.01053v2…