Technical Summary: LRAgent

Problem Statement

In multi-LLM agent systems, role specialization is frequently achieved through Multi-LoRA architectures, where multiple agents share a single pretrained backbone but utilize distinct, lightweight LoRA adapters for specific tasks (e.g., planning, action, reflection). While this design improves parameter efficiency, it introduces significant redundancy during inference.

When agents collaborate on long, tool-augmented trajectories, they often process overlapping context. However, existing systems typically maintain independent KV caches for each agent. This results in:

1. Memory Overhead: Each agent stores a full KV cache for the same shared context, leading to linear memory scaling with the number of agents.
2. Compute Overhead: Agents redundantly recompute hidden states and KV projections for context segments already processed by other agents.

Existing KV cache sharing methods (e.g., ICaRus, Cache2Cache, KVLink) generally focus on positional misalignment or require architectural modifications and additional training. Crucially, they do not explicitly exploit the Multi-LoRA structure, often failing to eliminate redundant computation for contexts that have already been processed by other agents in the same system.

Methodology

The authors propose LRAgent, a KV cache sharing framework specifically designed for Multi-LoRA agent systems. The core insight driving LRAgent is that for the same context, the base cache (activations from the shared pretrained backbone) remains highly similar across agents, while the adapter output (activations from the LoRA weights) acts as a small, largely decorrelated perturbation containing agent-specific information.

1. Cache Decomposition

LRAgent decomposes the value cache ($Y$) into two components:

• Base Cache ($Y_{base}$): Derived from the frozen pretrained weights ($W_0$). This component is shared across all agents.
• Adapter Output ($\Delta Y$): Derived from the LoRA weights ($A, B$). This component is specific to each agent.

The authors observe that the base cache exhibits high cosine similarity (>0.97) across agents, whereas the adapter output has low similarity. Consequently, sharing the base cache preserves most of the semantic information, while the adapter contribution can be stored efficiently.

2. Sharing Schemes

LRAgent introduces two specific sharing strategies:

• BaseShared:

  • Mechanism: Shares a single base cache across all agents. Each agent maintains its own LR cache (Low-Rank cache), which stores the intermediate activations after the LoRA down-projection ($X_i A_i$) in a low-rank dimension $r$.
  • Reconstruction: At runtime, the full adapter contribution is reconstructed by multiplying the LR cache by the LoRA up-projection matrix ($B_i$) only when needed.
  • Benefit: Reduces memory usage significantly (approx. $1/N$ of the non-shared baseline) by avoiding full-dimension storage of adapter outputs. However, it still requires recomputing the LR cache for the entire accumulated context when switching agents ("LR prefill").

• BaseLRShared:

  • Mechanism: Builds upon BaseShared by leveraging a Shared-A Multi-LoRA architecture, where the down-projection matrix $A$ is shared across all agents (while $B$ remains task-specific).
  • Innovation: Since $A$ is shared, the LR cache ($X_i A$) is also highly similar across agents. This scheme shares a single LR cache for the entire system in addition to the base cache.
  • Benefit: Eliminates the need for "LR prefill" recomputation. Agents can reuse both the base and LR caches for previously seen tokens, reducing computational complexity to be comparable to fully shared caching.

3. Flash-LoRA-Attention

To address the computational overhead of reconstructing the adapter contribution from the low-rank cache, the authors introduce Flash-LoRA-Attention.

• Problem: Naive reconstruction ($V_{lr} B$) expands the cache to full dimension $d_{out}$ before attention, incurring $O(L \cdot r \cdot d_{out})$ complexity.
• Solution: The kernel reorders matrix multiplications using associativity: $O = P(V_{base} + V_{lr}B) \rightarrow O = PV_{base} + (PV_{lr})B$.
• Effect: The attention-weighted multiplication is performed on the low-rank cache first ($P V_{lr}$), scaling with $O(L \cdot r)$. The up-projection ($B$) is applied only once per query block. This reduces the expansion overhead by a factor of $r/d_{out}$.

Key Results

The authors evaluated LRAgent on LLaMA-3.1-8B-Instruct and Ministral-8B-Instruct using HotpotQA and ScienceQA benchmarks with multi-hop agent trajectories (Plan, Action, Reflect roles).

• Accuracy:
  • BaseShared and BaseLRShared maintain accuracy very close to the Non-Shared baseline (average drop $\le 0.7\%$ and $\le 1.5\%$, respectively).
  • In contrast, FullShared (naive full cache sharing) and DroidSpeak (selective recomputation) suffer larger accuracy drops (up to $5.3\%$ and $2.6\%$, respectively), highlighting the importance of preserving agent-specific adapter contributions.
• Memory Efficiency:
  • LRAgent reduces KV cache memory usage by nearly 1/3 compared to the Non-Shared baseline, comparable to other cache-sharing methods and only marginally higher than FullShared.
  • The LR caches are negligible in size due to the low-rank dimension.
• Throughput and Latency:
  • With Flash-LoRA-Attention, BaseLRShared achieves throughput gains of up to 2.46$\times$ over Non-Shared, approaching the performance of FullShared.
  • Time-to-First-Token (TTFT) is reduced by up to 4.44$\times$ for BaseLRShared, significantly outperforming DroidSpeak.
  • BaseShared achieves up to 1.42$\times$ throughput gain, balancing memory efficiency with computational savings.

Significance and Claims

The paper claims that LRAgent is the first work to explicitly tailor KV cache sharing to Multi-LoRA agent settings. Its significance lies in:

1. Decoupling Shared vs. Specific: By separating the cache into a shared base component and a low-rank adapter component, LRAgent resolves the trade-off between memory/compute efficiency and role-specific accuracy.
2. Computational Elimination: Unlike prior methods that merely reuse memory but still recompute hidden states (e.g., DroidSpeak), BaseLRShared (under Shared-A) eliminates redundant computation for previously processed contexts, achieving efficiency close to full cache sharing without the accuracy penalty.
3. Kernel Optimization: The Flash-LoRA-Attention kernel demonstrates that low-rank cache expansion can be made computationally cheap through reordering, enabling practical deployment.

The authors conclude that LRAgent enables Multi-LoRA agent systems to achieve memory and inference efficiency close to fully shared caching while preserving the accuracy of non-shared baselines, making it a robust solution for scalable, multi-agent LLM applications.