Fix #2313: Merged KV caches report get_seq_length() == 0 and are silently discarded on tran - #2316
Open
Memtensor-AI wants to merge 2 commits into
Open
Conversation
`KVCacheMemory._concat_caches` built each merged layer via bare `layer_cls()` + direct `.keys` / `.values` assignment, which bypasses `DynamicLayer.lazy_initialization`. From transformers 4.57 onwards `DynamicLayer.get_seq_length()` short-circuits to 0 when `is_initialized` is False, so the merged cache reported length 0 and the first forward pass silently discarded every merged token — with no exception, warning, or log. Any user of `get_cache([...])` on an install inside the declared `transformers>=4.51.3,<5.0.0` range hit silent data loss. Fix: invoke `merged_layer.lazy_initialization(concat_keys)` before assigning the concatenated tensors, guarded by hasattr so older layouts and the legacy `key_cache` branch stay untouched. Also repair the pre-existing test fixture (`make_filled_cache` was still appending to `cache.key_cache`, which was removed in 4.57, so two of the four tests errored with AttributeError before this patch ever ran) and add a regression test that asserts `merged.get_seq_length() == L1 + L2` and `is_initialized is True`. Fixes MemTensor#2313
Memtensor-AI
added
ai:generated
Collaborator
Author
🤖 Open Code ReviewTarget: PR #2316 ✅ OpenCodeReview: Review complete: 0 finding(s) across 2 selected item(s). Generated by cloud-assistant via Open Code Review. |
Collaborator
Author
🔧 Open Code Review requested Agent fixOpen Code Review found 1 issue(s). I have resumed the development Agent to fix them.
The Agent will push a new commit to this PR branch. OCR will recheck after the commit is pushed. |
Address Open Code Review finding on the regression test added in c759b60. `merged.layers[0]` was indexed without first asserting the list is non-empty, so a regression in `_concat_caches` returning a DynamicCache with an empty `layers` list would raise an opaque IndexError instead of a descriptive AssertionError. Mirror the guard already used by test_get_cache_merge for a consistent, self-describing failure message.
Collaborator
Author
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fixes issue #2313: on
transformers >= 4.57,KVCacheMemory.get_cache([...])returned a mergedDynamicCachethat reportedget_seq_length() == 0because_concat_cachesbuilt each output layer via barelayer_cls()+ direct.keys/.valuesassignment, which bypassesDynamicLayer.lazy_initialization(). From 4.57 onwardDynamicLayer.get_seq_length()short-circuits to 0 whenis_initializedis False, so the attention mask was sized for the query alone andDynamicLayer.update()silently overwrote the merged tensors with empty ones on the first forward pass — every merged token vanished with no exception, warning, or log line. The declaredtransformers>=4.51.3,<5.0.0range inpyproject.tomladmits 4.57.x, so this affected a supported install matrix.The fix in
src/memos/memories/activation/kv.pyinvokesmerged_layer.lazy_initialization(concat_keys)before assigning the concatenated tensors on the newlayersbranch, guarded byhasattr(merged_layer, "lazy_initialization")andgetattr(merged_layer, "is_initialized", False)so the legacykey_cachebranch and older layer implementations stay untouched. This routes through the public API path that upstream uses, sodtype,device, andis_initializedare all set exactly astransformersexpects.tests/memories/activation/test_kv.pyalso had two pre-existing failures (AttributeError: 'DynamicCache' object has no attribute 'key_cache') becausemake_filled_cachestill calledcache.key_cache.append(...), an attribute removed in the same 4.57 refactor. The fixture is rewritten to useDynamicCache.update(...)with ahasattr(cache, "layers")fork for legacy fallback, and a new regression testtest_concat_layer_reports_initialized_and_full_seq_lengthassertsmerged.get_seq_length() == 2 * seq_len,keys.shape[-2] == 2 * seq_len, andis_initialized is True. This test fails against the pre-fix implementation and passes with the fix.Verification:
python3 -m pytest tests/memories/activation/test_kv.py -q→ 5 passed, 0 failed ontransformers==4.57.6. Ruff format and ruff check both clean on the two modified files. A standalone reproducer built directly from the issue also confirmsget_seq_length == 10(expected 10) andis_initialized == Trueafter the fix. Only two files changed (src/memos/memories/activation/kv.py,tests/memories/activation/test_kv.py);.ai-tasks/andopenspec/changes/were intentionally left untracked on the working branch and archived separately to the memos-autodev-specs repo.Related Issue (Required): Fixes #2313
Type of change
Please delete options that are not relevant.
How Has This Been Tested?
Not run; documentation-only change.
Checklist
@bittergreen please review this PR.
Reviewer Checklist