🤖 AI

KEEP: A KV-Cache-Centric Memory Management System for Efficient Embodied Planning

This paper introduces KEEP, a KV-cache-centric memory management system that employs static-dynamic construction, multi-hop re-computation, and layer-balanced loading to significantly accelerate embodied planning in memory-augmented LLMs while maintaining high accuracy.

Original authors: Zebin Yang, Tong Xie, Baotong Lu, Shaoshan Liu, Bo Yu, Meng Li

Published 2026-03-18

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zebin Yang, Tong Xie, Baotong Lu, Shaoshan Liu, Bo Yu, Meng Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a robot trying to cook a complex meal in a busy kitchen. You have a super-smart brain (a Large Language Model) that knows how to cook, but it has a short-term memory problem: every time you ask it to plan the next step, you have to read it a giant book containing your entire history, the current state of the kitchen, and the instructions.

If the book is too thick, your brain takes forever to read the beginning before it can even start thinking about the next move. This is the problem KEEP solves.

Here is a simple breakdown of how KEEP works, using everyday analogies:

The Problem: The "Giant Book" vs. The "Changing Kitchen"

In traditional robot planning, the robot keeps a text log of everything that happened.

The Issue: If you move a cup from the counter to the sink, the robot has to rewrite the entire history book to reflect that change. If the book is 10,000 pages long, rewriting even one sentence makes the whole book invalid for the next step.
The Result: The robot spends most of its time just "reading" and "rewriting" the book, leaving very little time to actually think about what to do next. It's like trying to drive a car while constantly rewriting the map in your head.

The Solution: KEEP (The Smart Librarian)

KEEP is a new system that manages the robot's memory not as text, but as "pre-calculated thought chunks" (called KV Caches). Think of these chunks as pre-cooked meals. Instead of cooking the whole meal from scratch every time, you just grab the pre-cooked parts and heat them up.

However, in a kitchen, things change fast. KEEP uses three clever tricks to handle this:

1. The "Static vs. Dynamic" Filing System

Imagine you have two types of files in your office:

Static Files: Things that rarely change, like the company's founding date or the layout of the building.
Dynamic Files: Things that change every minute, like the temperature of the oven or the location of the milk.

How KEEP works:
Instead of treating all files the same, KEEP puts them in different folders.

For Static files (like the building layout), it groups them together and calculates the "thought" for the whole group at once. It's efficient because these facts don't change often.
For Dynamic files (like the milk), it keeps them separate. If the milk moves, it only updates that one tiny file, not the whole group.
The Benefit: You don't have to rewrite the whole book just because the milk moved. You only update the "milk page."

2. The "Detective Chain" (Multi-hop Re-computation)

Sometimes, two pieces of information seem unrelated until you connect the dots.

Example: You need to "Unlock the door."
- Memory A: "The door is locked."
- Memory B: "The key is on the table."
- Memory C: "The table is in the kitchen."

If you just look at the door, you might forget the table. If you just look at the table, you might forget the door.
How KEEP works:
KEEP acts like a detective. It starts with the current goal ("Unlock the door") and asks, "What is important?" It finds the key. Then it asks, "Where is the key?" and finds the table. It follows this chain of logic (Multi-hop) to find the most critical connections.

The Trick: It only "re-cooks" (re-computes) the specific thoughts related to this chain. It ignores the irrelevant stuff (like the TV in the living room). This saves massive amounts of time while keeping the robot smart.

3. The "Assembly Line" (Layer-balanced Loading)

Imagine a factory assembly line where workers are building a car.

The Problem: In the old system, the first few workers (early layers of the brain) had very little to do, while the last few workers (later layers) were buried under a mountain of heavy parts to load. This caused the line to stop and start, creating "bubbles" of wasted time.
How KEEP works: KEEP acts as a smart manager. It sees that the first workers are idle, so it tells them to start loading parts for the later workers while the current work is being done.
The Benefit: The assembly line never stops. While one part of the brain is thinking, another part is already fetching the next memory. It keeps the pipeline full and smooth.

The Result: A Faster, Smarter Robot

By using these three tricks, KEEP allows the robot to:

Think faster: It cuts the time it takes to start planning (Time-to-First-Token) by nearly half compared to other methods.
Be more accurate: It doesn't forget important connections between memories, so it completes tasks (like cooking or cleaning) more successfully.
Save energy: It doesn't waste power recalculating things that haven't changed.

In short: KEEP is like upgrading a robot's brain from a chaotic notebook that needs constant rewriting to a smart, organized digital filing system that only updates what's necessary, connects the dots automatically, and keeps the workflow moving without ever stopping.

1. Problem Statement

Embodied planning tasks require agents to predict action sequences to achieve long-term goals in dynamic environments. Large Language Models (LLMs) are increasingly used for this, relying on memory (past experiences, environmental states, task history) to maintain a global view and avoid repetitive exploration.

However, existing memory management approaches face critical bottlenecks:

Text-Based Memory: Storing memory as raw text leads to excessively long prompts, causing high prefill latency (Time-To-First-Token, TTFT) and computational redundancy.
KV-Cache Inefficiency: While caching Key-Value (KV) pairs is a standard optimization in LLM inference, it fails in embodied planning due to frequent memory updates.
- Dynamic Updates: When an agent acts (e.g., "pick up milk"), the state of objects changes. In traditional KV caching, a single update invalidates the cache for all subsequent tokens in that block, forcing expensive recomputation.
- Granularity Mismatch: Existing methods use fixed-size blocks. Small updates invalidate large blocks (high recomputation cost), while smaller blocks sever cross-attention between memory segments (low accuracy).
- Static Recomputation: Current KV recomputation methods (e.g., CacheBlend, EPIC) use static heuristics (e.g., recomputing start/end tokens) that fail to adapt to the context-dependent importance of memories in dynamic scenarios.
- Pipeline Imbalance: Loading KV caches from CPU RAM to GPU creates pipeline bubbles because the cost of loading and computation is highly unbalanced across different LLM layers.

2. Methodology: KEEP System

KEEP introduces a KV-Cache-Centric memory management system designed to minimize recomputation while preserving memory interconnections. It consists of three core innovations:

A. Static-Dynamic Memory Construction

To address the granularity mismatch and invalidation issues:

Concept: Memory segments are clustered based on semantics and classified into Static and Dynamic groups based on their update frequency.
- Static Groups: Segments that have not changed for a threshold period ( $t$ steps). These are grouped together, and their KV caches are computed with full cross-attention to preserve rich contextual relationships.
- Dynamic Groups: Segments that change frequently. These are computed individually per segment to prevent a single update from invalidating the entire group's cache.
Mechanism: Groups transition between static and dynamic states based on stability. This hybrid approach minimizes KV invalidation while maintaining necessary interconnections.

B. Multi-hop Memory Re-computation

To address the issue of static recomputation failing to capture context-dependent importance:

Concept: Instead of recomputing fixed positions, KEEP dynamically identifies critical memory segments using importance propagation.
Mechanism:
1. Initialization: Calculate initial attention scores between the current query and all memory segments.
2. Propagation: Iteratively select top- $r$ important memories and update the importance scores of other memories based on their cross-attention to the selected set. This mimics a multi-hop reasoning chain (e.g., Query $\to$ Key $\to$ Table).
3. Selective Recomputation: Only the KV caches of the final stabilized set of important memories are recomputed in subsequent layers.
Granularity: Recomputation is performed at the memory segment level (not token level), allowing for contiguous loading and reducing I/O overhead compared to token-granularity methods.

C. Layer-balanced Memory Loading

To address pipeline bubbles caused by uneven KV loading costs across layers:

Observation: Early layers often recompute a large portion of memories (low load), while later layers recompute less (high load), creating idle time in the pipeline.
Mechanism: KEEP leverages the fact that if a memory segment is not recomputed in an early layer, it will not be recomputed in any subsequent layers.
- The system proactively pre-loads KV caches for these "stable" memories during the computation of earlier layers.
- This redistributes the loading burden from later layers to underutilized earlier layers, effectively eliminating pipeline bubbles and maximizing hardware utilization.

3. Key Contributions

Analysis: A detailed analysis of the fundamental differences between embodied planning (dynamic, frequent updates) and conventional LLM inference (static text).
Algorithmic Innovations:
- Static-Dynamic Memory Construction: Reduces KV cache invalidation by managing memory at mixed granularities.
- Multi-hop Memory Re-computation: Dynamically recovers critical cross-attention interactions via importance propagation.
- Layer-balanced Memory Loading: Optimizes hardware utilization by balancing KV loading and computation across layers.
System Implementation: A full implementation integrated into the vLLM serving engine, utilizing a pipelined schedule for efficient CPU-GPU data transfer.

4. Experimental Results

The system was evaluated on the ALFRED and WAH-NL benchmarks using Qwen-2.5 models (14B and 32B INT4).

Performance vs. Text-Based Methods: KEEP achieves a 2.68× speedup in TTFT compared to text-based memory methods with negligible accuracy loss.
Performance vs. KV Recomputation (CacheBlend):
- Accuracy: 4.13% improvement in Success Rate (SR) on ALFRED.
- Latency: 1.90× reduction in TTFT.
- On WAH-NL, KEEP showed a 3.37% improvement in Subgoal Success Rate and 1.91× TTFT reduction.
Ablation Studies:
- Removing Static-Dynamic Construction caused a 6.94% SR drop and 1.54× TTFT increase.
- Removing Multi-hop Recomputation caused a 2.52% SR drop.
- Removing Layer-balanced Loading caused a 1.20× TTFT increase.
Scalability: KEEP's latency scales much slower with the number of retrieved memory segments compared to full recomputation or prefix caching methods.

5. Significance

KEEP represents a significant advancement in Embodied AI by solving the "memory bottleneck" in LLM-based planning.

Efficiency: It enables the use of large context windows (thousands of tokens) without the prohibitive latency costs associated with traditional KV caching or text prompts.
Accuracy: By dynamically identifying and recomputing only the most relevant memory interactions, it maintains high planning accuracy even with aggressive caching.
Hardware Efficiency: The layer-balanced loading strategy provides a blueprint for optimizing memory-bound LLM workloads on heterogeneous hardware (CPU RAM + GPU VRAM).

In summary, KEEP bridges the gap between the theoretical efficiency of KV caching and the practical demands of dynamic, long-horizon embodied tasks, offering a scalable solution for next-generation intelligent agents.