memory: Hybrid context shift #17009
Open
memory: Hybrid context shift #17009
+8
−4
Conversation
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
The recurrent state is always assumed to be the state as of the last update from the final token in the sequence. When doing a partial erasure, if the range does not include the final token, the erasure can be considered a success since any memory used for the sequence prior to the final token (which is no memory) has been successfully removed. There is one potential case that this doesn't address which is the pruning of cache to remove sensitive data from the context. This wouldn't work for attention cache partial removal (in the middle) either since the KV state is linearly-dependent and states in later sequence positions would still be based on the state from the sensitive data, even if that data is no longer cached, so I don't think this is relevant, but it is worth noting that the semantics of this change for a partial erasure in the middle of the cache are essentially "my context is already compressed" and not "all trace of the removed tokens has been removed." ggml-org#16768 Branch: HybridContextShift-16768 Signed-off-by: Gabe Goodhart <[email protected]>
This prefix matching is explicitly attempting to remove the tokens at the end of the sequence that don't match. This is the operation that can't be performed on a recurrent cache due to the state being updated in place, so if this removal fails, we need to clear the whole cache. ggml-org#16768 Branch: HybridContextShift-16768 Signed-off-by: Gabe Goodhart <[email protected]>
compilade
reviewed
Nov 4, 2025
Signed-off-by: Gabe Goodhart <[email protected]> Co-authored-by: compilade <[email protected]>
ggerganov
reviewed
Nov 5, 2025
| // partial intersection is invalid | ||
| if ((0 < p0 && p0 < cell.pos) || (0 < p1 && p1 <= cell.pos)) { | ||
| // partial intersection is invalid if it includes the final pos | ||
| if ((0 < p0 && p0 <= cell.pos && p1 > cell.pos)) { |
There was a problem hiding this comment.
Suggested change
| if ((0 < p0 && p0 <= cell.pos && p1 > cell.pos)) { | |
| if (0 < p0 && p0 <= cell.pos && p1 > cell.pos) { |
| // partial intersection is invalid | ||
| if ((0 < p0 && p0 < cell.pos) || (0 < p1 && p1 <= cell.pos)) { | ||
| // partial intersection is invalid if it includes the final pos | ||
| if ((0 < p0 && p0 <= cell.pos && p1 > cell.pos)) { |
There was a problem hiding this comment.
Why do we check strictly larger than 0 rather than 0 <= p0?
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 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.
Closes #16768
cc @leok7v
Description
This PR addresses context shift failure caused when a hybrid-recurrent model hits its context limit and attempts to perform context shifting. The main change here is to loosen the restriction in the
llama_memory_recurrent::seq_rmto only refuse to do a partial erasure if the part being erased includes the final token in the sequence. Since recurrent states are fixed size, any partial erasure that does not include the final token can be considered a no-op.Testing
To validate the result, you can use the following which artificially limits the context length to force a context shift:
Without this fix, it will fail with
init_batch: failed to prepare attention ubatches, but with this fix, it will successfully continue generating and produce generated output that is relevant to the previous context.