Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ShardMemo: Masked MoE Routing for Sharded Agentic LLM Memory

Yang Zhao, Chengxiao Dai, Yue Xiu, Mengying Kou, Yuliang Zheng, Dusit Niyato

TL;DR

ShardMemo tackles memory bottlenecks in agentic LLMs by introducing a tiered memory service with per-agent working state (Tier A), sharded evidence memory (Tier B) employing scope-before-routing and masked MoE routing, and a versioned skill library (Tier C). The shard probing process is cast as masked MoE routing under explicit budgets, using Top-$B_{ ext{probe}}$ or adaptive Top-$P$ with cost-aware gating across heterogeneous shard families, trained from evidence-to-shard supervision when available. Empirical results on LoCoMo, HotpotQA, and ToolBench show consistent improvements in task quality and reductions in retrieval work and tail latency compared to baselines, including +5.11 to +6.82 F1 gains on LoCoMo and strong Tier C performance on ToolBench (Precision@3 = 0.97, StepRed = 1.94). This architecture enables scalable, cost-aware memory access for long-horizon agentic reasoning and multi-agent tool usage, with practical implications for production-grade LLM systems.

Abstract

Agentic large language model (LLM) systems rely on external memory for long-horizon state and concurrent multi-agent execution, but centralized indexes and heuristic partitions become bottlenecks as memory volume and parallel access grow. We present ShardMemo, a budgeted tiered memory service with Tier A per-agent working state, Tier B sharded evidence with shard-local approximate nearest neighbor (ANN) indexes, and Tier C, a versioned skill library. Tier B enforces scope-before-routing: structured eligibility constraints mask ineligible shards before routing or ANN search. We cast shard probing as masked mixture-of-experts (MoE) routing over eligible shards, probing up to $B_{\mathrm{probe}}$ shards via Top-$B_{\mathrm{probe}}$ or adaptive Top-$P$, and use cost-aware gating over profile/observation/session shard families; the router is trained from evidence-to-shard supervision. On LoCoMo, ShardMemo improves over the strongest baseline (GAM) by +5.11 to +6.82 F1 across question categories. Under a fixed-budget routing setting ($B_{\mathrm{probe}}=3$), ShardMemo improves over cosine-to-prototype shard routing by +6.87 F1 while reducing retrieval work (VecScan 521->414, -20.5%) and p95 latency (95->76 ms). On long-context HotpotQA, ShardMemo achieves 63.41/61.88/57.95 F1 at 56K/224K/448K tokens. On ToolBench, Tier C reaches 0.97 Precision@3 and 1.94 StepRed (+10.2% and +7.2% over embedding-similarity retrieval).

ShardMemo: Masked MoE Routing for Sharded Agentic LLM Memory

TL;DR

ShardMemo tackles memory bottlenecks in agentic LLMs by introducing a tiered memory service with per-agent working state (Tier A), sharded evidence memory (Tier B) employing scope-before-routing and masked MoE routing, and a versioned skill library (Tier C). The shard probing process is cast as masked MoE routing under explicit budgets, using Top- or adaptive Top- with cost-aware gating across heterogeneous shard families, trained from evidence-to-shard supervision when available. Empirical results on LoCoMo, HotpotQA, and ToolBench show consistent improvements in task quality and reductions in retrieval work and tail latency compared to baselines, including +5.11 to +6.82 F1 gains on LoCoMo and strong Tier C performance on ToolBench (Precision@3 = 0.97, StepRed = 1.94). This architecture enables scalable, cost-aware memory access for long-horizon agentic reasoning and multi-agent tool usage, with practical implications for production-grade LLM systems.

Abstract

Agentic large language model (LLM) systems rely on external memory for long-horizon state and concurrent multi-agent execution, but centralized indexes and heuristic partitions become bottlenecks as memory volume and parallel access grow. We present ShardMemo, a budgeted tiered memory service with Tier A per-agent working state, Tier B sharded evidence with shard-local approximate nearest neighbor (ANN) indexes, and Tier C, a versioned skill library. Tier B enforces scope-before-routing: structured eligibility constraints mask ineligible shards before routing or ANN search. We cast shard probing as masked mixture-of-experts (MoE) routing over eligible shards, probing up to shards via Top- or adaptive Top-, and use cost-aware gating over profile/observation/session shard families; the router is trained from evidence-to-shard supervision. On LoCoMo, ShardMemo improves over the strongest baseline (GAM) by +5.11 to +6.82 F1 across question categories. Under a fixed-budget routing setting (), ShardMemo improves over cosine-to-prototype shard routing by +6.87 F1 while reducing retrieval work (VecScan 521->414, -20.5%) and p95 latency (95->76 ms). On long-context HotpotQA, ShardMemo achieves 63.41/61.88/57.95 F1 at 56K/224K/448K tokens. On ToolBench, Tier C reaches 0.97 Precision@3 and 1.94 StepRed (+10.2% and +7.2% over embedding-similarity retrieval).
Paper Structure (20 sections, 14 equations, 4 figures, 4 tables, 2 algorithms)

This paper contains 20 sections, 14 equations, 4 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Problem Formulation
  4. Method
  5. ShardMemo Overview
  6. Write path.
  7. System Architecture
  8. Tier A (working memory)
  9. Tier B (evidence memory)
  10. Tier C (skill library)
  11. Tier B router training
  12. Experiment
  13. Experimental Setup
  14. Datasets
  15. Baselines
  16. ...and 5 more sections

Figures (4)

  • Figure 1: ShardMemo architecture overview.
  • Figure 2: Budget sweep on LoCoMo by varying the shard-probe cap $B_{\mathrm{probe}} \in\{1,2,4,8\}$ at fixed retrieval size $K = 10$.
  • Figure 3: Mean retrieval latency (ms) vs. skill budget $R\in\{1,2,3,4,5\}$ on ToolBench.
  • Figure 4: Case study of single-hop fact extraction on LoCoMo.