Support memdecode #4767
Open
lvhan028 wants to merge 44 commits into
Open
Conversation
Add separate memory state caches for linear-attention layers, reserve GPU memory for base and memory SSM pools, and load the dense memory model without base quantization. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds MemDecode support to the lmdeploy PyTorch engine by introducing a memory-model agent plus a base/memory logits fusion module, and wiring the feature through engine config parsing and runtime cache budgeting/lifecycle management.
Changes:
- Add
MemDecodeConfig,MemDecodeFusion(fixed + adaptive router), andMemDecodeAgentfor running a secondary “memory” model and fusing outputs. - Integrate MemDecode into the PyTorch engine/model agent lifecycle (build, forward path, cache swapping, reset/release) and executor cache sizing/validation.
- Add MemDecode-focused unit tests and an end-to-end smoke-test example script.
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
lmdeploy/pytorch/config.py |
Introduces MemDecodeConfig and threads it through MiscConfig. |
lmdeploy/pytorch/engine/config_builder.py |
Parses MemDecode options from hf_overrides and packaged memory_fusion/config.json. |
lmdeploy/pytorch/engine/engine.py |
Builds and attaches memdecode_config; validates incompatibility with speculative decoding. |
lmdeploy/pytorch/engine/model_agent/agent.py |
Builds/uses memdecode_agent in forward path; manages caches and lifecycle hooks. |
lmdeploy/pytorch/engine/executor/base.py |
Accounts for memory-model KV/state cache memory; adds MemDecode validation. |
lmdeploy/pytorch/memdecode/fusion.py |
New fusion module (fixed lambda + adaptive router) and vocab alignment helper. |
lmdeploy/pytorch/memdecode/agent.py |
New memory-model agent that builds/runs the memory model and performs fusion. |
lmdeploy/pytorch/memdecode/__init__.py |
Exposes MemDecode public API symbols. |
tests/pytorch/memdecode/test_fusion.py |
Adds unit tests for fusion alignment and adaptive routing behaviors. |
tests/pytorch/memdecode/test_agent.py |
Adds unit tests for MemDecodeAgent build/forward behavior and context handling. |
tests/pytorch/engine/test_model_agent.py |
Extends model-agent tests to cover MemDecode lifecycle and forward behavior. |
tests/pytorch/engine/test_executor_base.py |
Extends executor tests for MemDecode cache sizing and validation. |
lmdeploy/vl/model/qwen3_5.py |
Refactors InternS2 arch list handling for Qwen3.5 model wrapper. |
lmdeploy/pytorch/models/module_map.py |
Cleans up InternS2Preview module map updates. |
examples/memdec/main.py |
Adds a runnable MemDecode API server + chat completion smoke test script. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+156
to
+177
| base_logits = align_logits_to_base(base_logits, self.base_vocab_size) | ||
| memory_logits = align_logits_to_base(memory_logits, self.base_vocab_size) | ||
| base_log_probs = torch.log_softmax(base_logits, dim=-1) | ||
| memory_log_probs = torch.log_softmax(memory_logits, dim=-1) | ||
|
|
||
| if self.adaptive_router: | ||
| return self._adaptive_fusion( | ||
| base_log_probs, | ||
| memory_log_probs, | ||
| base_hidden_states=base_hidden_states, | ||
| memory_hidden_states=memory_hidden_states, | ||
| ) | ||
|
|
||
| if self.lambda_value == 0.0: | ||
| return base_log_probs | ||
| if self.lambda_value == 1.0: | ||
| return memory_log_probs | ||
|
|
||
| base_log_weight = base_log_probs.new_tensor(math.log1p(-self.lambda_value)) | ||
| memory_log_weight = memory_log_probs.new_tensor(math.log(self.lambda_value)) | ||
| fused = torch.logaddexp(base_log_probs + base_log_weight, memory_log_probs + memory_log_weight) | ||
| return fused |
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
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.
No description provided.