Skip to content

The stored activation DynamicCache is mutated in place by generation, so activation memory grows every turn #2301

Description

@jeojdi1

Re-checked 2026-08-28 against main at 185ebdb ("Dev v2.0.32", #2295). src/memos/memories/activation/kv.py is unchanged since #797 (2025-12-30) and _concat_caches is byte-identical on main and dev-v2.0.30, so everything below still applies at HEAD. The line numbers cite 9119efe, where it was measured.

Environment

MemOS at main / 9119efe; torch 2.13.0, transformers 5.15.1, CPU, Qwen/Qwen2.5-0.5B-Instruct.

Summary

mem_os/core.py:324-326 takes the stored cache object and passes it straight through:

kv_cache = next(iter(mem_cube.act_mem.get_all()), None)
past_key_values = kv_cache.memory if (kv_cache and hasattr(kv_cache, "memory")) else None
...
response = self.chat_llm.generate(current_messages, past_key_values=past_key_values)

HFLLM._prefill (llms/hf.py:305-323) hands that object to self.model(...), and
DynamicLayer.update appends in place (self.keys = torch.cat([self.keys, key_states], dim=-2)).
The stored KVCacheItem.memory therefore accumulates the prompt and the generated tokens of every
turn. _concat_caches returning caches[0] unchanged for a single id (kv.py:208-209) hands out the
same aliased object, and ActivationMemoryManager re-dumps it
(mem_scheduler/memory_manage_modules/activation_memory_manager.py:95), so the growth persists to
disk.

Minimal reproduction

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
M = "Qwen/Qwen2.5-0.5B-Instruct"
tok = AutoTokenizer.from_pretrained(M)
m = AutoModelForCausalLM.from_pretrained(M, dtype=torch.float32).eval()
ids = tok("The user lives in Berlin.", return_tensors="pt").input_ids
with torch.no_grad():
    stored = m(input_ids=ids, use_cache=True, return_dict=True).past_key_values
print("after build:", stored.get_seq_length())
for turn in range(1, 4):
    q = tok(f"\nQ{turn}: tell me something.", return_tensors="pt",
            add_special_tokens=False).input_ids
    with torch.no_grad():
        o = m(input_ids=q, use_cache=True, past_key_values=stored, return_dict=True)
        for _ in range(5):
            t = o.logits[:, -1, :].argmax(-1, keepdim=True)
            o = m(input_ids=t, use_cache=True, past_key_values=o.past_key_values, return_dict=True)
    print(f"after turn {turn}:", stored.get_seq_length(), "| same object:", o.past_key_values is stored)

Observed / Expected

after build: 6
after turn 1: 19 | same object: True
after turn 2: 32 | same object: True
after turn 3: 45 | same object: True

Expected: the stored activation memory stays at 6 — it represents a fixed set of memories, not the
conversation.

Why it matters

Activation memory silently stops being the thing it was built from after the first turn, and grows
without bound across a session.

Suggested fix

Pass a copy into generation (copy.deepcopy(kv_cache.memory), or DynamicCache + crop back to the
recorded length afterwards), and make _concat_caches not return the caller's stored object for
len(caches) == 1.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Labels

area:coreMOS 编排层 / 框架底座 / 跨模块问题status:needs-triageNeeds initial triage | 需要初步判断 & 问题复现

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions