Repository navigation
fix(activation): keep the stored KV cache read-only during generation (#2301) - #2552
RerankerGuo wants to merge 2 commits into
Conversation
Memtensor-AI
added
area:memory
🤖 Open Code ReviewTarget: PR #2552 🔍 OpenCodeReview found 4 issue(s) in this PR. 1.
|
|
575d34c to
5d4ce40
Compare
|
Thanks for the review — both findings are addressed in 1. Empty cache / layout selection — fixedThe reported mechanism is real for a cache that exposes the modern I could not reproduce it on the versions in the supported range, because That is exactly why the check is now made explicit instead of relying on layers = getattr(cache, "layers", None)
if layers is not None:
# The modern layout. The check is on the attribute itself and not on
# its content: an empty cache still has to be copied into an
# independent one, because the model appends its first layer to the
# cache it receives.
cloned = copy.copy(cache)
cloned.layers = [_clone_cache_layer(layer) for layer in layers]
return clonedThe modern layout now wins whenever the attribute exists, and the
2. Memory cost of the copy — documented, with the reasoningAgreed on the trade-off; it is now spelled out in
I kept the copy rather than the crop-back alternative for two reasons: If the activation memory ever grows large enough for the extra copy to Verification after the changepython -m pytest -q tests/memories/activation/ tests/llms/test_hf.py
# 26 passed, 1 warning in 4.92s
env HF_HUB_OFFLINE=1 python -m pytest -q tests/configs/test_base.py tests/test_utils_timing.py
# 38 passed in 4.38s
python -m pytest -q tests
# 2 failed, 882 passed, 9 skipped, 20 warnings in 349.59s (0:05:49)
# the same two pre-existing failures as before, unchanged by this PR
ruff check src/memos tests
# All checks passed!
ruff format --check src/memos tests
# 527 files already formatted |
5d4ce40 to
3827e0d
Compare
|
3827e0d to
6f4b12b
Compare
|
Final status of this PR (head Upstream CI: green. All 17 jobs of the The CI installs the AutoTest: "ENV ISSUE", not a test failure. The AutoTest runner has no Both modules already imported # The KV cache is a torch structure; skip the whole module where torch is
# not installed instead of failing the collection of the test suite.
torch = pytest.importorskip("torch")With |
|
6f4b12b to
59fb97a
Compare
|
Both findings are addressed in 1. Double clone on the
|
✅ Automated Test Results: PASSEDRequired tests passed (25/25 executed, 0 skipped). Duration: 26s Branch: Tested source: |
Memtensor-AI
added
status:ready
…MemTensor#2301) The activation memory stored in `KVCacheMemory` was handed to the model as the very same object, and `transformers` appends the newly generated key/value tensors to that object **in place** (`DynamicLayer.update` rebinds `keys`/`values` on the same layer object). So the stored activation memory grew with the prompt and the generated tokens of every turn, and `ActivationMemoryManager` wrote the growth back to disk on the next dump. The stored cache was supposed to represent a fixed set of memories, not the conversation. Three hops were involved: 1. `MemOS.chat` (`mem_os/core.py`), `mem_os/main.py`, `mem_chat/simple.py` and the scheduler analyzer pass the stored `kv_cache.memory` straight into `HFLLM.generate(past_key_values=...)`; 2. `HFLLM._prefill` forwards it to `model(past_key_values=...)`, and transformers appends the new tokens' K/V in place; 3. `KVCacheMemory._concat_caches` returns `caches[0]` unchanged for a single id, so the `get_cache()` merge path hands out the stored object too. This commit treats stored caches as read-only by construction, at the two boundaries where a cache leaves the store or enters the model: - `clone_dynamic_cache()` (new, `memories/activation/kv.py`): an independent copy with cloned K/V tensors. The modern layout (`layers`, transformers >= 4.55) is preferred whenever the attribute exists, so that the copy never shares its layer list with the original, even for an empty cache that the model fills with its first layer. The legacy layout (`key_cache`/`value_cache` lists, transformers < 4.55) is copied as well. On a layer object, `keys`/`values` and `key_cache`/`value_cache` are copied independently of each other, so a layer that populates only one side does not lose data. The layer objects are shallow-copied, so all non-tensor state (for example the initialization flag that `DynamicLayer.get_seq_length()` relies on) is preserved. - `KVCacheMemory._concat_caches`: the single-cache case now returns a clone instead of the stored object. - `HFLLM.generate` / `HFLLM.generate_stream`: clone the incoming `past_key_values` once at the boundary. This covers all four call sites without touching them; the cost is one cache copy per generate call, the same order as the prefill work itself. No production dependency, configuration, schema or API change. Tests: - `tests/llms/test_hf.py`: `test_generate_does_not_mutate_caller_cache` and `test_generate_stream_does_not_mutate_caller_cache` use a stand-in model that appends K/V in place, and assert that the model receives a copy and that the caller's cache still holds only the original tokens. Both fail before the fix (the model receives the caller's cache itself) and pass after it. - `tests/memories/activation/test_kv.py`: `test_clone_dynamic_cache_does_not_share_tensors`, `test_get_cache_single_item_returns_independent_copy`, `test_get_cache_merge_does_not_alias_inputs`, `test_clone_dynamic_cache_handles_empty_cache`, `test_clone_dynamic_cache_prefers_the_layers_layout`, `test_clone_dynamic_cache_handles_legacy_layout` and `test_clone_dynamic_cache_copies_layers_with_legacy_attribute_names` cover the store boundary and both cache layouts. - `tests/utils.py`: shared helpers to build a cache and to read it independently of the cache layout. The layout-independent helpers also replace the `key_cache`-only helper in `tests/memories/activation/test_kv.py`, which cannot build a populated cache at all with transformers >= 4.55. The two touched test modules now import torch through `pytest.importorskip("torch")`, so that a test environment without torch skips them instead of failing during collection. Fixes MemTensor#2301.
The activation memory hands out an owned copy of the stored cache (`get_cache()`), and `HFLLM.generate` copied that copy again before it was handed to the model. A caller that follows the public API therefore paid for two full copies of every K/V tensor and kept three sets alive during the generation: the stored cache, the copy of the copy, and the one the model writes into. Ownership is now explicit. A cache that leaves the activation memory is an *owned copy*: it shares no tensor with the stored cache, so it may be handed to a model that appends to it in place. `clone_dynamic_cache()` and the merge path of `_concat_caches()` mark the cache they return, and the model boundary skips the copy for a cache that is already owned. A cache that is stored back into the memory loses the marker, so the model still copies it before it appends to it. The stored cache stays read-only in every case: the model never writes into a cache that aliases it, and the number of K/V tensor sets alive during a generation drops from three to two. Tests: - `tests/llms/test_hf.py::test_generation_from_get_cache_copies_the_kv_tensors_only_once` walks the public path (`get_cache()` -> `generate()`), counts the copies (exactly one) and pins that the model received the cache the memory handed out instead of a second copy of it. - `tests/llms/test_hf.py::test_generation_from_get_cache_keeps_the_stored_cache_read_only` pins that the stored cache is not rewritten, also when an owned copy is stored back into the memory. - `tests/memories/activation/test_kv.py` pins the ownership marker on the single-cache and on the merge path, and that a stored cache is never an owned copy. Fixes MemTensor#2301.
59fb97a to
55d2d84
Compare
Memtensor-AI
added
status:in-progress
Description
Fixes #2301.
The activation memory stored in
KVCacheMemorywas handed to the model asthe very same object, and
transformersappends the newly generatedkey/value tensors to that object in place (
DynamicLayer.updaterebindskeys/valueson the same layer object).So the stored activation memory grew with the prompt and the generated tokens
of every turn, and
ActivationMemoryManagerwrote the growth back to disk onthe next dump. The stored cache is supposed to represent a fixed set of
memories, not the conversation.
Three hops were involved:
MemOS.chat(mem_os/core.py),mem_os/main.py,mem_chat/simple.pyandthe scheduler analyzer pass the stored
kv_cache.memorystraight intoHFLLM.generate(past_key_values=...);HFLLM._prefillforwards it tomodel(past_key_values=...), which appendsthe new tokens' K/V in place;
KVCacheMemory._concat_cachesreturnscaches[0]unchanged for a singleid, so the
get_cache()merge path hands out the stored object as well.This PR treats the stored caches as read-only by construction, at the two
boundaries where a cache leaves the store or enters the model:
clone_dynamic_cache()(new,memories/activation/kv.py): anindependent copy with cloned K/V tensors. The modern layout (
layers,transformers >= 4.55, the version pinned by
uv.lock) is preferred wheneverthe attribute exists, so that the copy never shares its layer list with the
original — not even for an empty cache, which the model fills with its first
layer. The legacy layout (
key_cache/value_cachelists, transformers< 4.55, the version pinned by
poetry.lock) is copied as well. On a layerobject,
keys/valuesandkey_cache/value_cacheare copiedindependently of each other, so a layer that populates only one side does
not lose data. The layer objects are shallow-copied, so all non-tensor state
(for example the initialization flag that
DynamicLayer.get_seq_length()relies on) is preserved.
KVCacheMemory._concat_caches: the single-cache case now returns a cloneinstead of the stored object.
HFLLM.generate/HFLLM.generate_stream: clone the incomingpast_key_valuesonce at the boundary. This covers all four call siteswithout touching them; the cost is one cache copy per generate call, the
same order as the prefill work itself. The trade-off is documented in
HFLLM._clone_cache.No production dependency, configuration, schema or API change.
Type of change
How Has This Been Tested?
Commands were run in a pinned CPU environment (Python 3.11.15) whose
dependency versions match the repository's
uv.lock(torch 2.14.1+cpu,transformers 4.57.6) and which was installed with a frozen, hash-checked
install:
New regression tests, and how they behave without the fix:
tests/llms/test_hf.py::test_generate_does_not_mutate_caller_cacheand::test_generate_stream_does_not_mutate_caller_cacheuse a stand-in modelthat appends K/V in place. They assert that the model receives a copy and
that the caller's cache still holds only the original tokens. Before the fix
both fail with
AssertionError: unexpectedly identical: DynamicCache(layers=[DynamicLayer])(2 failed, 6 passed); after the fix they pass.
tests/memories/activation/test_kv.py::test_get_cache_single_item_returns_independent_copyfails before the fix with
AssertionError: assert DynamicCache(layers=[DynamicLayer]) is not DynamicCache(layers=[DynamicLayer])and passes after it.
tests/memories/activation/test_kv.py::test_clone_dynamic_cache_does_not_share_tensors,::test_get_cache_merge_does_not_alias_inputs,::test_clone_dynamic_cache_handles_empty_cache,::test_clone_dynamic_cache_prefers_the_layers_layout,::test_clone_dynamic_cache_handles_legacy_layoutand::test_clone_dynamic_cache_copies_layers_with_legacy_attribute_namescoverthe store boundary and both cache layouts.
::test_clone_dynamic_cache_prefers_the_layers_layoutfails on theprevious implementation with
AssertionError: assert [] is not [](1 failed,10 passed) and passes after it.
The layout-independent helpers in
tests/utils.pyalso replace thekey_cache-only fixture intests/memories/activation/test_kv.py. Thatfixture cannot build a populated cache at all with transformers >= 4.55
(
AttributeError: 'DynamicCache' object has no attribute 'key_cache'), whichmade
test_get_cache_mergeandtest_delete_and_get_allfail on everycurrent install; they now pass on both cache layouts.
The two touched test modules now import torch through
pytest.importorskip("torch"), so that a test environment without torchskips them instead of failing during collection. With torch installed, as in
the repository CI, nothing changes.
Upstream CI
All 17 jobs of the
Python testsworkflow pass on the final head (Linux /Windows / macOS 14 / macOS 15 × Python 3.10–3.13, plus
sync); one job logsthe pytest summary
884 passed, 9 skipped, 37 warnings in 68.33s. The CIinstalls the
allextra, so the KV cache tests run there with thepoetry.lockversion oftransformers(4.53.2, the legacykey_cache/value_cachelayout), while the environment used above hastransformers 4.57.6fromuv.lock(the modernlayerslayout) — both cachelayouts are covered.
The AutoTest check reports
ENV ISSUErather than a test failure: thatrunner has no
torchinstalled at all, so the two test modules this PRtouches cannot run there — before the
pytest.importorskipchange theyfailed during collection, now they are skipped and the runner reports that
no test executed. Both modules already imported
torchat module levelbefore this PR; the repository CI installs the
allextra and runs them.Notes for the reviewer
CONTRIBUTING.mdasks for PRs againstdev, but thisrepository has no
devbranch — onlymainanddev-v2.0.28…dev-v2.0.36. This PR targetsmain, the default branch, andsrc/memos/memories/activation/kv.pyandsrc/memos/llms/hf.pyarebyte-identical on
mainanddev-v2.0.36, so retargeting is trivial.Happy to retarget to
dev-v2.0.36if that is preferred._concat_cacheshas a separate,still-open defect: the merged layers are built with a bare
layer_cls()anddirect
.keys/.valuesassignment, so on transformers >= 4.55 they arenever marked initialized and
get_seq_length()reports0, which makes themodel drop every merged token. That is issue Merged KV caches report get_seq_length() == 0 and are silently discarded on transformers >= 4.57 #2313 (auto-closed as
not_plannedafter inactivity) and it is deliberately not touched here,to keep this PR to one logical change.
Checklist
Reviewer Checklist
repository CI matrix is tracked separately on this PR.