AI Tech Inspire

Stop Clearing the KV Cache? Testing Persistent Attention as Chat Memory

What if a chat agent could carry forward its own attention states instead of calling out to a vector database on every turn? At AI Tech Inspire, we spotted a simple—but intriguing—experiment that keeps the KV cache between turns to act as a form of persistent, in-session memory. The concept is easy to test and, based on early numbers, surprisingly competitive with retrieval.

Quick experiment recap

  • Target use case: multi-turn conversation agents where later turns depend on earlier details.
  • Model: Llama 3 8B. Conversations: 40 turns; turn ~35 requires info from turn ~10.
  • Method: modify the transformers inference loop to keep past_key_values across generate() calls (persistent KV).
  • Manual evaluation: ~50 conversations.
  • Baseline 1: RAG with Chroma + basic embeddings → ~67% success.
  • Baseline 2: Better embeddings (E5-large) + reranking → ~78% success.
  • Persistent KV cache → ~84% success.
  • Latency: persistent KV avoids external retrieval, shaving ~150–200 ms per turn after the first few.
  • Memory cost: ~40 turns × ~200 tokens each → ~3–4 GB of KV cache; scales roughly linearly with tokens.
  • Related work: a GitHub project dubbed “EverMemOS” claims ~92% on a benchmark using compression (not independently verified here).

Why keeping the KV cache can work

In standard generation, a model constructs keys and values for each token and layer, then discards them after finishing the response. Many stacks rebuild context every turn by refeeding the whole prompt. Persistent KV flips that: preserve past_key_values so the model can directly attend to earlier tokens without rebuilding or doing external retrieval. This is different from classic RAG, where embeddings and a vector DB select snippets to re-inject into the prompt.

The upside is obvious: no lossy embedding step, no reranker misses, no extra RPCs. Once the cache is in memory, attention has O(1) access to earlier states. That explains the reported 84% success on long-thread references compared to RAG baselines (67% and 78%), along with lower per-turn latency after the conversation warms up.

Key takeaway: persistent KV trades external retrieval complexity for raw memory. It’s simple, fast after warm-up, and—at least in this test—more reliable on long-thread references.

But memory is the taxman

KV cache size grows with sequence length and model size. A back-of-the-envelope guide: memory ≈ 2 × layers × seq_len × hidden_dim × bytes_per_elem (the factor of 2 is for K and V). The exact footprint varies by architecture, GQA layout, and precision, but the trend is consistent: more tokens, more VRAM/DRAM.

At the reported scale (about 8k tokens total across turns), 3–4 GB is plausible for an 8B model. On a single GPU with CUDA, that can be fine—until multiple concurrent sessions spin up. Server operators will immediately ask: how many concurrent chats can I support before VRAM becomes the bottleneck?

Where persistent KV wins (and where it doesn’t)

  • Shallow knowledge within a session: Perfect when the reference is something the user already said 20 turns ago.
  • Latency-sensitive assistants: Removing vector DB round-trips can reclaim 150–200 ms per turn.
  • Deterministic context: No retrieval-time variance; the model’s attention is the ground truth.

However:

  • Scalability: Linear memory growth makes long sessions or many users expensive.
  • Cross-session memory: KV persistence is local to one session; it’s not a knowledge store.
  • External knowledge: KV can’t surface facts not present in the conversation. RAG still wins for docs, codebases, or product catalogs.
  • Editing the past: If you need to “forget” or correct earlier turns, a pure KV approach can’t surgically remove states without rebuilding.

How to try it (safely)

For those using Hugging Face Transformers with PyTorch, the experiment is approachable:

  • Generate with use_cache=True and capture past_key_values after each turn.
  • On the next user turn, prepend only the new tokens and pass the saved past_key_values back into model.generate().
  • Keep system and formatting tokens stable across turns. If you change the system prompt or template, the cached states won’t match.
  • Log per-turn latency. Expect faster second-plus turns because you skip rebuilding attention on old tokens.

Two practical tips:

  • If you must insert content in the “middle” (e.g., a tool result that chronologically belongs earlier), you’ll likely need to rebuild from raw tokens or switch to a hybrid approach.
  • Pin the cache in GPU memory when possible. Paging to host can negate latency wins.

Compression and memory tricks worth exploring

The GitHub project referenced as “EverMemOS” claims 92% with compression. Beyond that, several techniques are floating around production stacks and research:

  • Quantized KV cache (e.g., 8-bit or 4-bit). Some inference engines quantize KV without retraining.
  • Segmented or paged KV (think “paged attention”), retaining only recent or salient spans.
  • Heuristic pruning: keep states around entities or topics; drop filler turns.
  • Slimmer precision during idle intervals, restore to higher precision when revisiting old spans.
  • Hybrid “memory tokens”: snapshot summaries back into the sequence as learned tokens, then discard raw KV from very old turns.

If you already use frameworks like TensorFlow or have custom kernels on CUDA, you might experiment with kernel-level KV compression. The win is not just footprint—better cache locality can improve throughput.

When to prefer RAG anyway

Even with a promising 84% figure in this test, retrieval shines when:

  • Grounding is required: pulling exact passages from a regulated corpus.
  • Knowledge is large or volatile: product docs, source code, policies.
  • You need auditability: it’s easier to log and explain retrieved chunks than to “explain” attention over opaque KV states.
  • Cross-session continuity matters: a vector DB survives restarts and scales horizontally; KV is ephemeral.

For many assistants, the sweet spot may be hybrid: persistent KV for fluid, in-session coherence, plus RAG for facts and links to source material. This also curbs KV growth by avoiding re-stating large documents in the chat.

Comparisons and adjacent ideas

If you’ve used APIs like GPT, you’ve probably noticed that long prompts become expensive and slow. Persistent KV is a way to keep the cost of long threads in compute memory rather than re-tokenizing and re-encoding everything. On the tooling side, using Transformers with session-level KV feels like a natural extension of what many inference engines already do within a single generation call.

For model devs: architectural advances like grouped-query attention (GQA), sparse attention, and selective memory tokens can reduce KV growth. The ecosystem is also exploring learned summarization tokens and state distillation to compress old context. And in practical deployment, engines that prioritize paged KV and tenant-aware memory pools can make or break concurrency.

Open questions we’re watching

  • How far can KV compression go before quality drops? Is 92% reproducible across diverse tasks?
  • What’s the best policy for pruning? Recency, entities, or an attention-based salience score?
  • Can we combine KV persistence with lightweight, on-the-fly summaries to cap memory?
  • How does this behave with larger models or multimodal setups (e.g., image + text)?
  • What security model is needed so KV from one tenant can’t leak into another?

Try it, measure it, share the traces

Persistent KV is refreshingly simple: don’t press Ctrl+C on your attention states after each turn. It won’t replace RAG for document-heavy agents, but the early numbers suggest it can reduce latency and boost consistency for conversational references that live within the session. If you’re already comfortable with PyTorch and Transformers, it’s an evening experiment that could pay off in production.

As always, your mileage will vary by domain and model. Log per-turn latencies, track memory usage, and A/B against your current RAG pipeline. If you discover that a compressed or hybrid KV policy gets you most of the gains without blowing up VRAM, that’s a compelling path forward.

AI Tech Inspire will keep an eye on follow-up experiments and any open-source releases that package this approach cleanly. In the meantime, the question is simple: if your agent already has the attention it needs, why keep sending it on a retrieval detour?

Recommended Resources

As an Amazon Associate, I earn from qualifying purchases.